MemoryOS 1.0.0__py3-none-any.whl → 1.1.1__py3-none-any.whl

memos/mem_scheduler/general_modules/misc.py

@@ -1,9 +1,10 @@
 import json
+import os
 
 from contextlib import suppress
 from datetime import datetime
 from queue import Empty, Full, Queue
-from typing import TYPE_CHECKING, TypeVar
+from typing import TYPE_CHECKING, Any, Generic, TypeVar
 
 from pydantic import field_serializer
 
@@ -16,6 +17,75 @@ T = TypeVar("T")
 BaseModelType = TypeVar("T", bound="BaseModel")
 
 
+class EnvConfigMixin(Generic[T]):
+    """Abstract base class for environment variable configuration."""
+
+    ENV_PREFIX = "MEMSCHEDULER_"
+
+    @classmethod
+    def get_env_prefix(cls) -> str:
+        """Automatically generates environment variable prefix from class name.
+
+        Converts the class name to uppercase and appends an underscore.
+        If the class name ends with 'Config', that suffix is removed first.
+
+        Examples:
+            RabbitMQConfig -> "RABBITMQ_"
+            OpenAIConfig -> "OPENAI_"
+            GraphDBAuthConfig -> "GRAPH_DB_AUTH_"
+        """
+        class_name = cls.__name__
+        # Remove 'Config' suffix if present
+        if class_name.endswith("Config"):
+            class_name = class_name[:-6]
+        # Convert to uppercase and add trailing underscore
+
+        return f"{cls.ENV_PREFIX}{class_name.upper()}_"
+
+    @classmethod
+    def from_env(cls: type[T]) -> T:
+        """Creates a config instance from environment variables.
+
+        Reads all environment variables with the class-specific prefix and maps them
+        to corresponding configuration fields (converting to the appropriate types).
+
+        Returns:
+            An instance of the config class populated from environment variables.
+
+        Raises:
+            ValueError: If required environment variables are missing.
+        """
+        prefix = cls.get_env_prefix()
+        field_values = {}
+
+        for field_name, field_info in cls.model_fields.items():
+            env_var = f"{prefix}{field_name.upper()}"
+            field_type = field_info.annotation
+
+            if field_info.is_required() and env_var not in os.environ:
+                raise ValueError(f"Required environment variable {env_var} is missing")
+
+            if env_var in os.environ:
+                raw_value = os.environ[env_var]
+                field_values[field_name] = cls._parse_env_value(raw_value, field_type)
+            elif field_info.default is not None:
+                field_values[field_name] = field_info.default
+            else:
+                raise ValueError()
+        return cls(**field_values)
+
+    @classmethod
+    def _parse_env_value(cls, value: str, target_type: type) -> Any:
+        """Converts environment variable string to appropriate type."""
+        if target_type is bool:
+            return value.lower() in ("true", "1", "t", "y", "yes")
+        if target_type is int:
+            return int(value)
+        if target_type is float:
+            return float(value)
+        return value
+
+
 class DictConversionMixin:
     """
     Provides conversion functionality between Pydantic models and dictionaries,
@@ -44,6 +114,26 @@ class DictConversionMixin:
             dump_data["timestamp"] = self.serialize_datetime(self.timestamp, None)
         return dump_data
 
+    def to_json(self, **kwargs) -> str:
+        """
+        Convert model instance to a JSON string.
+        - Accepts the same kwargs as json.dumps (e.g., indent, ensure_ascii)
+        - Default settings make JSON human-readable and UTF-8 safe
+        """
+        return json.dumps(self.to_dict(), ensure_ascii=False, default=lambda o: str(o), **kwargs)
+
+    @classmethod
+    def from_json(cls: type[BaseModelType], json_str: str) -> BaseModelType:
+        """
+        Create model instance from a JSON string.
+        - Parses JSON into a dictionary and delegates to from_dict
+        """
+        try:
+            data = json.loads(json_str)
+        except json.JSONDecodeError as e:
+            raise ValueError(f"Invalid JSON string: {e}") from e
+        return cls.from_dict(data)
+
     @classmethod
     def from_dict(cls: type[BaseModelType], data: dict) -> BaseModelType:
         """
