feat: add strategy for fine search

Author: tangg555

examples/mem_scheduler/api_w_scheduler.py

@@ -15,6 +15,7 @@
 print(f"Queue maxsize: {getattr(mem_scheduler.memos_message_queue, 'maxsize', 'N/A')}")
 print("=====================================\n")
 
+mem_scheduler.memos_message_queue.debug_mode_on()
 queue = mem_scheduler.memos_message_queue
 queue.clear()
 

src/memos/mem_scheduler/base_scheduler.py

@@ -206,6 +206,9 @@ def initialize_modules(
 
         # start queue monitor if enabled and a bot is set later
 
+    def debug_mode_on(self):
+        self.memos_message_queue.debug_mode_on()
+
     def _cleanup_on_init_failure(self):
         """Clean up resources if initialization fails."""
         try:

src/memos/mem_scheduler/memory_manage_modules/retriever.py

@@ -11,6 +11,8 @@
 from memos.mem_scheduler.schemas.general_schemas import (
     DEFAULT_SCHEDULER_RETRIEVER_BATCH_SIZE,
     DEFAULT_SCHEDULER_RETRIEVER_RETRIES,
+    FINE_STRATEGY,
+    FineStrategy,
     TreeTextMemory_FINE_SEARCH_METHOD,
     TreeTextMemory_SEARCH_METHOD,
 )
@@ -93,9 +95,15 @@ def _build_enhancement_prompt(self, query_history: list[str], batch_texts: list[
                 if len(query_history) > 1
                 else query_history[0]
             )
-        text_memories = "\n".join([f"- {mem}" for i, mem in enumerate(batch_texts)])
+        # Include numbering for rewrite mode to help LLM reference original memory IDs
+        if FINE_STRATEGY == FineStrategy.REWRITE:
+            text_memories = "\n".join([f"- [{i}] {mem}" for i, mem in enumerate(batch_texts)])
+            prompt_name = "memory_rewrite_enhancement"
+        else:
+            text_memories = "\n".join([f"- {mem}" for i, mem in enumerate(batch_texts)])
+            prompt_name = "memory_recreate_enhancement"
         return self.build_prompt(
-            "memory_enhancement",
+            prompt_name,
             query_history=query_history,
             memories=text_memories,
         )
@@ -109,9 +117,11 @@ def _process_enhancement_batch(
     ) -> tuple[list[TextualMemoryItem], bool]:
         attempt = 0
         text_memories = [one.memory for one in memories]
+
         prompt = self._build_enhancement_prompt(
             query_history=query_history, batch_texts=text_memories
         )
+
         llm_response = None
         while attempt <= max(0, retries) + 1:
             try:
@@ -121,14 +131,51 @@ def _process_enhancement_batch(
                     # create new
                     enhanced_memories = []
                     user_id = memories[0].metadata.user_id
-                    for new_mem in processed_text_memories:
-                        enhanced_memories.append(
-                            TextualMemoryItem(
-                                memory=new_mem, metadata=TextualMemoryMetadata(user_id=user_id)
+                    if FINE_STRATEGY == FineStrategy.RECREATE:
+                        for new_mem in processed_text_memories:
+                            enhanced_memories.append(
+                                TextualMemoryItem(
+                                    memory=new_mem, metadata=TextualMemoryMetadata(user_id=user_id)
+                                )
                             )
-                        )
+                    elif FINE_STRATEGY == FineStrategy.REWRITE:
+                        # Parse index from each processed line and rewrite corresponding original memory
+                        def _parse_index_and_text(s: str) -> tuple[int | None, str]:
+                            import re
+
+                            s = (s or "").strip()
+                            # Preferred: [index] text
+                            m = re.match(r"^\s*\[(\d+)\]\s*(.+)$", s)
+                            if m:
+                                return int(m.group(1)), m.group(2).strip()
+                            # Fallback: index: text or index - text
+                            m = re.match(r"^\s*(\d+)\s*[:\-\)]\s*(.+)$", s)
+                            if m:
+                                return int(m.group(1)), m.group(2).strip()
+                            return None, s
+
+                        idx_to_original = dict(enumerate(memories))
+                        for j, item in enumerate(processed_text_memories):
+                            idx, new_text = _parse_index_and_text(item)
+                            if idx is not None and idx in idx_to_original:
+                                orig = idx_to_original[idx]
+                            else:
+                                # Fallback: align by order if index missing/invalid
+                                orig = memories[j] if j < len(memories) else None
+                            if not orig:
+                                continue
+                            enhanced_memories.append(
+                                TextualMemoryItem(
+                                    id=orig.id,
+                                    memory=new_text,
+                                    metadata=orig.metadata,
+                                )
+                            )
+                    else:
+                        logger.error(f"Fine search strategy {FINE_STRATEGY} not exists")
+
                     logger.info(
-                        f"[enhance_memories_with_query] ✅ done | prompt={prompt} | llm_response={llm_response}"
+                        f"[enhance_memories_with_query] ✅ done | Strategy={FINE_STRATEGY} | prompt={prompt} | llm_response={llm_response}"
                     )
                     return enhanced_memories, True
                 else:

src/memos/mem_scheduler/schemas/general_schemas.py

@@ -1,3 +1,5 @@
+import os
+
 from enum import Enum
 from pathlib import Path
 from typing import NewType
@@ -11,6 +13,13 @@ class SearchMode(str, Enum):
     MIXTURE = "mixture"
 
 
+class FineStrategy(str, Enum):
+    """Enumeration for fine strategies."""
+
+    REWRITE = "rewrite"
+    RECREATE = "recreate"
+
+
 FILE_PATH = Path(__file__).absolute()
 BASE_DIR = FILE_PATH.parent.parent.parent.parent.parent
 
@@ -74,3 +83,17 @@ class SearchMode(str, Enum):
 # new types
 UserID = NewType("UserID", str)
 MemCubeID = NewType("CubeID", str)
+
+# algorithm strategies
+DEFAULT_FINE_STRATEGY = FineStrategy.REWRITE
+
+# Read fine strategy from environment variable `FINE_STRATEGY`.
+# If provided and valid, use it; otherwise fall back to default.
+_env_fine_strategy = os.getenv("FINE_STRATEGY")
+if _env_fine_strategy:
+    try:
+        FINE_STRATEGY = FineStrategy(_env_fine_strategy)
+    except ValueError:
+        FINE_STRATEGY = DEFAULT_FINE_STRATEGY
+else:
+    FINE_STRATEGY = DEFAULT_FINE_STRATEGY

src/memos/mem_scheduler/task_schedule_modules/redis_queue.py

@@ -177,10 +177,6 @@ def get(
         try:
             stream_key = self.get_stream_key(user_id=user_id, mem_cube_id=mem_cube_id)
 
-            if stream_key not in self.seen_streams:
-                self.seen_streams.add(stream_key)
-                self._ensure_consumer_group(stream_key=stream_key)
-
             # Calculate timeout for Redis
             redis_timeout = None
             if block and timeout is not None:
@@ -204,6 +200,7 @@ def get(
                     logger.warning(
                         f"Consumer group or stream missing for '{stream_key}/{self.consumer_group}'. Attempting to create and retry."
                     )
+                    self._ensure_consumer_group(stream_key=stream_key)
                     messages = self._redis_conn.xreadgroup(
                         self.consumer_group,
                         self.consumer_name,
@@ -354,10 +351,9 @@ def clear(self) -> None:
 
             for stream_key in stream_keys:
                 # Delete the entire stream
-                self._redis_conn.delete(self.stream_key_prefix)
-                logger.info(f"Cleared Redis stream: {self.stream_key_prefix}")
-                # Recreate the consumer group
-                self._ensure_consumer_group(stream_key=stream_key)
+                self._redis_conn.delete(stream_key)
+                logger.info(f"Cleared Redis stream: {stream_key}")
+
         except Exception as e:
             logger.error(f"Failed to clear Redis queue: {e}")
 

src/memos/mem_scheduler/task_schedule_modules/task_queue.py

@@ -35,6 +35,11 @@ def __init__(
 
         self.disabled_handlers = disabled_handlers
 
+    def debug_mode_on(self):
+        self.memos_message_queue.stream_key_prefix = (
+            f"debug_mode:{self.memos_message_queue.stream_key_prefix}"
+        )
+
     def get_stream_keys(self) -> list[str]:
         if isinstance(self.memos_message_queue, SchedulerRedisQueue):
             return self.memos_message_queue.get_stream_keys()

src/memos/templates/mem_scheduler_prompts.py

@@ -390,7 +390,7 @@
 - Focus on whether the memories can fully answer the query without additional information
 """
 
-MEMORY_ENHANCEMENT_PROMPT = """
+MEMORY_RECREATE_ENHANCEMENT_PROMPT = """
 You are a knowledgeable and precise AI assistant.
 
 # GOAL
@@ -427,6 +427,49 @@
 Final Output:
 """
 
+# Rewrite version: return enhanced memories with original IDs
+MEMORY_REWRITE_ENHANCEMENT_PROMPT = """
+You are a knowledgeable and precise AI assistant.
+
+# GOAL
+Transform raw memories into clean, query-relevant facts — preserving timestamps and resolving ambiguities without inference. Return each enhanced fact with the ID of the original memory being modified.
+
+# RULES & THINKING STEPS
+1. Keep ONLY what’s relevant to the user’s query. Delete irrelevant memories entirely.
+2. Preserve ALL explicit timestamps (e.g., “on October 6”, “daily”, “after injury”).
+3. Resolve all ambiguities using only memory content:
+   - Pronouns → full name: “she” → “Melanie”
+   - Vague nouns → specific detail: “home” → “her childhood home in Guangzhou”
+   - “the user” → identity from context (e.g., “Melanie” if travel/running memories)
+4. Never invent, assume, or extrapolate.
+5. Each output line must be a standalone, clear, factual statement.
+6. Output format: one line per fact, starting with "- ", no extra text.
+
+# IMPORTANT FOR REWRITE
+- Each output line MUST include the original memory’s ID shown in the input list.
+- Use the index shown for each original memory (e.g., "[0]", "[1]") as the ID to reference which memory you are rewriting.
+- For every rewritten line, prefix with the corresponding index in square brackets.
+
+# OUTPUT FORMAT (STRICT)
+Return ONLY the following block, with **one enhanced memory per line**.
+Each line MUST start with "- " (dash + space) AND include index in square brackets.
+
+Wrap the final output inside:
+<answer>
+- [index] enhanced memory 1
+- [index] enhanced memory 2
+...
+</answer>
+
+## User Query
+{query_history}
+
+## Original Memories
+{memories}
+
+Final Output:
+"""
+
 # One-sentence prompt for recalling missing information to answer the query (English)
 ENLARGE_RECALL_PROMPT_ONE_SENTENCE = """
 You are a precise AI assistant. Your job is to analyze the user's query and the available memories to identify what specific information is missing to fully answer the query.
@@ -471,7 +514,8 @@
     "memory_redundancy_filtering": MEMORY_REDUNDANCY_FILTERING_PROMPT,
     "memory_combined_filtering": MEMORY_COMBINED_FILTERING_PROMPT,
     "memory_answer_ability_evaluation": MEMORY_ANSWER_ABILITY_EVALUATION_PROMPT,
-    "memory_enhancement": MEMORY_ENHANCEMENT_PROMPT,
+    "memory_recreate_enhancement": MEMORY_RECREATE_ENHANCEMENT_PROMPT,
+    "memory_rewrite_enhancement": MEMORY_REWRITE_ENHANCEMENT_PROMPT,
     "enlarge_recall": ENLARGE_RECALL_PROMPT_ONE_SENTENCE,
 }