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

memos/mem_scheduler/modules/redis_service.py

@@ -2,11 +2,9 @@ import asyncio
 import threading
 
 from collections.abc import Callable
+from typing import Any
 
-import redis
-
-from redis import Redis
-
+from memos.dependency import require_python_package
 from memos.log import get_logger
 from memos.mem_scheduler.modules.base import BaseSchedulerModule
 
@@ -15,6 +13,11 @@ logger = get_logger(__name__)
 
 
 class RedisSchedulerModule(BaseSchedulerModule):
+    @require_python_package(
+        import_name="redis",
+        install_command="pip install redis",
+        install_link="https://redis.readthedocs.io/en/stable/",
+    )
     def __init__(self):
         """
         intent_detector: Object used for intent recognition (such as the above IntentDetector)
@@ -35,23 +38,25 @@ class RedisSchedulerModule(BaseSchedulerModule):
         self._redis_listener_loop: asyncio.AbstractEventLoop | None = None
 
     @property
-    def redis(self) -> Redis:
+    def redis(self) -> Any:
         return self._redis_conn
 
     @redis.setter
-    def redis(self, value: Redis) -> None:
+    def redis(self, value: Any) -> None:
         self._redis_conn = value
 
     def initialize_redis(
         self, redis_host: str = "localhost", redis_port: int = 6379, redis_db: int = 0
     ):
+        import redis
+
         self.redis_host = redis_host
         self.redis_port = redis_port
         self.redis_db = redis_db
 
         try:
             logger.debug(f"Connecting to Redis at {redis_host}:{redis_port}/{redis_db}")
-            self._redis_conn = Redis(
+            self._redis_conn = redis.Redis(
                 host=self.redis_host, port=self.redis_port, db=self.redis_db, decode_responses=True
             )
             # test conn
@@ -63,21 +68,21 @@ class RedisSchedulerModule(BaseSchedulerModule):
         self._redis_conn.xtrim("user:queries:stream", self.query_list_capacity)
         return self._redis_conn
 
-    async def add_message_stream(self, message: dict):
+    async def redis_add_message_stream(self, message: dict):
         logger.debug(f"add_message_stream: {message}")
         return self._redis_conn.xadd("user:queries:stream", message)
 
-    async def consume_message_stream(self, message: dict):
+    async def redis_consume_message_stream(self, message: dict):
         logger.debug(f"consume_message_stream: {message}")
 
-    def _run_listener_async(self, handler: Callable):
+    def _redis_run_listener_async(self, handler: Callable):
         """Run the async listener in a separate thread"""
         self._redis_listener_loop = asyncio.new_event_loop()
         asyncio.set_event_loop(self._redis_listener_loop)
 
         async def listener_wrapper():
             try:
-                await self._listen_query_stream(handler)
+                await self.__redis_listen_query_stream(handler)
             except Exception as e:
                 logger.error(f"Listener thread error: {e}")
             finally:
@@ -85,8 +90,12 @@ class RedisSchedulerModule(BaseSchedulerModule):
 
         self._redis_listener_loop.run_until_complete(listener_wrapper())
 
-    async def _listen_query_stream(self, handler=None, last_id: str = "$", block_time: int = 2000):
+    async def __redis_listen_query_stream(
+        self, handler=None, last_id: str = "$", block_time: int = 2000
+    ):
         """Internal async stream listener"""
+        import redis
+
         self._redis_listener_running = True
         while self._redis_listener_running:
             try:
@@ -99,6 +108,7 @@ class RedisSchedulerModule(BaseSchedulerModule):
                     for _, stream_messages in messages:
                         for message_id, message_data in stream_messages:
                             try:
+                                print(f"deal with message_data {message_data}")
                                 await handler(message_data)
                                 last_id = message_id
                             except Exception as e:
@@ -112,17 +122,17 @@ class RedisSchedulerModule(BaseSchedulerModule):
                 logger.error(f"Unexpected error: {e}")
                 await asyncio.sleep(1)
 
-    def start_listening(self, handler: Callable | None = None):
+    def redis_start_listening(self, handler: Callable | None = None):
         """Start the Redis stream listener in a background thread"""
         if self._redis_listener_thread and self._redis_listener_thread.is_alive():
             logger.warning("Listener is already running")
             return
 
         if handler is None:
-            handler = self.consume_message_stream
+            handler = self.redis_consume_message_stream
 
         self._redis_listener_thread = threading.Thread(
-            target=self._run_listener_async,
+            target=self._redis_run_listener_async,
             args=(handler,),
             daemon=True,
             name="RedisListenerThread",
@@ -130,13 +140,7 @@ class RedisSchedulerModule(BaseSchedulerModule):
         self._redis_listener_thread.start()
         logger.info("Started Redis stream listener thread")
 
-    def close(self):
-        """Close Redis connection"""
-        if self._redis_conn is not None:
-            self._redis_conn.close()
-            self._redis_conn = None
-
-    def stop_listening(self):
+    def redis_stop_listening(self):
         """Stop the listener thread gracefully"""
         self._redis_listener_running = False
         if self._redis_listener_thread and self._redis_listener_thread.is_alive():
@@ -144,3 +148,9 @@ class RedisSchedulerModule(BaseSchedulerModule):
             if self._redis_listener_thread.is_alive():
                 logger.warning("Listener thread did not stop gracefully")
         logger.info("Redis stream listener stopped")
+
+    def redis_close(self):
+        """Close Redis connection"""
+        if self._redis_conn is not None:
+            self._redis_conn.close()
+            self._redis_conn = None

memos/mem_scheduler/modules/retriever.py

@@ -1,41 +1,268 @@
+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 (
+    TreeTextMemory_SEARCH_METHOD,
+)
+from memos.mem_scheduler.utils import (
+    extract_json_dict,
+    is_all_chinese,
+    is_all_english,
+    transform_name_to_key,
+)
+from memos.memories.textual.tree import TextualMemoryItem, TreeTextMemory
 
 
 logger = get_logger(__name__)
 
 
 class SchedulerRetriever(BaseSchedulerModule):
-    def __init__(self, chat_llm, context_window_size=5):
+    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
+
+    def search(
+        self, query: str, mem_cube: GeneralMemCube, top_k: int, method=TreeTextMemory_SEARCH_METHOD
+    ):
+        """Search in text memory with the given query.
+
+        Args:
+            query: The search query string
+            top_k: Number of top results to return
+            method: Search method to use
+
+        Returns:
+            Search results or None if not implemented
+        """
+        text_mem_base = mem_cube.text_mem
+        try:
+            if method == TreeTextMemory_SEARCH_METHOD:
+                assert isinstance(text_mem_base, TreeTextMemory)
+                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:
+                raise NotImplementedError(str(type(text_mem_base)))
+        except Exception as e:
+            logger.error(f"Fail to search. The exeption is {e}.", exc_info=True)
+            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]:
         """
