MemoryOS 0.2.1__py3-none-any.whl → 0.2.2__py3-none-any.whl

memos/mem_scheduler/general_scheduler.py

@@ -4,12 +4,16 @@ from memos.configs.mem_scheduler import GeneralSchedulerConfig
 from memos.log import get_logger
 from memos.mem_cube.general import GeneralMemCube
 from memos.mem_scheduler.base_scheduler import BaseScheduler
-from memos.mem_scheduler.modules.schemas import (
+from memos.mem_scheduler.schemas.general_schemas import (
     ADD_LABEL,
     ANSWER_LABEL,
+    DEFAULT_MAX_QUERY_KEY_WORDS,
     QUERY_LABEL,
-    ScheduleMessageItem,
+    MemCubeID,
+    UserID,
 )
+from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
+from memos.mem_scheduler.schemas.monitor_schemas import QueryMonitorItem
 from memos.memories.textual.tree import TextualMemoryItem, TreeTextMemory
 
 
@@ -29,6 +33,51 @@ class GeneralScheduler(BaseScheduler):
         }
         self.dispatcher.register_handlers(handlers)
 
+    # for evaluation
+    def search_for_eval(
+        self,
+        query: str,
+        user_id: UserID | str,
+        top_k: int,
+    ) -> list[str]:
+        query_keywords = self.monitor.extract_query_keywords(query=query)
+        logger.info(f'Extract keywords "{query_keywords}" from query "{query}"')
+
+        item = QueryMonitorItem(
+            query_text=query,
+            keywords=query_keywords,
+            max_keywords=DEFAULT_MAX_QUERY_KEY_WORDS,
+        )
+        self.monitor.query_monitors.put(item=item)
+        logger.debug(
+            f"Queries in monitor are {self.monitor.query_monitors.get_queries_with_timesort()}."
+        )
+
+        queries = [query]
+
+        # recall
+        cur_working_memory, new_candidates = self.process_session_turn(
+            queries=queries,
+            user_id=user_id,
+            mem_cube_id=self.current_mem_cube_id,
+            mem_cube=self.current_mem_cube,
+            top_k=self.top_k,
+        )
+        logger.info(f"Processed {queries} and get {len(new_candidates)} new candidate memories.")
+
+        # 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]
+
     def _query_message_consumer(self, messages: list[ScheduleMessageItem]) -> None:
         """
         Process and handle query trigger messages from the queue.
@@ -36,12 +85,12 @@ class GeneralScheduler(BaseScheduler):
         Args:
             messages: List of query messages to process
         """
-        logger.debug(f"Messages {messages} assigned to {QUERY_LABEL} handler.")
+        logger.info(f"Messages {messages} assigned to {QUERY_LABEL} handler.")
 
         # Process the query in a session turn
         grouped_messages = self.dispatcher.group_messages_by_user_and_cube(messages=messages)
 
-        self._validate_messages(messages=messages, label=QUERY_LABEL)
+        self.validate_schedule_messages(messages=messages, label=QUERY_LABEL)
 
         for user_id in grouped_messages:
             for mem_cube_id in grouped_messages[user_id]:
@@ -49,16 +98,50 @@ class GeneralScheduler(BaseScheduler):
                 if len(messages) == 0:
                     return
 
+                mem_cube = messages[0].mem_cube
+
                 # for status update
                 self._set_current_context_from_message(msg=messages[0])
 
