fix: prevent memory drift for referential queries in fast mode

src/memos/memories/textual/tree_text_memory/retrieve/task_goal_parser.py

@@ -14,6 +14,28 @@
 
 logger = get_logger(__name__)
 
+# Hints that indicate a query is referential / elliptical and needs
+# context augmentation from prior conversation.
+# Multi-char hints are safe for plain substring matching.
+_MULTI_CHAR_REFERENTIAL_HINTS = (
+    "其他",
+    "别的",
+    "这个",
+    "那个",
+    "继续",
+    "上面",
+    "刚说",
+    "other",
+    "another",
+    "it",
+    "this",
+    "that",
+    "more",
+)
+# Single-char hints are only matched at the *start* of the query to avoid
+# false positives like "还款计划" (contains "还") or "也门旅游" (contains "也").
+_SINGLE_CHAR_REFERENTIAL_HINTS = ("再", "也", "还", "他", "她", "它")
+
 
 class TaskGoalParser:
     """
@@ -44,38 +66,95 @@ def parse(
         """
 
         if mode == "fast":
-            return self._parse_fast(task_description, context=context, **kwargs)
+            return self._parse_fast(
+                task_description, context=context, conversation=conversation, **kwargs
+            )
         elif mode == "fine":
             if not self.llm:
                 raise ValueError("LLM not provided for slow mode.")
             return self._parse_fine(task_description, context, conversation, **kwargs)
         else:
             raise ValueError(f"Unknown mode: {mode}")
 
-    def _parse_fast(self, task_description: str, **kwargs) -> ParsedTaskGoal:
+    def _parse_fast(
+        self,
+        task_description: str,
+        *,
+        conversation: list[dict] | None = None,
+        **kwargs,
+    ) -> ParsedTaskGoal:
         """
         Fast mode: simple jieba word split.
+
+        If recent conversation history is available, a lightweight
+        context-augmented query is generated by prepending the most
+        recent user message content.  This prevents referential
+        queries (e.g. "再找找其他价格优惠的") from drifting to
+        unrelated older memories.
         """
         context = kwargs.get("context", "")
         use_fast_graph = kwargs.get("use_fast_graph", False)
+
+        # ─── Lightweight context expansion (no LLM call) ────────────────
+        # When conversation history exists, augment the query with the
+        # most recent user message so that embedding search has enough
+        # context to disambiguate referential / elliptical queries.
+        augmented_query = task_description
+        if conversation:
+            recent_user_msgs = [
+                msg["content"]
+                for msg in conversation
+                if isinstance(msg, dict)
+                and msg.get("role") == "user"
+                and isinstance(msg.get("content"), str)
+                and msg["content"].strip()
+            ]
+            if recent_user_msgs:
+                # The current query might already be the last item in the conversation
+                # history. If so, we want the *previous* user message as context.
+                if (
+                    recent_user_msgs[-1].strip() == task_description.strip()
+                    and len(recent_user_msgs) > 1
+                ):
+                    last_user_msg = recent_user_msgs[-2].strip()
+                elif recent_user_msgs[-1].strip() != task_description.strip():
+                    last_user_msg = recent_user_msgs[-1].strip()
+                else:
+                    last_user_msg = ""
+
+                # Only augment if the query is short or looks referential.
+                is_referential = (
+                    len(task_description) < 20
+                    or any(h in task_description for h in _MULTI_CHAR_REFERENTIAL_HINTS)
+                    or any(task_description.startswith(ch) for ch in _SINGLE_CHAR_REFERENTIAL_HINTS)
+                )
+                if is_referential and last_user_msg:
+                    # Truncate to avoid over-long embedding inputs.
+                    # Head truncation preserves the topic anchor at the beginning
+                    # of the prior message.
+                    max_ctx = 120
+                    if len(last_user_msg) > max_ctx:
+                        last_user_msg = last_user_msg[:max_ctx]
+                    augmented_query = f"{last_user_msg} {task_description}"
+
         if use_fast_graph:
-            desc_tokenized = self.tokenizer.tokenize_mixed(task_description)
+            desc_tokenized = self.tokenizer.tokenize_mixed(augmented_query)
             return ParsedTaskGoal(
-                memories=[task_description],
+                memories=[augmented_query],
                 keys=desc_tokenized,
                 tags=desc_tokenized,
                 goal_type="default",
-                rephrased_query=task_description,
+                rephrased_query=augmented_query,
                 internet_search=False,
                 context=context,
             )
         else:
             return ParsedTaskGoal(
-                memories=[task_description],
+                memories=[augmented_query],
                 keys=[],
                 tags=[],
                 goal_type="default",
-                rephrased_query=task_description,
+                rephrased_query=augmented_query,
                 internet_search=False,
                 context=context,
             )

tests/memories/textual/test_task_goal_parser_context.py

@@ -0,0 +1,99 @@
+"""Tests for context-aware query expansion in TaskGoalParser fast mode.
+
+Covers the fix for #1365: referential queries (e.g. "再找找其他价格优惠的")
+no longer drift to unrelated older memories because the most recent user
+message is prepended to the query for embedding search.
+"""
+
+from unittest.mock import MagicMock
+
+import pytest
+
+from memos.memories.textual.tree_text_memory.retrieve.task_goal_parser import (
+    TaskGoalParser,
+)
+
+
+@pytest.fixture
+def parser():
+    """Build a TaskGoalParser with a mock LLM (not used in fast mode)."""
+    return TaskGoalParser(llm=MagicMock())
+
+
+class TestParseFastContextExpansion:
+    """Tests that ``_parse_fast`` uses conversation history correctly."""
+
+    def test_no_conversation_keeps_original_query(self, parser):
+        """Without conversation, query is unchanged."""
+        result = parser._parse_fast("找商品", context="")
+        assert result.rephrased_query == "找商品"
+        assert result.memories == ["找商品"]
+
+    def test_referential_query_gets_augmented(self, parser):
+        """Short referential queries are prepended with last user message."""
+        conversation = [
+            {"role": "user", "content": "帮我找A商品"},
+            {"role": "assistant", "content": "找到了A商品"},
+            {"role": "user", "content": "再找找其他价格优惠的"},
+        ]
+        result = parser._parse_fast("再找找其他价格优惠的", conversation=conversation, context="")
+        assert "帮我找A商品" in result.rephrased_query
+        assert "再找找其他价格优惠的" in result.rephrased_query
+
+    def test_non_referential_long_query_not_augmented(self, parser):
+        """Long, self-contained queries are not augmented."""
+        conversation = [
+            {"role": "user", "content": "帮我找A商品"},
+            {"role": "assistant", "content": "找到了"},
+        ]
+        long_query = "请帮我搜索一下关于大型语言模型训练过程中梯度爆炸问题的解决方案"
+        result = parser._parse_fast(long_query, conversation=conversation, context="")
+        assert result.rephrased_query == long_query
+
+    def test_english_referential_query(self, parser):
+        """English referential words also trigger augmentation."""
+        conversation = [
+            {"role": "user", "content": "Find me a red dress"},
+            {"role": "assistant", "content": "Here are some red dresses"},
+        ]
+        result = parser._parse_fast("show me more", conversation=conversation, context="")
+        assert "Find me a red dress" in result.rephrased_query
+        assert "show me more" in result.rephrased_query
+
+    def test_empty_conversation(self, parser):
+        """Empty conversation list does not cause errors."""
+        result = parser._parse_fast("找商品", conversation=[], context="")
+        assert result.rephrased_query == "找商品"
+
+    def test_conversation_without_user_messages(self, parser):
+        """Conversation with only assistant messages does not augment."""
+        conversation = [
+            {"role": "assistant", "content": "Hello there"},
+        ]
+        result = parser._parse_fast("再找找", conversation=conversation, context="")
+        assert result.rephrased_query == "再找找"
+
+    def test_long_last_user_message_truncated(self, parser):
+        """Very long last user messages are truncated to 120 chars."""
+        long_msg = "A" * 300
+        conversation = [
+            {"role": "user", "content": long_msg},
+        ]
+        result = parser._parse_fast("再找找", conversation=conversation, context="")
+        # The augmented query should contain the truncated version (first 120 chars)
+        assert "A" * 120 in result.rephrased_query
+        assert "A" * 121 not in result.rephrased_query  # Verify truncation actually occurred
+
+    def test_parse_method_passes_conversation_in_fast_mode(self, parser):
+        """The public ``parse`` method passes conversation to ``_parse_fast``."""
+        conversation = [
+            {"role": "user", "content": "帮我找A商品"},
+            {"role": "assistant", "content": "好的"},
+        ]
+        result = parser.parse(
+            task_description="再找找",
+            context="",
+            conversation=conversation,
+            mode="fast",
+        )
+        assert "帮我找A商品" in result.rephrased_query