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

memos/mem_scheduler/general_scheduler.py

@@ -1,27 +1,16 @@
 import json
 
-from datetime import datetime, timedelta
-
 from memos.configs.mem_scheduler import GeneralSchedulerConfig
-from memos.llms.base import BaseLLM
 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.monitor import SchedulerMonitor
-from memos.mem_scheduler.modules.retriever import SchedulerRetriever
 from memos.mem_scheduler.modules.schemas import (
+    ADD_LABEL,
     ANSWER_LABEL,
-    DEFAULT_ACT_MEM_DUMP_PATH,
-    DEFAULT_ACTIVATION_MEM_SIZE,
-    NOT_INITIALIZED,
     QUERY_LABEL,
-    ScheduleLogForWebItem,
     ScheduleMessageItem,
-    TextMemory_SEARCH_METHOD,
-    TreeTextMemory_SEARCH_METHOD,
 )
 from memos.memories.textual.tree import TextualMemoryItem, TreeTextMemory
-from memos.templates.mem_scheduler_prompts import MEMORY_ASSEMBLY_TEMPLATE
 
 
 logger = get_logger(__name__)
@@ -31,275 +20,167 @@ class GeneralScheduler(BaseScheduler):
     def __init__(self, config: GeneralSchedulerConfig):
         """Initialize the scheduler with the given configuration."""
         super().__init__(config)
-        self.top_k = self.config.get("top_k", 10)
-        self.top_n = self.config.get("top_n", 5)
-        self.act_mem_update_interval = self.config.get("act_mem_update_interval", 300)
-        self.context_window_size = self.config.get("context_window_size", 5)
-        self.activation_mem_size = self.config.get(
-            "activation_mem_size", DEFAULT_ACTIVATION_MEM_SIZE
-        )
-        self.act_mem_dump_path = self.config.get("act_mem_dump_path", DEFAULT_ACT_MEM_DUMP_PATH)
-        self.search_method = TextMemory_SEARCH_METHOD
-        self._last_activation_mem_update_time = 0.0
-        self.query_list = []
 
         # register handlers
         handlers = {
-            QUERY_LABEL: self._query_message_consume,
-            ANSWER_LABEL: self._answer_message_consume,
+            QUERY_LABEL: self._query_message_consumer,
+            ANSWER_LABEL: self._answer_message_consumer,
+            ADD_LABEL: self._add_message_consumer,
         }
         self.dispatcher.register_handlers(handlers)
 
-    def initialize_modules(self, chat_llm: BaseLLM):
-        self.chat_llm = chat_llm
-        self.monitor = SchedulerMonitor(
-            chat_llm=self.chat_llm, activation_mem_size=self.activation_mem_size
-        )
-        self.retriever = SchedulerRetriever(chat_llm=self.chat_llm)
-        logger.debug("GeneralScheduler has been initialized")
-
-    def _answer_message_consume(self, messages: list[ScheduleMessageItem]) -> None:
+    def _query_message_consumer(self, messages: list[ScheduleMessageItem]) -> None:
         """
-        Process and handle answer trigger messages from the queue.
+        Process and handle query trigger messages from the queue.
 
         Args:
-          messages: List of answer messages to process
+            messages: List of query messages to process
         """
-        # TODO: This handler is not ready yet
-        logger.debug(f"Messages {messages} assigned to {ANSWER_LABEL} handler.")
-        for msg in messages:
-            if msg.label is not ANSWER_LABEL:
-                logger.error(f"_answer_message_consume is not designed for {msg.label}")
-                continue
-            answer = msg.content
-            self._current_user_id = msg.user_id
-            self._current_mem_cube_id = msg.mem_cube_id
-            self._current_mem_cube = msg.mem_cube
+        logger.debug(f"Messages {messages} assigned to {QUERY_LABEL} handler.")
 
-            # Get current activation memory items
-            current_activation_mem = [
-                item["memory"]
-                for item in self.monitor.activation_memory_freq_list
-                if item["memory"] is not None
-            ]
+        # Process the query in a session turn
+        grouped_messages = self.dispatcher.group_messages_by_user_and_cube(messages=messages)
 
-            # Update memory frequencies based on the answer
-            # TODO: not implemented
-            self.monitor.activation_memory_freq_list = self.monitor.update_freq(
-                answer=answer, activation_memory_freq_list=self.monitor.activation_memory_freq_list
-            )
+        self._validate_messages(messages=messages, label=QUERY_LABEL)
 
-            # Check if it's time to update activation memory
-            now = datetime.now()
-            if (now - self._last_activation_mem_update_time) >= timedelta(
-                seconds=self.act_mem_update_interval
-            ):
-                # TODO: not implemented
-                self.update_activation_memory(current_activation_mem)
-                self._last_activation_mem_update_time = now
+        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
 
