[Feat] Support resetting agent memory recursively or by keypath (#271)

* update `reset` method

* fix type hints

lagent/agents/agent.py

@@ -165,9 +165,26 @@ def register_hook(self, hook: Callable):
         self._hooks[handle.id] = hook
         return handle
 
-    def reset(self, session_id=0):
-        if self.memory:
-            self.memory.reset(session_id=session_id)
+    def reset(self,
+              session_id=0,
+              keypath: Optional[str] = None,
+              recursive: bool = False):
+        assert not (keypath and
+                    recursive), 'keypath and recursive can\'t be used together'
+        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]
+            agent.reset(session_id, recursive=False)
+        else:
+            if self.memory:
+                self.memory.reset(session_id=session_id)
+            if recursive:
+                for agent in getattr(self, '_agents', {}).values():
+                    agent.reset(session_id, recursive=True)
 
     def __repr__(self):
 

lagent/llms/openai.py

@@ -7,7 +7,7 @@
 from concurrent.futures import ThreadPoolExecutor
 from logging import getLogger
 from threading import Lock
-from typing import Dict, List, Optional, Union
+from typing import AsyncGenerator, Dict, List, Optional, Union
 
 import aiohttp
 import requests
@@ -701,7 +701,8 @@ async def _chat(self, messages: List[dict], **gen_params) -> str:
                            f'{max_num_retries} times. Check the logs for '
                            f'details. errmsg: {errmsg}')
 
-    async def _stream_chat(self, messages: List[dict], **gen_params) -> str:
+    async def _stream_chat(self, messages: List[dict],
+                           **gen_params) -> AsyncGenerator[str, None]:
         """Generate completion from a list of templates.
 
         Args: