Feat: update prompt and format need raw (#528)

* feat: update memos headers

* feat: headers add

* feat: update search agent

* feat: upadte mem story

* feat: update mem scehduler

* feat: update deepsearch mem code

* feat:  update deepsearch agent

* feat: update test code

* fix: remove dup config

* feat: dock search pipeline

* fix: code test

* feat: add test scripts

* feat: add test

* feat: update need_raw process

Author: fridayL

src/memos/mem_agent/deepsearch_agent.py

@@ -183,8 +183,6 @@ def run(self, query: str, **kwargs) -> str | list[TextualMemoryItem]:
             if search_results:
                 context_batch = [self._extract_context_from_memory(mem) for mem in search_results]
                 accumulated_context.extend(context_batch)
-                accumulated_memories.extend(search_results)
-
                 reflection_result = self.reflector.run(current_query, context_batch)
                 status = reflection_result.get("status", "sufficient")
                 reasoning = reflection_result.get("reasoning", "")
@@ -193,11 +191,14 @@ def run(self, query: str, **kwargs) -> str | list[TextualMemoryItem]:
 
                 if status == "sufficient":
                     logger.info("Sufficient information collected")
+                    accumulated_memories.extend(search_results)
                     break
                 elif status == "needs_raw":
                     logger.info("Need original sources, retrieving raw content")
+                    accumulated_memories.extend(self._set_source_from_memory(search_results))
                     break
                 elif status == "missing_info":
+                    accumulated_memories.extend(search_results)
                     missing_entities = reflection_result.get("missing_entities", [])
                     logger.info(f"Missing information: {missing_entities}")
                     current_query = reflection_result.get("new_search_query")
@@ -331,6 +332,22 @@ def _refine_query_for_missing_info(self, query: str, missing_entities: list[str]
 
         return refined_query
 
+    def _set_source_from_memory(
+        self, memory_items: list[TextualMemoryItem]
+    ) -> list[TextualMemoryItem]:
+        """set source from memory item"""
+        for memory_item in memory_items:
+            if not hasattr(memory_item.metadata, "sources"):
+                continue
+            chat_sources = [
+                f"{source.chat_time} {source.role}: {source.content}"
+                for source in memory_item.metadata.sources
+                if hasattr(source, "type") and source.type == "chat"
+            ]
+            if chat_sources:
+                memory_item.memory = "\n".join(chat_sources) + "\n"
+        return memory_items
+
     def _generate_final_answer(
         self,
         original_query: str,

src/memos/templates/mem_agent_prompts.py

@@ -35,7 +35,13 @@
 - "sufficient": Context fully answers the query
 - "missing_info": Key information is missing (e.g., specific dates, locations, details)
 - "needs_raw": Content is relevant but too summarized/vague, need original sources
-- "new_search_query": New search query to retrieve more information
+
+IMPORTANT for "new_search_query":
+- MUST preserve ALL specific entities from the original query (names, dates, times, locations, etc.)
+- DO NOT replace specific information with generic terms like "user", "person", "they", etc.
+- Keep the exact same subjects, time references, and key details as in the original query
+- Only modify the query to focus on the missing information while maintaining all original specifics
+- Example: If original query mentions "May 2024", keep "May 2024" in new query, don't change to "that month"
 
 Response:"""