-                self.process_session_turn(
-                    queries=[msg.content for msg in messages],
+                # update query monitors
+                for msg in messages:
+                    query = msg.content
+                    query_keywords = self.monitor.extract_query_keywords(query=query)
+                    logger.info(f'Extract keywords "{query_keywords}" from query "{query}"')
+
+                    item = QueryMonitorItem(
+                        query_text=query,
+                        keywords=query_keywords,
+                        max_keywords=DEFAULT_MAX_QUERY_KEY_WORDS,
+                    )
+                    self.monitor.query_monitors.put(item=item)
+                logger.debug(
+                    f"Queries in monitor are {self.monitor.query_monitors.get_queries_with_timesort()}."
+                )
+
+                queries = [msg.content for msg in messages]
+
+                # recall
+                cur_working_memory, new_candidates = self.process_session_turn(
+                    queries=queries,
                     user_id=user_id,
                     mem_cube_id=mem_cube_id,
-                    mem_cube=messages[0].mem_cube,
+                    mem_cube=mem_cube,
                     top_k=self.top_k,
                 )
+                logger.info(
+                    f"Processed {queries} and get {len(new_candidates)} new candidate memories."
+                )
+
+                # rerank
+                new_order_working_memory = self.replace_working_memory(
+                    user_id=user_id,
+                    mem_cube_id=mem_cube_id,
+                    mem_cube=mem_cube,
+                    original_memory=cur_working_memory,
+                    new_memory=new_candidates,
+                )
+                logger.info(f"size of new_order_working_memory: {len(new_order_working_memory)}")
 
     def _answer_message_consumer(self, messages: list[ScheduleMessageItem]) -> None:
         """
@@ -67,11 +150,11 @@ class GeneralScheduler(BaseScheduler):
         Args:
           messages: List of answer messages to process
         """
-        logger.debug(f"Messages {messages} assigned to {ANSWER_LABEL} handler.")
+        logger.info(f"Messages {messages} assigned to {ANSWER_LABEL} handler.")
         # Process the query in a session turn
         grouped_messages = self.dispatcher.group_messages_by_user_and_cube(messages=messages)
 
-        self._validate_messages(messages=messages, label=ANSWER_LABEL)
+        self.validate_schedule_messages(messages=messages, label=ANSWER_LABEL)
 
         for user_id in grouped_messages:
             for mem_cube_id in grouped_messages[user_id]:
@@ -82,8 +165,18 @@ class GeneralScheduler(BaseScheduler):
                 # for status update
                 self._set_current_context_from_message(msg=messages[0])
 
-                # update acivation memories
+                # update activation memories
                 if self.enable_act_memory_update:
+                    if (
+                        len(self.monitor.working_memory_monitors[user_id][mem_cube_id].memories)
+                        == 0
+                    ):
+                        self.initialize_working_memory_monitors(
+                            user_id=user_id,
+                            mem_cube_id=mem_cube_id,
+                            mem_cube=messages[0].mem_cube,
+                        )
+
                     self.update_activation_memory_periodically(
                         interval_seconds=self.monitor.act_mem_update_interval,
                         label=ANSWER_LABEL,
@@ -93,94 +186,121 @@ class GeneralScheduler(BaseScheduler):
                     )
 
     def _add_message_consumer(self, messages: list[ScheduleMessageItem]) -> None:
-        logger.debug(f"Messages {messages} assigned to {ADD_LABEL} handler.")
+        logger.info(f"Messages {messages} assigned to {ADD_LABEL} handler.")
         # Process the query in a session turn
         grouped_messages = self.dispatcher.group_messages_by_user_and_cube(messages=messages)
 
-        self._validate_messages(messages=messages, label=ADD_LABEL)
+        self.validate_schedule_messages(messages=messages, label=ADD_LABEL)
+        try:
+            for user_id in grouped_messages:
+                for mem_cube_id in grouped_messages[user_id]:
+                    messages = grouped_messages[user_id][mem_cube_id]
+                    if len(messages) == 0:
+                        return
 
-        for user_id in grouped_messages:
-            for mem_cube_id in grouped_messages[user_id]:
-                messages = grouped_messages[user_id][mem_cube_id]
-                if len(messages) == 0:
-                    return
+                    # for status update
+                    self._set_current_context_from_message(msg=messages[0])
 
-                # for status update
-                self._set_current_context_from_message(msg=messages[0])
+                    # submit logs
+                    for msg in messages:
+                        try:
+                            userinput_memory_ids = json.loads(msg.content)
+                        except Exception as e:
+                            logger.error(f"Error: {e}. Content: {msg.content}", exc_info=True)
+                            userinput_memory_ids = []
 
-                # submit logs
-                for msg in messages:
-                    user_inputs = json.loads(msg.content)
-                    self.log_adding_user_inputs(
-                        user_inputs=user_inputs,
-                        user_id=msg.user_id,
-                        mem_cube_id=msg.mem_cube_id,
-                        mem_cube=msg.mem_cube,
-                    )
+                        mem_cube = msg.mem_cube
+                        for memory_id in userinput_memory_ids:
+                            mem_item: TextualMemoryItem = mem_cube.text_mem.get(memory_id=memory_id)
+                            mem_type = mem_item.metadata.memory_type
+                            mem_content = mem_item.memory
 
-                # update acivation memories
-                if self.enable_act_memory_update:
-                    self.update_activation_memory_periodically(
-                        interval_seconds=self.monitor.act_mem_update_interval,
-                        label=ADD_LABEL,
-                        user_id=user_id,
-                        mem_cube_id=mem_cube_id,
-                        mem_cube=messages[0].mem_cube,
-                    )
+                            self.log_adding_memory(
+                                memory=mem_content,
+                                memory_type=mem_type,
+                                user_id=msg.user_id,
+                                mem_cube_id=msg.mem_cube_id,
+                                mem_cube=msg.mem_cube,
+                                log_func_callback=self._submit_web_logs,
+                            )
+
+                    # update activation memories
+                    if self.enable_act_memory_update:
+                        self.update_activation_memory_periodically(
+                            interval_seconds=self.monitor.act_mem_update_interval,
+                            label=ADD_LABEL,
+                            user_id=user_id,
+                            mem_cube_id=mem_cube_id,
+                            mem_cube=messages[0].mem_cube,
+                        )
+        except Exception as e:
+            logger.error(f"Error: {e}", exc_info=True)
 
     def process_session_turn(
         self,
         queries: str | list[str],
-        user_id: str,
-        mem_cube_id: str,
+        user_id: UserID | str,
+        mem_cube_id: MemCubeID | str,
         mem_cube: GeneralMemCube,
         top_k: int = 10,
-        query_history: list[str] | None = None,
-    ) -> None:
+    ) -> tuple[list[TextualMemoryItem], list[TextualMemoryItem]] | None:
         """
         Process a dialog turn:
         - If q_list reaches window size, trigger retrieval;
         - Immediately switch to the new memory if retrieval is triggered.
         """
