revision of mixture api: add conversation turn and reduce 2 stage ranking to 1 stage (#405)

* 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

* feat: Optimize mixture search and enhance API client

* feat: Add conversation_turn tracking for session-based memory search

    - Add conversation_turn field to APIMemoryHistoryEntryItem schema with default value 0
    - Implement session counter in OptimizedScheduler to track turn count per session_id
    - Update sync_search_data method to accept and store conversation_turn parameter
    - Maintain session history with LRU eviction (max 5 sessions)
    - Rename conversation_id to session_id for consistency with request object
    - Enable direct access to session_id from search requests

    This feature allows tracking conversation turns within the same session,
    providing better context for memory retrieval and search history management.

* adress time bug in monitor

* revise simple tree

* add mode to evaluation client; rewrite print to logger.info in db files

---------

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

Author: 3 people

evaluation/scripts/utils/client.py

@@ -181,7 +181,7 @@ def search(self, query, user_id, top_k):
                 "mem_cube_id": user_id,
                 "conversation_id": "",
                 "top_k": top_k,
-                "mode": "fast",
+                "mode": os.getenv("SEARCH_MODE", "fast"),
                 "handle_pref_mem": False,
             },
             ensure_ascii=False,
@@ -232,7 +232,7 @@ def search(self, query, user_id, top_k):
                 "query": query,
                 "user_id": user_id,
                 "memory_limit_number": top_k,
-                "mode": "mixture",
+                "mode": os.getenv("SEARCH_MODE", "fast"),
             }
         )
 

src/memos/graph_dbs/neo4j.py

@@ -1072,7 +1072,7 @@ def drop_database(self) -> None:
 
             with self.driver.session(database=self.system_db_name) as session:
                 session.run(f"DROP DATABASE {self.db_name} IF EXISTS")
-                print(f"Database '{self.db_name}' has been dropped.")
+                logger.info(f"Database '{self.db_name}' has been dropped.")
         else:
             raise ValueError(
                 f"Refusing to drop protected database: {self.db_name} in "

src/memos/graph_dbs/polardb.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/base_scheduler.py

@@ -8,24 +8,27 @@
     APISearchHistoryManager,
     TaskRunningStatus,
 )
-from memos.mem_scheduler.utils.db_utils import get_utc_now
 from memos.memories.textual.item import TextualMemoryItem
 
 
 logger = get_logger(__name__)
 
 
 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,
@@ -41,8 +44,12 @@ def sync_search_data(
         query: str,
         memories: list[TextualMemoryItem],
         formatted_memories: Any,
-        conversation_id: str | None = None,
+        session_id: str | None = None,
+        conversation_turn: int = 0,
     ) -> 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)
@@ -59,7 +66,7 @@ def sync_search_data(
                 query=query,
                 formatted_memories=formatted_memories,
                 task_status=TaskRunningStatus.COMPLETED,  # Use the provided running_status
-                conversation_id=conversation_id,
+                session_id=session_id,
                 memories=memories,
             )
 
@@ -69,18 +76,18 @@ def sync_search_data(
                 logger.warning(f"Failed to update entry with item_id: {item_id}")
         else:
             # Add new entry based on running_status
-            search_entry = APIMemoryHistoryEntryItem(
+            entry_item = APIMemoryHistoryEntryItem(
                 item_id=item_id,
                 query=query,
                 formatted_memories=formatted_memories,
                 memories=memories,
                 task_status=TaskRunningStatus.COMPLETED,
-                conversation_id=conversation_id,
-                created_time=get_utc_now(),
+                session_id=session_id,
+                conversation_turn=conversation_turn,
             )
 
             # Add directly to completed list as APIMemoryHistoryEntryItem instance
-            search_history.completed_entries.append(search_entry)
+            search_history.completed_entries.append(entry_item)
 
             # Maintain window size
             if len(search_history.completed_entries) > search_history.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/general_modules/api_misc.py

@@ -76,8 +76,8 @@ def __init__(
         ] = {}
 
         # Lifecycle monitor
-        self.last_activation_mem_update_time = datetime.min
-        self.last_query_consume_time = datetime.min
+        self.last_activation_mem_update_time = get_utc_now()
+        self.last_query_consume_time = get_utc_now()
 
         self._register_lock = Lock()
         self._process_llm = process_llm