MemoryOS 0.0.1__py3-none-any.whl → 0.1.12__py3-none-any.whl

memos/mem_scheduler/general_scheduler.py

@@ -0,0 +1,305 @@
+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 (
+    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__)
+
+
+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,
+        }
+        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:
+        """
+        Process and handle answer trigger messages from the queue.
+
+        Args:
+          messages: List of answer 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
+
+            # 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
+            ]
+
+            # 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
+            )
+
+            # 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
+
+            # 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)
+
+    def _query_message_consume(self, messages: list[ScheduleMessageItem]) -> None:
+        """
+        Process and handle query trigger messages from the queue.
+
+        Args:
+            messages: List of query 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)
+
+    def process_session_turn(
+        self,
+        query: str,
+        top_k: int = 10,
+        top_n: int = 5,
+    ) -> 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()
+        else:
+            logger.error("Not implemented!")
+            return
+        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
+        )
+        if intent_result["trigger_retrieval"]:
+            missing_evidence = intent_result["missing_evidence"]
+            num_evidence = len(missing_evidence)
+            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}")
+                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
+                    ]
+                )
+            )
+            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}")

memos/mem_scheduler/modules/base.py

@@ -0,0 +1,74 @@
+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.templates.mem_scheduler_prompts import PROMPT_MAPPING
+
+
+logger = get_logger(__name__)
+
+
+class BaseSchedulerModule:
+    def __init__(self):
+        """Initialize the scheduler with the given configuration."""
+        self.base_dir = Path(BASE_DIR)
+
+        self._chat_llm = 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:
+        if template_name not in PROMPT_MAPPING:
+            logger.error("Prompt template is not found!")
+        prompt = PROMPT_MAPPING[template_name]
+        return prompt
+
+    def build_prompt(self, template_name: str, **kwargs) -> str:
+        template = self.load_template(template_name)
+        if not template:
+            raise FileNotFoundError(f"Prompt template `{template_name}` not found.")
+        return template.format(**kwargs)
+
+    def _build_system_prompt(self, memories: list | None = None) -> str:
+        """Build system prompt with optional memories context."""
+        base_prompt = (
+            "You are a knowledgeable and helpful AI assistant. "
+            "You have access to conversation memories that help you provide more personalized responses. "
+            "Use the memories to understand the user's context, preferences, and past interactions. "
+            "If memories are provided, reference them naturally when relevant, but don't explicitly mention having memories."
+        )
+
+        if memories:
+            memory_context = "\n\n## Conversation Context:\n"
+            for i, memory in enumerate(memories, 1):
+                memory_context += f"{i}. {memory.memory}\n"
+            return base_prompt + memory_context
+
+        return base_prompt
+
+    def get_mem_cube(self, mem_cube_id: str) -> GeneralMemCube:
+        logger.error(f"mem_cube {mem_cube_id} does not exists.")
+        return self.mem_cubes.get(mem_cube_id, None)
+
+    @property
+    def chat_llm(self) -> BaseLLM:
+        """The memory cube associated with this MemChat."""
+        return self._chat_llm
+
+    @chat_llm.setter
+    def chat_llm(self, value: BaseLLM) -> None:
+        """The memory cube associated with this MemChat."""
+        self._chat_llm = value
+
+    @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

memos/mem_scheduler/modules/dispatcher.py

