Add support for agent memory

Author: RussellLuo

coagent/agents/__init__.py

@@ -8,6 +8,7 @@
     StructuredOutput,
     type_to_response_format_param,
 )
+from .memory import Memory, InMemMemory
 from .model import Model
 from .parallel import Aggregator, AggregationResult, Parallel
 from .sequential import Sequential

coagent/agents/chat_agent.py

@@ -25,6 +25,7 @@
     MCPTool,
     NamedMCPServer,
 )
+from .memory import Memory, NoMemory
 from .messages import ChatMessage, ChatHistory, StructuredOutput
 from .model import default_model, Model
 from .util import is_user_confirmed
@@ -222,6 +223,7 @@ def __init__(
         mcp_servers: list[NamedMCPServer] | None = None,
         mcp_server_agent_type: str = "mcp_server",
         model: Model = default_model,
+        memory: Memory | None = None,
         timeout: float = 300,
     ):
         super().__init__(timeout=timeout)
@@ -236,6 +238,7 @@ def __init__(
         self._swarm_client: Swarm = Swarm(self.model)
         self._swarm_agent: SwarmAgent | None = None
 
+        self._memory: Memory = memory or NoMemory()
         self._history: ChatHistory = ChatHistory(messages=[])
 
     @property
@@ -267,6 +270,10 @@ def mcp_server_agent_type(self) -> str:
     def model(self) -> Model:
         return self._model
 
+    @property
+    def memory(self) -> Memory:
+        return self._memory
+
     async def stopped(self) -> None:
         for server in self.mcp_servers:
             if server.connect is not None:
@@ -286,7 +293,7 @@ def get_swarm_client(self, extensions: dict) -> Swarm:
         if model_id:
             # We assume that non-empty model ID indicates the use of a dynamic model client.
             model = Model(
-                model=model_id,
+                id=model_id,
                 base_url=extensions.get("model_base_url", ""),
                 api_key=extensions.get("model_api_key", ""),
                 api_version=extensions.get("model_api_version", ""),
@@ -319,13 +326,15 @@ async def get_swarm_agent(self) -> SwarmAgent:
 
     async def agent(self, agent_type: str) -> AsyncIterator[ChatMessage]:
         """The candidate agent to delegate the conversation to."""
+        # TODO: Handle memory?
         async for chunk in Delegate(self, agent_type).handle(self._history):
             yield chunk
 
     @handler
     async def handle_history(
         self, msg: ChatHistory, ctx: Context
     ) -> AsyncIterator[ChatMessage]:
+        # TODO: Handle memory?
         response = self._handle_history(msg)
         async for resp in response:
             yield resp
@@ -334,15 +343,25 @@ async def handle_history(
     async def handle_message(
         self, msg: ChatMessage, ctx: Context
     ) -> AsyncIterator[ChatMessage]:
-        history = ChatHistory(messages=[msg])
+        existing = await self.memory.get_items()
+        history = ChatHistory(messages=existing + [msg])
+
         response = self._handle_history(history)
+        full_content = ""
         async for resp in response:
             yield resp
+            full_content += resp.content
+
+        await self.memory.add_items(
+            msg,  # input item
+            ChatMessage(role="assistant", content=full_content),  # output item
+        )
 
     @handler
     async def handle_structured_output(
         self, msg: StructuredOutput, ctx: Context
     ) -> AsyncIterator[ChatMessage]:
+        # TODO: Handle memory?
         match msg.input:
             case ChatMessage():
                 history = ChatHistory(messages=[msg.input])

coagent/agents/memory.py

@@ -0,0 +1,89 @@
+from typing import Protocol, runtime_checkable
+
+from .messages import ChatMessage
+
+
+@runtime_checkable
+class Memory(Protocol):
+    """Protocol for memory implementations.
+
+    Memory stores conversation history for a specific agent, allowing
+    agents to maintain context without requiring explicit manual memory management.
+    """
+
+    async def get_items(self, limit: int | None = None) -> list[ChatMessage]:
+        """Retrieve the conversation history for this memory.
+
+        Args:
+            limit: Maximum number of items to retrieve. If None, retrieves all items.
+                   When specified, returns the latest N items in chronological order.
+
+        Returns:
+            List of input items representing the conversation history
+        """
+        ...
+
+    async def add_items(self, *items: ChatMessage) -> None:
+        """Add new items to the conversation history.
+
+        Args:
+            items: List of input items to add to the history
+        """
+        ...
+
+    async def pop_item(self) -> ChatMessage | None:
+        """Remove and return the most recent item from the memory.
+
+        Returns:
+            The most recent item if it exists, None if the memory is empty
+        """
+        ...
+
+    async def clear_items(self) -> None:
+        """Clear all items for this memory."""
+        ...
+
+
+class NoMemory(list):
+    """Built-in memory implementation that stores no conversation history."""
+
+    async def get_items(self, limit: int | None = None) -> list[ChatMessage]:
+        """Retrieve the conversation history for this memory."""
+        return []
+
+    async def add_items(self, *items: ChatMessage) -> None:
+        """Add new items to the conversation history."""
+        return
+
+    async def pop_item(self) -> ChatMessage | None:
+        """Remove and return the most recent item from the memory."""
+        return
+
+    async def clear_items(self) -> None:
+        """Clear all items for this memory."""
+        return
+
+
+class InMemMemory(list):
+    """Built-in memory implementation that stores conversation history in memory."""
+
+    async def get_items(self, limit: int | None = None) -> list[ChatMessage]:
+        """Retrieve the conversation history for this memory."""
+        if limit is None:
+            return self
+        # Return the latest limit number of items.
+        return self[-limit:]
+
+    async def add_items(self, *items: ChatMessage) -> None:
+        """Add new items to the conversation history."""
+        self.extend(items)
+
+    async def pop_item(self) -> ChatMessage | None:
+        """Remove and return the most recent item from the memory."""
+        if not self:
+            return None
+        return self.pop()
+
+    async def clear_items(self) -> None:
+        """Clear all items for this memory."""
+        self.clear()