merge dev

Author: yuan.wang

examples/mem_scheduler/show_redis_status.py

@@ -0,0 +1,67 @@
+import time
+
+from memos.api.routers.server_router import mem_scheduler
+from memos.mem_scheduler.task_schedule_modules.redis_queue import SchedulerRedisQueue
+
+
+queue = mem_scheduler.memos_message_queue.memos_message_queue
+
+
+def fetch_status(queue: SchedulerRedisQueue) -> dict[str, dict[str, int]]:
+    """Fetch and print per-user Redis queue status using built-in API.
+
+    Returns a dict mapping user_id -> {"pending": int, "remaining": int}.
+    """
+    # This method will also print a summary and per-user counts.
+    return queue.show_task_status()
+
+
+def print_diff(prev: dict[str, dict[str, int]], curr: dict[str, dict[str, int]]) -> None:
+    """Print aggregated totals and per-user changes compared to previous snapshot."""
+    ts = time.strftime("%Y-%m-%d %H:%M:%S")
+    tot_p_prev = sum(v.get("pending", 0) for v in prev.values()) if prev else 0
+    tot_r_prev = sum(v.get("remaining", 0) for v in prev.values()) if prev else 0
+    tot_p_curr = sum(v.get("pending", 0) for v in curr.values())
+    tot_r_curr = sum(v.get("remaining", 0) for v in curr.values())
+
+    dp_tot = tot_p_curr - tot_p_prev
+    dr_tot = tot_r_curr - tot_r_prev
+
+    print(f"[{ts}] Total pending={tot_p_curr} ({dp_tot:+d}), remaining={tot_r_curr} ({dr_tot:+d})")
+
+    # Print per-user deltas (current counts are already printed by show_task_status)
+    all_uids = sorted(set(prev.keys()) | set(curr.keys()))
+    for uid in all_uids:
+        p_prev = prev.get(uid, {}).get("pending", 0)
+        r_prev = prev.get(uid, {}).get("remaining", 0)
+        p_curr = curr.get(uid, {}).get("pending", 0)
+        r_curr = curr.get(uid, {}).get("remaining", 0)
+        dp = p_curr - p_prev
+        dr = r_curr - r_prev
+        # Only print when there is any change to reduce noise
+        if dp != 0 or dr != 0:
+            print(f"  Δ {uid}: pending={dp:+d}, remaining={dr:+d}")
+
+
+# Note: queue.show_task_status() handles printing per-user counts internally.
+
+
+def main(interval_sec: float = 5.0) -> None:
+    prev: dict[str, dict[str, int]] = {}
+    while True:
+        try:
+            curr = fetch_status(queue)
+            print_diff(prev, curr)
+            print(f"stream_cache ({len(queue._stream_keys_cache)}): {queue._stream_keys_cache}")
+            prev = curr
+            time.sleep(interval_sec)
+        except KeyboardInterrupt:
+            print("Stopped.")
+            break
+        except Exception as e:
+            print(f"Error while fetching status: {e}")
+            time.sleep(interval_sec)
+
+
+if __name__ == "__main__":
+    main()

src/memos/llms/openai.py

@@ -28,7 +28,12 @@ def __init__(self, config: OpenAILLMConfig):
         )
         logger.info("OpenAI LLM instance initialized")
 
