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

memos/mem_scheduler/modules/monitor.py

@@ -6,18 +6,23 @@ from memos.llms.base import BaseLLM
 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.modules.misc import AutoDroppingQueue as Queue
-from memos.mem_scheduler.modules.schemas import (
+from memos.mem_scheduler.schemas.general_schemas import (
     DEFAULT_ACTIVATION_MEM_MONITOR_SIZE_LIMIT,
+    DEFAULT_WEIGHT_VECTOR_FOR_RANKING,
     DEFAULT_WORKING_MEM_MONITOR_SIZE_LIMIT,
     MONITOR_ACTIVATION_MEMORY_TYPE,
     MONITOR_WORKING_MEMORY_TYPE,
     MemCubeID,
-    MemoryMonitorManager,
     UserID,
 )
-from memos.mem_scheduler.utils import extract_json_dict
-from memos.memories.textual.tree import TextualMemoryItem, TreeTextMemory
+from memos.mem_scheduler.schemas.monitor_schemas import (
+    MemoryMonitorItem,
+    MemoryMonitorManager,
+    QueryMonitorItem,
+    QueryMonitorQueue,
+)
+from memos.mem_scheduler.utils.misc_utils import extract_json_dict
+from memos.memories.textual.tree import TreeTextMemory
 
 
 logger = get_logger(__name__)
@@ -31,7 +36,8 @@ class SchedulerMonitor(BaseSchedulerModule):
 
         # hyper-parameters
         self.config: BaseSchedulerConfig = config
-        self.act_mem_update_interval = self.config.get("act_mem_update_interval", 300)
+        self.act_mem_update_interval = self.config.get("act_mem_update_interval", 30)
+        self.query_trigger_interval = self.config.get("query_trigger_interval", 10)
 
         # Partial Retention Strategy
         self.partial_retention_number = 2
@@ -39,16 +45,39 @@ class SchedulerMonitor(BaseSchedulerModule):
         self.activation_mem_monitor_capacity = DEFAULT_ACTIVATION_MEM_MONITOR_SIZE_LIMIT
 
         # attributes
-        self.query_history = Queue(maxsize=self.config.context_window_size)
-        self.intent_history = Queue(maxsize=self.config.context_window_size)
+        # recording query_messages
+        self.query_monitors: QueryMonitorQueue[QueryMonitorItem] = QueryMonitorQueue(
+            maxsize=self.config.context_window_size
+        )
+
         self.working_memory_monitors: dict[UserID, dict[MemCubeID, MemoryMonitorManager]] = {}
         self.activation_memory_monitors: dict[UserID, dict[MemCubeID, MemoryMonitorManager]] = {}
 
         # Lifecycle monitor
-        self._last_activation_mem_update_time = datetime.min
+        self.last_activation_mem_update_time = datetime.min
+        self.last_query_consume_time = datetime.min
 
         self._process_llm = process_llm
 
+    def extract_query_keywords(self, query: str) -> list:
+        """Extracts core keywords from a user query based on specific semantic rules."""
+        prompt_name = "query_keywords_extraction"
+        prompt = self.build_prompt(
+            template_name=prompt_name,
+            query=query,
+        )
+        llm_response = self._process_llm.generate([{"role": "user", "content": prompt}])
+        try:
+            # Parse JSON output from LLM response
+            keywords = extract_json_dict(llm_response)
+            assert isinstance(keywords, list)
+        except Exception as e:
+            logger.error(
+                f"Failed to parse keywords from LLM response: {llm_response}. Error: {e!s}"
+            )
+            keywords = [query]
+        return keywords
+
     def register_memory_manager_if_not_exists(
         self,
         user_id: str,
@@ -90,13 +119,15 @@ class SchedulerMonitor(BaseSchedulerModule):
                 f"mem_cube_id={mem_cube_id} in the provided memory_monitors dictionary"
             )
 
-    def update_memory_monitors(self, user_id: str, mem_cube_id: str, mem_cube: GeneralMemCube):
+    def update_working_memory_monitors(
+        self,
+        new_working_memory_monitors: list[MemoryMonitorItem],
+        user_id: str,
+        mem_cube_id: str,
+        mem_cube: GeneralMemCube,
+    ):
         text_mem_base: TreeTextMemory = mem_cube.text_mem
-
-        if not isinstance(text_mem_base, TreeTextMemory):
-            logger.error("Not Implemented")
-            return
-
+        assert isinstance(text_mem_base, TreeTextMemory)
         self.working_mem_monitor_capacity = min(
             DEFAULT_WORKING_MEM_MONITOR_SIZE_LIMIT,
             (
@@ -105,17 +136,6 @@ class SchedulerMonitor(BaseSchedulerModule):
             ),
         )
 
-        self.update_working_memory_monitors(
-            user_id=user_id, mem_cube_id=mem_cube_id, mem_cube=mem_cube
-        )
-
-        self.update_activation_memory_monitors(
-            user_id=user_id, mem_cube_id=mem_cube_id, mem_cube=mem_cube
-        )
-
-    def update_working_memory_monitors(
-        self, user_id: str, mem_cube_id: str, mem_cube: GeneralMemCube
-    ):
         # register monitors
         self.register_memory_manager_if_not_exists(
             user_id=user_id,
@@ -124,14 +144,8 @@ class SchedulerMonitor(BaseSchedulerModule):
             max_capacity=self.working_mem_monitor_capacity,
         )
 
-        # === update working memory monitors ===
-        # Retrieve current working memory content
-        text_mem_base: TreeTextMemory = mem_cube.text_mem
-        working_memory: list[TextualMemoryItem] = text_mem_base.get_working_memory()
-        text_working_memory: list[str] = [w_m.memory for w_m in working_memory]
-
         self.working_memory_monitors[user_id][mem_cube_id].update_memories(
-            text_working_memories=text_working_memory,
+            new_memory_monitors=new_working_memory_monitors,
             partial_retention_number=self.partial_retention_number,
         )
 
@@ -149,16 +163,13 @@ class SchedulerMonitor(BaseSchedulerModule):
         # Sort by importance_score in descending order and take top k
         top_k_memories = sorted(
             self.working_memory_monitors[user_id][mem_cube_id].memories,
-            key=lambda m: m.get_score(),
+            key=lambda m: m.get_importance_score(weight_vector=DEFAULT_WEIGHT_VECTOR_FOR_RANKING),
             reverse=True,
         )[: self.activation_mem_monitor_capacity]
 
-        # Extract just the text from these memories
-        text_top_k_memories = [m.memory_text for m in top_k_memories]
-
         # Update the activation memory monitors with these important memories
         self.activation_memory_monitors[user_id][mem_cube_id].update_memories(
-            text_working_memories=text_top_k_memories,
+            new_memory_monitors=top_k_memories,
             partial_retention_number=self.partial_retention_number,
         )
 
@@ -206,10 +217,10 @@ class SchedulerMonitor(BaseSchedulerModule):
             )
             return []
 
