letta-nightly 0.6.39.dev20250313162623__py3-none-any.whl → 0.6.40.dev20250314173529__py3-none-any.whl

letta/server/server.py

@@ -19,16 +19,18 @@ import letta.system as system
 from letta.agent import Agent, save_agent
 from letta.config import LettaConfig
 from letta.data_sources.connectors import DataConnector, load_data
+from letta.dynamic_multi_agent import DynamicMultiAgent
 from letta.helpers.datetime_helpers import get_utc_time
 from letta.helpers.json_helpers import json_dumps, json_loads
 from letta.helpers.mcp_helpers import (
+    MCP_CONFIG_TOPLEVEL_KEY,
     BaseMCPClient,
-    LocalMCPClient,
-    LocalServerConfig,
     MCPServerType,
     MCPTool,
     SSEMCPClient,
     SSEServerConfig,
+    StdioMCPClient,
+    StdioServerConfig,
 )
 
 # TODO use custom interface
@@ -37,6 +39,7 @@ from letta.interface import CLIInterface  # for printing to terminal
 from letta.log import get_logger
 from letta.offline_memory_agent import OfflineMemoryAgent
 from letta.orm.errors import NoResultFound
+from letta.round_robin_multi_agent import RoundRobinMultiAgent
 from letta.schemas.agent import AgentState, AgentType, CreateAgent
 from letta.schemas.block import BlockUpdate
 from letta.schemas.embedding_config import EmbeddingConfig
@@ -44,12 +47,14 @@ from letta.schemas.embedding_config import EmbeddingConfig
 # openai schemas
 from letta.schemas.enums import JobStatus, MessageStreamStatus
 from letta.schemas.environment_variables import SandboxEnvironmentVariableCreate
+from letta.schemas.group import Group, ManagerType
 from letta.schemas.job import Job, JobUpdate
 from letta.schemas.letta_message import LegacyLettaMessage, LettaMessage, ToolReturnMessage
+from letta.schemas.letta_message_content import TextContent
 from letta.schemas.letta_response import LettaResponse
 from letta.schemas.llm_config import LLMConfig
 from letta.schemas.memory import ArchivalMemorySummary, ContextWindowOverview, Memory, RecallMemorySummary
