Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 19 additions & 2 deletions src/memos/mem_agent/deepsearch_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -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", "")
Expand All @@ -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")
Expand Down Expand Up @@ -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,
Expand Down
8 changes: 7 additions & 1 deletion src/memos/templates/mem_agent_prompts.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:"""

Expand Down
Loading