-        manager = monitor_dict[user_id][mem_cube_id]
+        manager: MemoryMonitorManager = monitor_dict[user_id][mem_cube_id]
         # Sort memories by recording_count in descending order and return top_k items
-        sorted_memories = sorted(manager.memories, key=lambda m: m.recording_count, reverse=True)
-        sorted_text_memories = [m.memory_text for m in sorted_memories[:top_k]]
+        sorted_memory_monitors = manager.get_sorted_mem_monitors(reverse=True)
+        sorted_text_memories = [m.memory_text for m in sorted_memory_monitors[:top_k]]
         return sorted_text_memories
 
     def get_monitors_info(self, user_id: str, mem_cube_id: str) -> dict[str, Any]:

memos/mem_scheduler/modules/rabbitmq_service.py

@@ -4,13 +4,13 @@ import threading
 import time
 
 from pathlib import Path
-from queue import Queue
 
 from memos.configs.mem_scheduler import AuthConfig, RabbitMQConfig
 from memos.dependency import require_python_package
 from memos.log import get_logger
 from memos.mem_scheduler.modules.base import BaseSchedulerModule
-from memos.mem_scheduler.modules.schemas import DIRECT_EXCHANGE_TYPE, FANOUT_EXCHANGE_TYPE
+from memos.mem_scheduler.modules.misc import AutoDroppingQueue
+from memos.mem_scheduler.schemas.general_schemas import DIRECT_EXCHANGE_TYPE, FANOUT_EXCHANGE_TYPE
 
 
 logger = get_logger(__name__)
