Merge branch 'dev' into feat/redis_scheduler

Author: glin93

src/memos/api/handlers/base_handler.py

@@ -161,6 +161,11 @@ def mos_server(self):
         """Get MOS server instance."""
         return self.deps.mos_server
 
+    @property
+    def deepsearch_agent(self):
+        """Get deepsearch agent instance."""
+        return self.deps.deepsearch_agent
+
     def _validate_dependencies(self, *required_deps: str) -> None:
         """
         Validate that required dependencies are available.

src/memos/api/handlers/chat_handler.py

@@ -402,9 +402,9 @@ def generate_chat_response() -> Generator[str, None, None]:
 
                     # Prepare preference markdown string
                     if chat_req.include_preference:
-                        pref_md_string = self._build_pref_md_string_for_playground(
-                            search_response.data["pref_mem"][0].get("memories", [])
-                        )
+                        pref_list = search_response.data.get("pref_mem") or []
+                        pref_memories = pref_list[0].get("memories", []) if pref_list else []
+                        pref_md_string = self._build_pref_md_string_for_playground(pref_memories)
                         yield f"data: {json.dumps({'type': 'pref_md_string', 'data': pref_md_string})}\n\n"
 
                     # Step 2: Build system prompt with memories
@@ -425,8 +425,6 @@ def generate_chat_response() -> Generator[str, None, None]:
                         f"current_system_prompt: {system_prompt}"
                     )
 
-                    yield f"data: {json.dumps({'type': 'status', 'data': '2'})}\n\n"
-
                     # Step 3: Generate streaming response from LLM
                     if (
                         chat_req.model_name_or_path
@@ -448,9 +446,11 @@ def generate_chat_response() -> Generator[str, None, None]:
                     for chunk in response_stream:
                         if chunk == "<think>":
                             in_think = True
+                            yield f"data: {json.dumps({'type': 'status', 'data': 'reasoning'})}\n\n"
                             continue
                         if chunk == "</think>":
                             in_think = False
+                            yield f"data: {json.dumps({'type': 'status', 'data': '2'})}\n\n"
                             continue
 
                         if in_think:
@@ -564,17 +564,17 @@ def _build_pref_md_string_for_playground(self, pref_mem_list: list[any]) -> str:
         explicit = []
         implicit = []
         for pref_mem in pref_mem_list:
-            if pref_mem["metadata"]["preference_type"] == "explicit":
+            if pref_mem["metadata"]["preference_type"] == "explicit_preference":
                 explicit.append(
                     {
-                        "content": pref_mem["preference"],
+                        "content": pref_mem["metadata"]["preference"],
                         "reasoning": pref_mem["metadata"]["reasoning"],
                     }
                 )
-            elif pref_mem["metadata"]["preference_type"] == "implicit":
+            elif pref_mem["metadata"]["preference_type"] == "implicit_preference":
                 implicit.append(
                     {
-                        "content": pref_mem["preference"],
+                        "content": pref_mem["metadata"]["preference"],
                         "reasoning": pref_mem["metadata"]["reasoning"],
                     }
                 )

src/memos/api/handlers/component_init.py

@@ -45,6 +45,7 @@
 
 if TYPE_CHECKING:
     from memos.memories.textual.tree import TreeTextMemory
+from memos.mem_agent.deepsearch_agent import DeepSearchMemAgent
 from memos.memories.textual.tree_text_memory.retrieve.internet_retriever_factory import (
     InternetRetrieverFactory,
 )
@@ -325,6 +326,10 @@ def init_server() -> dict[str, Any]:
         online_bot = get_online_bot_function() if dingding_enabled else None
         logger.info("DingDing bot is enabled")
 
+    deepsearch_agent = DeepSearchMemAgent(
+        llm=llm,
+        memory_retriever=tree_mem,
+    )
     # Return all components as a dictionary for easy access and extension
     return {
         "graph_db": graph_db,
@@ -349,4 +354,5 @@ def init_server() -> dict[str, Any]:
         "pref_mem": pref_mem,
         "online_bot": online_bot,
         "redis_client": redis_client,
+        "deepsearch_agent": deepsearch_agent,
     }

src/memos/api/handlers/search_handler.py

@@ -31,7 +31,9 @@ def __init__(self, dependencies: HandlerDependencies):
             dependencies: HandlerDependencies instance
         """
         super().__init__(dependencies)
