fix: resolve scheduler module import and Redis integration issues

Author: tangg555

src/memos/api/routers/server_router.py

@@ -1,3 +1,4 @@
+import json
 import os
 import traceback
 
@@ -29,7 +30,12 @@
 from memos.mem_reader.factory import MemReaderFactory
 from memos.mem_scheduler.orm_modules.base_model import BaseDBManager
 from memos.mem_scheduler.scheduler_factory import SchedulerFactory
-from memos.mem_scheduler.schemas.general_schemas import SearchMode
+from memos.mem_scheduler.schemas.general_schemas import (
+    API_MIX_SEARCH_LABEL,
+    SearchMode,
+)
+from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
+from memos.mem_scheduler.utils.db_utils import get_utc_now
 from memos.memories.textual.tree_text_memory.organize.manager import MemoryManager
 from memos.memories.textual.tree_text_memory.retrieve.internet_retriever_factory import (
     InternetRetrieverFactory,
@@ -101,6 +107,21 @@ def _get_default_memory_size(cube_config) -> dict[str, int]:
     }
 
 
+def _create_naive_mem_cube() -> NaiveMemCube:
+    """Create a NaiveMemCube instance with initialized components."""
+    naive_mem_cube = NaiveMemCube(
+        llm=llm,
+        embedder=embedder,
+        mem_reader=mem_reader,
+        graph_db=graph_db,
+        reranker=reranker,
+        internet_retriever=internet_retriever,
+        memory_manager=memory_manager,
+        default_cube_config=default_cube_config,
+    )
+    return naive_mem_cube
+
+
 def init_server():
     """Initialize server components and configurations."""
     # Get default cube configuration
@@ -152,6 +173,10 @@ def init_server():
     )
     mem_scheduler.start()
 
+    # Initialize SchedulerAPIModule
+    api_module = mem_scheduler.api_module
+
+    naive_mem_cube = _create_naive_mem_cube()
     return (
         graph_db,
         mem_reader,
@@ -163,6 +188,8 @@ def init_server():
         default_cube_config,
         mos_server,
         mem_scheduler,
+        naive_mem_cube,
+        api_module,
     )
 
 
@@ -178,24 +205,11 @@ def init_server():
     default_cube_config,
     mos_server,
     mem_scheduler,
+    naive_mem_cube,
+    api_module,
 ) = init_server()
 
 
-def _create_naive_mem_cube() -> NaiveMemCube:
-    """Create a NaiveMemCube instance with initialized components."""
-    naive_mem_cube = NaiveMemCube(
-        llm=llm,
-        embedder=embedder,
-        mem_reader=mem_reader,
-        graph_db=graph_db,
-        reranker=reranker,
-        internet_retriever=internet_retriever,
-        memory_manager=memory_manager,
-        default_cube_config=default_cube_config,
-    )
-    return naive_mem_cube
-
-
 def _format_memory_item(memory_data: Any) -> dict[str, Any]:
     """Format a single memory item for API response."""
     memory = memory_data.model_dump()
@@ -257,30 +271,99 @@ def mix_search_memories(
     search_req: APISearchRequest,
     user_context: UserContext,
 ):