@@ -38,7 +38,9 @@ class RabbitMQSchedulerModule(BaseSchedulerModule):
 
         # fixed params
         self.rabbitmq_message_cache_max_size = 10  # Max 10 messages
-        self.rabbitmq_message_cache = Queue(maxsize=self.rabbitmq_message_cache_max_size)
+        self.rabbitmq_message_cache = AutoDroppingQueue(
+            maxsize=self.rabbitmq_message_cache_max_size
+        )
         self.rabbitmq_connection_attempts = 3  # Max retry attempts on connection failure
         self.rabbitmq_retry_delay = 5  # Delay (seconds) between retries
         self.rabbitmq_heartbeat = 60  # Heartbeat interval (seconds) for connectio
@@ -214,12 +216,12 @@ class RabbitMQSchedulerModule(BaseSchedulerModule):
     def on_rabbitmq_message(self, channel, method, properties, body):
         """Handle incoming messages. Only for test."""
         try:
-            print(f"Received message: {body.decode()}")
-            self.rabbitmq_message_cache.put_nowait({"properties": properties, "body": body})
-            print(f"message delivery_tag: {method.delivery_tag}")
+            print(f"Received message: {body.decode()}\n")
+            self.rabbitmq_message_cache.put({"properties": properties, "body": body})
+            print(f"message delivery_tag: {method.delivery_tag}\n")
             channel.basic_ack(delivery_tag=method.delivery_tag)
         except Exception as e:
-            logger.error(f"Message handling failed: {e}")
+            logger.error(f"Message handling failed: {e}", exc_info=True)
 
     def wait_for_connection_ready(self):
         start_time = time.time()

memos/mem_scheduler/modules/retriever.py

@@ -1,20 +1,19 @@
-import logging
-
 from memos.configs.mem_scheduler import BaseSchedulerConfig