-from letta.schemas.message import Message, MessageCreate, MessageRole, MessageUpdate, TextContent
+from letta.schemas.message import Message, MessageCreate, MessageRole, MessageUpdate
 from letta.schemas.organization import Organization
 from letta.schemas.passage import Passage, PassageUpdate
 from letta.schemas.providers import (
@@ -80,6 +85,7 @@ from letta.server.rest_api.interface import StreamingServerInterface
 from letta.server.rest_api.utils import sse_async_generator
 from letta.services.agent_manager import AgentManager
 from letta.services.block_manager import BlockManager
+from letta.services.group_manager import GroupManager
 from letta.services.identity_manager import IdentityManager
 from letta.services.job_manager import JobManager
 from letta.services.message_manager import MessageManager
@@ -94,6 +100,7 @@ from letta.services.tool_execution_sandbox import ToolExecutionSandbox
 from letta.services.tool_manager import ToolManager
 from letta.services.user_manager import UserManager
 from letta.settings import model_settings, settings, tool_settings
+from letta.supervisor_multi_agent import SupervisorMultiAgent
 from letta.tracing import trace_method
 from letta.utils import get_friendly_error_msg
 
@@ -207,6 +214,7 @@ class SyncServer(Server):
         self.provider_manager = ProviderManager()
         self.step_manager = StepManager()
         self.identity_manager = IdentityManager()
+        self.group_manager = GroupManager()
 
         # Managers that interface with parallelism
         self.per_agent_lock_manager = PerAgentLockManager()
@@ -331,8 +339,8 @@ class SyncServer(Server):
         for server_name, server_config in mcp_server_configs.items():
             if server_config.type == MCPServerType.SSE:
                 self.mcp_clients[server_name] = SSEMCPClient()
-            elif server_config.type == MCPServerType.LOCAL:
-                self.mcp_clients[server_name] = LocalMCPClient()
+            elif server_config.type == MCPServerType.STDIO:
+                self.mcp_clients[server_name] = StdioMCPClient()
             else:
                 raise ValueError(f"Invalid MCP server config: {server_config}")
             try:
@@ -353,6 +361,8 @@ class SyncServer(Server):
         agent_lock = self.per_agent_lock_manager.get_lock(agent_id)
         with agent_lock:
             agent_state = self.agent_manager.get_agent_by_id(agent_id=agent_id, actor=actor)
+            if agent_state.multi_agent_group:
+                return self.load_multi_agent(agent_state.multi_agent_group, actor, interface, agent_state)
 
             interface = interface or self.default_interface_factory()
             if agent_state.agent_type == AgentType.memgpt_agent:
@@ -364,6 +374,46 @@ class SyncServer(Server):
 
             return agent
 
+    def load_multi_agent(
+        self, group: Group, actor: User, interface: Union[AgentInterface, None] = None, agent_state: Optional[AgentState] = None
+    ) -> Agent:
+        match group.manager_type:
+            case ManagerType.round_robin:
+                agent_state = agent_state or self.agent_manager.get_agent_by_id(agent_id=group.agent_ids[0], actor=actor)
+                return RoundRobinMultiAgent(
+                    agent_state=agent_state,
+                    interface=interface,
+                    user=actor,
+                    group_id=group.id,
+                    agent_ids=group.agent_ids,
+                    description=group.description,
+                    max_turns=group.max_turns,
+                )
+            case ManagerType.dynamic:
+                agent_state = agent_state or self.agent_manager.get_agent_by_id(agent_id=group.manager_agent_id, actor=actor)
+                return DynamicMultiAgent(
+                    agent_state=agent_state,
+                    interface=interface,
+                    user=actor,
+                    group_id=group.id,
+                    agent_ids=group.agent_ids,
+                    description=group.description,
+                    max_turns=group.max_turns,
+                    termination_token=group.termination_token,
+                )
+            case ManagerType.supervisor:
+                agent_state = agent_state or self.agent_manager.get_agent_by_id(agent_id=group.manager_agent_id, actor=actor)
+                return SupervisorMultiAgent(
+                    agent_state=agent_state,
+                    interface=interface,
+                    user=actor,
+                    group_id=group.id,
+                    agent_ids=group.agent_ids,
+                    description=group.description,
+                )
+            case _:
+                raise ValueError(f"Type {group.manager_type} is not supported.")
+
     def _step(
         self,
         actor: User,
@@ -690,7 +740,7 @@ class SyncServer(Server):
                     Message(
                         agent_id=agent_id,
                         role=message.role,
-                        content=[TextContent(text=message.content)],
+                        content=[TextContent(text=message.content)] if message.content else [],
                         name=message.name,
                         # assigned later?
                         model=None,
@@ -800,6 +850,9 @@ class SyncServer(Server):
         # TODO: @mindy look at moving this to agent_manager to avoid above extra call
         passages = self.passage_manager.insert_passage(agent_state=agent_state, agent_id=agent_id, text=memory_contents, actor=actor)
 
+        # rebuild agent system prompt - force since no archival change
+        self.agent_manager.rebuild_system_prompt(agent_id=agent_id, actor=actor, force=True)
+
         return passages
 
     def modify_archival_memory(self, agent_id: str, memory_id: str, passage: PassageUpdate, actor: User) -> List[Passage]:
@@ -809,10 +862,14 @@ class SyncServer(Server):
 
     def delete_archival_memory(self, memory_id: str, actor: User):
         # TODO check if it exists first, and throw error if not
-        # TODO: @mindy make this return the deleted passage instead
+        # TODO: need to also rebuild the prompt here
+        passage = self.passage_manager.get_passage_by_id(passage_id=memory_id, actor=actor)
+
+        # delete the passage
         self.passage_manager.delete_passage_by_id(passage_id=memory_id, actor=actor)
 
-        # TODO: return archival memory
+        # rebuild system prompt and force
+        self.agent_manager.rebuild_system_prompt(agent_id=passage.agent_id, actor=actor, force=True)
 
     def get_agent_recall(
         self,
@@ -931,6 +988,9 @@ class SyncServer(Server):
             new_passage_size = self.agent_manager.passage_size(actor=actor, agent_id=agent_id)
             assert new_passage_size >= curr_passage_size  # in case empty files are added
 
+            # rebuild system prompt and force
+            self.agent_manager.rebuild_system_prompt(agent_id=agent_id, actor=actor, force=True)
+
         return job
 
     def load_data(
@@ -1209,7 +1269,7 @@ class SyncServer(Server):
 
     # MCP wrappers
     # TODO support both command + SSE servers (via config)
-    def get_mcp_servers(self) -> dict[str, Union[SSEServerConfig, LocalServerConfig]]:
+    def get_mcp_servers(self) -> dict[str, Union[SSEServerConfig, StdioServerConfig]]:
         """List the MCP servers in the config (doesn't test that they are actually working)"""
         mcp_server_list = {}
 
@@ -1227,8 +1287,8 @@ class SyncServer(Server):
                 # Proper formatting is "mcpServers" key at the top level,
                 # then a dict with the MCP server name as the key,
                 # with the value being the schema from StdioServerParameters
-                if "mcpServers" in mcp_config:
-                    for server_name, server_params_raw in mcp_config["mcpServers"].items():
+                if MCP_CONFIG_TOPLEVEL_KEY in mcp_config:
+                    for server_name, server_params_raw in mcp_config[MCP_CONFIG_TOPLEVEL_KEY].items():
 
                         # No support for duplicate server names
                         if server_name in mcp_server_list:
@@ -1249,7 +1309,7 @@ class SyncServer(Server):
                         else:
                             # Attempt to parse the server params as a StdioServerParameters
                             try:
-                                server_params = LocalServerConfig(
+                                server_params = StdioServerConfig(
                                     server_name=server_name,
                                     command=server_params_raw["command"],
                                     args=server_params_raw.get("args", []),
@@ -1269,6 +1329,98 @@ class SyncServer(Server):
 
         return self.mcp_clients[mcp_server_name].list_tools()
 
+    def add_mcp_server_to_config(
+        self, server_config: Union[SSEServerConfig, StdioServerConfig], allow_upsert: bool = True
+    ) -> dict[str, Union[SSEServerConfig, StdioServerConfig]]:
+        """Add a new server config to the MCP config file"""
+
+        # If the config file doesn't exist, throw an error.
+        mcp_config_path = os.path.join(constants.LETTA_DIR, constants.MCP_CONFIG_NAME)
+        if not os.path.exists(mcp_config_path):
+            raise FileNotFoundError(f"MCP config file not found: {mcp_config_path}")
+
+        # If the file does exist, attempt to parse it get calling get_mcp_servers
+        try:
+            current_mcp_servers = self.get_mcp_servers()
+        except Exception as e:
+            # Raise an error telling the user to fix the config file
+            logger.error(f"Failed to parse MCP config file at {mcp_config_path}: {e}")
+            raise ValueError(f"Failed to parse MCP config file {mcp_config_path}")
+
+        # Check if the server name is already in the config
+        if server_config.server_name in current_mcp_servers and not allow_upsert:
+            raise ValueError(f"Server name {server_config.server_name} is already in the config file")
+
+        # Attempt to initialize the connection to the server
+        if server_config.type == MCPServerType.SSE:
+            new_mcp_client = SSEMCPClient()
+        elif server_config.type == MCPServerType.STDIO:
+            new_mcp_client = StdioMCPClient()
+        else:
+            raise ValueError(f"Invalid MCP server config: {server_config}")
+        try:
+            new_mcp_client.connect_to_server(server_config)
+        except:
+            logger.exception(f"Failed to connect to MCP server: {server_config.server_name}")
+            raise RuntimeError(f"Failed to connect to MCP server: {server_config.server_name}")
+        # Print out the tools that are connected
+        logger.info(f"Attempting to fetch tools from MCP server: {server_config.server_name}")
+        new_mcp_tools = new_mcp_client.list_tools()
+        logger.info(f"MCP tools connected: {', '.join([t.name for t in new_mcp_tools])}")
+        logger.debug(f"MCP tools: {', '.join([str(t) for t in new_mcp_tools])}")
+
+        # Now that we've confirmed the config is working, let's add it to the client list
+        self.mcp_clients[server_config.server_name] = new_mcp_client
+
+        # Add to the server file
+        current_mcp_servers[server_config.server_name] = server_config
+
+        # Write out the file, and make sure to in include the top-level mcpConfig
+        try:
+            new_mcp_file = {MCP_CONFIG_TOPLEVEL_KEY: {k: v.to_dict() for k, v in current_mcp_servers.items()}}
+            with open(mcp_config_path, "w") as f:
+                json.dump(new_mcp_file, f, indent=4)
+        except Exception as e:
+            logger.error(f"Failed to write MCP config file at {mcp_config_path}: {e}")
+            raise ValueError(f"Failed to write MCP config file {mcp_config_path}")
+
+        return list(current_mcp_servers.values())
+
+    def delete_mcp_server_from_config(self, server_name: str) -> dict[str, Union[SSEServerConfig, StdioServerConfig]]:
+        """Delete a server config from the MCP config file"""
+
+        # If the config file doesn't exist, throw an error.
+        mcp_config_path = os.path.join(constants.LETTA_DIR, constants.MCP_CONFIG_NAME)
+        if not os.path.exists(mcp_config_path):
+            raise FileNotFoundError(f"MCP config file not found: {mcp_config_path}")
+
+        # If the file does exist, attempt to parse it get calling get_mcp_servers
+        try:
+            current_mcp_servers = self.get_mcp_servers()
+        except Exception as e:
+            # Raise an error telling the user to fix the config file
+            logger.error(f"Failed to parse MCP config file at {mcp_config_path}: {e}")
+            raise ValueError(f"Failed to parse MCP config file {mcp_config_path}")
+
+        # Check if the server name is already in the config
+        # If it's not, throw an error
+        if server_name not in current_mcp_servers:
+            raise ValueError(f"Server name {server_name} not found in MCP config file")
+
+        # Remove from the server file
+        del current_mcp_servers[server_name]
+
+        # Write out the file, and make sure to in include the top-level mcpConfig
+        try:
+            new_mcp_file = {MCP_CONFIG_TOPLEVEL_KEY: {k: v.to_dict() for k, v in current_mcp_servers.items()}}
+            with open(mcp_config_path, "w") as f:
+                json.dump(new_mcp_file, f, indent=4)
+        except Exception as e:
+            logger.error(f"Failed to write MCP config file at {mcp_config_path}: {e}")
+            raise ValueError(f"Failed to write MCP config file {mcp_config_path}")
+
+        return list(current_mcp_servers.values())
+
     @trace_method
     async def send_message_to_agent(
         self,
@@ -1403,3 +1555,106 @@ class SyncServer(Server):
 
             traceback.print_exc()
             raise HTTPException(status_code=500, detail=f"{e}")
+
+    @trace_method
+    async def send_group_message_to_agent(
+        self,
+        group_id: str,
+        actor: User,
+        messages: Union[List[Message], List[MessageCreate]],
+        stream_steps: bool,
+        stream_tokens: bool,
+        chat_completion_mode: bool = False,
+        # Support for AssistantMessage
+        use_assistant_message: bool = True,
+        assistant_message_tool_name: str = constants.DEFAULT_MESSAGE_TOOL,
+        assistant_message_tool_kwarg: str = constants.DEFAULT_MESSAGE_TOOL_KWARG,
+        metadata: Optional[dict] = None,
+    ) -> Union[StreamingResponse, LettaResponse]:
+        include_final_message = True
+        if not stream_steps and stream_tokens:
+            raise HTTPException(status_code=400, detail="stream_steps must be 'true' if stream_tokens is 'true'")
+
+        try:
+            # fetch the group
+            group = self.group_manager.retrieve_group(group_id=group_id, actor=actor)
+            letta_multi_agent = self.load_multi_agent(group=group, actor=actor)
+
+            llm_config = letta_multi_agent.agent_state.llm_config
+            supports_token_streaming = ["openai", "anthropic", "deepseek"]
+            if stream_tokens and (
+                llm_config.model_endpoint_type not in supports_token_streaming or "inference.memgpt.ai" in llm_config.model_endpoint
+            ):
+                warnings.warn(
+                    f"Token streaming is only supported for models with type {' or '.join(supports_token_streaming)} in the model_endpoint: agent has endpoint type {llm_config.model_endpoint_type} and {llm_config.model_endpoint}. Setting stream_tokens to False."
+                )
+                stream_tokens = False
+
+            # Create a new interface per request
+            letta_multi_agent.interface = StreamingServerInterface(
+                use_assistant_message=use_assistant_message,
+                assistant_message_tool_name=assistant_message_tool_name,
+                assistant_message_tool_kwarg=assistant_message_tool_kwarg,
+                inner_thoughts_in_kwargs=(
+                    llm_config.put_inner_thoughts_in_kwargs if llm_config.put_inner_thoughts_in_kwargs is not None else False
+                ),
+            )
+            streaming_interface = letta_multi_agent.interface
+            if not isinstance(streaming_interface, StreamingServerInterface):
+                raise ValueError(f"Agent has wrong type of interface: {type(streaming_interface)}")
+            streaming_interface.streaming_mode = stream_tokens
+            streaming_interface.streaming_chat_completion_mode = chat_completion_mode
+            if metadata and hasattr(streaming_interface, "metadata"):
+                streaming_interface.metadata = metadata
+
+            streaming_interface.stream_start()
+            task = asyncio.create_task(
+                asyncio.to_thread(
+                    letta_multi_agent.step,
+                    messages=messages,
+                    chaining=self.chaining,
+                    max_chaining_steps=self.max_chaining_steps,
+                )
+            )
+
+            if stream_steps:
+                # return a stream
+                return StreamingResponse(
+                    sse_async_generator(
+                        streaming_interface.get_generator(),
+                        usage_task=task,
+                        finish_message=include_final_message,
+                    ),
+                    media_type="text/event-stream",
+                )
+
+            else:
+                # buffer the stream, then return the list
+                generated_stream = []
+                async for message in streaming_interface.get_generator():
+                    assert (
+                        isinstance(message, LettaMessage)
+                        or isinstance(message, LegacyLettaMessage)
+                        or isinstance(message, MessageStreamStatus)
+                    ), type(message)
+                    generated_stream.append(message)
+                    if message == MessageStreamStatus.done:
+                        break
+
+                # Get rid of the stream status messages
+                filtered_stream = [d for d in generated_stream if not isinstance(d, MessageStreamStatus)]
+                usage = await task
+
+                # By default the stream will be messages of type LettaMessage or LettaLegacyMessage
+                # If we want to convert these to Message, we can use the attached IDs
+                # NOTE: we will need to de-duplicate the Messsage IDs though (since Assistant->Inner+Func_Call)
+                # TODO: eventually update the interface to use `Message` and `MessageChunk` (new) inside the deque instead
+                return LettaResponse(messages=filtered_stream, usage=usage)
+        except HTTPException:
+            raise
+        except Exception as e:
+            print(e)
+            import traceback
+
+            traceback.print_exc()
+            raise HTTPException(status_code=500, detail=f"{e}")

letta/services/agent_manager.py

@@ -18,6 +18,7 @@ from letta.orm import Tool as ToolModel
 from letta.orm.enums import ToolType
 from letta.orm.errors import NoResultFound
 from letta.orm.sandbox_config import AgentEnvironmentVariable as AgentEnvironmentVariableModel
+from letta.orm.sqlalchemy_base import AccessType
 from letta.orm.sqlite_functions import adapt_array
 from letta.schemas.agent import AgentState as PydanticAgentState
 from letta.schemas.agent import AgentType, CreateAgent, UpdateAgent
@@ -35,10 +36,15 @@ from letta.schemas.tool_rule import ContinueToolRule as PydanticContinueToolRule
 from letta.schemas.tool_rule import TerminalToolRule as PydanticTerminalToolRule
 from letta.schemas.tool_rule import ToolRule as PydanticToolRule
 from letta.schemas.user import User as PydanticUser
-from letta.serialize_schemas import SerializedAgentSchema
-from letta.serialize_schemas.tool import SerializedToolSchema
+from letta.serialize_schemas import MarshmallowAgentSchema
+from letta.serialize_schemas.marshmallow_tool import SerializedToolSchema
+from letta.serialize_schemas.pydantic_agent_schema import AgentSchema
 from letta.services.block_manager import BlockManager
 from letta.services.helpers.agent_manager_helper import (
+    _apply_filters,
+    _apply_identity_filters,
+    _apply_pagination,
+    _apply_tag_filter,
     _process_relationship,
     _process_tags,
     check_supports_structured_output,
@@ -49,6 +55,7 @@ from letta.services.helpers.agent_manager_helper import (
 )
 from letta.services.identity_manager import IdentityManager
 from letta.services.message_manager import MessageManager
+from letta.services.passage_manager import PassageManager
 from letta.services.source_manager import SourceManager
 from letta.services.tool_manager import ToolManager
 from letta.settings import settings
@@ -70,6 +77,7 @@ class AgentManager:
         self.tool_manager = ToolManager()
         self.source_manager = SourceManager()
         self.message_manager = MessageManager()
+        self.passage_manager = PassageManager()
         self.identity_manager = IdentityManager()
 
     # ======================================================================================================================
@@ -326,39 +334,60 @@ class AgentManager:
             # Convert to PydanticAgentState and return
             return agent.to_pydantic()
 
-    @enforce_types
+    # TODO: Make this general and think about how to roll this into sqlalchemybase
     def list_agents(
         self,
         actor: PydanticUser,
+        name: Optional[str] = None,
+        tags: Optional[List[str]] = None,
+        match_all_tags: bool = False,
         before: Optional[str] = None,
         after: Optional[str] = None,
         limit: Optional[int] = 50,
-        tags: Optional[List[str]] = None,
-        match_all_tags: bool = False,
         query_text: Optional[str] = None,
-        identifier_keys: Optional[List[str]] = None,
+        project_id: Optional[str] = None,
+        template_id: Optional[str] = None,
+        base_template_id: Optional[str] = None,
         identity_id: Optional[str] = None,
-        **kwargs,
+        identifier_keys: Optional[List[str]] = None,
+        include_relationships: Optional[List[str]] = None,
     ) -> List[PydanticAgentState]:
         """
-        List agents that have the specified tags.
+        Retrieves agents with optimized filtering and optional field selection.
+
+        Args:
+            actor: The User requesting the list
+            name (Optional[str]): Filter by agent name.
+            tags (Optional[List[str]]): Filter agents by tags.
+            match_all_tags (bool): If True, only return agents that match ALL given tags.
+            before (Optional[str]): Cursor for pagination.
+            after (Optional[str]): Cursor for pagination.
+            limit (Optional[int]): Maximum number of agents to return.
+            query_text (Optional[str]): Search agents by name.
+            project_id (Optional[str]): Filter by project ID.
+            template_id (Optional[str]): Filter by template ID.
+            base_template_id (Optional[str]): Filter by base template ID.
+            identity_id (Optional[str]): Filter by identifier ID.
+            identifier_keys (Optional[List[str]]): Search agents by identifier keys.
+            include_relationships (Optional[List[str]]): List of fields to load for performance optimization.
+
+        Returns:
+            List[PydanticAgentState]: The filtered list of matching agents.
         """
         with self.session_maker() as session:
-            agents = AgentModel.list(
-                db_session=session,
-                before=before,
-                after=after,
-                limit=limit,
-                tags=tags,
-                match_all_tags=match_all_tags,
-                organization_id=actor.organization_id if actor else None,
-                query_text=query_text,
-                identifier_keys=identifier_keys,
-                identity_id=identity_id,
-                **kwargs,
-            )
+            query = select(AgentModel).distinct(AgentModel.created_at, AgentModel.id)
+            query = AgentModel.apply_access_predicate(query, actor, ["read"], AccessType.ORGANIZATION)
+
+            # Apply filters
+            query = _apply_filters(query, name, query_text, project_id, template_id, base_template_id)
+            query = _apply_identity_filters(query, identity_id, identifier_keys)
+            query = _apply_tag_filter(query, tags, match_all_tags)
+            query = _apply_pagination(query, before, after, session)
 
-            return [agent.to_pydantic() for agent in agents]
+            query = query.limit(limit)
+
+            agents = session.execute(query).scalars().all()
+            return [agent.to_pydantic(include_relationships=include_relationships) for agent in agents]
 
     @enforce_types
     def list_agents_matching_tags(
@@ -399,7 +428,7 @@ class AgentManager:
                 # Ensures agents match at least one tag in match_some
                 query = query.join(AgentsTags).where(AgentsTags.tag.in_(match_some))
 
-            query = query.group_by(AgentModel.id).limit(limit)
+            query = query.distinct(AgentModel.id).order_by(AgentModel.id).limit(limit)
 
             return list(session.execute(query).scalars())
 
@@ -434,29 +463,32 @@ class AgentManager:
         with self.session_maker() as session:
             # Retrieve the agent
             agent = AgentModel.read(db_session=session, identifier=agent_id, actor=actor)
+            # TODO check if it is managing a group
             agent.hard_delete(session)
 
     @enforce_types
-    def serialize(self, agent_id: str, actor: PydanticUser) -> dict:
+    def serialize(self, agent_id: str, actor: PydanticUser) -> AgentSchema:
         with self.session_maker() as session:
             # Retrieve the agent
             agent = AgentModel.read(db_session=session, identifier=agent_id, actor=actor)
-            schema = SerializedAgentSchema(session=session, actor=actor)
-            return schema.dump(agent)
+            schema = MarshmallowAgentSchema(session=session, actor=actor)
+            data = schema.dump(agent)
+            return AgentSchema(**data)
 
     @enforce_types
     def deserialize(
         self,
-        serialized_agent: dict,
+        serialized_agent: AgentSchema,
         actor: PydanticUser,
         append_copy_suffix: bool = True,
         override_existing_tools: bool = True,
         project_id: Optional[str] = None,
     ) -> PydanticAgentState:
+        serialized_agent = serialized_agent.model_dump()
         tool_data_list = serialized_agent.pop("tools", [])
 
         with self.session_maker() as session:
-            schema = SerializedAgentSchema(session=session, actor=actor)
+            schema = MarshmallowAgentSchema(session=session, actor=actor)
             agent = schema.load(serialized_agent, session=session)
             if append_copy_suffix:
                 agent.name += "_copy"
@@ -595,12 +627,17 @@ class AgentManager:
             # NOTE: a bit of a hack - we pull the timestamp from the message created_by
             memory_edit_timestamp = curr_system_message.created_at
 
+        num_messages = self.message_manager.size(actor=actor, agent_id=agent_id)
+        num_archival_memories = self.passage_manager.size(actor=actor, agent_id=agent_id)
+
         # update memory (TODO: potentially update recall/archival stats separately)
         new_system_message_str = compile_system_message(
             system_prompt=agent_state.system,
             in_context_memory=agent_state.memory,
             in_context_memory_last_edit=memory_edit_timestamp,
             recent_passages=self.list_passages(actor=actor, agent_id=agent_id, ascending=False, limit=10),
+            previous_message_count=num_messages,
+            archival_memory_size=num_archival_memories,
         )
 
         diff = united_diff(curr_system_message_openai["content"], new_system_message_str)