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
5 changes: 5 additions & 0 deletions src/memos/api/handlers/base_handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -161,6 +161,11 @@ def mos_server(self):
"""Get MOS server instance."""
return self.deps.mos_server

@property
def deepsearch_agent(self):
"""Get deepsearch agent instance."""
return self.deps.deepsearch_agent

def _validate_dependencies(self, *required_deps: str) -> None:
"""
Validate that required dependencies are available.
Expand Down
6 changes: 6 additions & 0 deletions src/memos/api/handlers/component_init.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@

if TYPE_CHECKING:
from memos.memories.textual.tree import TreeTextMemory
from memos.mem_agent.deepsearch_agent import DeepSearchMemAgent
from memos.memories.textual.tree_text_memory.retrieve.internet_retriever_factory import (
InternetRetrieverFactory,
)
Expand Down Expand Up @@ -307,6 +308,10 @@ def init_server() -> dict[str, Any]:
online_bot = get_online_bot_function() if dingding_enabled else None
logger.info("DingDing bot is enabled")

deepsearch_agent = DeepSearchMemAgent(
llm=llm,
memory_retriever=tree_mem,
)
# Return all components as a dictionary for easy access and extension
return {
"graph_db": graph_db,
Expand All @@ -330,4 +335,5 @@ def init_server() -> dict[str, Any]:
"text_mem": text_mem,
"pref_mem": pref_mem,
"online_bot": online_bot,
"deepsearch_agent": deepsearch_agent,
}
10 changes: 7 additions & 3 deletions src/memos/api/handlers/search_handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,9 @@ def __init__(self, dependencies: HandlerDependencies):
dependencies: HandlerDependencies instance
"""
super().__init__(dependencies)
self._validate_dependencies("naive_mem_cube", "mem_scheduler", "searcher")
self._validate_dependencies(
"naive_mem_cube", "mem_scheduler", "searcher", "deepsearch_agent"
)

def handle_search_memories(self, search_req: APISearchRequest) -> SearchResponse:
"""
Expand All @@ -52,10 +54,10 @@ def handle_search_memories(self, search_req: APISearchRequest) -> SearchResponse

results = cube_view.search_memories(search_req)

self.logger.info(f"[AddHandler] Final add results count={len(results)}")
self.logger.info(f"[SearchHandler] Final search results count={len(results)}")

return SearchResponse(
message="Memory searched successfully",
message="Search completed successfully",
data=results,
)

Expand Down Expand Up @@ -83,6 +85,7 @@ def _build_cube_view(self, search_req: APISearchRequest) -> MemCubeView:
mem_scheduler=self.mem_scheduler,
logger=self.logger,
searcher=self.searcher,
deepsearch_agent=self.deepsearch_agent,
)
else:
single_views = [
Expand All @@ -93,6 +96,7 @@ def _build_cube_view(self, search_req: APISearchRequest) -> MemCubeView:
mem_scheduler=self.mem_scheduler,
logger=self.logger,
searcher=self.searcher,
deepsearch_agent=self.deepsearch_agent,
)
for cube_id in cube_ids
]
Expand Down
11 changes: 5 additions & 6 deletions src/memos/mem_agent/deepsearch_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,8 @@
if TYPE_CHECKING:
from memos.types import MessageList

logger = get_logger(__name__)


class JSONResponseParser:
"""Elegant JSON response parser for LLM outputs"""
Expand All @@ -48,9 +50,6 @@ def parse(response: str) -> dict[str, Any]:
raise ValueError(f"Cannot parse JSON response: {response[:100]}...")


logger = get_logger(__name__)


class QueryRewriter(BaseMemAgent):
"""Specialized agent for rewriting queries based on conversation history"""

Expand Down Expand Up @@ -141,7 +140,7 @@ def __init__(
memory_retriever: Memory retrieval interface (e.g., naive_mem_cube.text_mem)
config: Configuration for deep search behavior
"""
self.config = config or DeepSearchAgentConfig()
self.config = config or DeepSearchAgentConfig(agent_name="DeepSearchMemAgent")
self.max_iterations = self.config.max_iterations
self.timeout = self.config.timeout
self.llm: BaseLLM = llm
Expand Down Expand Up @@ -219,7 +218,7 @@ def run(self, query: str, **kwargs) -> str | list[TextualMemoryItem]:
return self._remove_duplicate_memories(accumulated_memories)
else:
return self._generate_final_answer(
query, accumulated_memories, accumulated_context, "", history
query, accumulated_memories, accumulated_context, history
)

def _remove_duplicate_memories(
Expand Down Expand Up @@ -248,9 +247,9 @@ def _generate_final_answer(
original_query: str,
search_results: list[TextualMemoryItem],
context: list[str],
missing_info: str = "",
history: list[str] | None = None,
sources: list[str] | None = None,
missing_info: str | None = None,
) -> str:
"""
Generate the final answer.
Expand Down
8 changes: 6 additions & 2 deletions src/memos/multi_mem_cube/single_cube.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ class SingleCubeView(MemCubeView):
mem_scheduler: Any
logger: Any
searcher: Any
deepsearch_agent: Any

def add_memories(self, add_req: APIADDRequest) -> list[dict[str, Any]]:
"""
Expand Down Expand Up @@ -247,8 +248,11 @@ def _fast_search(
def _deep_search(
self, search_req: APISearchRequest, user_context: UserContext, max_thinking_depth: int
) -> list:
logger.error("waiting to be implemented")
return []
deepsearch_results = self.deepsearch_agent.run(
search_req.query, user_id=user_context.mem_cube_id
)
formatted_memories = [format_memory_item(data) for data in deepsearch_results]
return formatted_memories

def _fine_search(
self,
Expand Down
4 changes: 3 additions & 1 deletion src/memos/templates/mem_agent_prompts.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,12 +22,14 @@
{context}

Analyze the context and determine the next step. Return your response in JSON format with the following structure:
{{
```json
{{
"status": "sufficient|missing_info|needs_raw",
"reasoning": "Brief explanation of your decision",
"missing_entities": ["entity1", "entity2"],
"new_search_query": "new search query",
}}
```

Status definitions:
- "sufficient": Context fully answers the query
Expand Down
1 change: 1 addition & 0 deletions tests/api/test_server_router.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ def mock_init_server():
"pref_mem": None,
"online_bot": None,
"chat_llms": Mock(),
"deepsearch_agent": Mock(),
}

with patch("memos.api.handlers.init_server", return_value=mock_components):
Expand Down
Loading
Loading