@@ -102,3 +192,11 @@ class AutoDroppingQueue(Queue[T]):
     def get_queue_content_without_pop(self) -> list[T]:
         """Return a copy of the queue's contents without modifying it."""
         return list(self.queue)
+
+    def clear(self) -> None:
+        """Remove all items from the queue.
+
+        This operation is thread-safe.
+        """
+        with self.mutex:
+            self.queue.clear()

memos/mem_scheduler/general_modules/scheduler_logger.py

@@ -69,7 +69,7 @@ class SchedulerLoggerModule(BaseSchedulerModule):
                 and mem_cube_id in self.monitor.activation_memory_monitors[user_id]
             ):
                 activation_monitor = self.monitor.activation_memory_monitors[user_id][mem_cube_id]
-                transformed_act_memory_size = len(activation_monitor.memories)
+                transformed_act_memory_size = len(activation_monitor.obj.memories)
                 logger.info(
                     f'activation_memory_monitors currently has "{transformed_act_memory_size}" transformed memory size'
                 )
@@ -98,6 +98,7 @@ class SchedulerLoggerModule(BaseSchedulerModule):
         )
         return log_message
 
+    # TODO: 日志打出来数量不对
     @log_exceptions(logger=logger)
     def log_working_memory_replacement(
         self,
@@ -125,6 +126,7 @@ class SchedulerLoggerModule(BaseSchedulerModule):
         added_memories = list(new_set - original_set)  # Present in new but not original
 
         # recording messages
+        log_messages = []
         for memory in added_memories:
             normalized_mem = transform_name_to_key(name=memory)
             if normalized_mem not in memory_type_map:
@@ -145,11 +147,13 @@ class SchedulerLoggerModule(BaseSchedulerModule):
                 mem_cube_id=mem_cube_id,
                 mem_cube=mem_cube,
             )
