add tool search

Author: yuan.wang

src/memos/api/handlers/formatters_handler.py

@@ -90,3 +90,37 @@ def post_process_pref_mem(
         memories_result["pref_note"] = pref_note
 
     return memories_result
+
+
+def post_process_textual_mem(
+    memories_result: dict[str, Any],
+    text_formatted_mem: list[dict[str, Any]],
+    mem_cube_id: str,
+) -> dict[str, Any]:
+    """
+    Post-process text and tool memory results.
+    """
+    fact_mem = [
+        mem
+        for mem in text_formatted_mem
+        if mem["metadata"]["memory_type"] not in ["ToolSchemaMemory", "ToolTrajectoryMemory"]
+    ]
+    tool_mem = [
+        mem
+        for mem in text_formatted_mem
+        if mem["metadata"]["memory_type"] in ["ToolSchemaMemory", "ToolTrajectoryMemory"]
+    ]
+
+    memories_result["text_mem"].append(
+        {
+            "cube_id": mem_cube_id,
+            "memories": fact_mem,
+        }
+    )
+    memories_result["tool_mem"].append(
+        {
+            "cube_id": mem_cube_id,
+            "memories": tool_mem,
+        }
+    )
+    return memories_result

src/memos/api/product_models.py

@@ -326,6 +326,21 @@ class APISearchRequest(BaseRequest):
         ),
     )
 
+    search_tool_memory: bool = Field(
+        True,
+        description=(
+            "Whether to retrieve tool memories along with general memories. "
+            "If enabled, the system will automatically recall tool memories "
+            "relevant to the query. Default: True."
+        ),
+    )
+
+    tool_mem_top_k: int = Field(
+        6,
+        ge=0,
+        description="Number of tool memories to retrieve (top-K). Default: 6.",
+    )
+
     # ==== Filter conditions ====
     # TODO: maybe add detailed description later
     filter: dict[str, Any] | None = Field(

src/memos/mem_reader/read_multi_modal/tool_parser.py

@@ -135,91 +135,6 @@ def rebuild_from_source(
     ) -> ChatCompletionToolMessageParam:
         """Rebuild tool message from SourceMessage."""
 
-        # Priority 1: Use original_part if available
-        if hasattr(source, "original_part") and source.original_part:
-            original = source.original_part
-            # If it's a content part, wrap it in a message
-            if isinstance(original, dict) and "type" in original:
-                return {
-                    "role": source.role or "user",
-                    "tool_call_id": source.tool_call_id or "",
-                    "content": [original],
-                    "chat_time": source.chat_time,
-                    "message_id": source.message_id,
-                }
-            # If it's already a full message, return it
-            if isinstance(original, dict) and "role" in original:
-                return original
-
-        # Priority 2: Rebuild from source fields
-        if source.type == "text":
-            return {
-                "role": source.role or "tool",
-                "content": [
-                    {
-                        "type": "text",
-                        "text": source.content or "",
-                    }
-                ],
-                "chat_time": source.chat_time,
-                "message_id": source.message_id,
-            }
-        elif source.type == "file":
-            return {
-                "role": source.role or "tool",
-                "content": [
-                    {
-                        "type": "file",
-                        "file": {
-                            "file_id": source.file_id or "",
-                            "filename": source.filename or "",
-                            "file_data": source.content or "",
-                        },
-                    }
-                ],
-                "chat_time": source.chat_time,
-                "message_id": source.message_id,
-            }
-        elif source.type == "image_url":
-            return {
-                "role": source.role or "tool",
-                "content": [
-                    {
-                        "type": "image_url",
-                        "image_url": {
-                            "url": source.content or "",
-                            "detail": source.detail or "auto",
-                        },
-                    }
-                ],
-                "chat_time": source.chat_time,
-                "message_id": source.message_id,
-            }
-        elif source.type == "input_audio":
-            return {
-                "role": source.role or "tool",
-                "content": [
-                    {
-                        "type": "input_audio",
-                        "input_audio": {
-                            "data": source.content or "",
-                            "format": source.format or "wav",
-                        },
-                    }
-                ],
-                "chat_time": source.chat_time,
-                "message_id": source.message_id,
-            }
-
-        # Simple text message
-        return {
-            "role": "tool",
-            "content": source.content or "",
-            "tool_call_id": source.message_id or "",
-            "chat_time": source.chat_time,
-            "message_id": source.message_id,
-        }
-
     def parse_fast(
         self,
         message: ChatCompletionToolMessageParam,
@@ -261,25 +176,5 @@ def parse_fine(
         info: dict[str, Any],
         **kwargs,
     ) -> list[TextualMemoryItem]:
-        content = message.get("content", "")
-        if isinstance(content, list):
-            part_type = content[0].get("type", "")
-            if part_type == "text":
-                # text will fine parse in full chat content, no need to parse specially
-                return []
-            elif part_type == "file":
-                # use file content parser to parse file content, no need to parse here
-                return []
-            elif part_type == "image_url":
-                # TODO: use multi-modal llm to generate mem by image url
-                content = content[0].get("image_url", {}).get("url", "")
-                return []
-            elif part_type == "input_audio":
-                # TODO: unsupport audio for now
-                return []
-            else:
-                logger.warning(f"[ToolParser] Unsupported part type: {part_type}")
-                return []
-        else:
-            # simple string content message, fine parse in full chat content, no need to parse specially
-            return []
+        # tool message no special multimodal handling is required in fine mode.
+        return []

src/memos/mem_scheduler/optimized_scheduler.py

@@ -159,6 +159,8 @@ def mix_search_memories(
             search_filter=search_filter,
             search_priority=search_priority,
             info=info,
+            search_tool_memory=search_req.search_tool_memory,
+            tool_mem_top_k=search_req.tool_mem_top_k,
         )
 
         # Try to get pre-computed memories if available
@@ -182,6 +184,8 @@ def mix_search_memories(
             top_k=search_req.top_k,
             user_name=user_context.mem_cube_id,
             info=info,
+            search_tool_memory=search_req.search_tool_memory,
+            tool_mem_top_k=search_req.tool_mem_top_k,
         )
         memories = merged_memories[: search_req.top_k]
 

src/memos/memories/textual/tree.py

@@ -166,6 +166,8 @@ def search(
         search_priority: dict | None = None,
         search_filter: dict | None = None,
         user_name: str | None = None,
+        search_tool_memory: bool = False,
+        tool_mem_top_k: int = 6,
         **kwargs,
     ) -> list[TextualMemoryItem]:
         """Search for memories based on a query.
@@ -223,6 +225,8 @@ def search(
             search_priority,
             user_name=user_name,
             plugin=kwargs.get("plugin", False),
+            search_tool_memory=search_tool_memory,
+            tool_mem_top_k=tool_mem_top_k,
         )
 
     def get_relevant_subgraph(

src/memos/memories/textual/tree_text_memory/retrieve/recall.py

@@ -59,7 +59,13 @@ def retrieve(
         Returns:
             list: Combined memory items.
         """
-        if memory_scope not in ["WorkingMemory", "LongTermMemory", "UserMemory"]:
+        if memory_scope not in [
+            "WorkingMemory",
+            "LongTermMemory",
+            "UserMemory",
+            "ToolSchemaMemory",
+            "ToolTrajectoryMemory",
+        ]:
             raise ValueError(f"Unsupported memory scope: {memory_scope}")
 
         if memory_scope == "WorkingMemory":