-        self._validate_dependencies("naive_mem_cube", "mem_scheduler", "searcher")
+        self._validate_dependencies(
+            "naive_mem_cube", "mem_scheduler", "searcher", "deepsearch_agent"
+        )
 
     def handle_search_memories(self, search_req: APISearchRequest) -> SearchResponse:
         """
@@ -52,9 +54,10 @@ def handle_search_memories(self, search_req: APISearchRequest) -> SearchResponse
 
         results = cube_view.search_memories(search_req)
 
-        self.logger.info(f"[AddHandler] Final add results count={len(results)}")
+        self.logger.info(f"[SearchHandler] Final search results count={len(results)}")
+
         return SearchResponse(
-            message="Memory searched successfully",
+            message="Search completed successfully",
             data=results,
         )
 
@@ -82,6 +85,7 @@ def _build_cube_view(self, search_req: APISearchRequest) -> MemCubeView:
                 mem_scheduler=self.mem_scheduler,
                 logger=self.logger,
                 searcher=self.searcher,
+                deepsearch_agent=self.deepsearch_agent,
             )
         else:
             single_views = [
@@ -92,6 +96,7 @@ def _build_cube_view(self, search_req: APISearchRequest) -> MemCubeView:
                     mem_scheduler=self.mem_scheduler,
                     logger=self.logger,
                     searcher=self.searcher,
+                    deepsearch_agent=self.deepsearch_agent,
                 )
                 for cube_id in cube_ids
             ]

src/memos/mem_agent/deepsearch_agent.py

@@ -26,6 +26,8 @@
 if TYPE_CHECKING:
     from memos.types import MessageList
 
+logger = get_logger(__name__)
+
 
 class JSONResponseParser:
     """Elegant JSON response parser for LLM outputs"""
@@ -48,9 +50,6 @@ def parse(response: str) -> dict[str, Any]:
         raise ValueError(f"Cannot parse JSON response: {response[:100]}...")
 
 
-logger = get_logger(__name__)
-
-
 class QueryRewriter(BaseMemAgent):
     """Specialized agent for rewriting queries based on conversation history"""
 
@@ -141,7 +140,7 @@ def __init__(
             memory_retriever: Memory retrieval interface (e.g., naive_mem_cube.text_mem)
             config: Configuration for deep search behavior
         """
-        self.config = config or DeepSearchAgentConfig()
+        self.config = config or DeepSearchAgentConfig(agent_name="DeepSearchMemAgent")
         self.max_iterations = self.config.max_iterations
         self.timeout = self.config.timeout
         self.llm: BaseLLM = llm
@@ -219,7 +218,7 @@ def run(self, query: str, **kwargs) -> str | list[TextualMemoryItem]:
             return self._remove_duplicate_memories(accumulated_memories)
         else:
             return self._generate_final_answer(
-                query, accumulated_memories, accumulated_context, "", history
+                query, accumulated_memories, accumulated_context, history
             )
 
     def _remove_duplicate_memories(
@@ -248,9 +247,9 @@ def _generate_final_answer(
         original_query: str,
         search_results: list[TextualMemoryItem],
         context: list[str],
-        missing_info: str = "",
         history: list[str] | None = None,
         sources: list[str] | None = None,
+        missing_info: str | None = None,
     ) -> str:
         """
         Generate the final answer.

src/memos/multi_mem_cube/single_cube.py

@@ -45,6 +45,7 @@ class SingleCubeView(MemCubeView):
     mem_scheduler: Any
     logger: Any
     searcher: Any
+    deepsearch_agent: Any
 
     def add_memories(self, add_req: APIADDRequest) -> list[dict[str, Any]]:
         """
@@ -202,6 +203,15 @@ def _deep_search(
         formatted_memories = [format_memory_item(data) for data in enhanced_memories]
         return formatted_memories
 
+    def _deep_search(
+        self, search_req: APISearchRequest, user_context: UserContext, max_thinking_depth: int
+    ) -> list:
+        deepsearch_results = self.deepsearch_agent.run(
+            search_req.query, user_id=user_context.mem_cube_id
+        )
+        formatted_memories = [format_memory_item(data) for data in deepsearch_results]
+        return formatted_memories
+
     def _fine_search(
         self,
         search_req: APISearchRequest,

src/memos/templates/mem_agent_prompts.py

@@ -22,12 +22,14 @@
 {context}
 
 Analyze the context and determine the next step. Return your response in JSON format with the following structure:
-{{
+ ```json
+ {{
     "status": "sufficient|missing_info|needs_raw",
     "reasoning": "Brief explanation of your decision",
     "missing_entities": ["entity1", "entity2"],
     "new_search_query": "new search query",
 }}
+```
 
 Status definitions:
 - "sufficient": Context fully answers the query

tests/api/test_server_router.py

@@ -49,6 +49,7 @@ def mock_init_server():
         "online_bot": None,
         "chat_llms": Mock(),
         "redis_client": Mock(),
+        "deepsearch_agent": Mock(),
     }
 
     with patch("memos.api.handlers.init_server", return_value=mock_components):