addressed all conflicts

Author: tangg555

evaluation/scripts/utils/client.py

@@ -311,7 +311,7 @@ def add(self, messages, user_id, iso_date):
                 agent_name=self.agent_id,
                 session_date=iso_date,
             )
-            self.wait_for_completion(response.task_id)
+            self.wait_for_completion(response.item_id)
         except Exception as error:
             print("❌ Error saving conversation:", error)
 

src/memos/api/config.py

@@ -334,8 +334,8 @@ def get_scheduler_config() -> dict[str, Any]:
                 "thread_pool_max_workers": int(
                     os.getenv("MOS_SCHEDULER_THREAD_POOL_MAX_WORKERS", "10")
                 ),
-                "consume_interval_seconds": int(
-                    os.getenv("MOS_SCHEDULER_CONSUME_INTERVAL_SECONDS", "3")
+                "consume_interval_seconds": float(
+                    os.getenv("MOS_SCHEDULER_CONSUME_INTERVAL_SECONDS", "0.01")
                 ),
                 "enable_parallel_dispatch": os.getenv(
                     "MOS_SCHEDULER_ENABLE_PARALLEL_DISPATCH", "true"

src/memos/api/routers/server_router.py

@@ -1,9 +1,8 @@
-import json
 import os
 import traceback
 
 from concurrent.futures import ThreadPoolExecutor
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from fastapi import APIRouter, HTTPException
 
@@ -33,11 +32,8 @@
 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 (
-    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.prefer_text_memory.config import (
     AdderConfigFactory,
     ExtractorConfigFactory,
@@ -54,6 +50,10 @@
 )
 from memos.reranker.factory import RerankerFactory
 from memos.templates.instruction_completion import instruct_completion
+
+
+if TYPE_CHECKING:
+    from memos.mem_scheduler.optimized_scheduler import OptimizedScheduler
 from memos.types import MOSSearchResult, UserContext
 from memos.vec_dbs.factory import VecDBFactory
 
@@ -208,36 +208,53 @@ def init_server():
         online_bot=False,
     )
 
+    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,
+        vector_db=vector_db,
+        pref_extractor=pref_extractor,
+        pref_adder=pref_adder,
+        pref_retriever=pref_retriever,
+    )
+
     # Initialize Scheduler
     scheduler_config_dict = APIConfig.get_scheduler_config()
     scheduler_config = SchedulerConfigFactory(
         backend="optimized_scheduler", config=scheduler_config_dict
     )
-    mem_scheduler = SchedulerFactory.from_config(scheduler_config)
+    mem_scheduler: OptimizedScheduler = SchedulerFactory.from_config(scheduler_config)
     mem_scheduler.initialize_modules(
         chat_llm=llm,
         process_llm=mem_reader.llm,
         db_engine=BaseDBManager.create_default_sqlite_engine(),
     )
+    mem_scheduler.current_mem_cube = naive_mem_cube
     mem_scheduler.start()
 
     # Initialize SchedulerAPIModule
     api_module = mem_scheduler.api_module
 
-    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,
-        vector_db=vector_db,
-        pref_extractor=pref_extractor,
-        pref_adder=pref_adder,
-        pref_retriever=pref_retriever,
+    # Initialize Scheduler
+    scheduler_config_dict = APIConfig.get_scheduler_config()
+    scheduler_config = SchedulerConfigFactory(
+        backend="optimized_scheduler", config=scheduler_config_dict
+    )
+    mem_scheduler = SchedulerFactory.from_config(scheduler_config)
+    mem_scheduler.initialize_modules(
+        chat_llm=llm,
+        process_llm=mem_reader.llm,
+        db_engine=BaseDBManager.create_default_sqlite_engine(),
     )
+    mem_scheduler.start()
+
+    # Initialize SchedulerAPIModule
+    api_module = mem_scheduler.api_module
 
     return (
         graph_db,
@@ -398,96 +415,12 @@ def mix_search_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
+    formatted_memories = mem_scheduler.mix_search_memories(
+        search_req=search_req,
+        user_context=user_context,
+    )
+    return formatted_memories
 
 
 def fine_search_memories(

src/memos/mem_scheduler/base_scheduler.py

@@ -502,7 +502,7 @@ def update_activation_memory_periodically(
         except Exception as e:
             logger.error(f"Error in update_activation_memory_periodically: {e}", exc_info=True)
 
-    async def submit_messages(self, messages: ScheduleMessageItem | list[ScheduleMessageItem]):
+    def submit_messages(self, messages: ScheduleMessageItem | list[ScheduleMessageItem]):
         """Submit messages to the message queue (either local queue or Redis)."""
         if isinstance(messages, ScheduleMessageItem):
             messages = [messages]  # transform single message to list
@@ -519,7 +519,7 @@ async def submit_messages(self, messages: ScheduleMessageItem | list[ScheduleMes
 
             if self.use_redis_queue:
                 # Use Redis stream for message queue
-                await self.redis_add_message_stream(message.to_dict())
+                self.redis_add_message_stream(message.to_dict())
                 logger.info(f"Submitted message to Redis: {message.label} - {message.content}")
             else:
                 # Use local queue
@@ -774,34 +774,6 @@ def unregister_handlers(self, labels: list[str]) -> dict[str, bool]:
         return self.dispatcher.unregister_handlers(labels)
 
     def get_running_tasks(self, filter_func: Callable | None = None) -> dict[str, dict]:
-        """
-        Get currently running tasks, optionally filtered by a custom function.
-
-        This method delegates to the dispatcher's get_running_tasks method.
-
-        Args:
-            filter_func: Optional function to filter tasks. Should accept a RunningTaskItem
-                        and return True if the task should be included in results.
-
-        Returns:
-            dict[str, dict]: Dictionary mapping task IDs to task information dictionaries.
-                           Each task dict contains: item_id, user_id, mem_cube_id, task_info,
-                           task_name, start_time, end_time, status, result, error_message, messages
-
-        Examples:
-            # Get all running tasks
-            all_tasks = scheduler.get_running_tasks()
-
-            # Get tasks for specific user
-            user_tasks = scheduler.get_running_tasks(
-                filter_func=lambda task: task.user_id == "user123"
-            )
-
-            # Get tasks with specific status
-            active_tasks = scheduler.get_running_tasks(
-                filter_func=lambda task: task.status == "running"
-            )
-        """
         if not self.dispatcher:
             logger.warning("Dispatcher is not initialized, returning empty tasks dict")
             return {}