Skip to content

Commit 0a4a935

Browse files
authored
Feat: add deepsearch agent for memos (#517)
* 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
1 parent 1519e33 commit 0a4a935

File tree

6 files changed

+752
-0
lines changed

6 files changed

+752
-0
lines changed
Lines changed: 191 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,191 @@
1+
"""
2+
DeepSearch Agent Usage Examples - Simplified Version
3+
4+
This example demonstrates simplified initialization of DeepSearchMemAgent without
5+
external config builders, using APIConfig methods directly.
6+
"""
7+
8+
import os
9+
10+
from typing import Any
11+
12+
from memos.api.config import APIConfig
13+
from memos.configs.embedder import EmbedderConfigFactory
14+
from memos.configs.graph_db import GraphDBConfigFactory
15+
from memos.configs.internet_retriever import InternetRetrieverConfigFactory
16+
from memos.configs.llm import LLMConfigFactory
17+
from memos.configs.mem_agent import MemAgentConfigFactory
18+
from memos.configs.mem_reader import MemReaderConfigFactory
19+
from memos.configs.reranker import RerankerConfigFactory
20+
from memos.embedders.factory import EmbedderFactory
21+
from memos.graph_dbs.factory import GraphStoreFactory
22+
from memos.llms.factory import LLMFactory
23+
from memos.log import get_logger
24+
from memos.mem_agent.deepsearch_agent import DeepSearchMemAgent
25+
from memos.mem_agent.factory import MemAgentFactory
26+
from memos.mem_cube.navie import NaiveMemCube
27+
from memos.mem_reader.factory import MemReaderFactory
28+
from memos.memories.textual.simple_tree import SimpleTreeTextMemory
29+
from memos.memories.textual.tree_text_memory.organize.manager import MemoryManager
30+
from memos.memories.textual.tree_text_memory.retrieve.internet_retriever_factory import (
31+
InternetRetrieverFactory,
32+
)
33+
from memos.reranker.factory import RerankerFactory
34+
35+
36+
logger = get_logger(__name__)
37+
38+
39+
def build_minimal_components():
40+
"""
41+
Build minimal components for DeepSearchMemAgent with simplified configuration.
42+
43+
This function creates all necessary components using APIConfig methods,
44+
similar to config_builders.py but inline for easier customization.
45+
"""
46+
logger.info("Initializing simplified MemOS components...")
47+
48+
# Build component configurations using APIConfig methods (like config_builders.py)
49+
50+
# Graph DB configuration - using APIConfig.get_nebular_config()
51+
graph_db_backend = os.getenv("NEO4J_BACKEND", "polardb").lower()
52+
graph_db_backend_map = {
53+
"polardb": APIConfig.get_polardb_config(),
54+
}
55+
graph_db_config = GraphDBConfigFactory.model_validate(
56+
{
57+
"backend": graph_db_backend,
58+
"config": graph_db_backend_map[graph_db_backend],
59+
}
60+
)
61+
62+
# LLM configuration - using APIConfig.get_openai_config()
63+
llm_config = LLMConfigFactory.model_validate(
64+
{
65+
"backend": "openai",
66+
"config": APIConfig.get_openai_config(),
67+
}
68+
)
69+
70+
# Embedder configuration - using APIConfig.get_embedder_config()
71+
embedder_config = EmbedderConfigFactory.model_validate(APIConfig.get_embedder_config())
72+
73+
# Memory reader configuration - using APIConfig.get_product_default_config()
74+
mem_reader_config = MemReaderConfigFactory.model_validate(
75+
APIConfig.get_product_default_config()["mem_reader"]
76+
)
77+
78+
# Reranker configuration - using APIConfig.get_reranker_config()
79+
reranker_config = RerankerConfigFactory.model_validate(APIConfig.get_reranker_config())
80+
81+
# Internet retriever configuration - using APIConfig.get_internet_config()
82+
internet_retriever_config = InternetRetrieverConfigFactory.model_validate(
83+
APIConfig.get_internet_config()
84+
)
85+
86+
logger.debug("Component configurations built successfully")
87+
88+
# Create component instances
89+
graph_db = GraphStoreFactory.from_config(graph_db_config)
90+
llm = LLMFactory.from_config(llm_config)
91+
embedder = EmbedderFactory.from_config(embedder_config)
92+
mem_reader = MemReaderFactory.from_config(mem_reader_config)
93+
reranker = RerankerFactory.from_config(reranker_config)
94+
internet_retriever = InternetRetrieverFactory.from_config(
95+
internet_retriever_config, embedder=embedder
96+
)
97+
98+
logger.debug("Core components instantiated")
99+
100+
# Get default cube configuration like component_init.py
101+
default_cube_config = APIConfig.get_default_cube_config()
102+
103+
# Get default memory size from cube config (like component_init.py)
104+
def get_memory_size_from_config(cube_config):
105+
return getattr(cube_config.text_mem.config, "memory_size", None) or {
106+
"WorkingMemory": 20,
107+
"LongTermMemory": 1500,
108+
"UserMemory": 480,
109+
}
110+
111+
memory_size = get_memory_size_from_config(default_cube_config)
112+
is_reorganize = getattr(default_cube_config.text_mem.config, "reorganize", False)
113+
114+
# Initialize memory manager with config from APIConfig
115+
memory_manager = MemoryManager(
116+
graph_db,
117+
embedder,
118+
llm,
119+
memory_size=memory_size,
120+
is_reorganize=is_reorganize,
121+
)
122+
text_memory_config = default_cube_config.text_mem.config
123+
text_mem = SimpleTreeTextMemory(
124+
llm=llm,
125+
embedder=embedder,
126+
mem_reader=mem_reader,
127+
graph_db=graph_db,
128+
reranker=reranker,
129+
memory_manager=memory_manager,
130+
config=text_memory_config,
131+
internet_retriever=internet_retriever,
132+
)
133+
134+
naive_mem_cube = NaiveMemCube(
135+
text_mem=text_mem,
136+
pref_mem=None, # Simplified: no preference memory
137+
act_mem=None,
138+
para_mem=None,
139+
)
140+
141+
return {
142+
"llm": llm,
143+
"naive_mem_cube": naive_mem_cube,
144+
"embedder": embedder,
145+
"graph_db": graph_db,
146+
"mem_reader": mem_reader,
147+
}
148+
149+
150+
def factory_initialization() -> tuple[DeepSearchMemAgent, dict[str, Any]]:
151+
# Build necessary components with simplified setup
152+
components = build_minimal_components()
153+
llm = components["llm"]
154+
naive_mem_cube = components["naive_mem_cube"]
155+
156+
# Create configuration Factory with simplified config
157+
agent_config_factory = MemAgentConfigFactory(
158+
backend="deep_search",
159+
config={
160+
"agent_name": "SimplifiedDeepSearchAgent",
161+
"description": "Simplified intelligent agent for deep search",
162+
"max_iterations": 3, # Maximum number of iterations
163+
"timeout": 60, # Timeout in seconds
164+
},
165+
)
166+
167+
# Create Agent using Factory
168+
# Pass text_mem as memory_retriever, it provides search method
169+
deep_search_agent = MemAgentFactory.from_config(
170+
config_factory=agent_config_factory, llm=llm, memory_retriever=naive_mem_cube.text_mem
171+
)
172+
173+
logger.info("✓ DeepSearchMemAgent created successfully")
174+
logger.info(f" - Agent name: {deep_search_agent.config.agent_name}")
175+
logger.info(f" - Max iterations: {deep_search_agent.max_iterations}")
176+
logger.info(f" - Timeout: {deep_search_agent.timeout} seconds")
177+
178+
return deep_search_agent, components
179+
180+
181+
def main():
182+
agent_factory, components_factory = factory_initialization()
183+
results = agent_factory.run(
184+
"Caroline met up with friends, family, and mentors in early July 2023.",
185+
user_id="locomo_exp_user_0_speaker_b_ct-1118",
186+
)
187+
print(results)
188+
189+
190+
if __name__ == "__main__":
191+
main()

src/memos/configs/mem_agent.py

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
from typing import Any, ClassVar
2+
3+
from pydantic import Field, field_validator, model_validator
4+
5+
from memos.configs.base import BaseConfig
6+
7+
8+
class BaseAgentConfig(BaseConfig):
9+
"""Base configuration class for agents."""
10+
11+
agent_name: str = Field(..., description="Name of the agent")
12+
description: str | None = Field(default=None, description="Description of the agent")
13+
14+
15+
class SimpleAgentConfig(BaseAgentConfig):
16+
"""Simple agent configuration class."""
17+
18+
max_iterations: int = Field(
19+
default=10, description="Maximum number of iterations for the agent"
20+
)
21+
timeout: int = Field(default=30, description="Timeout in seconds for agent execution")
22+
23+
24+
class DeepSearchAgentConfig(BaseAgentConfig):
25+
"""Deep search agent configuration class."""
26+
27+
max_iterations: int = Field(default=3, description="Maximum number of iterations for the agent")
28+
timeout: int = Field(default=30, description="Timeout in seconds for agent execution")
29+
30+
31+
class MemAgentConfigFactory(BaseConfig):
32+
"""Factory class for creating agent configurations."""
33+
34+
backend: str = Field(..., description="Backend for agent")
35+
config: dict[str, Any] = Field(..., description="Configuration for the agent backend")
36+
37+
backend_to_class: ClassVar[dict[str, Any]] = {
38+
"simple": SimpleAgentConfig,
39+
"deep_search": DeepSearchAgentConfig,
40+
}
41+
42+
@field_validator("backend")
43+
@classmethod
44+
def validate_backend(cls, backend: str) -> str:
45+
"""Validate the backend field."""
46+
if backend not in cls.backend_to_class:
47+
raise ValueError(f"Invalid backend: {backend}")
48+
return backend
49+
50+
@model_validator(mode="after")
51+
def create_config(self) -> "MemAgentConfigFactory":
52+
config_class = self.backend_to_class[self.backend]
53+
self.config = config_class(**self.config)
54+
return self

src/memos/mem_agent/base.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
from abc import ABC, abstractmethod
2+
3+
from memos.configs.mem_agent import BaseAgentConfig
4+
5+
6+
class BaseMemAgent(ABC):
7+
"""
8+
Base class for all agents.
9+
"""
10+
11+
def __init__(self, config: BaseAgentConfig):
12+
"""Initialize the BaseMemAgent with the given configuration."""
13+
self.config = config
14+
15+
@abstractmethod
16+
def run(self, input: str) -> str:
17+
"""
18+
Run the agent.
19+
"""

0 commit comments

Comments
 (0)