-            # recording messages
-            log_message = self.create_autofilled_log_item(
-                log_title="memos answer triggers scheduling...",
-                label=ANSWER_LABEL,
-                log_content="activation_memory has been updated",
-            )
-            self._submit_web_logs(messages=log_message)
+                # for status update
+                self._set_current_context_from_message(msg=messages[0])
 
-    def _query_message_consume(self, messages: list[ScheduleMessageItem]) -> None:
+                self.process_session_turn(
+                    queries=[msg.content for msg in messages],
+                    user_id=user_id,
+                    mem_cube_id=mem_cube_id,
+                    mem_cube=messages[0].mem_cube,
+                    top_k=self.top_k,
+                )
+
+    def _answer_message_consumer(self, messages: list[ScheduleMessageItem]) -> None:
         """
-        Process and handle query trigger messages from the queue.
+        Process and handle answer trigger messages from the queue.
 
         Args:
-            messages: List of query messages to process
+          messages: List of answer messages to process
         """
-        logger.debug(f"Messages {messages} assigned to {QUERY_LABEL} handler.")
-        for msg in messages:
-            if msg.label is not QUERY_LABEL:
-                logger.error(f"_query_message_consume is not designed for {msg.label}")
-                continue
-            # Process the query in a session turn
-            self._current_user_id = msg.user_id
-            self._current_mem_cube_id = msg.mem_cube_id
-            self._current_mem_cube = msg.mem_cube
-            self.process_session_turn(query=msg.content, top_k=self.top_k, top_n=self.top_n)
+        logger.debug(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)
+
+        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])
+
+                # update acivation memories
+                if self.enable_act_memory_update:
+                    self.update_activation_memory_periodically(
+                        interval_seconds=self.monitor.act_mem_update_interval,
+                        label=ANSWER_LABEL,
+                        user_id=user_id,
+                        mem_cube_id=mem_cube_id,
+                        mem_cube=messages[0].mem_cube,
+                    )
+
+    def _add_message_consumer(self, messages: list[ScheduleMessageItem]) -> None:
+        logger.debug(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)
+
+        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])
+
+                # 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,
+                    )
+
+                # 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,
+                    )
 
     def process_session_turn(
         self,
-        query: str,
+        queries: str | list[str],
+        user_id: str,
+        mem_cube_id: str,
+        mem_cube: GeneralMemCube,
         top_k: int = 10,
-        top_n: int = 5,
+        query_history: list[str] | None = None,
     ) -> None:
         """
         Process a dialog turn:
         - If q_list reaches window size, trigger retrieval;
         - Immediately switch to the new memory if retrieval is triggered.
         """
-        q_list = [query]
-        self.query_list.append(query)
-        text_mem_base = self.mem_cube.text_mem
-        if isinstance(text_mem_base, TreeTextMemory):
-            working_memory: list[TextualMemoryItem] = text_mem_base.get_working_memory()
+        if isinstance(queries, str):
+            queries = [queries]
+
+        if query_history is None:
+            query_history = queries
         else:
-            logger.error("Not implemented!")
+            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]
         intent_result = self.monitor.detect_intent(
-            q_list=q_list, text_working_memory=text_working_memory
+            q_list=query_history, text_working_memory=text_working_memory
         )
+
         if intent_result["trigger_retrieval"]:
-            missing_evidence = intent_result["missing_evidence"]
-            num_evidence = len(missing_evidence)
+            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_evidence:
-                logger.debug(f"missing_evidence: {item}")
-                results = self.search(query=item, top_k=k_per_evidence, method=self.search_method)
-                logger.debug(f"search results for {missing_evidence}: {results}")
+            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)
 
