fix bugs in memory (#753)

vinci-grape · web-flow · commit ff32df1dba14 · 2025-10-24T17:12:22.000+08:00
diff --git a/ms_agent/agent/llm_agent.py b/ms_agent/agent/llm_agent.py
@@ -14,7 +14,7 @@
 from ms_agent.llm.llm import LLM
 from ms_agent.llm.utils import Message
 from ms_agent.memory import Memory, memory_mapping
-from ms_agent.memory.mem0ai import Mem0Memory
+from ms_agent.memory.mem0ai import Mem0Memory, SharedMemoryManager
 from ms_agent.rag.base import RAG
 from ms_agent.rag.utils import rag_mapping
 from ms_agent.tools import ToolManager
@@ -342,19 +342,19 @@ async def load_memory(self):
                     service = self.config.llm.service
                     config_dict = {
                         'model':
-                        self.config.llm.model,
+                        _memory.summary_model,
                         'provider':
                         'openai',
                         'openai_base_url':
                         getattr(self.config.llm, f'{service}_base_url', None),
                         'openai_api_key':
                         getattr(self.config.llm, f'{service}_api_key', None),
+                        'max_tokens':
+                        _memory.max_tokens,
                     }
                     llm_config_obj = OmegaConf.create(config_dict)
                     setattr(_memory, 'llm', llm_config_obj)
-
                 if memory_type == 'mem0':
-                    from ms_agent.memory.mem0ai import SharedMemoryManager
                     shared_memory = SharedMemoryManager.get_shared_memory(
                         _memory)
                     self.memory_tools.append(shared_memory)
@@ -638,7 +638,7 @@ async def run_loop(self, messages: Union[List[Message], str],
                 async for messages in self.step(messages):
                     yield messages
                 self.runtime.round += 1
-                # save history
+                # save memory and history
                 self.save_memory(messages)
                 self.save_history(messages)
 
diff --git a/ms_agent/memory/mem0ai.py b/ms_agent/memory/mem0ai.py
@@ -131,10 +131,7 @@ def _initialize_memory(self):
                                 self.config.llm.openai_base_url),
                         'api_key':
                         getattr(self.config, 'embedder_api_key',
-                                self.config.llm.openai_api_key),
-                        'max_tokens':
-                        getattr(self.config, 'max_tokens',
-                                self.config.llm.max_tokens),
+                                self.config.llm.openai_api_key)
                     }
                 },
                 'llm': {
@@ -147,7 +144,10 @@ def _initialize_memory(self):
                                 self.config.llm.openai_base_url),
                         'api_key':
                         getattr(self.config, 'llm_api_key',
-                                self.config.llm.openai_api_key)
+                                self.config.llm.openai_api_key),
+                        'max_tokens':
+                        getattr(self.config, 'max_tokens',
+                                self.config.llm.max_tokens),
                     }
                 },
             }
diff --git a/ms_agent/tools/filesystem_tool.py b/ms_agent/tools/filesystem_tool.py
@@ -68,16 +68,21 @@ async def get_tools(self):
                 Tool(
                     tool_name='read_file',
                     server_name='file_system',
-                    description='Read the content of a file',
+                    description='Read the content of file(s)',
                     parameters={
                         'type': 'object',
                         'properties': {
-                            'path': {
-                                'type': 'string',
-                                'description': 'The relative path of the file',
+                            'paths': {
+                                'type':
+                                'array',
+                                'items': {
+                                    'type': 'string'
+                                },
+                                'description':
+                                'List of relative file path(s) to read',
                             }
                         },
-                        'required': ['path'],
+                        'required': ['paths'],
                         'additionalProperties': False
                     }),
                 Tool(
@@ -154,20 +159,23 @@ async def write_file(self, path: str, content: str):
         except Exception as e:
             return f'Write file <{path}> failed, error: ' + str(e)
 
-    async def read_file(self, path: str):
-        """Read the content of a file.
+    async def read_file(self, paths: list[str]):
+        """Read the content of file(s).
 
         Args:
-            path(`path`): The relative file path to read, a prefix dir will be automatically concatenated.
+            paths(`list[str]`): List of relative file path(s) to read, a prefix dir will be automatically concatenated.
 
         Returns:
-            The file content or error message.
+            Dictionary mapping file path(s) to their content or error messages.
         """
-        try:
-            with open(os.path.join(self.output_dir, path), 'r') as f:
-                return f.read()
-        except Exception as e:
-            return f'Read file <{path}> failed, error: ' + str(e)
+        results = {}
+        for path in paths:
+            try:
+                with open(os.path.join(self.output_dir, path), 'r') as f:
+                    results[path] = f.read()
+            except Exception as e:
+                results[path] = f'Read file <{path}> failed, error: ' + str(e)
+        return str(results)
 
     async def list_files(self, path: str = None):
         """List all files in a directory.
diff --git a/projects/code_scratch/coding.yaml b/projects/code_scratch/coding.yaml
@@ -4,6 +4,7 @@ llm:
   openai_api_key:
   openai_base_url: https://dashscope.aliyuncs.com/compatible-mode/v1
 
+
 generation_config:
   temperature: 0.2
   top_k: 20
