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

memos/mem_scheduler/general_scheduler.py

@@ -1,27 +1,20 @@
 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 (
+from memos.mem_scheduler.schemas.general_schemas import (
+    ADD_LABEL,
     ANSWER_LABEL,
-    DEFAULT_ACT_MEM_DUMP_PATH,
-    DEFAULT_ACTIVATION_MEM_SIZE,
-    NOT_INITIALIZED,
+    DEFAULT_MAX_QUERY_KEY_WORDS,
     QUERY_LABEL,
-    ScheduleLogForWebItem,
-    ScheduleMessageItem,
-    TextMemory_SEARCH_METHOD,
-    TreeTextMemory_SEARCH_METHOD,
+    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
-from memos.templates.mem_scheduler_prompts import MEMORY_ASSEMBLY_TEMPLATE
 
 
 logger = get_logger(__name__)
@@ -31,275 +24,283 @@ 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
+    # 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,
         )
-        self.retriever = SchedulerRetriever(chat_llm=self.chat_llm)
-        logger.debug("GeneralScheduler has been initialized")
+        new_order_working_memory = new_order_working_memory[:top_k]
+        logger.info(f"size of new_order_working_memory: {len(new_order_working_memory)}")
 
-    def _answer_message_consume(self, messages: list[ScheduleMessageItem]) -> None:
+        return [m.memory for m in new_order_working_memory]
+
+    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
-
-            # 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
-            )
+        logger.info(f"Messages {messages} assigned to {QUERY_LABEL} handler.")
 
-            # 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)
+        # Process the query in a session turn
+        grouped_messages = self.dispatcher.group_messages_by_user_and_cube(messages=messages)
+
+        self.validate_schedule_messages(messages=messages, label=QUERY_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
+
+                mem_cube = messages[0].mem_cube
 
-    def _query_message_consume(self, messages: list[ScheduleMessageItem]) -> None:
+                # for status update
+                self._set_current_context_from_message(msg=messages[0])
+
+                # 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=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:
         """
-        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.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_schedule_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 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,
+                        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.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_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 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 = []
+
+                        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
+
+                            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,
-        query: str,
+        queries: str | list[str],
+        user_id: UserID | str,
+        mem_cube_id: MemCubeID | str,
+        mem_cube: GeneralMemCube,
         top_k: int = 10,
-        top_n: int = 5,
-    ) -> 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.
         """
-        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
+        text_mem_base = mem_cube.text_mem
+        if not isinstance(text_mem_base, TreeTextMemory):
+            logger.error("Not implemented!", exc_info=True)
+            return
 
-    @property
-    def mem_cube(self) -> GeneralMemCube:
-        """The memory cube associated with this MemChat."""
-        return self._current_mem_cube
+        logger.info(f"Processing {len(queries)} queries.")
 
-    @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
+        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
+        )
 
-    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}"
-                    )
+        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
 
-            text_mem_base.replace_working_memory(new_order_memory[top_n:])
-            new_order_memory = new_order_memory[:top_n]
+        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"The working memory has been replaced with {len(new_order_memory)} new memories."
+                f'Query schedule triggered for user "{user_id}" and mem_cube "{mem_cube_id}".'
+                f" Missing evidences: {intent_result['missing_evidences']}"
             )
-        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"
+        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
             )
-            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.")
+            logger.info(
+                f"search results for {missing_evidences}: {[one.memory for one in results]}"
+            )
+            new_candidates.extend(results)
 
-        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
-                    ]
-                )
+        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}"
             )
-            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}")
+            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
 
 
@@ -16,8 +16,9 @@ class BaseSchedulerModule:
         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._process_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:
@@ -63,12 +64,20 @@ 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."""
-        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