-from memos.dependency import require_python_package
 from memos.llms.base import BaseLLM
 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.modules.schemas import (
+from memos.mem_scheduler.schemas.general_schemas import (
     TreeTextMemory_SEARCH_METHOD,
 )
-from memos.mem_scheduler.utils import (
-    extract_json_dict,
-    is_all_chinese,
-    is_all_english,
+from memos.mem_scheduler.utils.filter_utils import (
+    filter_similar_memories,
+    filter_too_short_memories,
     transform_name_to_key,
 )
+from memos.mem_scheduler.utils.misc_utils import (
+    extract_json_dict,
+)
 from memos.memories.textual.tree import TextualMemoryItem, TreeTextMemory
 
 
@@ -25,19 +24,16 @@ class SchedulerRetriever(BaseSchedulerModule):
     def __init__(self, process_llm: BaseLLM, config: BaseSchedulerConfig):
         super().__init__()
 
-        self.config: BaseSchedulerConfig = config
-        self.process_llm = process_llm
-
         # hyper-parameters
         self.filter_similarity_threshold = 0.75
         self.filter_min_length_threshold = 6
 
-        # log function callbacks
-        self.log_working_memory_replacement = None
+        self.config: BaseSchedulerConfig = config
+        self.process_llm = process_llm
 
     def search(
         self, query: str, mem_cube: GeneralMemCube, top_k: int, method=TreeTextMemory_SEARCH_METHOD
-    ):
+    ) -> list[TextualMemoryItem]:
         """Search in text memory with the given query.
 
         Args:
@@ -66,203 +62,124 @@ class SchedulerRetriever(BaseSchedulerModule):
             results = []
         return results
 
-    @require_python_package(
-        import_name="sklearn",
-        install_command="pip install scikit-learn",
-        install_link="https://scikit-learn.org/stable/install.html",
-    )
-    def filter_similar_memories(
-        self, text_memories: list[str], similarity_threshold: float = 0.75
-    ) -> list[str]:
+    def rerank_memories(
+        self,
+        queries: list[str],
+        original_memories: list[str],
+        top_k: int,
+    ) -> (list[str], bool):
         """
-        Filters out low-quality or duplicate memories based on text similarity.
+        Rerank memories based on relevance to given queries using LLM.
 
         Args:
-            text_memories: List of text memories to filter
-            similarity_threshold: Threshold for considering memories duplicates (0.0-1.0)
-                                Higher values mean stricter filtering
+            queries: List of query strings to determine relevance
+            original_memories: List of memory strings to be reranked
+            top_k: Number of top memories to return after reranking
 
         Returns:
-            List of filtered memories with duplicates removed
-        """
-        from sklearn.feature_extraction.text import TfidfVectorizer
-        from sklearn.metrics.pairwise import cosine_similarity
-
-        if not text_memories:
-            logging.warning("Received empty memories list - nothing to filter")
-            return []
-
-        for idx in range(len(text_memories)):
-            if not isinstance(text_memories[idx], str):
-                logger.error(
-                    f"{text_memories[idx]} in memories is not a string,"
-                    f" and now has been transformed to be a string."
-                )
-                text_memories[idx] = str(text_memories[idx])
-
-        try:
-            # Step 1: Vectorize texts using TF-IDF
-            vectorizer = TfidfVectorizer()
-            tfidf_matrix = vectorizer.fit_transform(text_memories)
-
-            # Step 2: Calculate pairwise similarity matrix
-            similarity_matrix = cosine_similarity(tfidf_matrix)
-
-            # Step 3: Identify duplicates
-            to_keep = []
-            removal_reasons = {}
-
-            for current_idx in range(len(text_memories)):
-                is_duplicate = False
-
-                # Compare with already kept memories
-                for kept_idx in to_keep:
-                    similarity_score = similarity_matrix[current_idx, kept_idx]
-
-                    if similarity_score > similarity_threshold:
-                        is_duplicate = True
-                        # Generate removal reason with sample text
-                        removal_reasons[current_idx] = (
-                            f"Memory too similar (score: {similarity_score:.2f}) to kept memory #{kept_idx}. "
-                            f"Kept: '{text_memories[kept_idx][:100]}...' | "
-                            f"Removed: '{text_memories[current_idx][:100]}...'"
-                        )
-                        logger.info(removal_reasons)
-                        break
-
-                if not is_duplicate:
-                    to_keep.append(current_idx)
-
-            # Return filtered memories
-            return [text_memories[i] for i in sorted(to_keep)]
-
-        except Exception as e:
-            logging.error(f"Error filtering memories: {e!s}")
-            return text_memories  # Return original list if error occurs
-
-    def filter_too_short_memories(
-        self, text_memories: list[str], min_length_threshold: int = 20
-    ) -> list[str]:
-        """
-        Filters out text memories that fall below the minimum length requirement.
-        Handles both English (word count) and Chinese (character count) differently.
-
-        Args:
-            text_memories: List of text memories to be filtered
-            min_length_threshold: Minimum length required to keep a memory.
-                                For English: word count, for Chinese: character count.
+            List of reranked memory strings (length <= top_k)
 
-        Returns:
-            List of filtered memories meeting the length requirement
+        Note:
+            If LLM reranking fails, falls back to original order (truncated to top_k)
         """
-        if not text_memories:
-            logging.debug("Empty memories list received in short memory filter")
-            return []
+        success_flag = False
 
-        filtered_memories = []
-        removed_count = 0
+        logger.info(f"Starting memory reranking for {len(original_memories)} memories")
 
-        for memory in text_memories:
-            stripped_memory = memory.strip()
-            if not stripped_memory:  # Skip empty/whitespace memories
-                removed_count += 1
-                continue
-
-            # Determine measurement method based on language
-            if is_all_english(stripped_memory):
-                length = len(stripped_memory.split())  # Word count for English
-            elif is_all_chinese(stripped_memory):
-                length = len(stripped_memory)  # Character count for Chinese
-            else:
-                logger.debug(
-                    f"Mixed-language memory, using character count: {stripped_memory[:50]}..."
-                )
-                length = len(stripped_memory)  # Default to character count
+        # Build LLM prompt for memory reranking
+        prompt = self.build_prompt(
+            "memory_reranking",
+            queries=[f"[0] {queries[0]}"],
+            current_order=[f"[{i}] {mem}" for i, mem in enumerate(original_memories)],
+        )
+        logger.debug(f"Generated reranking prompt: {prompt[:200]}...")  # Log first 200 chars
 
-            if length >= min_length_threshold:
-                filtered_memories.append(memory)
-            else:
-                removed_count += 1
+        # Get LLM response
+        response = self.process_llm.generate([{"role": "user", "content": prompt}])
+        logger.debug(f"Received LLM response: {response[:200]}...")  # Log first 200 chars
 
-        if removed_count > 0:
+        try:
+            # Parse JSON response
+            response = extract_json_dict(response)
+            new_order = response["new_order"][:top_k]
+            text_memories_with_new_order = [original_memories[idx] for idx in new_order]
             logger.info(
-                f"Filtered out {removed_count} short memories "
-                f"(below {min_length_threshold} units). "
-                f"Total remaining: {len(filtered_memories)}"
+                f"Successfully reranked memories. Returning top {len(text_memories_with_new_order)} items;"
+                f"Ranking reasoning: {response['reasoning']}"
             )
+            success_flag = True
+        except Exception as e:
+            logger.error(
+                f"Failed to rerank memories with LLM. Exception: {e}. Raw response: {response} ",
+                exc_info=True,
+            )
+            text_memories_with_new_order = original_memories[:top_k]
+            success_flag = False
+        return text_memories_with_new_order, success_flag
 
-        return filtered_memories
-
-    def replace_working_memory(
+    def process_and_rerank_memories(
         self,
         queries: list[str],
-        user_id: str,
-        mem_cube_id: str,
-        mem_cube: GeneralMemCube,
         original_memory: list[TextualMemoryItem],
         new_memory: list[TextualMemoryItem],
         top_k: int = 10,
-    ) -> None | list[TextualMemoryItem]:
-        """Replace working memory with new memories after reranking."""
-        memories_with_new_order = None
-        text_mem_base = mem_cube.text_mem
-        if isinstance(text_mem_base, TreeTextMemory):
-            text_mem_base: TreeTextMemory = text_mem_base
-            combined_memory = original_memory + new_memory
-            memory_map = {
-                transform_name_to_key(name=mem_obj.memory): mem_obj for mem_obj in combined_memory
-            }
-            combined_text_memory = [transform_name_to_key(name=m.memory) for m in combined_memory]
-
-            # apply filters
-            filtered_combined_text_memory = self.filter_similar_memories(
-                text_memories=combined_text_memory,
-                similarity_threshold=self.filter_similarity_threshold,
-            )
-
-            filtered_combined_text_memory = self.filter_too_short_memories(
-                text_memories=filtered_combined_text_memory,
-                min_length_threshold=self.filter_min_length_threshold,
-            )
+    ) -> list[TextualMemoryItem] | None:
+        """
+        Process and rerank memory items by combining original and new memories,
+        applying filters, and then reranking based on relevance to queries.
 
-            unique_memory = list(dict.fromkeys(filtered_combined_text_memory))
+        Args:
+            queries: List of query strings to rerank memories against
+            original_memory: List of original TextualMemoryItem objects
+            new_memory: List of new TextualMemoryItem objects to merge
+            top_k: Maximum number of memories to return after reranking
 
-            try:
-                prompt = self.build_prompt(
-                    "memory_reranking",
-                    queries=queries,
-                    current_order=unique_memory,
-                    staging_buffer=[],
+        Returns:
+            List of reranked TextualMemoryItem objects, or None if processing fails
+        """
+        # Combine original and new memories into a single list
+        combined_memory = original_memory + new_memory
+
+        # Create a mapping from normalized text to memory objects
+        memory_map = {
+            transform_name_to_key(name=mem_obj.memory): mem_obj for mem_obj in combined_memory
+        }
+
+        # Extract normalized text representations from all memory items
+        combined_text_memory = [m.memory for m in combined_memory]
+
+        # Apply similarity filter to remove overly similar memories
+        filtered_combined_text_memory = filter_similar_memories(
+            text_memories=combined_text_memory,
+            similarity_threshold=self.filter_similarity_threshold,
+        )
+
+        # Apply length filter to remove memories that are too short
+        filtered_combined_text_memory = filter_too_short_memories(
+            text_memories=filtered_combined_text_memory,
+            min_length_threshold=self.filter_min_length_threshold,
+        )
+
+        # Ensure uniqueness of memory texts using dictionary keys (preserves order)
+        unique_memory = list(dict.fromkeys(filtered_combined_text_memory))
+
+        # Rerank the filtered memories based on relevance to the queries
+        text_memories_with_new_order, success_flag = self.rerank_memories(
+            queries=queries,
+            original_memories=unique_memory,
+            top_k=top_k,
+        )
+
+        # Map reranked text entries back to their original memory objects
+        memories_with_new_order = []
+        for text in text_memories_with_new_order:
+            normalized_text = transform_name_to_key(name=text)
+            if normalized_text in memory_map:  # Ensure correct key matching
+                memories_with_new_order.append(memory_map[normalized_text])
+            else:
+                logger.warning(
+                    f"Memory text not found in memory map. text: {text};\n"
+                    f"Keys of memory_map: {memory_map.keys()}"
                 )
-                response = self.process_llm.generate([{"role": "user", "content": prompt}])
-                response = extract_json_dict(response)
-                text_memories_with_new_order = response.get("new_order", [])[:top_k]
-            except Exception as e:
-                logger.error(f"Fail to rerank with LLM, Exeption: {e}.", exc_info=True)
-                text_memories_with_new_order = unique_memory[:top_k]
-
-            memories_with_new_order = []
-            for text in text_memories_with_new_order:
-                normalized_text = transform_name_to_key(name=text)
-                if text in memory_map:
-                    memories_with_new_order.append(memory_map[normalized_text])
-                else:
-                    logger.warning(
-                        f"Memory text not found in memory map. text: {text}; keys of memory_map: {memory_map.keys()}"
-                    )
-
-            text_mem_base.replace_working_memory(memories_with_new_order)
-            logger.info(
-                f"The working memory has been replaced with {len(memories_with_new_order)} new memories."
-            )
-            self.log_working_memory_replacement(
-                original_memory=original_memory,
-                new_memory=memories_with_new_order,
-                user_id=user_id,
-                mem_cube_id=mem_cube_id,
-                mem_cube=mem_cube,
-            )
-        else:
-            logger.error("memory_base is not supported")
 
-        return memories_with_new_order
+        return memories_with_new_order, success_flag