Feat/tool memory

src/memos/api/routers/server_router.py

@@ -90,6 +90,7 @@
 status_tracker = TaskStatusTracker(redis_client=redis_client)
 embedder = components["embedder"]
 graph_db = components["graph_db"]
+vector_db = components["vector_db"]
 
 
 # =============================================================================
@@ -359,6 +360,13 @@ def get_user_names_by_memory_ids(request: GetUserNamesByMemoryIdsRequest):
             ),
         )
     result = graph_db.get_user_names_by_memory_ids(memory_ids=request.memory_ids)
+    if vector_db:
+        prefs = []
+        for collection_name in ["explicit_preference", "implicit_preference"]:
+            prefs.extend(
+                vector_db.get_by_ids(collection_name=collection_name, ids=request.memory_ids)
+            )
+        result.update({pref.id: pref.payload.get("mem_cube_id", None) for pref in prefs})
     return GetUserNamesByMemoryIdsResponse(
         code=200,
         message="Successfully",

src/memos/mem_reader/multi_modal_struct.py

@@ -1,5 +1,6 @@
 import concurrent.futures
 import json
+import re
 import traceback
 
 from typing import Any
@@ -547,7 +548,11 @@ def _process_tool_trajectory_fine(
         for fast_item in fast_memory_items:
             # Extract memory text (string content)
             mem_str = fast_item.memory or ""
-            if not mem_str.strip() or "tool:" not in mem_str:
+            if not mem_str.strip() or (
+                "tool:" not in mem_str
+                and "[tool_calls]:" not in mem_str
+                and not re.search(r"<tool_schema>.*?</tool_schema>", mem_str, re.DOTALL)
+            ):
                 continue
             try:
                 resp = self._get_llm_tool_trajectory_response(mem_str)
@@ -563,6 +568,8 @@ def _process_tool_trajectory_fine(
                         value=m.get("trajectory", ""),
                         info=info,
                         memory_type=memory_type,
+                        correctness=m.get("correctness", ""),
+                        experience=m.get("experience", ""),
                         tool_used_status=m.get("tool_used_status", []),
                     )
                     fine_memory_items.append(node)
@@ -606,16 +613,22 @@ def _process_multi_modal_data(
         if mode == "fast":
             return fast_memory_items
         else:
-            # Part A: call llm
+            # Part A: call llm in parallel using thread pool
             fine_memory_items = []
-            fine_memory_items_string_parser = self._process_string_fine(
-                fast_memory_items, info, custom_tags
-            )
-            fine_memory_items.extend(fine_memory_items_string_parser)
 
-            fine_memory_items_tool_trajectory_parser = self._process_tool_trajectory_fine(
-                fast_memory_items, info
-            )
+            with ContextThreadPoolExecutor(max_workers=2) as executor:
+                future_string = executor.submit(
+                    self._process_string_fine, fast_memory_items, info, custom_tags
+                )
+                future_tool = executor.submit(
+                    self._process_tool_trajectory_fine, fast_memory_items, info
+                )
+
+                # Collect results
+                fine_memory_items_string_parser = future_string.result()
+                fine_memory_items_tool_trajectory_parser = future_tool.result()
+
+            fine_memory_items.extend(fine_memory_items_string_parser)
             fine_memory_items.extend(fine_memory_items_tool_trajectory_parser)
 
             # Part B: get fine multimodal items
@@ -658,13 +671,18 @@ def _process_transfer_multi_modal_data(
         }
 
         fine_memory_items = []
-        # Part A: call llm
-        fine_memory_items_string_parser = self._process_string_fine([raw_node], info, custom_tags)
-        fine_memory_items.extend(fine_memory_items_string_parser)
+        # Part A: call llm in parallel using thread pool
+        with ContextThreadPoolExecutor(max_workers=2) as executor:
+            future_string = executor.submit(
+                self._process_string_fine, [raw_node], info, custom_tags
+            )
+            future_tool = executor.submit(self._process_tool_trajectory_fine, [raw_node], info)
 
-        fine_memory_items_tool_trajectory_parser = self._process_tool_trajectory_fine(
-            [raw_node], info
-        )
+            # Collect results
+            fine_memory_items_string_parser = future_string.result()
+            fine_memory_items_tool_trajectory_parser = future_tool.result()
+
+        fine_memory_items.extend(fine_memory_items_string_parser)
         fine_memory_items.extend(fine_memory_items_tool_trajectory_parser)
 
         # Part B: get fine multimodal items

src/memos/mem_reader/read_multi_modal/system_parser.py

@@ -42,9 +42,10 @@ def create_source(
         info: dict[str, Any],
     ) -> SourceMessage:
         """Create SourceMessage from system message."""
-        content = message["content"]
+
+        content = message.get("content", "")
         if isinstance(content, dict):
-            content = content["text"]
+            content = content.get("text", "")
 
         content_wo_tool_schema = re.sub(
             r"<tool_schema>(.*?)</tool_schema>",
@@ -84,17 +85,146 @@ def parse_fast(
         info: dict[str, Any],
         **kwargs,
     ) -> list[TextualMemoryItem]:
-        content = message["content"]
+        content = message.get("content", "")
         if isinstance(content, dict):
-            content = content["text"]
+            content = content.get("text", "")
 
-        # Replace tool_schema content with "omitted" in remaining content
-        content_wo_tool_schema = re.sub(
-            r"<tool_schema>(.*?)</tool_schema>",
-            r"<tool_schema>omitted</tool_schema>",
-            content,
-            flags=re.DOTALL,
-        )
+        # Find first tool_schema block
+        tool_schema_pattern = r"<tool_schema>(.*?)</tool_schema>"
+        match = re.search(tool_schema_pattern, content, flags=re.DOTALL)
+
+        if match:
+            original_text = match.group(0)  # Complete <tool_schema>...</tool_schema> block
+            schema_content = match.group(1)  # Content between the tags
+
+            # Parse tool schema
+            try:
+                tool_schema = json.loads(schema_content)
+                assert isinstance(tool_schema, list), "Tool schema must be a list[dict]"
+            except json.JSONDecodeError:
+                try:
+                    tool_schema = ast.literal_eval(schema_content)
+                    assert isinstance(tool_schema, list), "Tool schema must be a list[dict]"
+                except (ValueError, SyntaxError, AssertionError):
+                    logger.warning(
+                        f"[SystemParser] Failed to parse tool schema with both JSON and ast.literal_eval: {schema_content[:100]}..."
+                    )
+                    tool_schema = None
+            except AssertionError:
+                logger.warning(
+                    f"[SystemParser] Tool schema must be a list[dict]: {schema_content[:100]}..."
+                )
+                tool_schema = None
+
+            # Process and replace
+            if tool_schema is not None:
+
+                def remove_descriptions(obj):
+                    """Recursively remove all 'description' keys from a nested dict/list structure."""
+                    if isinstance(obj, dict):
+                        return {
+                            k: remove_descriptions(v) for k, v in obj.items() if k != "description"
+                        }
+                    elif isinstance(obj, list):
+                        return [remove_descriptions(item) for item in obj]
+                    else:
+                        return obj
+
+                def keep_first_layer_params(obj):
+                    """Only keep first layer parameter information, remove nested parameters."""
+                    if isinstance(obj, list):
+                        return [keep_first_layer_params(item) for item in obj]
+                    elif isinstance(obj, dict):
+                        result = {}
+                        for k, v in obj.items():
+                            if k == "properties" and isinstance(v, dict):
+                                # For properties, only keep first layer parameter names and types
+                                first_layer_props = {}
+                                for param_name, param_info in v.items():
+                                    if isinstance(param_info, dict):
+                                        # Only keep type and basic info, remove nested properties
+                                        first_layer_props[param_name] = {
+                                            key: val
+                                            for key, val in param_info.items()
+                                            if key in ["type", "enum", "required"]
+                                            and key != "properties"
+                                        }
+                                    else:
+                                        first_layer_props[param_name] = param_info
+                                result[k] = first_layer_props
+                            elif k == "parameters" and isinstance(v, dict):
+                                # Process parameters object but only keep first layer
+                                result[k] = keep_first_layer_params(v)
+                            elif isinstance(v, dict | list) and k != "properties":
+                                result[k] = keep_first_layer_params(v)
+                            else:
+                                result[k] = v
+                        return result
+                    else:
+                        return obj
+
+                def format_tool_schema_readable(tool_schema):
+                    """Convert tool schema to readable format: tool_name: [param1 (type1), ...](required: ...)"""
+                    lines = []
+                    for tool in tool_schema:
+                        if not tool:
+                            continue
+
+                        # Handle both new format and old-style OpenAI function format
+                        if tool.get("type") == "function" and "function" in tool:
+                            tool_info = tool.get("function")
+                            if not tool_info:
+                                continue
+                        else:
+                            tool_info = tool
+
+                        tool_name = tool_info.get("name", "unknown")
+                        params_obj = tool_info.get("parameters", {})
+                        properties = params_obj.get("properties", {})
+                        required = params_obj.get("required", [])
+
+                        # Format parameters
+                        param_strs = []
+                        for param_name, param_info in properties.items():
+                            if isinstance(param_info, dict):
+                                param_type = param_info.get("type", "any")
+                                # Handle enum
+                                if "enum" in param_info and param_info["enum"] is not None:
+                                    # Ensure all enum values are strings
+                                    enum_values = [str(v) for v in param_info["enum"]]
+                                    param_type = f"{param_type}[{', '.join(enum_values)}]"
+                                param_strs.append(f"{param_name} ({param_type})")
+                            else:
+                                param_strs.append(f"{param_name} (any)")
+
+                        # Format required parameters
+                        # Ensure all required parameter names are strings
+                        required_strs = [str(r) for r in required] if required else []
+                        required_str = (
+                            f"(required: {', '.join(required_strs)})" if required_strs else ""
+                        )
+
+                        # Construct the line
+                        params_part = f"[{', '.join(param_strs)}]" if param_strs else "[]"
+                        line = f"{tool_name}: {params_part}{required_str}"
+                        lines.append(line)
+
+                    return "\n".join(lines)
+
+                # First keep only first layer params, then remove descriptions
+                simple_tool_schema = keep_first_layer_params(tool_schema)
+                simple_tool_schema = remove_descriptions(simple_tool_schema)
+                # change to readable format
+                readable_schema = format_tool_schema_readable(simple_tool_schema)
+
+                processed_text = f"<tool_schema>{readable_schema}</tool_schema>"
+                content = content.replace(original_text, processed_text, 1)
+
+        parts = ["system: "]
+        if message.get("chat_time"):
+            parts.append(f"[{message.get('chat_time')}]: ")
+        prefix = "".join(parts)
+        msg_line = f"{prefix}{content}\n"
 
         source = self.create_source(message, info)
 
@@ -104,7 +234,7 @@ def parse_fast(
         session_id = info_.pop("session_id", "")
 
         # Split parsed text into chunks
-        content_chunks = self._split_text(content_wo_tool_schema)
+        content_chunks = self._split_text(msg_line)
 
         memory_items = []
         for _chunk_idx, chunk_text in enumerate(content_chunks):
@@ -132,9 +262,9 @@ def parse_fine(
         info: dict[str, Any],
         **kwargs,
     ) -> list[TextualMemoryItem]:
-        content = message["content"]
+        content = message.get("content", "")
         if isinstance(content, dict):
-            content = content["text"]
+            content = content.get("text", "")
         try:
             tool_schema = json.loads(content)
             assert isinstance(tool_schema, list), "Tool schema must be a list[dict]"

src/memos/memories/textual/preference.py

@@ -248,7 +248,7 @@ def get_all(self) -> list[TextualMemoryItem]:
         Returns:
             list[TextualMemoryItem]: List of all memories.
         """
-        all_collections = self.vector_db.list_collections()
+        all_collections = ["explicit_preference", "implicit_preference"]
         all_memories = {}
         for collection_name in all_collections:
             items = self.vector_db.get_all(collection_name)

src/memos/memories/textual/simple_preference.py

@@ -90,7 +90,7 @@ def get_with_collection_name(
                 return None
             return TextualMemoryItem(
                 id=res.id,
-                memory=res.payload.get("dialog_str", ""),
+                memory=res.memory,
                 metadata=PreferenceTextualMemoryMetadata(**res.payload),
             )
         except Exception as e:
@@ -116,7 +116,7 @@ def get_by_ids_with_collection_name(
             return [
                 TextualMemoryItem(
                     id=memo.id,
-                    memory=memo.payload.get("dialog_str", ""),
+                    memory=memo.memory,
                     metadata=PreferenceTextualMemoryMetadata(**memo.payload),
                 )
                 for memo in res
@@ -132,14 +132,14 @@ def get_all(self) -> list[TextualMemoryItem]:
         Returns:
             list[TextualMemoryItem]: List of all memories.
         """
-        all_collections = self.vector_db.list_collections()
+        all_collections = ["explicit_preference", "implicit_preference"]
         all_memories = {}
         for collection_name in all_collections:
             items = self.vector_db.get_all(collection_name)
             all_memories[collection_name] = [
                 TextualMemoryItem(
                     id=memo.id,
-                    memory=memo.payload.get("dialog_str", ""),
+                    memory=memo.memory,
                     metadata=PreferenceTextualMemoryMetadata(**memo.payload),
                 )
                 for memo in items