-        monitor: Object used to acquire monitoring information
-        mem_cube: Object/interface for querying the underlying database
-        context_window_size: Size of the context window for conversation history
+        Filters out low-quality or duplicate memories based on text similarity.
+
+        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
+
+        Returns:
+            List of filtered memories with duplicates removed
         """
-        super().__init__()
+        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 = {}
 
-        self.monitors = {}
-        self.context_window_size = context_window_size
+            for current_idx in range(len(text_memories)):
+                is_duplicate = False
 
-        self._chat_llm = chat_llm
-        self._current_mem_cube = None
+                # Compare with already kept memories
+                for kept_idx in to_keep:
+                    similarity_score = similarity_matrix[current_idx, kept_idx]
 
-    @property
-    def memory_texts(self) -> list[str]:
-        """The memory cube associated with this MemChat."""
-        return self._memory_text_list
+                    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
 
-    @memory_texts.setter
-    def memory_texts(self, value: list[str]) -> None:
-        """The memory cube associated with this MemChat."""
-        self._memory_text_list = value
+                if not is_duplicate:
+                    to_keep.append(current_idx)
 
-    def fetch_context(self):
+            # 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]:
         """
-        Extract the context window from the current conversation
-        conversation_history: a list (in chronological order)
+        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.
+
+        Returns:
+            List of filtered memories meeting the length requirement
         """
-        return self._memory_text_list[-self.context_window_size :]
+        if not text_memories:
+            logging.debug("Empty memories list received in short memory filter")
+            return []
+
+        filtered_memories = []
+        removed_count = 0
+
+        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
+
+            if length >= min_length_threshold:
+                filtered_memories.append(memory)
+            else:
+                removed_count += 1
+
+        if removed_count > 0:
+            logger.info(
+                f"Filtered out {removed_count} short memories "
+                f"(below {min_length_threshold} units). "
+                f"Total remaining: {len(filtered_memories)}"
+            )
+
+        return filtered_memories
+
+    def replace_working_memory(
+        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,
+            )
+
+            unique_memory = list(dict.fromkeys(filtered_combined_text_memory))
+
+            try:
+                prompt = self.build_prompt(
+                    "memory_reranking",
+                    queries=queries,
+                    current_order=unique_memory,
+                    staging_buffer=[],
+                )
+                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")
 
-    def retrieve(self, query: str, memory_texts: list[str], top_k: int = 5) -> list[str]:
-        return None
+        return memories_with_new_order