-            log_func_callback([log_message])
-            logger.info(
-                f"{len(added_memories)} {LONG_TERM_MEMORY_TYPE} memorie(s) "
-                f"transformed to {WORKING_MEMORY_TYPE} memories."
-            )
+            log_messages.append(log_message)
+
+        logger.info(
+            f"{len(added_memories)} {LONG_TERM_MEMORY_TYPE} memorie(s) "
+            f"transformed to {WORKING_MEMORY_TYPE} memories."
+        )
+        log_func_callback(log_messages)
 
     @log_exceptions(logger=logger)
     def log_activation_memory_update(
@@ -170,6 +174,7 @@ class SchedulerLoggerModule(BaseSchedulerModule):
         added_memories = list(new_set - original_set)  # Present in new but not original
 
         # recording messages
+        log_messages = []
         for mem in added_memories:
             log_message_a = self.create_autofilled_log_item(
                 log_content=mem,
@@ -194,12 +199,13 @@ class SchedulerLoggerModule(BaseSchedulerModule):
                 mem_cube_id=mem_cube_id,
                 mem_cube=mem_cube,
             )
-            logger.info(
-                f"{len(added_memories)} {ACTIVATION_MEMORY_TYPE} memorie(s) "
-                f"transformed to {PARAMETER_MEMORY_TYPE} memories."
-            )
 
-            log_func_callback([log_message_a, log_message_b])
+            log_messages.extend([log_message_a, log_message_b])
+        logger.info(
+            f"{len(added_memories)} {ACTIVATION_MEMORY_TYPE} memorie(s) "
+            f"transformed to {PARAMETER_MEMORY_TYPE} memories."
+        )
+        log_func_callback(log_messages)
 
     @log_exceptions(logger=logger)
     def log_adding_memory(

memos/mem_scheduler/general_scheduler.py

@@ -27,6 +27,8 @@ class GeneralScheduler(BaseScheduler):
         """Initialize the scheduler with the given configuration."""
         super().__init__(config)
 
+        self.query_key_words_limit = self.config.get("query_key_words_limit", 20)
+
         # register handlers
         handlers = {
             QUERY_LABEL: self._query_message_consumer,
@@ -35,78 +37,6 @@ class GeneralScheduler(BaseScheduler):
         }
         self.dispatcher.register_handlers(handlers)
 
-    # for evaluation
-    def search_for_eval(
-        self, query: str, user_id: UserID | str, top_k: int, scheduler_flag: bool = True
-    ) -> (list[str], bool):
-        self.monitor.register_query_monitor_if_not_exists(
-            user_id=user_id, mem_cube_id=self.current_mem_cube_id
-        )
-
-        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,
-        )
-        query_monitor = self.monitor.query_monitors[user_id][self.current_mem_cube_id]
-        query_monitor.put(item=item)
-        logger.debug(f"Queries in monitor are {query_monitor.get_queries_with_timesort()}.")
-
-        queries = [query]
-
-        # recall
-        mem_cube = self.current_mem_cube
-        text_mem_base = mem_cube.text_mem
-
-        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
-        )
-
-        if not scheduler_flag:
-            return text_working_memory, intent_result["trigger_retrieval"]
-        else:
-            if intent_result["trigger_retrieval"]:
-                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,
-                    )
-                    logger.info(
-                        f"search results for {missing_evidences}: {[one.memory for one in results]}"
-                    )
-                    new_candidates.extend(results)
-                print(
-                    f"missing_evidences: {missing_evidences} and get {len(new_candidates)} new candidate memories."
-                )
-            else:
-                new_candidates = []
-                print(f"intent_result: {intent_result}. not triggered")
-
-            # 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,
-            )
-            new_order_working_memory = new_order_working_memory[:top_k]
-            logger.info(f"size of new_order_working_memory: {len(new_order_working_memory)}")
-
-            return [m.memory for m in new_order_working_memory], intent_result["trigger_retrieval"]
-
     def _query_message_consumer(self, messages: list[ScheduleMessageItem]) -> None:
         """
         Process and handle query trigger messages from the queue.
@@ -140,7 +70,9 @@ class GeneralScheduler(BaseScheduler):
 
                     query = msg.content
                     query_keywords = self.monitor.extract_query_keywords(query=query)
-                    logger.info(f'Extract keywords "{query_keywords}" from query "{query}"')
+                    logger.info(
+                        f'Extracted keywords "{query_keywords}" from query "{query}" for user_id={user_id}'
+                    )
 
                     if len(query_keywords) == 0:
                         stripped_query = query.strip()
@@ -155,21 +87,26 @@ class GeneralScheduler(BaseScheduler):
                             )
                             words = stripped_query  # Default to character count
 
-                        query_keywords = list(set(words[:20]))
+                        query_keywords = list(set(words[: self.query_key_words_limit]))
                         logger.error(
-                            f"Keyword extraction failed for query. Using fallback keywords: {query_keywords[:10]}... (truncated)"
+                            f"Keyword extraction failed for query '{query}' (user_id={user_id}). Using fallback keywords: {query_keywords[:10]}... (truncated)",
+                            exc_info=True,
                         )
 
                     item = QueryMonitorItem(
+                        user_id=user_id,
+                        mem_cube_id=mem_cube_id,
                         query_text=query,
                         keywords=query_keywords,
                         max_keywords=DEFAULT_MAX_QUERY_KEY_WORDS,
                     )
 
-                    self.monitor.query_monitors[user_id][mem_cube_id].put(item=item)
+                    query_db_manager = self.monitor.query_monitors[user_id][mem_cube_id]
+                    query_db_manager.obj.put(item=item)
+                    # Sync with database after adding new item
+                    query_db_manager.sync_with_orm()
                 logger.debug(
-                    f"Queries in monitor are "
-                    f"{self.monitor.query_monitors[user_id][mem_cube_id].get_queries_with_timesort()}."
+                    f"Queries in monitor for user_id={user_id}, mem_cube_id={mem_cube_id}: {query_db_manager.obj.get_queries_with_timesort()}"
                 )
 
                 queries = [msg.content for msg in messages]
@@ -183,7 +120,7 @@ class GeneralScheduler(BaseScheduler):
                     top_k=self.top_k,
                 )
                 logger.info(
-                    f"Processed {queries} and get {len(new_candidates)} new candidate memories."
+                    f"Processed {len(queries)} queries {queries} and retrieved {len(new_candidates)} new candidate memories for user_id={user_id}"
                 )
 
                 # rerank
@@ -194,7 +131,9 @@ class GeneralScheduler(BaseScheduler):
                     original_memory=cur_working_memory,
                     new_memory=new_candidates,
                 )
-                logger.info(f"size of new_order_working_memory: {len(new_order_working_memory)}")
+                logger.info(
+                    f"Final working memory size: {len(new_order_working_memory)} memories for user_id={user_id}"
+                )
 
                 # update activation memories
                 logger.info(
@@ -293,10 +232,17 @@ class GeneralScheduler(BaseScheduler):
 
         text_mem_base = mem_cube.text_mem
         if not isinstance(text_mem_base, TreeTextMemory):
-            logger.error("Not implemented!", exc_info=True)
+            logger.error(
+                f"Not implemented! Expected TreeTextMemory but got {type(text_mem_base).__name__} "
+                f"for mem_cube_id={mem_cube_id}, user_id={user_id}. "
+                f"text_mem_base value: {text_mem_base}",
+                exc_info=True,
+            )
             return
 
-        logger.info(f"Processing {len(queries)} queries.")
+        logger.info(
+            f"Processing {len(queries)} queries for user_id={user_id}, mem_cube_id={mem_cube_id}"
+        )
 
         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]
@@ -312,16 +258,20 @@ class GeneralScheduler(BaseScheduler):
             time_trigger_flag = True
 
         if (not intent_result["trigger_retrieval"]) and (not time_trigger_flag):
-            logger.info(f"Query schedule not triggered. Intent_result: {intent_result}")
+            logger.info(
+                f"Query schedule not triggered for user_id={user_id}, mem_cube_id={mem_cube_id}. 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")
+            logger.info(
+                f"Query schedule forced to trigger due to time ticker for user_id={user_id}, mem_cube_id={mem_cube_id}"
+            )
             intent_result["trigger_retrieval"] = True
             intent_result["missing_evidences"] = queries
         else:
             logger.info(
-                f'Query schedule triggered for user "{user_id}" and mem_cube "{mem_cube_id}".'
-                f" Missing evidences: {intent_result['missing_evidences']}"
+                f"Query schedule triggered for user_id={user_id}, mem_cube_id={mem_cube_id}. "
+                f"Missing evidences: {intent_result['missing_evidences']}"
             )
 
         missing_evidences = intent_result["missing_evidences"]
@@ -329,7 +279,9 @@ class GeneralScheduler(BaseScheduler):
         k_per_evidence = max(1, top_k // max(1, num_evidence))
         new_candidates = []
         for item in missing_evidences:
-            logger.info(f"missing_evidences: {item}")
+            logger.info(
+                f"Searching for missing evidence: '{item}' with top_k={k_per_evidence} for user_id={user_id}"
+            )
             info = {
                 "user_id": user_id,
                 "session_id": "",
@@ -343,7 +295,7 @@ class GeneralScheduler(BaseScheduler):
                 info=info,
             )
             logger.info(
-                f"search results for {missing_evidences}: {[one.memory for one in results]}"
+                f"Search results for missing evidence '{item}': {[one.memory for one in results]}"
             )
             new_candidates.extend(results)
         return cur_working_memory, new_candidates

memos/mem_scheduler/memory_manage_modules/__init__.py

@@ -0,0 +1,5 @@
+from .memory_filter import MemoryFilter
+from .retriever import SchedulerRetriever
+
+
+__all__ = ["MemoryFilter", "SchedulerRetriever"]