feat: redis for sync history memories and new api of mixture search (#398)

* debug an error function name

* feat: Add DynamicCache compatibility for different transformers versions

- Fix build_kv_cache method in hf.py to handle both old and new DynamicCache structures
  - Support new 'layers' attribute with key_cache/value_cache or keys/values
  - Maintain backward compatibility with direct key_cache/value_cache attributes
  - Add comprehensive error handling and logging for unsupported structures

- Update move_dynamic_cache_htod function in kv.py for cross-version compatibility
  - Handle layers-based structure in newer transformers versions
  - Support alternative attribute names (keys/values vs key_cache/value_cache)
  - Preserve original functionality for older transformers versions

- Add comprehensive tests for DynamicCache compatibility
  - Test activation memory update with mock DynamicCache layers
  - Verify layers attribute access across different transformers versions
  - Fix scheduler logger mock to include memory_manager attribute
This resolves AttributeError issues when using different versions of the
transformers library and ensures robust handling of DynamicCache objects.

debug

* feat: implement APIAnalyzerForScheduler for memory operations

- Add APIAnalyzerForScheduler class with search/add operations
- Support requests and http.client with connection reuse
- Include comprehensive error handling and dynamic configuration
- Add English test suite with real-world conversation scenarios

* feat: Add search_ws API endpoint and enhance API analyzer functionality

- Add search_ws endpoint in server_router.py for scheduler-enabled search
- Fix missing imports: time module, SearchRequest class, and get_mos_product_instance function
- Implement search_ws method in api_analyzer.py with HTTP client support
- Add _search_ws_with_requests and _search_ws_with_http_client private methods
- Include search_ws usage example in demonstration code
- Enhance scheduler and dispatcher capabilities for improved memory management
- Expand test coverage to ensure functionality stability

This update primarily strengthens the memory scheduling system's search capabilities,
providing users with more flexible API interface options.

* fix: resolve test failures and warnings in test suite

- Fix Pydantic serialization warning in test_memos_chen_tang_hello_world
  * Add warnings filter to suppress UserWarning from Pydantic serialization

- Fix KeyError: 'past_key_values' in test_build_kv_cache_and_generation
  * Update mock configuration to properly return forward_output with past_key_values
  * Add DynamicCache version compatibility handling in test mocks
  * Support both old and new transformers versions with layers/key_cache attributes
  * Improve assertion logic to check all model calls for required parameters

- Update base_scheduler.py to use centralized DEFAULT_MAX_INTERNAL_MESSAGE_QUEUE_SIZE constant
  * Add import for DEFAULT_MAX_INTERNAL_MESSAGE_QUEUE_SIZE from general_schemas
  * Replace hardcoded value 100 with configurable constant (1000)

All tests now pass successfully with proper version compatibility handling.

* feat: add a test_robustness execution to test thread pool execution

* feat: optimize scheduler configuration and API search functionality

- Add DEFAULT_TOP_K and DEFAULT_CONTEXT_WINDOW_SIZE global constants in general_schemas.py
- Update base_scheduler.py to use global default values instead of hardcoded numbers
- Fix SchedulerConfigFactory initialization issue by using keyword argument expansion
- Resolve UnboundLocalError variable conflict in search_memories_ws function
- Fix indentation and parameter issues in OptimizedScheduler search_for_api method
- Improve code standardization and maintainability

* feat: Add Redis auto-initialization with fallback strategies

- Add auto_initialize_redis() with config/env/local fallback
- Move Redis logic from dispatcher_monitor to redis_service
- Update base_scheduler to use auto initialization
- Add proper resource cleanup and error handling

* feat: add database connection management to ORM module

- Add MySQL engine loading from environment variables in BaseDBManager
- Add Redis connection loading from environment variables in BaseDBManager
- Enhance database configuration validation and error handling
- Complete database adapter infrastructure for ORM module
- Provide unified database connection management interface

This update provides comprehensive database connection management capabilities
for the mem_scheduler module, supporting dynamic MySQL and Redis configuration
loading from environment variables, establishing reliable data persistence
foundation for scheduling services and API services.

* remove part of test

* feat: add Redis-based ORM with multiprocess synchronization

- Add RedisDBManager and RedisLockableORM classes
- Implement atomic locking mechanism for concurrent access
- Add merge functionality for different object types
- Include comprehensive test suite and examples
- Fix Redis key type conflicts in lock operations

* fix: resolve scheduler module import and Redis integration issues

* revise naive memcube creation in server router

* remove long-time tests in test_scheduler

* remove redis test which needs .env

* refactor all codes about mixture search with scheduler

* fix: resolve Redis API synchronization issues and implement search API with reranker

- Fix running_entries to running_task_ids migration across codebase
- Update sync_search_data method to properly handle TaskRunningStatus
- Correct variable naming and logic in API synchronization flow
- Implement search API endpoint with reranker functionality
- Update test files to reflect new running_task_ids convention
- Ensure proper Redis state management for concurrent tasks

* remove a test for api module

* revise to pass the test suite

* address some bugs to make mix_search normally running

* modify codes according to evaluation logs

---------

Co-authored-by: CaralHsi <[email protected]>

Author: tangg555

evaluation/scripts/utils/client.py

@@ -183,6 +183,7 @@ def search(self, query, user_id, top_k):
                 "mem_cube_id": user_id,
                 "conversation_id": "",
                 "top_k": top_k,
+                "mode": "mixture",
             },
             ensure_ascii=False,
         )
@@ -230,6 +231,7 @@ def search(self, query, user_id, top_k):
                 "query": query,
                 "user_id": user_id,
                 "memory_limit_number": top_k,
+                "mode": "mixture",
             }
         )
 
@@ -311,7 +313,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

@@ -361,8 +361,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
 
@@ -154,7 +154,6 @@ def init_server():
 
     # Build component configurations
     graph_db_config = _build_graph_db_config()
-    print(graph_db_config)
     llm_config = _build_llm_config()
     embedder_config = _build_embedder_config()
     mem_reader_config = _build_mem_reader_config()
@@ -209,22 +208,6 @@ def init_server():
         online_bot=False,
     )
 
-    # 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
-
     naive_mem_cube = NaiveMemCube(
         llm=llm,
         embedder=embedder,
@@ -240,6 +223,23 @@ def init_server():
         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: 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
+
     return (
         graph_db,
         mem_reader,
@@ -400,96 +400,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/configs/mem_scheduler.py

@@ -15,6 +15,7 @@
     DEFAULT_CONSUME_INTERVAL_SECONDS,
     DEFAULT_CONTEXT_WINDOW_SIZE,
     DEFAULT_MAX_INTERNAL_MESSAGE_QUEUE_SIZE,
+    DEFAULT_MULTI_TASK_RUNNING_TIMEOUT,
     DEFAULT_THREAD_POOL_MAX_WORKERS,
     DEFAULT_TOP_K,
     DEFAULT_USE_REDIS_QUEUE,
@@ -59,6 +60,10 @@ class BaseSchedulerConfig(BaseConfig):
         default=DEFAULT_MAX_INTERNAL_MESSAGE_QUEUE_SIZE,
         description="Maximum size of internal message queue when not using Redis",
     )
+    multi_task_running_timeout: int = Field(
+        default=DEFAULT_MULTI_TASK_RUNNING_TIMEOUT,
+        description="Default timeout for multi-task running operations in seconds",
+    )
 
 
 class GeneralSchedulerConfig(BaseSchedulerConfig):