-        if isinstance(queries, str):
-            queries = [queries]
-
-        if query_history is None:
-            query_history = queries
-        else:
-            query_history.extend(queries)
 
         text_mem_base = mem_cube.text_mem
         if not isinstance(text_mem_base, TreeTextMemory):
             logger.error("Not implemented!", exc_info=True)
             return
 
-        working_memory: list[TextualMemoryItem] = text_mem_base.get_working_memory()
-        text_working_memory: list[str] = [w_m.memory for w_m in working_memory]
+        logger.info(f"Processing {len(queries)} queries.")
+
+        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=query_history, text_working_memory=text_working_memory
+            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.debug(f"missing_evidences: {item}")
-                results = self.retriever.search(
-                    query=item, mem_cube=mem_cube, top_k=k_per_evidence, method=self.search_method
-                )
-                logger.debug(f"search results for {missing_evidences}: {results}")
-                new_candidates.extend(results)
-
-            new_order_working_memory = self.retriever.replace_working_memory(
-                queries=queries,
-                user_id=user_id,
-                mem_cube_id=mem_cube_id,
-                mem_cube=mem_cube,
-                original_memory=working_memory,
-                new_memory=new_candidates,
-                top_k=top_k,
+        time_trigger_flag = False
+        if self.monitor.timed_trigger(
+            last_time=self.monitor.last_query_consume_time,
+            interval_seconds=self.monitor.query_trigger_interval,
+        ):
+            time_trigger_flag = True
+
+        if (not intent_result["trigger_retrieval"]) and (not time_trigger_flag):
+            logger.info(f"Query schedule not triggered. Intent_result: {intent_result}")
+            return
+        elif (not intent_result["trigger_retrieval"]) and time_trigger_flag:
+            logger.info("Query schedule is forced to trigger due to time ticker")
+            intent_result["trigger_retrieval"] = True
+            intent_result["missing_evidences"] = queries
+        else:
+            logger.info(
+                f'Query schedule triggered for user "{user_id}" and mem_cube "{mem_cube_id}".'
+                f" Missing evidences: {intent_result['missing_evidences']}"
+            )
+
+        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)
+
+        if len(new_candidates) == 0:
+            logger.warning(
+                f"As new_candidates is empty, new_candidates is set same to working_memory.\n"
+                f"time_trigger_flag: {time_trigger_flag}; intent_result: {intent_result}"
             )