-    @timed_with_status(log_prefix="OpenAI LLM", log_args=["model_name_or_path"])
+    @timed_with_status(
+        log_prefix="OpenAI LLM",
+        log_extra_args=lambda self, messages, **kwargs: {
+            "model_name_or_path": kwargs.get("model_name_or_path", self.config.model_name_or_path)
+        },
+    )
     def generate(self, messages: MessageList, **kwargs) -> str:
         """Generate a response from OpenAI LLM, optionally overriding generation params."""
         response = self.client.chat.completions.create(
@@ -55,7 +60,12 @@ def generate(self, messages: MessageList, **kwargs) -> str:
             return reasoning_content + response_content
         return response_content
 
-    @timed_with_status(log_prefix="OpenAI LLM", log_args=["model_name_or_path"])
+    @timed_with_status(
+        log_prefix="OpenAI LLM",
+        log_extra_args=lambda self, messages, **kwargs: {
+            "model_name_or_path": self.config.model_name_or_path
+        },
+    )
     def generate_stream(self, messages: MessageList, **kwargs) -> Generator[str, None, None]:
         """Stream response from OpenAI LLM with optional reasoning support."""
         if kwargs.get("tools"):

src/memos/mem_reader/simple_struct.py

@@ -453,7 +453,7 @@ def get_memory(
     @staticmethod
     def _parse_hallucination_filter_response(text: str) -> tuple[bool, dict[int, dict]]:
         """Parse index-keyed JSON from hallucination filter response.
-        Expected shape: { "0": {"delete": bool, "rewritten": str, "reason": str}, ... }
+        Expected shape: { "0": {"need_rewrite": bool, "rewritten_suffix": str, "reason": str}, ... }
         Returns (success, parsed_dict) with int keys.
         """
         try:
@@ -476,27 +476,33 @@ def _parse_hallucination_filter_response(text: str) -> tuple[bool, dict[int, dic
                     continue
             if not isinstance(v, dict):
                 continue
-            delete_flag = v.get("delete")
-            rewritten = v.get("rewritten", "")
+            need_rewrite = v.get("need_rewrite")
+            rewritten_suffix = v.get("rewritten_suffix", "")
             reason = v.get("reason", "")
             if (
-                isinstance(delete_flag, bool)
-                and isinstance(rewritten, str)
+                isinstance(need_rewrite, bool)
+                and isinstance(rewritten_suffix, str)
                 and isinstance(reason, str)
             ):
-                result[idx] = {"delete": delete_flag, "rewritten": rewritten, "reason": reason}
+                result[idx] = {
+                    "need_rewrite": need_rewrite,
+                    "rewritten_suffix": rewritten_suffix,
+                    "reason": reason,
+                }
 
         return (len(result) > 0), result
 
     def filter_hallucination_in_memories(
-        self, user_messages: list[str], memory_list: list[TextualMemoryItem]
+        self, messages: list[dict], memory_list: list[TextualMemoryItem]
     ) -> list[TextualMemoryItem]:
-        flat_memories = [one.memory for one in memory_list]
+        # Build input objects with memory text and metadata (timestamps, sources, etc.)
         template = PROMPT_MAPPING["hallucination_filter"]
         prompt_args = {
-            "user_messages_inline": "\n".join([f"- {memory}" for memory in user_messages]),
+            "messages_inline": "\n".join(
+                [f"- [{message['role']}]: {message['content']}" for message in messages]
+            ),
             "memories_inline": json.dumps(
-                {str(i): memory for i, memory in enumerate(flat_memories)},
+                {idx: mem.memory for idx, mem in enumerate(memory_list)},
                 ensure_ascii=False,
                 indent=2,
             ),
@@ -511,40 +517,31 @@ def filter_hallucination_in_memories(
                 f"[filter_hallucination_in_memories] Hallucination filter parsed successfully: {success}"
             )
             if success:
+                new_mem_list = []
                 logger.info(f"Hallucination filter result: {parsed}")
-                total = len(memory_list)
-                keep_flags = [True] * total
+                assert len(parsed) == len(memory_list)
                 for mem_idx, content in parsed.items():
-                    # Validate index bounds
-                    if not isinstance(mem_idx, int) or mem_idx < 0 or mem_idx >= total:
-                        logger.warning(
-                            f"[filter_hallucination_in_memories] Ignoring out-of-range index: {mem_idx}"
-                        )
-                        continue
-
-                    delete_flag = content.get("delete", False)
-                    rewritten = content.get("rewritten", None)
+                    need_rewrite = content.get("need_rewrite", False)
+                    rewritten_suffix = content.get("rewritten_suffix", "")
                     reason = content.get("reason", "")
 
-                    logger.info(
-                        f"[filter_hallucination_in_memories] index={mem_idx}, delete={delete_flag}, rewritten='{(rewritten or '')[:100]}', reason='{reason[:120]}'"
-                    )
+                    # Append a new memory item instead of replacing the original
+                    if (
+                        need_rewrite
+                        and isinstance(rewritten_suffix, str)
+                        and len(rewritten_suffix.strip()) > 0
+                    ):
+                        original_text = memory_list[mem_idx].memory
+
+                        logger.info(
+                            f"[filter_hallucination_in_memories] index={mem_idx}, need_rewrite={need_rewrite}, rewritten_suffix='{rewritten_suffix}', reason='{reason}', original memory='{original_text}', action='append_suffix'"
+                        )
 
-                    if delete_flag is True and rewritten is not None:
-                        # Mark for deletion
-                        keep_flags[mem_idx] = False
+                        # Append only the suffix to the original memory text
+                        memory_list[mem_idx].memory = original_text + rewritten_suffix
+                        new_mem_list.append(memory_list[mem_idx])
                     else:
-                        # Apply rewrite if provided (safe-by-default: keep item when not mentioned or delete=False)
-                        try:
-                            if isinstance(rewritten, str):
-                                memory_list[mem_idx].memory = rewritten
-                        except Exception as e:
-                            logger.warning(
-                                f"[filter_hallucination_in_memories] Failed to apply rewrite for index {mem_idx}: {e}"
-                            )
-
-                # Build result, preserving original order; keep items not mentioned by LLM by default
-                new_mem_list = [memory_list[i] for i in range(total) if keep_flags[i]]
+                        new_mem_list.append(memory_list[mem_idx])
                 return new_mem_list
             else:
                 logger.warning("Hallucination filter parsing failed or returned empty result.")
@@ -602,11 +599,8 @@ def _read_memory(
             # Build inputs
             new_memory_list = []
             for unit_messages, unit_memory_list in zip(messages, memory_list, strict=False):
-                unit_user_messages = [
-                    msg["content"] for msg in unit_messages if msg["role"] == "user"
-                ]
                 unit_memory_list = self.filter_hallucination_in_memories(
-                    user_messages=unit_user_messages, memory_list=unit_memory_list
+                    messages=unit_messages, memory_list=unit_memory_list
                 )
                 new_memory_list.append(unit_memory_list)
             memory_list = new_memory_list

src/memos/mem_scheduler/general_scheduler.py

@@ -126,7 +126,10 @@ def long_memory_update_process(
             top_k=self.top_k,
         )
         logger.info(
-            f"[long_memory_update_process] Processed {len(queries)} queries {queries} and retrieved {len(new_candidates)} new candidate memories for user_id={user_id}"
+            # Build the candidate preview string outside the f-string to avoid backslashes in expression
+            f"[long_memory_update_process] Processed {len(queries)} queries {queries} and retrieved {len(new_candidates)} "
+            f"new candidate memories for user_id={user_id}: "
+            + ("\n- " + "\n- ".join([f"{one.id}: {one.memory}" for one in new_candidates]))
         )
 
         # rerank
@@ -141,10 +144,14 @@ def long_memory_update_process(
             f"[long_memory_update_process] Final working memory size: {len(new_order_working_memory)} memories for user_id={user_id}"
         )
 
-        old_memory_texts = [mem.memory for mem in cur_working_memory]
-        new_memory_texts = [mem.memory for mem in new_order_working_memory]
+        old_memory_texts = "\n- " + "\n- ".join(
+            [f"{one.id}: {one.memory}" for one in cur_working_memory]
+        )
+        new_memory_texts = "\n- " + "\n- ".join(
+            [f"{one.id}: {one.memory}" for one in new_order_working_memory]
+        )
 
-        logger.debug(
+        logger.info(
             f"[long_memory_update_process] For user_id='{user_id}', mem_cube_id='{mem_cube_id}': "
             f"Scheduler replaced working memory based on query history {queries}. "
             f"Old working memory ({len(old_memory_texts)} items): {old_memory_texts}. "
@@ -1413,20 +1420,21 @@ def process_session_turn(
             logger.info(
                 f"[process_session_turn] Searching for missing evidence: '{item}' with top_k={k_per_evidence} for user_id={user_id}"
             )
-            info = {
-                "user_id": user_id,
-                "session_id": "",
-            }
 
+            search_args = {}
             results: list[TextualMemoryItem] = self.retriever.search(
                 query=item,
+                user_id=user_id,
+                mem_cube_id=mem_cube_id,
                 mem_cube=mem_cube,
                 top_k=k_per_evidence,
                 method=self.search_method,
-                info=info,
+                search_args=search_args,
             )
+
             logger.info(
-                f"[process_session_turn] Search results for missing evidence '{item}': {[one.memory for one in results]}"
+                f"[process_session_turn] Search results for missing evidence '{item}': "
+                + ("\n- " + "\n- ".join([f"{one.id}: {one.memory}" for one in results]))
             )
             new_candidates.extend(results)
         return cur_working_memory, new_candidates

src/memos/mem_scheduler/memory_manage_modules/retriever.py

@@ -22,7 +22,11 @@
 from memos.mem_scheduler.utils.misc_utils import extract_json_obj, extract_list_items_in_answer
 from memos.memories.textual.item import TextualMemoryMetadata
 from memos.memories.textual.tree import TextualMemoryItem, TreeTextMemory
-from memos.types.general_types import FINE_STRATEGY, FineStrategy
+from memos.types.general_types import (
+    FINE_STRATEGY,
+    FineStrategy,
+    SearchMode,
+)
 
 # Extract JSON response
 from .memory_filter import MemoryFilter
@@ -237,10 +241,12 @@ def recall_for_missing_memories(
     def search(
         self,
         query: str,
+        user_id: str,
+        mem_cube_id: str,
         mem_cube: GeneralMemCube,
         top_k: int,
         method: str = TreeTextMemory_SEARCH_METHOD,
-        info: dict | None = None,
+        search_args: dict | None = None,
     ) -> list[TextualMemoryItem]:
         """Search in text memory with the given query.
 
@@ -253,22 +259,67 @@ def search(
             Search results or None if not implemented
         """
         text_mem_base = mem_cube.text_mem
+        # Normalize default for mutable argument
+        search_args = search_args or {}
         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": ""}
+                session_id = search_args.get("session_id", "default_session")
+                target_session_id = session_id
+                search_priority = (
+                    {"session_id": target_session_id} if "session_id" in search_args else None
+                )
+                search_filter = search_args.get("filter")
+                search_source = search_args.get("source")
+                plugin = bool(search_source is not None and search_source == "plugin")
+                user_name = search_args.get("user_name", mem_cube_id)
+                internet_search = search_args.get("internet_search", False)
+                chat_history = search_args.get("chat_history")
+                search_tool_memory = search_args.get("search_tool_memory", False)
+                tool_mem_top_k = search_args.get("tool_mem_top_k", 6)
+                playground_search_goal_parser = search_args.get(
+                    "playground_search_goal_parser", False
+                )
 
-                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, info=info
+                info = search_args.get(
+                    "info",
+                    {
+                        "user_id": user_id,
+                        "session_id": target_session_id,
+                        "chat_history": chat_history,
+                    },
                 )
-                results_user = text_mem_base.search(
-                    query=query, top_k=top_k, memory_type="UserMemory", mode=mode, info=info
+
+                results_long_term = mem_cube.text_mem.search(
+                    query=query,
+                    user_name=user_name,
+                    top_k=top_k,
+                    mode=SearchMode.FAST,
+                    manual_close_internet=not internet_search,
+                    memory_type="LongTermMemory",
+                    search_filter=search_filter,
+                    search_priority=search_priority,
+                    info=info,
+                    plugin=plugin,
+                    search_tool_memory=search_tool_memory,
+                    tool_mem_top_k=tool_mem_top_k,
+                    playground_search_goal_parser=playground_search_goal_parser,
+                )
+
+                results_user = mem_cube.text_mem.search(
+                    query=query,
+                    user_name=user_name,
+                    top_k=top_k,
+                    mode=SearchMode.FAST,
+                    manual_close_internet=not internet_search,
+                    memory_type="UserMemory",
+                    search_filter=search_filter,
+                    search_priority=search_priority,
+                    info=info,
+                    plugin=plugin,
+                    search_tool_memory=search_tool_memory,
+                    tool_mem_top_k=tool_mem_top_k,
+                    playground_search_goal_parser=playground_search_goal_parser,
                 )
                 results = results_long_term + results_user
             else: