revision of mixture api: add conversation turn and reduce 2 stage ranking to 1 stage

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/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

@@ -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/monitors/general_monitor.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