-            logger.debug(f"size of new_order_working_memory: {len(new_order_working_memory)}")
+            new_candidates = cur_working_memory
+        return cur_working_memory, new_candidates

memos/mem_scheduler/modules/base.py

@@ -3,7 +3,7 @@ from pathlib import Path
 from memos.llms.base import BaseLLM
 from memos.log import get_logger
 from memos.mem_cube.general import GeneralMemCube
-from memos.mem_scheduler.modules.schemas import BASE_DIR
+from memos.mem_scheduler.schemas.general_schemas import BASE_DIR
 from memos.templates.mem_scheduler_prompts import PROMPT_MAPPING
 
 
@@ -17,8 +17,8 @@ class BaseSchedulerModule:
 
         self._chat_llm = None
         self._process_llm = None
-        self._current_mem_cube_id: str | None = None
-        self._current_mem_cube: GeneralMemCube | None = None
+        self.current_mem_cube_id: str | None = None
+        self.current_mem_cube: GeneralMemCube | None = None
         self.mem_cubes: dict[str, GeneralMemCube] = {}
 
     def load_template(self, template_name: str) -> str:
@@ -75,9 +75,9 @@ class BaseSchedulerModule:
     @property
     def mem_cube(self) -> GeneralMemCube:
         """The memory cube associated with this MemChat."""
-        return self._current_mem_cube
+        return self.current_mem_cube
 
     @mem_cube.setter
     def mem_cube(self, value: GeneralMemCube) -> None:
         """The memory cube associated with this MemChat."""
-        self._current_mem_cube = value
+        self.current_mem_cube = value

memos/mem_scheduler/modules/dispatcher.py

@@ -4,7 +4,7 @@ from concurrent.futures import ThreadPoolExecutor
 
 from memos.log import get_logger
 from memos.mem_scheduler.modules.base import BaseSchedulerModule
-from memos.mem_scheduler.modules.schemas import ScheduleMessageItem
+from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
 
 
 logger = get_logger(__name__)
@@ -22,7 +22,7 @@ class SchedulerDispatcher(BaseSchedulerModule):
     - Bulk handler registration
     """
 
-    def __init__(self, max_workers=3, enable_parallel_dispatch=False):
+    def __init__(self, max_workers=30, enable_parallel_dispatch=False):
         super().__init__()
         # Main dispatcher thread pool
         self.max_workers = max_workers
@@ -128,16 +128,13 @@ class SchedulerDispatcher(BaseSchedulerModule):
             else:
                 handler = self.handlers[label]
             # dispatch to different handler
-            logger.debug(f"Dispatch {len(msgs)} messages to {label} handler.")
+            logger.debug(f"Dispatch {len(msgs)} message(s) to {label} handler.")
             if self.enable_parallel_dispatch and self.dispatcher_executor is not None:
                 # Capture variables in lambda to avoid loop variable issues
-                # TODO check this
-                future = self.dispatcher_executor.submit(handler, msgs)
-                logger.debug(f"Dispatched {len(msgs)} messages as future task")
-                return future
+                self.dispatcher_executor.submit(handler, msgs)
+                logger.info(f"Dispatched {len(msgs)} message(s) as future task")
             else:
                 handler(msgs)
-                return None
 
     def join(self, timeout: float | None = None) -> bool:
         """Wait for all dispatched tasks to complete.
@@ -159,7 +156,7 @@ class SchedulerDispatcher(BaseSchedulerModule):
         if self.dispatcher_executor is not None:
             self.dispatcher_executor.shutdown(wait=True)
         self._running = False
-        logger.info("Dispatcher has been shutdown")
+        logger.info("Dispatcher has been shutdown.")
 
     def __enter__(self):
         self._running = True

