Merge branch 'feat/demo-product' into feat/demo-product

Author: CaralHsi

src/memos/mem_scheduler/base_scheduler.py

@@ -17,7 +17,6 @@
 from memos.mem_scheduler.modules.retriever import SchedulerRetriever
 from memos.mem_scheduler.modules.schemas import (
     ACTIVATION_MEMORY_TYPE,
-    ACTIVATION_MEMORY_VLLM_BACKEND,
     ADD_LABEL,
     DEFAULT_ACT_MEM_DUMP_PATH,
     DEFAULT_CONSUME_INTERVAL_SECONDS,
@@ -33,7 +32,8 @@
     TreeTextMemory_SEARCH_METHOD,
 )
 from memos.mem_scheduler.utils import normalize_name
-from memos.memories.activation.kv import KVCacheItem, KVCacheMemory
+from memos.memories.activation.kv import KVCacheMemory
+from memos.memories.activation.vllmkv import VLLMKVCacheItem, VLLMKVCacheMemory
 from memos.memories.textual.tree import TextualMemoryItem, TreeTextMemory
 from memos.templates.mem_scheduler_prompts import MEMORY_ASSEMBLY_TEMPLATE
 
@@ -88,7 +88,6 @@ def __init__(self, config: BaseSchedulerConfig):
         self.auth_config_path: str | Path | None = self.config.get("auth_config_path", None)
         self.auth_config = None
         self.rabbitmq_config = None
-        self.act_mem_backend = ACTIVATION_MEMORY_VLLM_BACKEND
 
     def initialize_modules(self, chat_llm: BaseLLM, process_llm: BaseLLM | None = None):
         if process_llm is None:
@@ -186,8 +185,13 @@ def update_activation_memory(
             logger.error("Not Implemented.")
 
         try:
-            assert isinstance(mem_cube.act_mem, KVCacheMemory)
-            act_mem: KVCacheMemory = mem_cube.act_mem
+            if isinstance(mem_cube.act_mem, VLLMKVCacheMemory):
+                act_mem: VLLMKVCacheMemory = mem_cube.act_mem
+            elif isinstance(mem_cube.act_mem, KVCacheMemory):
+                act_mem: KVCacheMemory = mem_cube.act_mem
+            else:
+                logger.error("Not Implemented.")
+                return
 
             text_memory = MEMORY_ASSEMBLY_TEMPLATE.format(
                 memory_text="".join(
@@ -200,14 +204,17 @@ def update_activation_memory(
             )
 
             # huggingface or vllm kv cache
-            original_cache_items: list[KVCacheItem] = act_mem.get_all()
-            pre_cache_item: KVCacheItem = original_cache_items[-1]
-            original_text_memories = pre_cache_item.records.text_memories
-            act_mem.delete_all()
-            cache_item: KVCacheItem = act_mem.extract(text_memory)
+            original_cache_items: list[VLLMKVCacheItem] = act_mem.get_all()
+            original_text_memories = []
+            if len(original_cache_items) > 0:
+                pre_cache_item: VLLMKVCacheItem = original_cache_items[-1]
+                original_text_memories = pre_cache_item.records.text_memories
+                act_mem.delete_all()
+
+            cache_item = act_mem.extract(text_memory)
             cache_item.records.text_memories = new_text_memories
 
-            act_mem.add(cache_item)
+            act_mem.add([cache_item])
             act_mem.dump(self.act_mem_dump_path)
 
             self.log_activation_memory_update(
@@ -219,10 +226,11 @@ def update_activation_memory(
             )
 
         except Exception as e:
-            logger.warning(f"MOS-based activation memory update failed: {e}")
+            logger.warning(f"MOS-based activation memory update failed: {e}", exc_info=True)
 
     def update_activation_memory_periodically(
         self,
+        interval_seconds: int,
         user_id: str,
         mem_cube_id: str,
         mem_cube: GeneralMemCube,
@@ -231,8 +239,10 @@ def update_activation_memory_periodically(
 
         if self.monitor.timed_trigger(
             last_time=self.monitor._last_activation_mem_update_time,
-            interval_seconds=self.monitor.act_mem_update_interval,
+            interval_seconds=interval_seconds,
         ):
+            logger.info(f"Updating activation memory for user {user_id} and mem_cube {mem_cube_id}")
+
             self.monitor.update_memory_monitors(
                 user_id=user_id, mem_cube_id=mem_cube_id, mem_cube=mem_cube
             )
@@ -241,10 +251,24 @@ def update_activation_memory_periodically(
                 m.memory_text for m in self.monitor.activation_memory_monitors[user_id][mem_cube_id]
             ]
 
+            logger.info(
+                f"Collected {len(new_activation_memories)} new memory entries for processing"
+            )
+
             self.update_activation_memory(new_memories=new_activation_memories, mem_cube=mem_cube)
 
             self.monitor._last_activation_mem_update_time = datetime.now()
 
+            logger.debug(
+                f"Activation memory update completed at {self.monitor._last_activation_mem_update_time}"
+            )
+        else:
+            logger.info(
+                f"Skipping update - {interval_seconds} second interval not yet reached. "
+                f"Last update time is {self.monitor._last_activation_mem_update_time} and now is"
+                f"{datetime.now()}"
+            )
+
     def submit_messages(self, messages: ScheduleMessageItem | list[ScheduleMessageItem]):
         """Submit multiple messages to the message queue."""
         if isinstance(messages, ScheduleMessageItem):

src/memos/mem_scheduler/general_scheduler.py

@@ -82,8 +82,10 @@ def _answer_message_consumer(self, messages: list[ScheduleMessageItem]) -> None:
                 # for status update
                 self._set_current_context_from_message(msg=messages[0])
 
+                # update acivation memories
                 if self.enable_act_memory_update:
                     self.update_activation_memory_periodically(
+                        interval_seconds=self.monitor.act_mem_update_interval,
                         user_id=user_id,
                         mem_cube_id=mem_cube_id,
                         mem_cube=messages[0].mem_cube,
@@ -105,13 +107,6 @@ def _add_message_consumer(self, messages: list[ScheduleMessageItem]) -> None:
                 # for status update
                 self._set_current_context_from_message(msg=messages[0])
 
-                if self.enable_act_memory_update:
-                    self.update_activation_memory_periodically(
-                        user_id=user_id,
-                        mem_cube_id=mem_cube_id,
-                        mem_cube=messages[0].mem_cube,
-                    )
-
                 # submit logs
                 for msg in messages:
                     user_inputs = json.loads(msg.content)
@@ -122,6 +117,15 @@ def _add_message_consumer(self, messages: list[ScheduleMessageItem]) -> None:
                         mem_cube=msg.mem_cube,
                     )
 
+                # update acivation memories
+                if self.enable_act_memory_update:
+                    self.update_activation_memory_periodically(
+                        interval_seconds=self.monitor.act_mem_update_interval,
+                        user_id=user_id,
+                        mem_cube_id=mem_cube_id,
+                        mem_cube=messages[0].mem_cube,
+                    )
+
     def process_session_turn(
         self,
         queries: str | list[str],
@@ -146,7 +150,7 @@ def process_session_turn(
 
         text_mem_base = mem_cube.text_mem
         if not isinstance(text_mem_base, TreeTextMemory):
-            logger.error("Not implemented!")
+            logger.error("Not implemented!", exc_info=True)
             return
 
         working_memory: list[TextualMemoryItem] = text_mem_base.get_working_memory()

src/memos/mem_scheduler/modules/retriever.py

@@ -59,7 +59,7 @@ def search(
             else:
                 raise NotImplementedError(str(type(text_mem_base)))
         except Exception as e:
-            logger.error(f"Fail to search. The exeption is {e}.")
+            logger.error(f"Fail to search. The exeption is {e}.", exc_info=True)
             results = []
         return results
 
@@ -142,6 +142,7 @@ def filter_too_short_memories(
         Returns:
             List of memories with short entries removed
         """
+        # TODO: Chinese
         if not text_memories:
             logging.debug("Received empty memories list in short memory filter")
             return []
@@ -209,7 +210,7 @@ def replace_working_memory(
                 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}.")
+                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 = []

src/memos/mem_scheduler/modules/schemas.py

@@ -47,9 +47,6 @@
 MONITOR_WORKING_MEMORY_TYPE = "MonitorWorkingMemoryType"
 MONITOR_ACTIVATION_MEMORY_TYPE = "MonitorActivationMemoryType"
 
-# Backends
-ACTIVATION_MEMORY_VLLM_BACKEND = "VLLMForActivationMemory"
-ACTIVATION_MEMORY_HF_BACKEND = "HugggingaceForActivationMemory"
 
 # new types
 UserID = NewType("UserID", str)