feat: add agent types

letta/client/client.py

@@ -9,7 +9,7 @@
 from letta.data_sources.connectors import DataConnector
 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 import AgentState, CreateAgent, UpdateAgentState, AgentType
 from letta.schemas.block import (
     Block,
     CreateBlock,
@@ -68,6 +68,7 @@ def agent_exists(self, agent_id: Optional[str] = None, agent_name: Optional[str]
     def create_agent(
         self,
         name: Optional[str] = None,
+        agent_type: Optional[AgentType] = 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)),
@@ -319,6 +320,8 @@ def agent_exists(self, agent_id: str) -> bool:
     def create_agent(
         self,
         name: Optional[str] = None,
+        # agent config
+        agent_type: Optional[AgentType] = None,
         # model configs
         embedding_config: EmbeddingConfig = None,
         llm_config: LLMConfig = None,
@@ -381,6 +384,7 @@ def create_agent(
             memory=memory,
             tools=tool_names,
             system=system,
+            agent_type=agent_type if agent_type else AgentType.base_agent,
             llm_config=llm_config if llm_config else self._default_llm_config,
             embedding_config=embedding_config if embedding_config else self._default_embedding_config,
         )
@@ -1462,6 +1466,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_type: Optional[AgentType] = None,
         # model configs
         embedding_config: EmbeddingConfig = None,
         llm_config: LLMConfig = None,
@@ -1524,6 +1530,7 @@ def create_agent(
                 memory=memory,
                 tools=tool_names,
                 system=system,
+                agent_type=agent_type if agent_type else AgentType.base_agent,
                 llm_config=llm_config if llm_config else self._default_llm_config,
                 embedding_config=embedding_config if embedding_config else self._default_embedding_config,
             ),

letta/metadata.py

@@ -218,6 +218,7 @@ class AgentModel(Base):
     tools = Column(JSON)
 
     # configs
+    agent_type = Column(String)
     llm_config = Column(LLMConfigColumn)
     embedding_config = Column(EmbeddingConfigColumn)
 
@@ -243,6 +244,7 @@ def to_record(self) -> AgentState:
             memory=Memory.load(self.memory),  # load dictionary
             system=self.system,
             tools=self.tools,
+            agent_type=self.agent_type,
             llm_config=self.llm_config,
             embedding_config=self.embedding_config,
             metadata_=self.metadata_,

letta/schemas/agent.py

@@ -1,5 +1,6 @@
 import uuid
 from datetime import datetime
+from enum import Enum
 from typing import Dict, List, Optional, Union
 
 from pydantic import BaseModel, Field, field_validator
@@ -21,6 +22,15 @@ class BaseAgent(LettaBase, validate_assignment=True):
     user_id: Optional[str] = Field(None, description="The user id of the agent.")
 
 
+class AgentType(str, Enum):
+    """
+    Enum to represent the type of agent.
+    """
+
+    base_agent = "base_agent"
+    split_thread_agent = "split_thread_agent"
+
+
 class AgentState(BaseAgent):
     """
     Representation of an agent's state. This is the state of the agent at a given time, and is persisted in the DB backend. The state has all the information needed to recreate a persisted agent.
@@ -52,6 +62,9 @@ class AgentState(BaseAgent):
     # system prompt
     system: str = Field(..., description="The system prompt used by the agent.")
 
+    # agent configuration
+    agent_type: AgentType = Field(..., description="The type of 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 +77,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_type: Optional[AgentType] = Field(None, description="The type of 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/server/server.py

@@ -11,6 +11,7 @@
 from fastapi import HTTPException
 
 import letta.constants as constants
+from letta.schemas.agent import AgentType
 import letta.server.utils as server_utils
 import letta.system as system
 from letta.agent import Agent, save_agent
@@ -335,7 +336,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_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}")
@@ -787,6 +791,7 @@ def create_agent(
                 name=request.name,
                 user_id=user_id,
                 tools=request.tools if request.tools else [],
+                agent_type=request.agent_type or AgentType.base_agent,
                 llm_config=llm_config,
                 embedding_config=embedding_config,
                 system=request.system,

tests/test_agent_type.py

@@ -0,0 +1,42 @@
+import json
+import pytest
+
+from letta.schemas.agent import AgentType
+from letta.client.client import create_client
+from letta.schemas.embedding_config import EmbeddingConfig
+from letta.schemas.llm_config import LLMConfig
+
+
+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_type=agent_type,
+            llm_config=LLMConfig.default_config("gpt-4"),
+            embedding_config=EmbeddingConfig.default_config(provider="openai"),
+        )
+    else:
+        agent_state = client.create_agent(
+            llm_config=LLMConfig.default_config("gpt-4"),
+            embedding_config=EmbeddingConfig.default_config(provider="openai"),
+        )
+
+    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)