-    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
-    naive_mem_cube = _create_naive_mem_cube()
-    search_results = naive_mem_cube.text_mem.search(
-        query=search_req.query,
-        user_name=user_context.mem_cube_id,
-        top_k=search_req.top_k,
-        mode=search_req.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,
-        },
-    )
-    formatted_memories = [_format_memory_item(data) for data in search_results]
-
-    return formatted_memories
+    """
+    Mix search memories: fast search + async fine search
+    """
+    # Get fast memories first
+    fast_memories = fast_search_memories(search_req, user_context)
+
+    # Check if scheduler and dispatcher are available for async execution
+    if mem_scheduler and hasattr(mem_scheduler, "dispatcher") and mem_scheduler.dispatcher:
+        try:
+            # Create message for async fine search
+            message_content = {
+                "search_req": {
+                    "query": search_req.query,
+                    "user_id": search_req.user_id,
+                    "session_id": search_req.session_id,
+                    "top_k": search_req.top_k,
+                    "internet_search": search_req.internet_search,
+                    "moscube": search_req.moscube,
+                    "chat_history": search_req.chat_history,
+                },
+                "user_context": {"mem_cube_id": user_context.mem_cube_id},
+            }
+
+            message = ScheduleMessageItem(
+                item_id=f"mix_search_{search_req.user_id}_{get_utc_now().timestamp()}",
+                user_id=search_req.user_id,
+                mem_cube_id=user_context.mem_cube_id,
+                label=API_MIX_SEARCH_LABEL,
+                mem_cube=naive_mem_cube,
+                content=json.dumps(message_content),
+                timestamp=get_utc_now(),
+            )
+
+            # Submit async task
+            mem_scheduler.dispatcher.submit_message(message)
+            logger.info(f"Submitted async fine search task for user {search_req.user_id}")
+
+            # Try to get pre-computed fine memories if available
+            try:
+                pre_fine_memories = api_module.get_pre_fine_memories(
+                    user_id=search_req.user_id, mem_cube_id=user_context.mem_cube_id
+                )
+                if pre_fine_memories:
+                    # Merge fast and pre-computed fine memories
+                    all_memories = fast_memories + pre_fine_memories
+                    # Remove duplicates based on content
+                    seen_contents = set()
+                    unique_memories = []
+                    for memory in all_memories:
+                        content_key = memory.get("content", "")
+                        if content_key not in seen_contents:
+                            seen_contents.add(content_key)
+                            unique_memories.append(memory)
+                    return unique_memories
+            except Exception as e:
+                logger.warning(f"Failed to get pre-computed fine memories: {e}")
+
+        except Exception as e:
+            logger.error(f"Failed to submit async fine search task: {e}")
+            # Fall back to synchronous execution
+
+    # Fallback: synchronous fine search
+    try:
+        fine_memories = fine_search_memories(search_req, user_context)
+
+        # Merge fast and fine memories
+        all_memories = fast_memories + fine_memories
+
+        # Remove duplicates based on content
+        seen_contents = set()
+        unique_memories = []
+        for memory in all_memories:
+            content_key = memory.get("content", "")
+            if content_key not in seen_contents:
+                seen_contents.add(content_key)
+                unique_memories.append(memory)
+
+        # Sync search data to Redis
+        try:
+            api_module.sync_search_data(
+                user_id=search_req.user_id,
+                mem_cube_id=user_context.mem_cube_id,
+                query=search_req.query,
+                formatted_memories=unique_memories,
+            )
+        except Exception as e:
+            logger.error(f"Failed to sync search data: {e}")
+
+        return unique_memories
+
+    except Exception as e:
+        logger.error(f"Fine search failed: {e}")
+        return fast_memories
 
 
 def fine_search_memories(
@@ -293,12 +376,11 @@ def fine_search_memories(
     search_filter = {"session_id": search_req.session_id} if search_req.session_id else None
 
     # Create MemCube and perform search
-    naive_mem_cube = _create_naive_mem_cube()
     search_results = naive_mem_cube.text_mem.search(
         query=search_req.query,
         user_name=user_context.mem_cube_id,
         top_k=search_req.top_k,
-        mode=search_req.mode,
+        mode=SearchMode.FINE,
         manual_close_internet=not search_req.internet_search,
         moscube=search_req.moscube,
         search_filter=search_filter,
@@ -323,12 +405,11 @@ def fast_search_memories(
     search_filter = {"session_id": search_req.session_id} if search_req.session_id else None
 
     # Create MemCube and perform search
-    naive_mem_cube = _create_naive_mem_cube()
     search_results = naive_mem_cube.text_mem.search(
         query=search_req.query,
         user_name=user_context.mem_cube_id,
         top_k=search_req.top_k,
-        mode=search_req.mode,
+        mode=SearchMode.FAST,
         manual_close_internet=not search_req.internet_search,
         moscube=search_req.moscube,
         search_filter=search_filter,

src/memos/mem_scheduler/general_modules/api_misc.py

@@ -0,0 +1,115 @@
+import threading
+
+from typing import Any
+
+from memos.log import get_logger
+from memos.mem_scheduler.general_modules.base import BaseSchedulerModule
+from memos.mem_scheduler.orm_modules.redis_model import RedisDBManager, SimpleListManager
+
+
+logger = get_logger(__name__)
+
+
+class SchedulerAPIModule(BaseSchedulerModule):
+    def __init__(self):
+        super().__init__()
+
+        self.search_history_managers: dict[str, RedisDBManager] = {}
+
+    def get_search_history_manager(self, user_id: str, mem_cube_id: str) -> RedisDBManager:
+        """Get or create a Redis manager for search history."""
+        key = f"search_history:{user_id}:{mem_cube_id}"
+        if key not in self.search_history_managers:
+            self.search_history_managers[key] = RedisDBManager(
+                user_id=user_id, mem_cube_id=mem_cube_id
+            )
+        return self.search_history_managers[key]
+
+    def sync_search_data(
+        self, user_id: str, mem_cube_id: str, query: str, formatted_memories: Any
+    ) -> None:
+        """
+        Sync search data to Redis, maintaining a list of size 5.
+
+        Args:
+            user_id: User identifier
+            mem_cube_id: Memory cube identifier
+            query: Search query string
+            formatted_memories: Formatted search results
+        """
+        try:
+            # Get the search history manager
+            manager = self.get_search_history_manager(user_id, mem_cube_id)
+
+            # Create search data entry
+            search_entry = {
+                "query": query,
+                "formatted_memories": formatted_memories,
+                "timestamp": threading.current_thread().ident,  # Use thread ID as simple timestamp
+            }
+
+            # Load existing search history
+            existing_data = manager.load_from_db()
+
+            if existing_data is None:
+                search_history = SimpleListManager([])
+            else:
+                # If existing data is a SimpleListManager, use it; otherwise create new one
+                if isinstance(existing_data, SimpleListManager):
+                    search_history = existing_data
+                else:
+                    search_history = SimpleListManager([])
+
+            # Add new entry and keep only latest 5
+            search_history.add_item(str(search_entry))
+            if len(search_history) > 5:
+                # Keep only the latest 5 items
+                search_history.items = search_history.items[-5:]
+
+            # Save back to Redis
+            manager.save_to_db(search_history)
+
+            logger.info(
+                f"Synced search data for user {user_id}, mem_cube {mem_cube_id}. History size: {len(search_history)}"
+            )
+
+        except Exception as e:
+            logger.error(f"Failed to sync search data: {e}", exc_info=True)
+
+    def get_pre_fine_memories(self, user_id: str, mem_cube_id: str) -> list:
+        """
+        Get the most recent pre-computed fine memories from search history.
+
+        Args:
+            user_id: User identifier
+            mem_cube_id: Memory cube identifier
+
+        Returns:
+            List of formatted memories from the most recent search, or empty list if none found
+        """
+        try:
+            manager = self.get_search_history_manager(user_id, mem_cube_id)
+            search_history_key = "search_history_list"
+            existing_data = manager.load_from_db(search_history_key)
+
+            if existing_data is None:
+                return []
+
+            search_history = (
+                existing_data.obj_instance
+                if hasattr(existing_data, "obj_instance")
+                else existing_data
+            )
+
+            if not search_history or len(search_history) == 0:
+                return []
+
+            # Return the formatted_memories from the most recent search
+            latest_entry = search_history[-1]
+            return (
+                latest_entry.get("formatted_memories", []) if isinstance(latest_entry, dict) else []
+            )
+
+        except Exception as e:
+            logger.error(f"Failed to get pre-computed fine memories: {e}", exc_info=True)
+            return []