From 7f103b3cb7b6745029312a9a79b66d38aab88e35 Mon Sep 17 00:00:00 2001
From: Wenqiang Wei <46308778+endxxxx@users.noreply.github.com>
Date: Thu, 31 Jul 2025 19:21:50 +0800
Subject: [PATCH] fix: avoid over-rewriting query and wrong reference display

---
 .../textual/tree_text_memory/retrieve/searcher.py    | 12 ++++++++++++
 .../textual/tree_text_memory/retrieve/utils.py       |  2 +-
 src/memos/templates/mos_prompts.py                   |  2 +-
 3 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/src/memos/memories/textual/tree_text_memory/retrieve/searcher.py b/src/memos/memories/textual/tree_text_memory/retrieve/searcher.py
index fde0b9a58..61e152338 100644
--- a/src/memos/memories/textual/tree_text_memory/retrieve/searcher.py
+++ b/src/memos/memories/textual/tree_text_memory/retrieve/searcher.py
@@ -75,6 +75,18 @@ def search(
             context = [related_node["memory"] for related_node in related_nodes]
             context = list(set(context))
 
+            # add some knowledge retrieved from internet to the context to avoid misunderstanding while parsing the task goal.
+            if self.internet_retriever:
+                supplyment_memory_items = self.internet_retriever.retrieve_from_internet(
+                    query=query, top_k=3
+                )
+                context.extend(
+                    [
+                        each_supplyment_item.memory.partition("\nContent: ")[-1]
+                        for each_supplyment_item in supplyment_memory_items
+                    ]
+                )
+
         # Step 1a: Parse task structure into topic, concept, and fact levels
         parsed_goal = self.task_goal_parser.parse(
             task_description=query,
diff --git a/src/memos/memories/textual/tree_text_memory/retrieve/utils.py b/src/memos/memories/textual/tree_text_memory/retrieve/utils.py
index 9c2d6e8fb..de389ef28 100644
--- a/src/memos/memories/textual/tree_text_memory/retrieve/utils.py
+++ b/src/memos/memories/textual/tree_text_memory/retrieve/utils.py
@@ -5,7 +5,7 @@
 2. Tags: thematic tags to help categorize and retrieve related memories.
 3. Goal Type: retrieval | qa | generation
 4. Rephrased instruction: Give a rephrased task instruction based on the former conversation to make it less confusing to look alone. If you think the task instruction is easy enough to understand, or there is no former conversation, set "rephrased_instruction" to an empty string.
-5. Need for internet search: If you think you need to search the internet to finish the rephrased/original user task instruction, set "internet_search" to True. Otherwise, set it to False.
+5. Need for internet search: If the user's task instruction only involves objective facts or can be completed without introducing external knowledge, set "internet_search" to False. Otherwise, set it to True.
 6. Memories: Provide 2–5 short semantic expansions or rephrasings of the rephrased/original user task instruction. These are used for improved embedding search coverage. Each should be clear, concise, and meaningful for retrieval.
 
 Task description:
diff --git a/src/memos/templates/mos_prompts.py b/src/memos/templates/mos_prompts.py
index ddf576ee2..77427060e 100644
--- a/src/memos/templates/mos_prompts.py
+++ b/src/memos/templates/mos_prompts.py
@@ -92,7 +92,7 @@
   - `memoriesID` is the specific memory ID from the available memories list
 
 ### Reference Examples
-- Correct: `[1:abc123]`, `[2:def456]`, `[3:ghi789]`, `[4:jkl101]`, `[5:mno112]`
+- Correct: `[1:abc123]`, `[2:def456]`, `[3:ghi789]`, `[4:jkl101][5:mno112]` (concatenate reference annotation directly while citing multiple memories)
 - Incorrect: `[1:abc123,2:def456]` (do not use connected format)
 
 ## Response Guidelines