feat: add agent types

letta/client/client.py

@@ -10,6 +10,7 @@
 from letta.functions.functions import parse_source_code
 from letta.memory import get_memory_functions
 from letta.schemas.agent import AgentState, CreateAgent, UpdateAgentState
+from letta.schemas.agent_config import AgentConfig
 from letta.schemas.block import (
     Block,
     CreateBlock,
@@ -68,6 +69,7 @@ def agent_exists(self, agent_id: Optional[str] = None, agent_name: Optional[str]
     def create_agent(
         self,
         name: Optional[str] = None,
+        agent_config: Optional[AgentConfig] = None,
         embedding_config: Optional[EmbeddingConfig] = None,
         llm_config: Optional[LLMConfig] = None,
         memory: Memory = ChatMemory(human=get_human_text(DEFAULT_HUMAN), persona=get_persona_text(DEFAULT_PERSONA)),
@@ -314,6 +316,8 @@ def agent_exists(self, agent_id: str) -> bool:
     def create_agent(
         self,
         name: Optional[str] = None,
+        # agent config
+        agent_config: Optional[AgentConfig] = None,
         # model configs
         embedding_config: Optional[EmbeddingConfig] = None,
         llm_config: Optional[LLMConfig] = None,
@@ -372,6 +376,7 @@ def create_agent(
             memory=memory,
             tools=tool_names,
             system=system,
+            agent_config=agent_config,
             llm_config=llm_config,
             embedding_config=embedding_config,
         )
@@ -1404,6 +1409,8 @@ def agent_exists(self, agent_id: Optional[str] = None, agent_name: Optional[str]
     def create_agent(
         self,
         name: Optional[str] = None,
+        # agent config
+        agent_config: Optional[AgentConfig] = None,
         # model configs
         embedding_config: Optional[EmbeddingConfig] = None,
         llm_config: Optional[LLMConfig] = None,
@@ -1462,6 +1469,7 @@ def create_agent(
                 memory=memory,
                 tools=tool_names,
                 system=system,
+                agent_config=agent_config,
                 llm_config=llm_config,
                 embedding_config=embedding_config,
             ),

letta/metadata.py

@@ -21,6 +21,7 @@
 
 from letta.config import LettaConfig
 from letta.schemas.agent import AgentState
+from letta.schemas.agent_config import AgentConfig
 from letta.schemas.api_key import APIKey
 from letta.schemas.block import Block, Human, Persona
 from letta.schemas.embedding_config import EmbeddingConfig
@@ -39,6 +40,28 @@
 Base = declarative_base()
 
 
+class AgentConfigColumn(TypeDecorator):
+    """Custom type for storing AgentConfig as JSON"""
+
+    impl = JSON
+    cache_ok = True
+
+    def load_dialect_impl(self, dialect):
+        return dialect.type_descriptor(JSON())
+
+    def process_bind_param(self, value, dialect):
+        if value:
+            # return vars(value)
+            if isinstance(value, AgentConfig):
+                return value.model_dump()
+        return value
+
+    def process_result_value(self, value, dialect):
+        if value:
+            return AgentConfig(**value)
+        return value
+
+
 class LLMConfigColumn(TypeDecorator):
     """Custom type for storing LLMConfig as JSON"""
 
@@ -206,6 +229,7 @@ class AgentModel(Base):
     tools = Column(JSON)
 
     # configs
+    agent_config = Column(AgentConfigColumn)
     llm_config = Column(LLMConfigColumn)
     embedding_config = Column(EmbeddingConfigColumn)
 
@@ -231,6 +255,7 @@ def to_record(self) -> AgentState:
             memory=Memory.load(self.memory),  # load dictionary
             system=self.system,
             tools=self.tools,
+            agent_config=self.agent_config,
             llm_config=self.llm_config,
             embedding_config=self.embedding_config,
             metadata_=self.metadata_,

letta/schemas/agent.py

@@ -4,6 +4,7 @@
 
 from pydantic import BaseModel, Field, field_validator
 
+from letta.schemas.agent_config import AgentConfig
 from letta.schemas.embedding_config import EmbeddingConfig
 from letta.schemas.letta_base import LettaBase
 from letta.schemas.llm_config import LLMConfig
@@ -52,6 +53,9 @@ class AgentState(BaseAgent):
     # system prompt
     system: str = Field(..., description="The system prompt used by the agent.")
 
+    # agent config
+    agent_config: AgentConfig = Field(..., description="The agent configuration used by the agent.")
+
     # llm information
     llm_config: LLMConfig = Field(..., description="The LLM configuration used by the agent.")
     embedding_config: EmbeddingConfig = Field(..., description="The embedding configuration used by the agent.")
@@ -64,6 +68,7 @@ class CreateAgent(BaseAgent):
     memory: Optional[Memory] = Field(None, description="The in-context memory of the agent.")
     tools: Optional[List[str]] = Field(None, description="The tools used by the agent.")
     system: Optional[str] = Field(None, description="The system prompt used by the agent.")
+    agent_config: Optional[AgentConfig] = Field(None, description="The agent configuration used by the agent.")
     llm_config: Optional[LLMConfig] = Field(None, description="The LLM configuration used by the agent.")
     embedding_config: Optional[EmbeddingConfig] = Field(None, description="The embedding configuration used by the agent.")
 

letta/schemas/agent_config.py

@@ -0,0 +1,24 @@
+from typing import Optional
+from enum import Enum
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class AgentType(str, Enum):
+    base_agent = "base_agent"
+    split_thread_agent = "split_thread_agent"
+
+
+class AgentConfig(BaseModel):
+    """
+    Configuration for a Letta agent. This object specifies all the information necessary to access an agent to usage with Letta, except for secret keys.
+
+    Attributes:
+        agent_type (str): The type of agent.
+    """
+
+    agent_type: AgentType = Field(..., description="The type of agent.")
+
+    @classmethod
+    def default_config(cls):
+        return cls(agent_type=AgentType.base_agent)

letta/server/server.py

@@ -11,6 +11,7 @@
 from fastapi import HTTPException
 
 import letta.constants as constants
+from letta.schemas.agent_config import AgentType
 import letta.server.utils as server_utils
 import letta.system as system
 from letta.agent import Agent, save_agent
@@ -259,6 +260,7 @@ def __init__(
 
         # Generate default LLM/Embedding configs for the server
         # TODO: we may also want to do the same thing with default persona/human/etc.
+        self.server_agent_config = settings.agent_config
         self.server_llm_config = settings.llm_config
         self.server_embedding_config = settings.embedding_config
         # self.server_llm_config = LLMConfig(
@@ -357,7 +359,10 @@ def _load_agent(self, user_id: str, agent_id: str, interface: Union[AgentInterfa
             # Make sure the memory is a memory object
             assert isinstance(agent_state.memory, Memory)
 
-            letta_agent = Agent(agent_state=agent_state, interface=interface, tools=tool_objs)
+            if agent_state.agent_config.agent_type == AgentType.base_agent:
+                letta_agent = Agent(agent_state=agent_state, interface=interface, tools=tool_objs)
+            else:
+                raise NotImplementedError("Only base agents are supported as of right now!")
 
             # Add the agent to the in-memory store and return its reference
             logger.debug(f"Adding agent to the agent cache: user_id={user_id}, agent_id={agent_id}")
@@ -768,6 +773,7 @@ def create_agent(
             # model configuration
             llm_config = request.llm_config if request.llm_config else self.server_llm_config
             embedding_config = request.embedding_config if request.embedding_config else self.server_embedding_config
+            agent_config = request.agent_config if request.agent_config else self.server_agent_config
 
             # get tools + make sure they exist
             tool_objs = []
@@ -809,6 +815,7 @@ def create_agent(
                 name=request.name,
                 user_id=user_id,
                 tools=request.tools if request.tools else [],
+                agent_config=agent_config,
                 llm_config=llm_config,
                 embedding_config=embedding_config,
                 system=request.system,

letta/settings.py

@@ -5,6 +5,7 @@
 from pydantic import Field
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
+from letta.schemas.agent_config import AgentConfig
 from letta.schemas.embedding_config import EmbeddingConfig
 from letta.schemas.llm_config import LLMConfig
 from letta.utils import printd
@@ -25,6 +26,9 @@ class Settings(BaseSettings):
     pg_port: Optional[int] = None
     pg_uri: Optional[str] = None  # option to specifiy full uri
 
+    # agent configuration
+    agent_type: Optional[str] = None
+
     # llm configuration
     llm_endpoint: Optional[str] = None
     llm_endpoint_type: Optional[str] = None
@@ -38,6 +42,13 @@ class Settings(BaseSettings):
     embedding_model: Optional[str] = None
     embedding_chunk_size: int = 300
 
+    @property
+    def agent_config(self):
+        if self.agent_type:
+            return AgentConfig(agent_type=self.agent_type)
+        else:
+            return AgentConfig.default_config()
+
     @property
     def llm_config(self):
 

tests/test_agent_config.py

@@ -0,0 +1,33 @@
+import json
+import pytest
+
+from letta.schemas.agent_config import AgentConfig, AgentType
+from letta.client.client import create_client
+
+
+def test_agent_creation(agent_type):
+    client = create_client()
+    assert client is not None
+
+    if agent_type is not None:
+        agent_state = client.create_agent(agent_config=AgentConfig(agent_type=agent_type))
+    else:
+        agent_state = client.create_agent()
+
+    agent = client.get_agent(agent_id=agent_state.id)
+    assert agent is not None
+
+    response = client.user_message(agent_id=agent_state.id, message="My name is Vivek.")
+    assert response is not None
+
+    print(f"Successfully created a agent of type {agent_type}!")
+
+
+if __name__ == "__main__":
+    # Test normal agent creation
+    test_agent_creation(None)
+
+    # Test agent creation with agent type
+    test_agent_creation(AgentType.base_agent)
+    with pytest.raises(NotImplementedError):
+        test_agent_creation(AgentType.split_thread_agent)