-            # recording messages
-            log_message = self.create_autofilled_log_item(
-                log_title="user query triggers scheduling...",
-                label=QUERY_LABEL,
-                log_content=f"search new candidates for working memory: {len(new_candidates)}",
-            )
-            self._submit_web_logs(messages=log_message)
-            new_order_working_memory = self.replace_working_memory(
-                original_memory=working_memory, new_memory=new_candidates, top_k=top_k, top_n=top_n
-            )
-            self.update_activation_memory(new_order_working_memory)
-
-    def create_autofilled_log_item(
-        self, log_title: str, log_content: str, label: str
-    ) -> ScheduleLogForWebItem:
-        # TODO: create the log iterm with real stats
-        text_mem_base: TreeTextMemory = self.mem_cube.text_mem
-        current_memory_sizes = {
-            "long_term_memory_size": NOT_INITIALIZED,
-            "user_memory_size": NOT_INITIALIZED,
-            "working_memory_size": NOT_INITIALIZED,
-            "transformed_act_memory_size": NOT_INITIALIZED,
-            "parameter_memory_size": NOT_INITIALIZED,
-        }
-        memory_capacities = {
-            "long_term_memory_capacity": text_mem_base.memory_manager.memory_size["LongTermMemory"],
-            "user_memory_capacity": text_mem_base.memory_manager.memory_size["UserMemory"],
-            "working_memory_capacity": text_mem_base.memory_manager.memory_size["WorkingMemory"],
-            "transformed_act_memory_capacity": NOT_INITIALIZED,
-            "parameter_memory_capacity": NOT_INITIALIZED,
-        }
-
-        log_message = ScheduleLogForWebItem(
-            user_id=self._current_user_id,
-            mem_cube_id=self._current_mem_cube_id,
-            label=label,
-            log_title=log_title,
-            log_content=log_content,
-            current_memory_sizes=current_memory_sizes,
-            memory_capacities=memory_capacities,
-        )
-        return log_message
-
-    @property
-    def mem_cube(self) -> GeneralMemCube:
-        """The memory cube associated with this MemChat."""
-        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.retriever.mem_cube = value
-
-    def replace_working_memory(
-        self,
-        original_memory: list[TextualMemoryItem],
-        new_memory: list[TextualMemoryItem],
-        top_k: int = 10,
-        top_n: int = 5,
-    ) -> None | list[TextualMemoryItem]:
-        new_order_memory = None
-        text_mem_base = self.mem_cube.text_mem
-        if isinstance(text_mem_base, TreeTextMemory):
-            text_mem_base: TreeTextMemory = text_mem_base
-            combined_text_memory = [new_m.memory for new_m in original_memory] + [
-                new_m.memory for new_m in new_memory
-            ]
-            combined_memory = original_memory + new_memory
-            memory_map = {mem_obj.memory: mem_obj for mem_obj in combined_memory}
-
-            unique_memory = list(dict.fromkeys(combined_text_memory))
-            prompt = self.build_prompt(
-                "memory_reranking", query="", current_order=unique_memory, staging_buffer=[]
-            )
-            response = self.chat_llm.generate([{"role": "user", "content": prompt}])
-            response = json.loads(response)
-            new_order_text_memory = response.get("new_order", [])[: top_n + top_k]
-
-            new_order_memory = []
-            for text in new_order_text_memory:
-                if text in memory_map:
-                    new_order_memory.append(memory_map[text])
-                else:
-                    logger.warning(
-                        f"Memory text not found in memory map. text: {text}; memory_map: {memory_map}"
-                    )
-
-            text_mem_base.replace_working_memory(new_order_memory[top_n:])
-            new_order_memory = new_order_memory[:top_n]
-            logger.info(
-                f"The working memory has been replaced with {len(new_order_memory)} new memories."
-            )
-        else:
-            logger.error("memory_base is not supported")
-
-        return new_order_memory
-
-    def search(self, query: str, top_k: int, method=TreeTextMemory_SEARCH_METHOD):
-        text_mem_base = self.mem_cube.text_mem
-        if isinstance(text_mem_base, TreeTextMemory) and method == TextMemory_SEARCH_METHOD:
-            results_long_term = text_mem_base.search(
-                query=query, top_k=top_k, memory_type="LongTermMemory"
-            )
-            results_user = text_mem_base.search(query=query, top_k=top_k, memory_type="UserMemory")
-            results = results_long_term + results_user
-        else:
-            logger.error("Not implemented.")
-            results = None
-        return results
-
-    def update_activation_memory(self, new_memory: list[str | TextualMemoryItem]) -> None:
-        """
-        Update activation memory by extracting KVCacheItems from new_memory (list of str),
-        add them to a KVCacheMemory instance, and dump to disk.
-        """
-        # TODO: The function of update activation memory is waiting to test
-        if len(new_memory) == 0:
-            logger.error("update_activation_memory: new_memory is empty.")
-            return
-        if isinstance(new_memory[0], TextualMemoryItem):
-            new_text_memory = [mem.memory for mem in new_memory]
-        elif isinstance(new_memory[0], str):
-            new_text_memory = new_memory
-        else:
-            logger.error("Not Implemented.")
-
-        try:
-            act_mem = self.mem_cube.act_mem
-
-            text_memory = MEMORY_ASSEMBLY_TEMPLATE.format(
-                memory_text="".join(
-                    [
-                        f"{i + 1}. {sentence.strip()}\n"
-                        for i, sentence in enumerate(new_text_memory)
-                        if sentence.strip()  # Skip empty strings
-                    ]
-                )
+            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,
             )
