feat: Optimize mixture search and enhance API client

Author: tangg555

src/memos/mem_scheduler/base_scheduler.py

@@ -6,6 +6,7 @@
 from collections.abc import Callable
 from datetime import datetime
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 from sqlalchemy.engine import Engine
 
@@ -50,6 +51,10 @@
 from memos.templates.mem_scheduler_prompts import MEMORY_ASSEMBLY_TEMPLATE
 
 
+if TYPE_CHECKING:
+    from memos.mem_cube.base import BaseMemCube
+
+
 logger = get_logger(__name__)
 
 
@@ -124,7 +129,7 @@ def __init__(self, config: BaseSchedulerConfig):
         self._context_lock = threading.Lock()
         self.current_user_id: UserID | str | None = None
         self.current_mem_cube_id: MemCubeID | str | None = None
-        self.current_mem_cube: GeneralMemCube | None = None
+        self.current_mem_cube: BaseMemCube | None = None
         self.auth_config_path: str | Path | None = self.config.get("auth_config_path", None)
         self.auth_config = None
         self.rabbitmq_config = None

src/memos/mem_scheduler/general_modules/api_misc.py

@@ -16,16 +16,20 @@
 
 
 class SchedulerAPIModule(BaseSchedulerModule):
-    def __init__(self, window_size=5):
+    def __init__(self, window_size: int | None = None, history_memory_turns: int | None = None):
         super().__init__()
         self.window_size = window_size
+        self.history_memory_turns = history_memory_turns
         self.search_history_managers: dict[str, APIRedisDBManager] = {}
-        self.pre_memory_turns = 5
 
     def get_search_history_manager(self, user_id: str, mem_cube_id: str) -> APIRedisDBManager:
         """Get or create a Redis manager for search history."""
+        logger.info(
+            f"Getting search history manager for user_id: {user_id}, mem_cube_id: {mem_cube_id}"
+        )
         key = f"search_history:{user_id}:{mem_cube_id}"
         if key not in self.search_history_managers:
+            logger.info(f"Creating new search history manager for key: {key}")
             self.search_history_managers[key] = APIRedisDBManager(
                 user_id=user_id,
                 mem_cube_id=mem_cube_id,
@@ -43,6 +47,9 @@ def sync_search_data(
         formatted_memories: Any,
         conversation_id: str | None = None,
     ) -> Any:
+        logger.info(
+            f"Syncing search data for item_id: {item_id}, user_id: {user_id}, mem_cube_id: {mem_cube_id}"
+        )
         # Get the search history manager
         manager = self.get_search_history_manager(user_id, mem_cube_id)
         manager.sync_with_redis(size_limit=self.window_size)
@@ -101,37 +108,22 @@ def sync_search_data(
         manager.sync_with_redis(size_limit=self.window_size)
         return manager
 
-    def get_pre_memories(self, user_id: str, mem_cube_id: str) -> list:
-        """
-        Get pre-computed memories from the most recent completed search entry.
-
-        Args:
-            user_id: User identifier
-            mem_cube_id: Memory cube identifier
-
-        Returns:
-            List of TextualMemoryItem objects from the most recent completed search
-        """
-        manager = self.get_search_history_manager(user_id, mem_cube_id)
-
-        existing_data = manager.load_from_db()
-        if existing_data is None:
-            return []
-
-        search_history: APISearchHistoryManager = existing_data
-
-        # Get memories from the most recent completed entry
-        history_memories = search_history.get_history_memories(turns=self.pre_memory_turns)
-        return history_memories
-
-    def get_history_memories(self, user_id: str, mem_cube_id: str, n: int) -> list:
+    def get_history_memories(
+        self, user_id: str, mem_cube_id: str, turns: int | None = None
+    ) -> list:
         """Get history memories for backward compatibility with tests."""
+        logger.info(
+            f"Getting history memories for user_id: {user_id}, mem_cube_id: {mem_cube_id}, turns: {turns}"
+        )
         manager = self.get_search_history_manager(user_id, mem_cube_id)
         existing_data = manager.load_from_db()
 
         if existing_data is None:
             return []
 
+        if turns is None:
+            turns = self.history_memory_turns
+
         # Handle different data formats
         if isinstance(existing_data, APISearchHistoryManager):
             search_history = existing_data
@@ -142,4 +134,4 @@ def get_history_memories(self, user_id: str, mem_cube_id: str, n: int) -> list:
             except Exception:
                 return []
 
-        return search_history.get_history_memories(turns=n)
+        return search_history.get_history_memories(turns=turns)

src/memos/mem_scheduler/optimized_scheduler.py

@@ -1,11 +1,13 @@
 import json
+import os
 
 from typing import TYPE_CHECKING
 
 from memos.api.product_models import APISearchRequest
 from memos.configs.mem_scheduler import GeneralSchedulerConfig
 from memos.log import get_logger
 from memos.mem_cube.general import GeneralMemCube
+from memos.mem_cube.navie import NaiveMemCube
 from memos.mem_scheduler.general_modules.api_misc import SchedulerAPIModule
 from memos.mem_scheduler.general_scheduler import GeneralScheduler
 from memos.mem_scheduler.schemas.general_schemas import (
@@ -23,6 +25,7 @@
 
 if TYPE_CHECKING:
     from memos.mem_scheduler.schemas.monitor_schemas import MemoryMonitorItem
+    from memos.memories.textual.tree_text_memory.retrieve.searcher import Searcher
     from memos.reranker.http_bge import HTTPBGEReranker
 
 
@@ -34,43 +37,19 @@ class OptimizedScheduler(GeneralScheduler):
 
     def __init__(self, config: GeneralSchedulerConfig):
         super().__init__(config)
-        self.api_module = SchedulerAPIModule()
+        self.window_size = int(os.getenv("API_SEARCH_WINDOW_SIZE", 5))
+        self.history_memory_turns = int(os.getenv("API_SEARCH_HISTORY_TURNS", 5))
+
+        self.api_module = SchedulerAPIModule(
+            window_size=self.window_size,
+            history_memory_turns=self.history_memory_turns,
+        )
         self.register_handlers(
             {
                 API_MIX_SEARCH_LABEL: self._api_mix_search_message_consumer,
             }
         )
 
-    def search_memories(
-        self,
-        search_req: APISearchRequest,
-        user_context: UserContext,
-        mem_cube: GeneralMemCube,
-        mode: SearchMode,
-    ):
-        """Fine search memories function copied from server_router to avoid circular import"""
-        target_session_id = search_req.session_id
-        if not target_session_id:
-            target_session_id = "default_session"
-        search_filter = {"session_id": search_req.session_id} if search_req.session_id else None
-
-        # Create MemCube and perform search
-        search_results = mem_cube.text_mem.search(
-            query=search_req.query,
-            user_name=user_context.mem_cube_id,
-            top_k=search_req.top_k,
-            mode=mode,
-            manual_close_internet=not search_req.internet_search,
-            moscube=search_req.moscube,
-            search_filter=search_filter,
-            info={
-                "user_id": search_req.user_id,
-                "session_id": target_session_id,
-                "chat_history": search_req.chat_history,
-            },
-        )
-        return search_results
-
     def submit_memory_history_async_task(
         self,
         search_req: APISearchRequest,
@@ -110,6 +89,36 @@ def submit_memory_history_async_task(
         logger.info(f"Submitted async fine search task for user {search_req.user_id}")
         return async_task_id
 
+    def search_memories(
+        self,
+        search_req: APISearchRequest,
+        user_context: UserContext,
+        mem_cube: NaiveMemCube,
+        mode: SearchMode,
+    ):
+        """Fine search memories function copied from server_router to avoid circular import"""
+        target_session_id = search_req.session_id
+        if not target_session_id:
+            target_session_id = "default_session"
+        search_filter = {"session_id": search_req.session_id} if search_req.session_id else None
+
+        # Create MemCube and perform search
+        search_results = mem_cube.text_mem.search(
+            query=search_req.query,
+            user_name=user_context.mem_cube_id,
+            top_k=search_req.top_k,
+            mode=mode,
+            manual_close_internet=not search_req.internet_search,
+            moscube=search_req.moscube,
+            search_filter=search_filter,
+            info={
+                "user_id": search_req.user_id,
+                "session_id": target_session_id,
+                "chat_history": search_req.chat_history,
+            },
+        )
+        return search_results
+
     def mix_search_memories(
         self,
         search_req: APISearchRequest,
@@ -122,12 +131,33 @@ def mix_search_memories(
         # Get mem_cube for fast search
         mem_cube = self.current_mem_cube
 
-        # Perform fast search
-        fast_memories = self.search_memories(
-            search_req=search_req,
-            user_context=user_context,
-            mem_cube=mem_cube,
+        target_session_id = search_req.session_id
+        if not target_session_id:
+            target_session_id = "default_session"
+        search_filter = {"session_id": search_req.session_id} if search_req.session_id else None
+
+        text_mem: TreeTextMemory = mem_cube.text_mem
+        searcher: Searcher = text_mem.get_searcher(
+            manual_close_internet=not search_req.internet_search,
+            moscube=False,
+        )
+        # Rerank Memories - reranker expects TextualMemoryItem objects
+        reranker: HTTPBGEReranker = text_mem.reranker
+        info = {
+            "user_id": search_req.user_id,
+            "session_id": target_session_id,
+            "chat_history": search_req.chat_history,
+        }
+
+        fast_retrieved_memories = searcher.retrieve(
+            query=search_req.query,
+            user_name=user_context.mem_cube_id,
+            top_k=search_req.top_k,
             mode=SearchMode.FAST,
+            manual_close_internet=not search_req.internet_search,
+            moscube=search_req.moscube,
+            search_filter=search_filter,
+            info=info,
         )
 
         self.submit_memory_history_async_task(
@@ -136,76 +166,69 @@ def mix_search_memories(
         )
 
         # Try to get pre-computed fine memories if available
-        pre_fine_memories = self.api_module.get_pre_memories(
-            user_id=search_req.user_id, mem_cube_id=user_context.mem_cube_id
+        history_memories = self.api_module.get_history_memories(
+            user_id=search_req.user_id,
+            mem_cube_id=user_context.mem_cube_id,
+            turns=self.history_memory_turns,
         )
-        if not pre_fine_memories:
+        if not history_memories:
+            fast_memories = searcher.post_retrieve(
+                retrieved_results=fast_retrieved_memories,
+                top_k=search_req.top_k,
+                user_name=user_context.mem_cube_id,
+                info=info,
+            )
             # Format fast memories for return
             formatted_memories = [format_textual_memory_item(data) for data in fast_memories]
             return formatted_memories
 
-        # Merge fast and pre-computed fine memories (both are TextualMemoryItem objects)
-        combined_memories = fast_memories + pre_fine_memories
-        # Remove duplicates based on memory content
-        seen_contents = set()
-        unique_memories = []
-        for memory in combined_memories:
-            # Both fast_memories and pre_fine_memories are TextualMemoryItem objects
-            content_key = memory.memory  # Use .memory attribute instead of .get("content", "")
-            if content_key not in seen_contents:
-                seen_contents.add(content_key)
-                unique_memories.append(memory)
-
-        # Rerank Memories - reranker expects TextualMemoryItem objects
-        reranker: HTTPBGEReranker = mem_cube.text_mem.reranker
-
-        # Use search_req parameters for reranking
-        target_session_id = search_req.session_id
-        if not target_session_id:
-            target_session_id = "default_session"
-        search_filter = {"session_id": search_req.session_id} if search_req.session_id else None
-
-        sorted_results = reranker.rerank(
+        sorted_history_memories = reranker.rerank(
             query=search_req.query,  # Use search_req.query instead of undefined query
-            graph_results=unique_memories,  # Pass TextualMemoryItem objects directly
+            graph_results=history_memories,  # Pass TextualMemoryItem objects directly
             top_k=search_req.top_k,  # Use search_req.top_k instead of undefined top_k
             search_filter=search_filter,
         )
 
+        sorted_results = fast_retrieved_memories + sorted_history_memories
+        final_results = searcher.post_retrieve(
+            retrieved_results=sorted_results,
+            top_k=search_req.top_k,
+            user_name=user_context.mem_cube_id,
+            info=info,
+        )
+
         formatted_memories = [
-            format_textual_memory_item(item) for item, score in sorted_results[: search_req.top_k]
+            format_textual_memory_item(item) for item in final_results[: search_req.top_k]
         ]
 
         return formatted_memories
 
     def update_search_memories_to_redis(
         self,
-        user_id: str,
-        mem_cube_id: str,
         messages: list[ScheduleMessageItem],
     ):
-        mem_cube = messages[0].mem_cube
+        mem_cube: NaiveMemCube = self.current_mem_cube
 
         for msg in messages:
             content_dict = json.loads(msg.content)
             search_req = content_dict["search_req"]
             user_context = content_dict["user_context"]
 
-            fine_memories: list[TextualMemoryItem] = self.search_memories(
+            memories: list[TextualMemoryItem] = self.search_memories(
                 search_req=APISearchRequest(**content_dict["search_req"]),
                 user_context=UserContext(**content_dict["user_context"]),
                 mem_cube=mem_cube,
-                mode=SearchMode.FINE,
+                mode=SearchMode.FAST,
             )
-            formatted_memories = [format_textual_memory_item(data) for data in fine_memories]
+            formatted_memories = [format_textual_memory_item(data) for data in memories]
 
             # Sync search data to Redis
             self.api_module.sync_search_data(
                 item_id=msg.item_id,
                 user_id=search_req["user_id"],
                 mem_cube_id=user_context["mem_cube_id"],
                 query=search_req["query"],
-                memories=fine_memories,
+                memories=memories,
                 formatted_memories=formatted_memories,
             )
 
@@ -228,9 +251,7 @@ def _api_mix_search_message_consumer(self, messages: list[ScheduleMessageItem])
                 messages = grouped_messages[user_id][mem_cube_id]
                 if len(messages) == 0:
                     return
-                self.update_search_memories_to_redis(
-                    user_id=user_id, mem_cube_id=mem_cube_id, messages=messages
-                )
+                self.update_search_memories_to_redis(messages=messages)
 
     def replace_working_memory(
         self,