memos/mem_scheduler/modules/misc.py

@@ -1,20 +1,85 @@
-import threading
+import json
 
+from contextlib import suppress
+from datetime import datetime
 from queue import Empty, Full, Queue
-from typing import TypeVar
+from typing import TYPE_CHECKING, TypeVar
 
+from pydantic import field_serializer
+
+
+if TYPE_CHECKING:
+    from pydantic import BaseModel
 
 T = TypeVar("T")
 
+BaseModelType = TypeVar("T", bound="BaseModel")
+
+
+class DictConversionMixin:
+    """
+    Provides conversion functionality between Pydantic models and dictionaries,
+    including datetime serialization handling.
+    """
+
+    @field_serializer("timestamp", check_fields=False)
+    def serialize_datetime(self, dt: datetime | None, _info) -> str | None:
+        """
+        Custom datetime serialization logic.
+        - Supports timezone-aware datetime objects
+        - Compatible with models without timestamp field (via check_fields=False)
+        """
+        if dt is None:
+            return None
+        return dt.isoformat()
+
+    def to_dict(self) -> dict:
+        """
+        Convert model instance to dictionary.
+        - Uses model_dump to ensure field consistency
+        - Prioritizes custom serializer for timestamp handling
+        """
+        dump_data = self.model_dump()
+        if hasattr(self, "timestamp") and self.timestamp is not None:
+            dump_data["timestamp"] = self.serialize_datetime(self.timestamp, None)
+        return dump_data
+
+    @classmethod
+    def from_dict(cls: type[BaseModelType], data: dict) -> BaseModelType:
+        """
+        Create model instance from dictionary.
+        - Automatically converts timestamp strings to datetime objects
+        """
+        data_copy = data.copy()  # Avoid modifying original dictionary
+        if "timestamp" in data_copy and isinstance(data_copy["timestamp"], str):
+            try:
+                data_copy["timestamp"] = datetime.fromisoformat(data_copy["timestamp"])
+            except ValueError:
+                # Handle invalid time formats - adjust as needed (e.g., log warning or set to None)
+                data_copy["timestamp"] = None
+
+        return cls(**data_copy)
+
+    def __str__(self) -> str:
+        """
+        Convert to formatted JSON string.
+        - Used for user-friendly display in print() or str() calls
+        """
+        return json.dumps(
+            self.to_dict(),
+            indent=4,
+            ensure_ascii=False,
+            default=lambda o: str(o),  # Handle other non-serializable objects
+        )
+
 
 class AutoDroppingQueue(Queue[T]):
     """A thread-safe queue that automatically drops the oldest item when full."""
 
     def __init__(self, maxsize: int = 0):
         super().__init__(maxsize=maxsize)
-        self._lock = threading.Lock()  # Additional lock to prevent race conditions
 
-    def put(self, item: T, block: bool = True, timeout: float | None = None) -> None:
+    def put(self, item: T, block: bool = False, timeout: float | None = None) -> None:
         """Put an item into the queue.
 
         If the queue is full, the oldest item will be automatically removed to make space.
@@ -25,15 +90,15 @@ class AutoDroppingQueue(Queue[T]):
             block: Ignored (kept for compatibility with Queue interface)
             timeout: Ignored (kept for compatibility with Queue interface)
         """
-        with self._lock:  # Ensure atomic operation
-            try:
-                # First try non-blocking put
-                super().put(item, block=False)
-            except Full:
-                # If queue is full, remove the oldest item
-                from contextlib import suppress
-
-                with suppress(Empty):
-                    self.get_nowait()  # Remove oldest item
-                # Retry putting the new item
-                super().put(item, block=False)
+        try:
+            # First try non-blocking put
+            super().put(item, block=block, timeout=timeout)
+        except Full:
+            with suppress(Empty):
+                self.get_nowait()  # Remove oldest item
+            # Retry putting the new item
+            super().put(item, block=block, timeout=timeout)
+
+    def get_queue_content_without_pop(self) -> list[T]:
+        """Return a copy of the queue's contents without modifying it."""
+        return list(self.queue)