Feat/tool memory (#583)

* function call supoort

* add tool parser

* rename multi model to modal

* rename multi modal

* tool mem support

* modify multi-modal code

* pref support multi-modal messages

* modify bug in chat handle

* fix pre commit

* modify code

* add tool search

* tool search

* add split chunck in system and tool

---------

Co-authored-by: yuan.wang <[email protected]>

Author: Wang-Daoji

docker/requirements.txt

@@ -160,3 +160,4 @@ xlrd==2.0.2
 xlsxwriter==3.2.5
 prometheus-client==0.23.1
 pymilvus==2.5.12
+langchain-text-splitters==1.0.0

src/memos/api/handlers/chat_handler.py

@@ -142,7 +142,9 @@ def handle_chat_complete(self, chat_req: APIChatCompleteRequest) -> dict[str, An
 
             # Step 2: Build system prompt
             system_prompt = self._build_system_prompt(
-                filtered_memories, search_response.data["pref_string"], chat_req.system_prompt
+                filtered_memories,
+                search_response.data.get("pref_string", ""),
+                chat_req.system_prompt,
             )
 
             # Prepare message history
@@ -257,7 +259,7 @@ def generate_chat_response() -> Generator[str, None, None]:
                     # Step 2: Build system prompt with memories
                     system_prompt = self._build_system_prompt(
                         filtered_memories,
-                        search_response.data["pref_string"],
+                        search_response.data.get("pref_string", ""),
                         chat_req.system_prompt,
                     )
 
@@ -449,7 +451,7 @@ def generate_chat_response() -> Generator[str, None, None]:
 
                     # Step 2: Build system prompt with memories
                     system_prompt = self._build_enhance_system_prompt(
-                        filtered_memories, search_response.data["pref_string"]
+                        filtered_memories, search_response.data.get("pref_string", "")
                     )
 
                     # Prepare messages

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

@@ -6,7 +6,7 @@
 
 # Import message types from core types module
 from memos.log import get_logger
-from memos.types import MessageDict, PermissionDict, SearchMode
+from memos.types import MessageList, MessagesType, PermissionDict, SearchMode
 
 
 logger = get_logger(__name__)
@@ -56,7 +56,7 @@ class Message(BaseModel):
 
 class MemoryCreate(BaseRequest):
     user_id: str = Field(..., description="User ID")
-    messages: list | None = Field(None, description="List of messages to store.")
+    messages: MessageList | None = Field(None, description="List of messages to store.")
     memory_content: str | None = Field(None, description="Content to store as memory")
     doc_path: str | None = Field(None, description="Path to document to store")
     mem_cube_id: str | None = Field(None, description="ID of the memory cube")
@@ -83,7 +83,7 @@ class ChatRequest(BaseRequest):
     writable_cube_ids: list[str] | None = Field(
         None, description="List of cube IDs user can write for multi-cube chat"
     )
-    history: list | None = Field(None, description="Chat history")
+    history: MessageList | None = Field(None, description="Chat history")
     mode: SearchMode = Field(SearchMode.FAST, description="search mode: fast, fine, or mixture")
     system_prompt: str | None = Field(None, description="Base system prompt to use for chat")
     top_k: int = Field(10, description="Number of results to return")
@@ -165,7 +165,7 @@ class ChatCompleteRequest(BaseRequest):
     user_id: str = Field(..., description="User ID")
     query: str = Field(..., description="Chat query message")
     mem_cube_id: str | None = Field(None, description="Cube ID to use for chat")
-    history: list | None = Field(None, description="Chat history")
+    history: MessageList | None = Field(None, description="Chat history")
     internet_search: bool = Field(False, description="Whether to use internet search")
     system_prompt: str | None = Field(None, description="Base prompt to use for chat")
     top_k: int = Field(10, description="Number of results to return")
@@ -251,7 +251,7 @@ class MemoryCreateRequest(BaseRequest):
     """Request model for creating memories."""
 
     user_id: str = Field(..., description="User ID")
-    messages: str | list | None = Field(None, description="List of messages to store.")
+    messages: str | MessagesType | None = Field(None, description="List of messages to store.")
     memory_content: str | None = Field(None, description="Memory content to store")
     doc_path: str | None = Field(None, description="Path to document to store")
     mem_cube_id: str | None = Field(None, description="Cube ID")
@@ -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(
@@ -360,7 +375,7 @@ class APISearchRequest(BaseRequest):
     )
 
     # ==== Context ====
-    chat_history: list | None = Field(
+    chat_history: MessageList | None = Field(
         None,
         description=(
             "Historical chat messages used internally by algorithms. "
@@ -490,7 +505,7 @@ class APIADDRequest(BaseRequest):
     )
 
     # ==== Input content ====
-    messages: str | list | None = Field(
+    messages: MessagesType | None = Field(
         None,
         description=(
             "List of messages to store. Supports: "
@@ -506,7 +521,7 @@ class APIADDRequest(BaseRequest):
     )
 
     # ==== Chat history ====
-    chat_history: list | None = Field(
+    chat_history: MessageList | None = Field(
         None,
         description=(
             "Historical chat messages used internally by algorithms. "
@@ -636,7 +651,7 @@ class APIFeedbackRequest(BaseRequest):
         "default_session", description="Session ID for soft-filtering memories"
     )
     task_id: str | None = Field(None, description="Task ID for monitering async tasks")
-    history: list[MessageDict] | None = Field(..., description="Chat history")
+    history: MessageList | None = Field(..., description="Chat history")
     retrieved_memory_ids: list[str] | None = Field(
         None, description="Retrieved memory ids at last turn"
     )
@@ -670,7 +685,7 @@ class APIChatCompleteRequest(BaseRequest):
     writable_cube_ids: list[str] | None = Field(
         None, description="List of cube IDs user can write for multi-cube chat"
     )
-    history: list | None = Field(None, description="Chat history")
+    history: MessageList | None = Field(None, description="Chat history")
     mode: SearchMode = Field(SearchMode.FAST, description="search mode: fast, fine, or mixture")
     system_prompt: str | None = Field(None, description="Base system prompt to use for chat")
     top_k: int = Field(10, description="Number of results to return")
@@ -739,7 +754,7 @@ class SuggestionRequest(BaseRequest):
     user_id: str = Field(..., description="User ID")
     mem_cube_id: str = Field(..., description="Cube ID")
     language: Literal["zh", "en"] = Field("zh", description="Language for suggestions")
-    message: list | None = Field(None, description="List of messages to store.")
+    message: MessagesType | None = Field(None, description="List of messages to store.")
 
 
 # ─── MemOS Client Response Models ──────────────────────────────────────────────

src/memos/mem_reader/multi_modal_struct.py

@@ -1,4 +1,5 @@
 import concurrent.futures
+import json
 import traceback
 
 from typing import Any
@@ -7,8 +8,9 @@
 from memos.configs.mem_reader import MultiModalStructMemReaderConfig
 from memos.context.context import ContextThreadPoolExecutor
 from memos.mem_reader.read_multi_modal import MultiModalParser
-from memos.mem_reader.simple_struct import SimpleStructMemReader
+from memos.mem_reader.simple_struct import SimpleStructMemReader, detect_lang
 from memos.memories.textual.item import TextualMemoryItem
+from memos.templates.tool_mem_prompts import TOOL_TRAJECTORY_PROMPT_EN, TOOL_TRAJECTORY_PROMPT_ZH
 from memos.types import MessagesType
 from memos.utils import timed
 
@@ -297,6 +299,61 @@ def _process_string_fine(
 
         return fine_memory_items
 
+    def _get_llm_tool_trajectory_response(self, mem_str: str) -> dict:
+        """
+        Generete tool trajectory experience item by llm.
+        """
+        try:
+            lang = detect_lang(mem_str)
+            template = TOOL_TRAJECTORY_PROMPT_ZH if lang == "zh" else TOOL_TRAJECTORY_PROMPT_EN
+            prompt = template.replace("{messages}", mem_str)
+            rsp = self.llm.generate([{"role": "user", "content": prompt}])
+            rsp = rsp.replace("```json", "").replace("```", "")
+            return json.loads(rsp)
+        except Exception as e:
+            logger.error(f"[MultiModalFine] Error calling LLM for tool trajectory: {e}")
+            return []
+
+    def _process_tool_trajectory_fine(
+        self,
+        fast_memory_items: list[TextualMemoryItem],
+        info: dict[str, Any],
+    ) -> list[TextualMemoryItem]:
+        """
+        Process tool trajectory memory items through LLM to generate fine mode memories.
+        """
+        if not fast_memory_items:
+            return []
+
+        fine_memory_items = []
+
+        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:
+                continue
+            try:
+                resp = self._get_llm_tool_trajectory_response(mem_str)
+            except Exception as e:
+                logger.error(f"[MultiModalFine] Error calling LLM for tool trajectory: {e}")
+                continue
+            for m in resp:
+                try:
+                    # Normalize memory_type (same as simple_struct)
+                    memory_type = "ToolTrajectoryMemory"
+
+                    node = self._make_memory_item(
+                        value=m.get("trajectory", ""),
+                        info=info,
+                        memory_type=memory_type,
+                        tool_used_status=m.get("tool_used_status", []),
+                    )
+                    fine_memory_items.append(node)
+                except Exception as e:
+                    logger.error(f"[MultiModalFine] parse error for tool trajectory: {e}")
+
+        return fine_memory_items
+
     @timed
     def _process_multi_modal_data(
         self, scene_data_info: MessagesType, info, mode: str = "fine", **kwargs
@@ -339,6 +396,11 @@ def _process_multi_modal_data(
             )
             fine_memory_items.extend(fine_memory_items_string_parser)
 
+            fine_memory_items_tool_trajectory_parser = self._process_tool_trajectory_fine(
+                fast_memory_items, info
+            )
+            fine_memory_items.extend(fine_memory_items_tool_trajectory_parser)
+
             # Part B: get fine multimodal items
             for fast_item in fast_memory_items:
                 sources = fast_item.metadata.sources
@@ -377,6 +439,12 @@ def _process_transfer_multi_modal_data(
         # 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)
+
+        fine_memory_items_tool_trajectory_parser = self._process_tool_trajectory_fine(
+            [raw_node], info
+        )
+        fine_memory_items.extend(fine_memory_items_tool_trajectory_parser)
+
         # Part B: get fine multimodal items
         for source in sources:
             items = self.multi_modal_parser.process_transfer(