Improve Triage to support memory

Author: RussellLuo

coagent/agents/triage.py

@@ -23,6 +23,7 @@
 
 from .aswarm import Agent as SwarmAgent, Swarm
 from .chat_agent import ChatHistory, ChatMessage, Delegate
+from .memory import Memory, NoMemory
 from .model import default_model, Model
 
 
@@ -47,6 +48,7 @@ def __init__(
         name: str = "",
         system: str = "",
         model: Model = default_model,
+        memory: Memory | None = None,
         static_agents: list[str] | None = None,
         dynamic_agents: list[DiscoveryQuery] | None = None,
         timeout: float = 300,
@@ -64,6 +66,7 @@ def __init__(
         self._sub_agents: dict[str, Schema] = {}
         self._swarm_agent: SwarmAgent | None = None
 
+        self._memory: Memory = memory or NoMemory()
         self._history: ChatHistory = ChatHistory(messages=[])
 
     @property
@@ -83,6 +86,10 @@ def system(self) -> str:
     def model(self) -> Model:
         return self._model
 
+    @property
+    def memory(self) -> Memory:
+        return self._memory
+
     @property
     def static_agents(self) -> list[str] | None:
         return self._static_agents
@@ -131,6 +138,7 @@ async def _update_swarm_agent(self) -> None:
 
     def _transfer_to_agent(self, agent_type: str):
         async def run() -> AsyncIterator[ChatMessage]:
+            # TODO: Handle memory?
             async for chunk in Delegate(self, agent_type).handle(self._history):
                 yield chunk
 
@@ -203,6 +211,7 @@ async def deregister_sub_agents(
     async def handle_history(
         self, msg: ChatHistory, ctx: Context
     ) -> AsyncIterator[ChatMessage]:
+        # TODO: Handle memory?
         response = self._handle_history(msg, ctx)
         async for resp in response:
             yield resp
@@ -211,10 +220,19 @@ 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, ctx)
+        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
+        )
 
     async def _handle_history(
         self, msg: ChatHistory, ctx: Context