MemoryOS 1.0.0__py3-none-any.whl → 1.1.1__py3-none-any.whl

memos/mem_scheduler/analyzer/scheduler_for_eval.py

@@ -0,0 +1,280 @@
+from __future__ import annotations
+
+import time
+
+from functools import wraps
+from typing import TYPE_CHECKING, Any, ClassVar
+
+from memos.log import get_logger
+from memos.mem_scheduler.general_scheduler import GeneralScheduler
+from memos.mem_scheduler.schemas.general_schemas import (
+    DEFAULT_MAX_QUERY_KEY_WORDS,
+    UserID,
+)
+from memos.mem_scheduler.schemas.monitor_schemas import QueryMonitorItem
+
+
+if TYPE_CHECKING:
+    from memos.memories.textual.tree import TextualMemoryItem
+
+
+logger = get_logger(__name__)
+
+
+class SchedulerForEval(GeneralScheduler):
+    """
+    A scheduler class that inherits from GeneralScheduler and provides evaluation-specific functionality.
+    This class extends GeneralScheduler with evaluation methods.
+    """
+
+    # Class variable to store timing information for all instances
+    timer_cache: ClassVar[dict[str, dict[str, Any]]] = {}
+
+    def __init__(self, config):
+        """
+        Initialize the SchedulerForEval with the same configuration as GeneralScheduler.
+
+        Args:
+            config: Configuration object for the scheduler
+        """
+        super().__init__(config)
+        # Initialize instance timer_cache
+        self.timer_cache = {}
+
+    @staticmethod
+    def time_it(func_name: str | None = None):
+        """
+        Static method decorator to measure function execution time and store in timer_cache.
+
+        Args:
+            func_name: Custom name for the function in timer_cache. If None, uses function.__name__
+        """
+
+        def decorator(func):
+            @wraps(func)
+            def wrapper(self, *args, **kwargs):
+                # Get function name
+                name = func_name or func.__name__
+
+                # Start timing
+                start_time = time.time()
+                result = func(self, *args, **kwargs)
+                end_time = time.time()
+
+                # Calculate execution time
+                exec_time = end_time - start_time
+
+                # Format time as HH:MM:SS.mmm
+                hours = int(exec_time // 3600)
+                minutes = int((exec_time % 3600) // 60)
+                seconds = exec_time % 60
+
+                if hours > 0:
+                    time_str = f"{hours:02d}:{minutes:02d}:{seconds:06.3f}"
+                else:
+                    time_str = f"{minutes:02d}:{seconds:06.3f}"
+
+                # Store in timer_cache
+                if not hasattr(self, "timer_cache"):
+                    self.timer_cache = {}
+
+                self.timer_cache[name] = {
+                    "time_str": time_str,
+                    "seconds": exec_time,
+                }
+
+                logger.info(f"{name} executed in {time_str}")
+                return result
+
+            return wrapper
+
+        return decorator
+
+    def get_timer_summary(self) -> str:
+        """
+        Get a summary of all timed functions.
+
+        Returns:
+            Formatted string with timing information
+        """
+        if not self.timer_cache:
+            return "No timing data available."
+
+        summary = "=== Timing Summary ===\n"
+        for func_name, data in self.timer_cache.items():
+            summary += f"{func_name}: {data['time_str']} (at {data['timestamp']})\n"
+
+        return summary
+
+    def clear_timer_cache(self):
+        """Clear the timer cache."""
+        self.timer_cache.clear()
+
+    @time_it("update_working_memory")
+    def update_working_memory_for_eval(
+        self, query: str, user_id: UserID | str, top_k: int
+    ) -> list[str]:
+        """
+        Update working memory based on query and return the updated memory list.
+
+        Args:
+            query: The query string
+            user_id: User identifier
+            top_k: Number of top memories to return
+
+        Returns:
+            List of memory strings from updated working memory
+        """
+        self.monitor.register_query_monitor_if_not_exists(
+            user_id=user_id, mem_cube_id=self.current_mem_cube_id
+        )
+
+        query_keywords = self.monitor.extract_query_keywords(query=query)
+        logger.info(f'Extract keywords "{query_keywords}" from query "{query}"')
+
+        item = QueryMonitorItem(
+            user_id=user_id,
+            mem_cube_id=self.current_mem_cube_id,
+            query_text=query,
+            keywords=query_keywords,
+            max_keywords=DEFAULT_MAX_QUERY_KEY_WORDS,
+        )
+        query_db_manager = self.monitor.query_monitors[user_id][self.current_mem_cube_id]
+        query_db_manager.obj.put(item=item)
+        # Sync with database after adding new item
+        query_db_manager.sync_with_orm()
+        logger.debug(f"Queries in monitor are {query_db_manager.obj.get_queries_with_timesort()}.")
+
+        queries = [query]
+
+        # recall
+        mem_cube = self.current_mem_cube
+        text_mem_base = mem_cube.text_mem
+
+        cur_working_memory: list[TextualMemoryItem] = text_mem_base.get_working_memory()
+        text_working_memory: list[str] = [w_m.memory for w_m in cur_working_memory]
+        intent_result = self.monitor.detect_intent(
+            q_list=queries, text_working_memory=text_working_memory
+        )
+
+        if intent_result["trigger_retrieval"]:
+            missing_evidences = intent_result["missing_evidences"]
+            num_evidence = len(missing_evidences)
+            k_per_evidence = max(1, top_k // max(1, num_evidence))
+            new_candidates = []
+            for item in missing_evidences:
+                logger.info(f"missing_evidences: {item}")
+                results: list[TextualMemoryItem] = self.retriever.search(
+                    query=item,
+                    mem_cube=mem_cube,
+                    top_k=k_per_evidence,
+                    method=self.search_method,
+                )
+                logger.info(
+                    f"search results for {missing_evidences}: {[one.memory for one in results]}"
+                )
+                new_candidates.extend(results)
+            logger.info(
+                f"missing_evidences: {missing_evidences} and get {len(new_candidates)} new candidate memories."
+            )
+        else:
+            new_candidates = []
+            logger.info(f"intent_result: {intent_result}. not triggered")
+
+        # rerank
+        new_order_working_memory = self.replace_working_memory(
+            user_id=user_id,
+            mem_cube_id=self.current_mem_cube_id,
+            mem_cube=self.current_mem_cube,
+            original_memory=cur_working_memory,
+            new_memory=new_candidates,
+        )
+        new_order_working_memory = new_order_working_memory[:top_k]
+        logger.info(f"size of new_order_working_memory: {len(new_order_working_memory)}")
+
+        return [m.memory for m in new_order_working_memory]
+
+    @time_it("memory_answer_ability")
+    def evaluate_memory_answer_ability(
+        self, query: str, memory_texts: list[str], top_k: int = 100
+    ) -> bool:
+        """
+        Use LLM to evaluate whether the given memories can answer the query.
+
+        Args:
+            query: The query string to evaluate
+            memory_texts: List of memory texts to check against
+            top_k: Maximum number of memories to consider for evaluation
+
+        Returns:
+            Boolean indicating whether the memories can answer the query
+        """
+        # Limit the number of memories to evaluate
+        limited_memories = memory_texts[:top_k] if memory_texts else []
+
+        # Build prompt using the template
+        prompt = self.monitor.build_prompt(
+            template_name="memory_answer_ability_evaluation",
+            query=query,
+            memory_list="\n".join([f"- {memory}" for memory in limited_memories])
+            if limited_memories
+            else "No memories available",
+        )
+
+        # Use the process LLM to generate response
+        response = self.monitor._process_llm.generate([{"role": "user", "content": prompt}])
+
+        try:
+            # Extract JSON response
+            from memos.mem_scheduler.utils.misc_utils import extract_json_dict
+
+            result = extract_json_dict(response)
+
+            # Validate response structure
+            if "result" in result:
+                logger.info(
+                    f"Memory answer ability evaluation result: {result['result']}, reason: {result.get('reason', 'No reason provided')}"
+                )
+                return result["result"]
+            else:
+                logger.warning(f"Invalid response structure from LLM: {result}")
+                return False
+
+        except Exception as e:
+            logger.error(
+                f"Failed to parse LLM response for memory answer ability evaluation: {response}. Error: {e}"
+            )
+            # Fallback: return False if we can't determine answer ability
+            return False
+
+    @time_it("search_for_eval")
+    def search_for_eval(
+        self, query: str, user_id: UserID | str, top_k: int, scheduler_flag: bool = True
+    ) -> list[str]:
+        """
+        Original search_for_eval function refactored to use the new decomposed functions.
+
+        Args:
+            query: The query string
+            user_id: User identifier
+            top_k: Number of top memories to return
+            scheduler_flag: Whether to update working memory or just evaluate
+
+        Returns:
+            Tuple of (memory_list, can_answer_boolean)
+        """
+        if not scheduler_flag:
+            # Get current working memory without updating
+            mem_cube = self.current_mem_cube
+            text_mem_base = mem_cube.text_mem
+            cur_working_memory: list[TextualMemoryItem] = text_mem_base.get_working_memory()
+            text_working_memory: list[str] = [w_m.memory for w_m in cur_working_memory]
+
+            return text_working_memory
+        else:
+            # Update working memory and get the result
+            updated_memories = self.update_working_memory_for_eval(
+                query=query, user_id=user_id, top_k=top_k
+            )
+
+            return updated_memories

memos/mem_scheduler/base_scheduler.py

@@ -5,16 +5,16 @@ import time
 from datetime import datetime
 from pathlib import Path
 
+from sqlalchemy.engine import Engine
+
 from memos.configs.mem_scheduler import AuthConfig, BaseSchedulerConfig
 from memos.llms.base import BaseLLM
 from memos.log import get_logger
 from memos.mem_cube.general import GeneralMemCube
 from memos.mem_scheduler.general_modules.dispatcher import SchedulerDispatcher
 from memos.mem_scheduler.general_modules.misc import AutoDroppingQueue as Queue
-from memos.mem_scheduler.general_modules.rabbitmq_service import RabbitMQSchedulerModule
-from memos.mem_scheduler.general_modules.redis_service import RedisSchedulerModule
-from memos.mem_scheduler.general_modules.retriever import SchedulerRetriever
 from memos.mem_scheduler.general_modules.scheduler_logger import SchedulerLoggerModule
+from memos.mem_scheduler.memory_manage_modules.retriever import SchedulerRetriever
 from memos.mem_scheduler.monitors.dispatcher_monitor import SchedulerDispatcherMonitor
 from memos.mem_scheduler.monitors.general_monitor import SchedulerGeneralMonitor
 from memos.mem_scheduler.schemas.general_schemas import (
@@ -33,6 +33,8 @@ from memos.mem_scheduler.schemas.monitor_schemas import MemoryMonitorItem
 from memos.mem_scheduler.utils.filter_utils import (
     transform_name_to_key,
 )
+from memos.mem_scheduler.webservice_modules.rabbitmq_service import RabbitMQSchedulerModule
+from memos.mem_scheduler.webservice_modules.redis_service import RedisSchedulerModule
 from memos.memories.activation.kv import KVCacheMemory
 from memos.memories.activation.vllmkv import VLLMKVCacheItem, VLLMKVCacheMemory
 from memos.memories.textual.tree import TextualMemoryItem, TreeTextMemory
@@ -62,6 +64,7 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule, SchedulerLogg
         )
 
         self.retriever: SchedulerRetriever | None = None
+        self.db_engine: Engine | None = None
         self.monitor: SchedulerGeneralMonitor | None = None
         self.dispatcher_monitor: SchedulerDispatcherMonitor | None = None
         self.dispatcher = SchedulerDispatcher(
@@ -70,12 +73,15 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule, SchedulerLogg
         )
 
         # internal message queue
-        self.max_internal_messae_queue_size = 100
+        self.max_internal_message_queue_size = self.config.get(
+            "max_internal_message_queue_size", 100
+        )
         self.memos_message_queue: Queue[ScheduleMessageItem] = Queue(
-            maxsize=self.max_internal_messae_queue_size
+            maxsize=self.max_internal_message_queue_size
         )
+        self.max_web_log_queue_size = self.config.get("max_web_log_queue_size", 50)
         self._web_log_message_queue: Queue[ScheduleLogForWebItem] = Queue(
-            maxsize=self.max_internal_messae_queue_size
+            maxsize=self.max_web_log_queue_size
         )
         self._consumer_thread = None  # Reference to our consumer thread
         self._running = False
@@ -92,34 +98,57 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule, SchedulerLogg
         self.auth_config = None
         self.rabbitmq_config = None
 
-    def initialize_modules(self, chat_llm: BaseLLM, process_llm: BaseLLM | None = None):
+    def initialize_modules(
+        self,
+        chat_llm: BaseLLM,
+        process_llm: BaseLLM | None = None,
+        db_engine: Engine | None = None,
+    ):
         if process_llm is None:
             process_llm = chat_llm
 
-        # initialize submodules
-        self.chat_llm = chat_llm
-        self.process_llm = process_llm
-        self.monitor = SchedulerGeneralMonitor(process_llm=self.process_llm, config=self.config)
-        self.dispatcher_monitor = SchedulerDispatcherMonitor(config=self.config)
-        self.retriever = SchedulerRetriever(process_llm=self.process_llm, config=self.config)
+        try:
+            # initialize submodules
+            self.chat_llm = chat_llm
+            self.process_llm = process_llm
+            self.db_engine = db_engine
+            self.monitor = SchedulerGeneralMonitor(
+                process_llm=self.process_llm, config=self.config, db_engine=self.db_engine
+            )
+            self.db_engine = self.monitor.db_engine
+            self.dispatcher_monitor = SchedulerDispatcherMonitor(config=self.config)
+            self.retriever = SchedulerRetriever(process_llm=self.process_llm, config=self.config)
+
+            if self.enable_parallel_dispatch:
+                self.dispatcher_monitor.initialize(dispatcher=self.dispatcher)
+                self.dispatcher_monitor.start()
+
+            # initialize with auth_config
+            if self.auth_config_path is not None and Path(self.auth_config_path).exists():
+                self.auth_config = AuthConfig.from_local_config(config_path=self.auth_config_path)
+            elif AuthConfig.default_config_exists():
+                self.auth_config = AuthConfig.from_local_config()
+            else:
+                self.auth_config = AuthConfig.from_local_env()
 
-        if self.enable_parallel_dispatch:
-            self.dispatcher_monitor.initialize(dispatcher=self.dispatcher)
-            self.dispatcher_monitor.start()
-
-        # initialize with auth_cofig
-        if self.auth_config_path is not None and Path(self.auth_config_path).exists():
-            self.auth_config = AuthConfig.from_local_config(config_path=self.auth_config_path)
-        elif AuthConfig.default_config_exists():
-            self.auth_config = AuthConfig.from_local_config()
-        else:
-            self.auth_config = None
+            if self.auth_config is not None:
+                self.rabbitmq_config = self.auth_config.rabbitmq
+                self.initialize_rabbitmq(config=self.rabbitmq_config)
 
-        if self.auth_config is not None:
-            self.rabbitmq_config = self.auth_config.rabbitmq
-            self.initialize_rabbitmq(config=self.rabbitmq_config)
+            logger.debug("GeneralScheduler has been initialized")
+        except Exception as e:
+            logger.error(f"Failed to initialize scheduler modules: {e}", exc_info=True)
+            # Clean up any partially initialized resources
+            self._cleanup_on_init_failure()
+            raise
 
-        logger.debug("GeneralScheduler has been initialized")
+    def _cleanup_on_init_failure(self):
+        """Clean up resources if initialization fails."""
+        try:
+            if hasattr(self, "dispatcher_monitor") and self.dispatcher_monitor is not None:
+                self.dispatcher_monitor.stop()
+        except Exception as e:
+            logger.warning(f"Error during cleanup: {e}")
 
     @property
     def mem_cube(self) -> GeneralMemCube:
@@ -200,8 +229,11 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule, SchedulerLogg
             text_mem_base: TreeTextMemory = text_mem_base
 
             # process rerank memories with llm
-            query_monitor = self.monitor.query_monitors[user_id][mem_cube_id]
-            query_history = query_monitor.get_queries_with_timesort()
+            query_db_manager = self.monitor.query_monitors[user_id][mem_cube_id]
+            # Sync with database to get latest query history
+            query_db_manager.sync_with_orm()
+
+            query_history = query_db_manager.obj.get_queries_with_timesort()
             memories_with_new_order, rerank_success_flag = (
                 self.retriever.process_and_rerank_memories(
                     queries=query_history,
@@ -211,8 +243,27 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule, SchedulerLogg
                 )
             )
 
-            # update working memory monitors
-            query_keywords = query_monitor.get_keywords_collections()
+            # Filter completely unrelated memories according to query_history
+            logger.info(f"Filtering memories based on query history: {len(query_history)} queries")
+            filtered_memories, filter_success_flag = self.retriever.filter_unrelated_memories(
+                query_history=query_history,
+                memories=memories_with_new_order,
+            )
+
+            if filter_success_flag:
+                logger.info(
+                    f"Memory filtering completed successfully. "
+                    f"Filtered from {len(memories_with_new_order)} to {len(filtered_memories)} memories"
+                )
+                memories_with_new_order = filtered_memories
+            else:
+                logger.warning(
+                    "Memory filtering failed - keeping all memories as fallback. "
+                    f"Original count: {len(memories_with_new_order)}"
+                )
+
+            # Update working memory monitors
+            query_keywords = query_db_manager.obj.get_keywords_collections()
             logger.info(
                 f"Processing {len(memories_with_new_order)} memories with {len(query_keywords)} query keywords"
             )
@@ -235,7 +286,7 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule, SchedulerLogg
 
             mem_monitors: list[MemoryMonitorItem] = self.monitor.working_memory_monitors[user_id][
                 mem_cube_id
-            ].get_sorted_mem_monitors(reverse=True)
+            ].obj.get_sorted_mem_monitors(reverse=True)
             new_working_memories = [mem_monitor.tree_memory_item for mem_monitor in mem_monitors]
 
             text_mem_base.replace_working_memory(memories=new_working_memories)
@@ -278,6 +329,7 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule, SchedulerLogg
             new_text_memories = new_memories
         else:
             logger.error("Not Implemented.")
+            return
 
         try:
             if isinstance(mem_cube.act_mem, VLLMKVCacheMemory):
@@ -333,7 +385,9 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule, SchedulerLogg
             )
 
         except Exception as e:
-            logger.warning(f"MOS-based activation memory update failed: {e}", exc_info=True)
+            logger.error(f"MOS-based activation memory update failed: {e}", exc_info=True)
+            # Re-raise the exception if it's critical for the operation
+            # For now, we'll continue execution but this should be reviewed
 
     def update_activation_memory_periodically(
         self,
@@ -358,7 +412,8 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule, SchedulerLogg
                 if (
                     user_id not in self.monitor.working_memory_monitors
                     or mem_cube_id not in self.monitor.working_memory_monitors[user_id]
-                    or len(self.monitor.working_memory_monitors[user_id][mem_cube_id].memories) == 0
+                    or len(self.monitor.working_memory_monitors[user_id][mem_cube_id].obj.memories)
+                    == 0
                 ):
                     logger.warning(
                         "No memories found in working_memory_monitors, activation memory update is skipped"
@@ -369,9 +424,13 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule, SchedulerLogg
                     user_id=user_id, mem_cube_id=mem_cube_id, mem_cube=mem_cube
                 )
 
+                # Sync with database to get latest activation memories
+                activation_db_manager = self.monitor.activation_memory_monitors[user_id][
+                    mem_cube_id
+                ]
+                activation_db_manager.sync_with_orm()
                 new_activation_memories = [
-                    m.memory_text
-                    for m in self.monitor.activation_memory_monitors[user_id][mem_cube_id].memories
+                    m.memory_text for m in activation_db_manager.obj.memories
                 ]
 
                 logger.info(
@@ -412,6 +471,11 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule, SchedulerLogg
             messages = [messages]  # transform single message to list
 
         for message in messages:
+            if not isinstance(message, ScheduleMessageItem):
+                error_msg = f"Invalid message type: {type(message)}, expected ScheduleMessageItem"
+                logger.error(error_msg)
+                raise TypeError(error_msg)
+
             self.memos_message_queue.put(message)
             logger.info(f"Submitted message: {message.label} - {message.content}")
 
@@ -427,6 +491,11 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule, SchedulerLogg
             messages = [messages]  # transform single message to list
 
         for message in messages:
+            if not isinstance(message, ScheduleLogForWebItem):
+                error_msg = f"Invalid message type: {type(message)}, expected ScheduleLogForWebItem"
+                logger.error(error_msg)
+                raise TypeError(error_msg)
+
             self._web_log_message_queue.put(message)
             message_info = message.debug_info()
             logger.debug(f"Submitted Scheduling log for web: {message_info}")
@@ -461,25 +530,26 @@ class BaseScheduler(RabbitMQSchedulerModule, RedisSchedulerModule, SchedulerLogg
         """
         while self._running:  # Use a running flag for graceful shutdown
             try:
-                # Check if queue has messages (non-blocking)
-                if not self.memos_message_queue.empty():
-                    # Get all available messages at once
-                    messages = []
-                    while not self.memos_message_queue.empty():
-                        try:
-                            messages.append(self.memos_message_queue.get_nowait())
-                        except queue.Empty:
-                            break
-
-                    if messages:
-                        try:
-                            self.dispatcher.dispatch(messages)
-                        except Exception as e:
-                            logger.error(f"Error dispatching messages: {e!s}")
-                        finally:
-                            # Mark all messages as processed
-                            for _ in messages:
-                                self.memos_message_queue.task_done()
+                # Get all available messages at once (thread-safe approach)
+                messages = []
+                while True:
+                    try:
+                        # Use get_nowait() directly without empty() check to avoid race conditions
+                        message = self.memos_message_queue.get_nowait()
+                        messages.append(message)
+                    except queue.Empty:
+                        # No more messages available
+                        break
+
+                if messages:
+                    try:
+                        self.dispatcher.dispatch(messages)
+                    except Exception as e:
+                        logger.error(f"Error dispatching messages: {e!s}")
+                    finally:
+                        # Mark all messages as processed
+                        for _ in messages:
+                            self.memos_message_queue.task_done()
 
                 # Sleep briefly to prevent busy waiting
                 time.sleep(self._consume_interval)  # Adjust interval as needed

memos/mem_scheduler/general_modules/dispatcher.py

@@ -2,8 +2,8 @@ import concurrent
 
 from collections import defaultdict
 from collections.abc import Callable
-from concurrent.futures import ThreadPoolExecutor
 
+from memos.context.context import ContextThreadPoolExecutor
 from memos.log import get_logger
 from memos.mem_scheduler.general_modules.base import BaseSchedulerModule
 from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
@@ -33,7 +33,7 @@ class SchedulerDispatcher(BaseSchedulerModule):
         self.enable_parallel_dispatch = enable_parallel_dispatch
         self.thread_name_prefix = "dispatcher"
         if self.enable_parallel_dispatch:
-            self.dispatcher_executor = ThreadPoolExecutor(
+            self.dispatcher_executor = ContextThreadPoolExecutor(
                 max_workers=self.max_workers, thread_name_prefix=self.thread_name_prefix
             )
         else: