fix bugs & new feat: fix bugs in mem_scheduler examples, and remove initialize working memories (logically uneccessary). change the function parameters of search as the function input info as an addition

Author: tangg555

examples/mem_scheduler/memos_w_scheduler.py

@@ -124,8 +124,7 @@ def run_with_scheduler_init():
         query = item["question"]
         print(f"Query:\n {query}\n")
         response = mos.chat(query=query, user_id=user_id)
-        print(f"Answer:\n {response}")
-        print("===== Chat End =====")
+        print(f"Answer:\n {response}\n")
 
     show_web_logs(mem_scheduler=mos.mem_scheduler)
 

examples/mem_scheduler/try_schedule_modules.py

@@ -151,6 +151,12 @@ def show_web_logs(mem_scheduler: GeneralScheduler):
         mos_config.mem_reader.config.llm.config.api_base = auth_config.openai.base_url
 
         mem_cube_config.text_mem.config.graph_db.config.uri = auth_config.graph_db.uri
+        mem_cube_config.text_mem.config.graph_db.config.user = auth_config.graph_db.user
+        mem_cube_config.text_mem.config.graph_db.config.password = auth_config.graph_db.password
+        mem_cube_config.text_mem.config.graph_db.config.db_name = auth_config.graph_db.db_name
+        mem_cube_config.text_mem.config.graph_db.config.auto_create = (
+            auth_config.graph_db.auto_create
+        )
 
     # Initialization
     mos = MOSForTestScheduler(mos_config)

src/memos/mem_scheduler/base_scheduler.py

@@ -132,7 +132,7 @@ def _set_current_context_from_message(self, msg: ScheduleMessageItem) -> None:
             self.current_mem_cube_id = msg.mem_cube_id
             self.current_mem_cube = msg.mem_cube
 
-    def transform_memories_to_monitors(
+    def transform_working_memories_to_monitors(
         self, query_keywords, memories: list[TextualMemoryItem]
     ) -> list[MemoryMonitorItem]:
         """
@@ -193,9 +193,8 @@ def replace_working_memory(
             text_mem_base: TreeTextMemory = text_mem_base
 
             # process rerank memories with llm
-            query_history = self.monitor.query_monitors[user_id][
-                mem_cube_id
-            ].get_queries_with_timesort()
+            query_monitor = self.monitor.query_monitors[user_id][mem_cube_id]
+            query_history = query_monitor.get_queries_with_timesort()
             memories_with_new_order, rerank_success_flag = (
                 self.retriever.process_and_rerank_memories(
                     queries=query_history,
@@ -206,13 +205,11 @@ def replace_working_memory(
             )
 
             # update working memory monitors
-            query_keywords = self.monitor.query_monitors[user_id][
-                mem_cube_id
-            ].get_keywords_collections()
+            query_keywords = query_monitor.get_keywords_collections()
             logger.debug(
                 f"Processing {len(memories_with_new_order)} memories with {len(query_keywords)} query keywords"
             )
-            new_working_memory_monitors = self.transform_memories_to_monitors(
+            new_working_memory_monitors = self.transform_working_memories_to_monitors(
                 query_keywords=query_keywords,
                 memories=memories_with_new_order,
             )
@@ -252,25 +249,6 @@ def replace_working_memory(
 
         return memories_with_new_order
 
-    def initialize_working_memory_monitors(
-        self,
-        user_id: UserID | str,
-        mem_cube_id: MemCubeID | str,
-        mem_cube: GeneralMemCube,
-    ):
-        text_mem_base: TreeTextMemory = mem_cube.text_mem
-        working_memories = text_mem_base.get_working_memory()
-
-        working_memory_monitors = self.transform_memories_to_monitors(
-            memories=working_memories,
-        )
-        self.monitor.update_working_memory_monitors(
-            new_working_memory_monitors=working_memory_monitors,
-            user_id=user_id,
-            mem_cube_id=mem_cube_id,
-            mem_cube=mem_cube,
-        )
-
     def update_activation_memory(
         self,
         new_memories: list[str | TextualMemoryItem],
@@ -374,13 +352,9 @@ def update_activation_memory_periodically(
                     or len(self.monitor.working_memory_monitors[user_id][mem_cube_id].memories) == 0
                 ):
                     logger.warning(
-                        "No memories found in working_memory_monitors, initializing from current working_memories"
-                    )
-                    self.initialize_working_memory_monitors(
-                        user_id=user_id,
-                        mem_cube_id=mem_cube_id,
-                        mem_cube=mem_cube,
+                        "No memories found in working_memory_monitors, activation memory update is skipped"
                     )
+                    return
 
                 self.monitor.update_activation_memory_monitors(
                     user_id=user_id, mem_cube_id=mem_cube_id, mem_cube=mem_cube

src/memos/mem_scheduler/general_scheduler.py

@@ -326,8 +326,17 @@ def process_session_turn(
         new_candidates = []
         for item in missing_evidences:
             logger.info(f"missing_evidences: {item}")
+            info = {
+                "user_id": user_id,
+                "session_id": "",
+            }
+
             results: list[TextualMemoryItem] = self.retriever.search(
-                query=item, mem_cube=mem_cube, top_k=k_per_evidence, method=self.search_method
+                query=item,
+                mem_cube=mem_cube,
+                top_k=k_per_evidence,
+                method=self.search_method,
+                info=info,
             )
             logger.info(
                 f"search results for {missing_evidences}: {[one.memory for one in results]}"

src/memos/mem_scheduler/modules/retriever.py

@@ -33,7 +33,12 @@ def __init__(self, process_llm: BaseLLM, config: BaseSchedulerConfig):
         self.process_llm = process_llm
 
     def search(
-        self, query: str, mem_cube: GeneralMemCube, top_k: int, method=TreeTextMemory_SEARCH_METHOD
+        self,
+        query: str,
+        mem_cube: GeneralMemCube,
+        top_k: int,
+        method: str = TreeTextMemory_SEARCH_METHOD,
+        info: dict | None = None,
     ) -> list[TextualMemoryItem]:
         """Search in text memory with the given query.
 
@@ -49,12 +54,19 @@ def search(
         try:
             if method in [TreeTextMemory_SEARCH_METHOD, TreeTextMemory_FINE_SEARCH_METHOD]:
                 assert isinstance(text_mem_base, TreeTextMemory)
+                if info is None:
+                    logger.warning(
+                        "Please input 'info' when use tree.search so that "
+                        "the database would store the consume history."
+                    )
+                    info = {"user_id": "", "session_id": ""}
+
                 mode = "fast" if method == TreeTextMemory_SEARCH_METHOD else "fine"
                 results_long_term = text_mem_base.search(
-                    query=query, top_k=top_k, memory_type="LongTermMemory", mode=mode
+                    query=query, top_k=top_k, memory_type="LongTermMemory", mode=mode, info=info
                 )
                 results_user = text_mem_base.search(
-                    query=query, top_k=top_k, memory_type="UserMemory", mode=mode
+                    query=query, top_k=top_k, memory_type="UserMemory", mode=mode, info=info
                 )
                 results = results_long_term + results_user
             else:

src/memos/mem_scheduler/mos_for_test_scheduler.py

@@ -81,7 +81,13 @@ def chat(self, query: str, user_id: str | None = None) -> str:
 
                 # from mem_cube
                 memories = mem_cube.text_mem.search(
-                    query, top_k=self.config.top_k - topk_for_scheduler
+                    query,
+                    top_k=self.config.top_k - topk_for_scheduler,
+                    info={
+                        "user_id": target_user_id,
+                        "session_id": self.session_id,
+                        "chat_history": chat_history.chat_history,
+                    },
                 )
                 text_memories = [m.memory for m in memories]
                 print(f"Search results with new working memories: {text_memories}")

src/memos/memories/textual/base.py

@@ -36,7 +36,7 @@ def update(self, memory_id: str, new_memory: TextualMemoryItem | dict[str, Any])
         """Update a memory by memory_id."""
 
     @abstractmethod
-    def search(self, query: str, top_k: int, info=None) -> list[TextualMemoryItem]:
+    def search(self, query: str, top_k: int, info=None, **kwargs) -> list[TextualMemoryItem]:
         """Search for memories based on a query.
         Args:
             query (str): The query to search for.

src/memos/memories/textual/general.py

@@ -114,7 +114,7 @@ def update(self, memory_id: str, new_memory: TextualMemoryItem | dict[str, Any])
 
         self.vector_db.update(memory_id, vec_db_item)
 
-    def search(self, query: str, top_k: int) -> list[TextualMemoryItem]:
+    def search(self, query: str, top_k: int, info=None, **kwargs) -> list[TextualMemoryItem]:
         """Search for memories based on a query.
         Args:
             query (str): The query to search for.