Refactor get_message interface (#314)

refactor `get_message` interface

Author: braisedpork1964

lagent/agents/agent.py

@@ -168,26 +168,27 @@ def reset(self, keypath: Optional[str] = None, recursive: bool = False):
                 for agent in getattr(self, '_agents', {}).values():
                     agent.reset(recursive=True)
 
-    def get_messages(self, keypath: Optional[str] = None) -> List[dict]:
-        """Get OpenAI format messages from memory.
+    def get_messages(self, prefix='', destination=None) -> Dict[str, List[dict]]:
+        """Get OpenAI format messages from all agents recursively, similar to state_dict.
 
         Args:
-            keypath (Optional[str]): The keypath of the sub-agent to get messages from. Default is None.
+            prefix (str): The prefix to prepend to the key. Default is ''.
+            destination (Optional[Dict]): The destination dict to store the messages. Default is None.
 
         Returns:
-            List[dict]: The messages from the memory including the sub-agent's system prompt.
+            Dict[str, List[dict]]: A dict mapping agent keypaths to their OpenAI format message lists.
         """
-        if keypath:
-            keys, agent = keypath.split('.'), self
-            for key in keys:
-                agents = getattr(agent, '_agents', {})
-                if key not in agents:
-                    raise KeyError(f'No sub-agent named {key} in {agent}')
-                agent = agents[key]
-            return agent.get_messages()
+        if destination is None:
+            destination = {}
         if self.aggregator:
-            return self.aggregator.aggregate(self.memory, self.name, self.output_format, self.template)
-        raise ValueError(f'{self.name} has no aggregator to get messages')
+            messages = self.aggregator.aggregate(self.memory, self.name, self.output_format, self.template)
+            if isinstance(messages, tuple):
+                messages, _ = messages
+            destination[prefix + 'messages'] = messages
+        for name, agent in getattr(self, '_agents', {}).items():
+            if isinstance(agent, Agent):
+                agent.get_messages(destination=destination, prefix=prefix + name + '.')
+        return destination
 
     def _scroll_buffer(self, message, hash_func=lambda m: m.uid):
         if not self.memory:

lagent/serving/sandbox/daemon.py

@@ -28,7 +28,6 @@
 """
 
 from __future__ import annotations
-
 import argparse
 import asyncio
 import inspect
@@ -96,9 +95,7 @@ async def start(self) -> None:
         """Start listening. Removes stale socket file if present."""
         if os.path.exists(self.sock_path):
             os.unlink(self.sock_path)
-        self._server = await asyncio.start_unix_server(
-            self._handle_client, path=self.sock_path
-        )
+        self._server = await asyncio.start_unix_server(self._handle_client, path=self.sock_path)
         os.chmod(self.sock_path, 0o777)
         logger.info("%s listening on %s", self.__class__.__name__, self.sock_path)
         await self._server.serve_forever()
@@ -139,10 +136,12 @@ async def _dispatch(self, request: dict) -> dict:
         if cmd == "ping":
             return {"status": "ok", "type": self.daemon_type}
         if cmd == "shutdown":
+
             async def _delayed_close():
                 await asyncio.sleep(0.1)
                 if self._server:
                     self._server.close()
+
             asyncio.create_task(_delayed_close())
             return {"status": "shutting_down"}
         return {"error": f"Unknown command: {cmd}"}
@@ -197,18 +196,22 @@ async def _dispatch(self, request: dict) -> dict:
         name = request.get("name")
         parameters = request.get("parameters", {})
         if not name:
-            return dataclass2dict(ActionReturn(
-                errmsg="Missing 'name' in request",
-                state=ActionStatusCode.ARGS_ERROR,
-            ))
+            return dataclass2dict(
+                ActionReturn(
+                    errmsg="Missing 'name' in request",
+                    state=ActionStatusCode.ARGS_ERROR,
+                )
+            )
 
         try:
             action_return = await self.executor.forward(name, parameters)
         except Exception as e:
             logger.exception("Action %s failed", name)
             action_return = ActionReturn(
-                args=parameters, type=name,
-                errmsg=str(e), state=ActionStatusCode.API_ERROR,
+                args=parameters,
+                type=name,
+                errmsg=str(e),
+                state=ActionStatusCode.API_ERROR,
             )
         return dataclass2dict(action_return)
 
@@ -349,6 +352,12 @@ async def _dispatch(self, request: dict) -> dict:
             except Exception as e:
                 return {"error": str(e)}
 
+        if cmd == 'get_messages':
+            try:
+                return self.agent.get_messages()
+            except Exception as e:
+                return {"error": str(e)}
+
         return {"error": f"Unknown command: {cmd}"}
 
     @staticmethod
@@ -404,11 +413,14 @@ def main():
     # -- start --
     p_start = sub.add_parser("start", help="Start the daemon")
     p_start.add_argument(
-        "--sock", default="/tmp/lagent_action.sock",
+        "--sock",
+        default="/tmp/lagent_action.sock",
         help="Unix socket path",
     )
     p_start.add_argument(
-        "--mode", choices=["actions", "agent"], default="actions",
+        "--mode",
+        choices=["actions", "agent"],
+        default="actions",
         help="'actions' for Level 1 (ActionDaemon), 'agent' for Level 2 (AgentDaemon)",
     )
     p_start.add_argument(