-            act_mem.delete_all()
-            cache_item = act_mem.extract(text_memory)
-            act_mem.add(cache_item)
-            act_mem.dump(self.act_mem_dump_path)
-        except Exception as e:
-            logger.warning(f"MOS-based activation memory update failed: {e}")
+            logger.debug(f"size of new_order_working_memory: {len(new_order_working_memory)}")

memos/mem_scheduler/modules/base.py

@@ -16,6 +16,7 @@ class BaseSchedulerModule:
         self.base_dir = Path(BASE_DIR)
 
         self._chat_llm = None
+        self._process_llm = None
         self._current_mem_cube_id: str | None = None
         self._current_mem_cube: GeneralMemCube | None = None
         self.mem_cubes: dict[str, GeneralMemCube] = {}
@@ -63,6 +64,14 @@ class BaseSchedulerModule:
         """The memory cube associated with this MemChat."""
         self._chat_llm = value
 
+    @property
+    def process_llm(self) -> BaseLLM:
+        return self._process_llm
+
+    @process_llm.setter
+    def process_llm(self, value: BaseLLM) -> None:
+        self._process_llm = value
+
     @property
     def mem_cube(self) -> GeneralMemCube:
         """The memory cube associated with this MemChat."""

memos/mem_scheduler/modules/dispatcher.py

@@ -73,6 +73,38 @@ class SchedulerDispatcher(BaseSchedulerModule):
     def _default_message_handler(self, messages: list[ScheduleMessageItem]) -> None:
         logger.debug(f"Using _default_message_handler to deal with messages: {messages}")
 
+    def group_messages_by_user_and_cube(
+        self, messages: list[ScheduleMessageItem]
+    ) -> dict[str, dict[str, list[ScheduleMessageItem]]]:
+        """
+        Groups messages into a nested dictionary structure first by user_id, then by mem_cube_id.
+
+        Args:
+            messages: List of ScheduleMessageItem objects to be grouped
+
+        Returns:
+            A nested dictionary with the structure:
+            {
+                "user_id_1": {
+                    "mem_cube_id_1": [msg1, msg2, ...],
+                    "mem_cube_id_2": [msg3, msg4, ...],
+                    ...
+                },
+                "user_id_2": {
+                    ...
+                },
+                ...
+            }
+            Where each msg is the original ScheduleMessageItem object
+        """
+        grouped_dict = defaultdict(lambda: defaultdict(list))
+
+        for msg in messages:
+            grouped_dict[msg.user_id][msg.mem_cube_id].append(msg)
+
+        # Convert defaultdict to regular dict for cleaner output
+        return {user_id: dict(cube_groups) for user_id, cube_groups in grouped_dict.items()}
+
     def dispatch(self, msg_list: list[ScheduleMessageItem]):
         """
         Dispatch a list of messages to their respective handlers.
@@ -98,6 +130,40 @@ class SchedulerDispatcher(BaseSchedulerModule):
             # dispatch to different handler
             logger.debug(f"Dispatch {len(msgs)} messages to {label} handler.")
             if self.enable_parallel_dispatch and self.dispatcher_executor is not None:
-                self.dispatcher_executor.submit(handler, msgs)
+                # 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
             else:
-                handler(msgs)  # Direct serial execution
+                handler(msgs)
+                return None
+
+    def join(self, timeout: float | None = None) -> bool:
+        """Wait for all dispatched tasks to complete.
+
+        Args:
+            timeout: Maximum time to wait in seconds. None means wait forever.
+
+        Returns:
+            bool: True if all tasks completed, False if timeout occurred.
+        """
+        if not self.enable_parallel_dispatch or self.dispatcher_executor is None:
+            return True  # 串行模式无需等待
+
+        self.dispatcher_executor.shutdown(wait=True, timeout=timeout)
+        return True
+
+    def shutdown(self) -> None:
+        """Gracefully shutdown the dispatcher."""
+        if self.dispatcher_executor is not None:
+            self.dispatcher_executor.shutdown(wait=True)
+        self._running = False
+        logger.info("Dispatcher has been shutdown")
+
+    def __enter__(self):
+        self._running = True
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.shutdown()

memos/mem_scheduler/modules/misc.py

@@ -0,0 +1,39 @@
+import threading
+
+from queue import Empty, Full, Queue
+from typing import TypeVar
+
+
+T = TypeVar("T")
+
+
+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:
+        """Put an item into the queue.
+
+        If the queue is full, the oldest item will be automatically removed to make space.
+        This operation is thread-safe.
+
+        Args:
+            item: The item to be put into the queue
+            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)