@@ -0,0 +1,103 @@
+from collections import defaultdict
+from collections.abc import Callable
+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
+
+
+logger = get_logger(__name__)
+
+
+class SchedulerDispatcher(BaseSchedulerModule):
+    """
+    Thread pool-based message dispatcher that routes messages to dedicated handlers
+    based on their labels.
+
+    Features:
+    - Dedicated thread pool per message label
+    - Batch message processing
+    - Graceful shutdown
+    - Bulk handler registration
+    """
+
+    def __init__(self, max_workers=3, enable_parallel_dispatch=False):
+        super().__init__()
+        # Main dispatcher thread pool
+        self.max_workers = max_workers
+        # Only initialize thread pool if in parallel mode
+        self.enable_parallel_dispatch = enable_parallel_dispatch
+        if self.enable_parallel_dispatch:
+            self.dispatcher_executor = ThreadPoolExecutor(
+                max_workers=self.max_workers, thread_name_prefix="dispatcher"
+            )
+        else:
+            self.dispatcher_executor = None
+        logger.info(f"enable_parallel_dispatch is set to {self.enable_parallel_dispatch}")
+        # Registered message handlers
+        self.handlers: dict[str, Callable] = {}
+        # Dispatcher running state
+        self._running = False
+
+    def register_handler(self, label: str, handler: Callable[[list[ScheduleMessageItem]], None]):
+        """
+        Register a handler function for a specific message label.
+
+        Args:
+            label: Message label to handle
+            handler: Callable that processes messages of this label
+        """
+        self.handlers[label] = handler
+
+    def register_handlers(
+        self, handlers: dict[str, Callable[[list[ScheduleMessageItem]], None]]
+    ) -> None:
+        """
+        Bulk register multiple handlers from a dictionary.
+
+        Args:
+            handlers: Dictionary mapping labels to handler functions
+                      Format: {label: handler_callable}
+        """
+        for label, handler in handlers.items():
+            if not isinstance(label, str):
+                logger.error(f"Invalid label type: {type(label)}. Expected str.")
+                continue
+            if not callable(handler):
+                logger.error(f"Handler for label '{label}' is not callable.")
+                continue
+            self.register_handler(label=label, handler=handler)
+        logger.info(f"Registered {len(handlers)} handlers in bulk")
+
+    def _default_message_handler(self, messages: list[ScheduleMessageItem]) -> None:
+        logger.debug(f"Using _default_message_handler to deal with messages: {messages}")
+
+    def dispatch(self, msg_list: list[ScheduleMessageItem]):
+        """
+        Dispatch a list of messages to their respective handlers.
+
+        Args:
+            msg_list: List of ScheduleMessageItem objects to process
+        """
+
+        # Group messages by their labels
+        label_groups = defaultdict(list)
+
+        # Organize messages by label
+        for message in msg_list:
+            label_groups[message.label].append(message)
+
+        # Process each label group
+        for label, msgs in label_groups.items():
+            if label not in self.handlers:
+                logger.error(f"No handler registered for label: {label}")
+                handler = self._default_message_handler
+            else:
+                handler = self.handlers[label]
+            # 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)
+            else:
+                handler(msgs)  # Direct serial execution

memos/mem_scheduler/modules/monitor.py

@@ -0,0 +1,82 @@
+import json
+
+from typing import Any
+
+from memos.log import get_logger
+from memos.mem_cube.general import GeneralMemCube
+from memos.mem_scheduler.modules.base import BaseSchedulerModule
+from memos.mem_scheduler.utils import extract_json_dict
+from memos.memories.textual.tree import TreeTextMemory
+
+
+logger = get_logger(__name__)
+
+
+class SchedulerMonitor(BaseSchedulerModule):
+    def __init__(self, chat_llm, activation_mem_size=5):
+        super().__init__()
+        self.statistics = {}
+        self.intent_history: list[str] = []
+        self.activation_mem_size = activation_mem_size
+        self.activation_memory_freq_list = [
+            {"memory": None, "count": 0} for _ in range(self.activation_mem_size)
+        ]
+
+        self._chat_llm = chat_llm
+
+    def update_stats(self, mem_cube):
+        self.statistics["activation_mem_size"] = self.activation_mem_size
+        mem_cube_info = self.get_mem_cube_info(mem_cube)
+        self.statistics.update(mem_cube_info)
+
+    def get_mem_cube_info(self, mem_cube: GeneralMemCube):
+        mem_cube_info = {}
+
+        text_mem = mem_cube.text_mem
+        if isinstance(text_mem, TreeTextMemory):
+            memory_size_dict = text_mem.memory_manager.memory_size
+            mem_cube_info["text_mem"] = memory_size_dict
+        else:
+            logger.error("Not Implemented")
+
+        return mem_cube_info
+
+    def detect_intent(
+        self,
+        q_list: list[str],
+        text_working_memory: list[str],
+        prompt_name="intent_recognizing",
+    ) -> dict[str, Any]:
+        """
+        Detect the intent of the user input.
+        """
+        prompt = self.build_prompt(
+            template_name=prompt_name,
+            q_list=q_list,
+            working_memory_list=text_working_memory,
+        )
+        response = self._chat_llm.generate([{"role": "user", "content": prompt}])
+        response = extract_json_dict(response)
+        return response
+
+    def update_freq(
+        self,
+        answer: str,
+        activation_memory_freq_list: list[dict],
+        prompt_name="freq_detecting",
+    ) -> list[dict]:
+        """
+        Use LLM to detect which memories in activation_memory_freq_list appear in the answer,
+        increment their count by 1, and return the updated list.
+        """
+        prompt = self.build_prompt(
+            template_name=prompt_name,
+            answer=answer,
+            activation_memory_freq_list=activation_memory_freq_list,
+        )
+        response = self._chat_llm.generate([{"role": "user", "content": prompt}])
+        try:
+            result = json.loads(response)
+        except Exception:
+            result = activation_memory_freq_list
+        return result