feat: integrate Mem0 memory management into code_scratch (#730)

Co-authored-by: 班扬 <[email protected]>

Author: vinci-grape

ms_agent/agent/llm_agent.py

@@ -235,8 +235,14 @@ async def _prepare_memory(self):
                 assert _memory.name in memory_mapping, (
                     f'{_memory.name} not in memory_mapping, '
                     f'which supports: {list(memory_mapping.keys())}')
-                self.memory_tools.append(memory_mapping[_memory.name](
-                    self.config))
+                if _memory.name == 'mem0':
+                    from .memory.mem0ai import SharedMemoryManager
+                    shared_memory = SharedMemoryManager.get_shared_memory(
+                        _memory)
+                    self.memory_tools.append(shared_memory)
+                else:
+                    self.memory_tools.append(
+                        memory_mapping[_memory.name](_memory))
 
     async def _prepare_planer(self):
         """Load and initialize the planer component from the config."""
@@ -446,7 +452,21 @@ def _save_history(self, messages: List[Message], **kwargs):
             messages (List[Message]): Current message history to save.
         """
         query = messages[1].content
-        if not query or not self.task or self.task == 'subtask':
+        if not query:
+            return
+
+        if self.memory_tools:
+            user_id = None
+            if hasattr(self.config, 'memory') and self.config.memory:
+                for memory_config in self.config.memory:
+                    if hasattr(memory_config,
+                               'user_id') and memory_config.user_id:
+                        user_id = memory_config.user_id
+                        break
+            for memory_tool in self.memory_tools:
+                memory_tool._add_memories_from_conversation(messages, user_id)
+
+        if not self.task or self.task == 'subtask':
             return
         config: DictConfig = deepcopy(self.config)  # noqa
         config.runtime = self.runtime.to_dict()
@@ -456,6 +476,21 @@ def _save_history(self, messages: List[Message], **kwargs):
             config=config,
             messages=messages)
 
+    def _save_memory(self, messages: List[Message], **kwargs):
+        """
+        Save memories to disk for future resuming.
+
+        Args:
+            messages (List[Message]): Current message history to save.
+        """
+
+        if self.memory_tools:
+            agent_id = self.tag
+            for memory_tool in self.memory_tools:
+                memory_tool._add_memories_from_procedural(
+                    messages, 'subagent', agent_id, 'procedural_memory')
+        return
+
     async def _run(self, messages: Union[List[Message], str],
                    **kwargs) -> AsyncGenerator[Any, Any]:
         """Run the agent, mainly contains a llm calling and tool calling loop.
@@ -512,6 +547,9 @@ async def _run(self, messages: Union[List[Message], str],
                 # save history
                 self._save_history(messages, **kwargs)
 
+            # save memory
+            self._save_memory(messages, **kwargs)
+
             await self._loop_callback('on_task_end', messages)
             await self._cleanup_tools()
             yield messages