diff --git a/src/memos/api/config.py b/src/memos/api/config.py index 61e73f9ae..b911610ca 100644 --- a/src/memos/api/config.py +++ b/src/memos/api/config.py @@ -115,6 +115,48 @@ def get_embedder_config() -> dict[str, Any]: }, } + @staticmethod + def get_internet_config() -> dict[str, Any]: + """Get embedder configuration.""" + return { + "backend": "xinyu", + "config": { + "api_key": os.getenv("XINYU_API_KEY"), + "search_engine_id": os.getenv("XINYU_SEARCH_ENGINE_ID"), + "max_results": 15, + "num_per_request": 10, + "reader": { + "backend": "simple_struct", + "config": { + "llm": { + "backend": "openai", + "config": { + "model_name_or_path": os.getenv("MEMRADER_MODEL"), + "temperature": 0.6, + "max_tokens": 5000, + "top_p": 0.95, + "top_k": 20, + "api_key": "EMPTY", + "api_base": os.getenv("MEMRADER_API_BASE"), + "remove_think_prefix": True, + "extra_body": {"chat_template_kwargs": {"enable_thinking": False}}, + }, + }, + "embedder": APIConfig.get_embedder_config(), + "chunker": { + "backend": "sentence", + "config": { + "tokenizer_or_token_counter": "gpt2", + "chunk_size": 512, + "chunk_overlap": 128, + "min_sentences_per_chunk": 1, + }, + }, + }, + }, + }, + } + @staticmethod def get_neo4j_community_config(user_id: str | None = None) -> dict[str, Any]: """Get Neo4j community configuration.""" @@ -340,7 +382,6 @@ def create_user_config(user_name: str, user_id: str) -> tuple[MOSConfig, General "top_k": 30, "max_turns_window": 20, } - # Add scheduler configuration if enabled if APIConfig.is_scheduler_enabled(): config_dict["mem_scheduler"] = APIConfig.get_scheduler_config() @@ -352,7 +393,11 @@ def create_user_config(user_name: str, user_id: str) -> tuple[MOSConfig, General neo4j_community_config = APIConfig.get_neo4j_community_config(user_id) neo4j_config = APIConfig.get_neo4j_config(user_id) - + internet_config = ( + APIConfig.get_internet_config() + if os.getenv("ENABLE_INTERNET", "false").lower() == "true" + else None + ) graph_db_backend_map = { "neo4j-community": neo4j_community_config, "neo4j": neo4j_config, @@ -360,6 +405,7 @@ def create_user_config(user_name: str, user_id: str) -> tuple[MOSConfig, General graph_db_backend = os.getenv("NEO4J_BACKEND", "neo4j-community").lower() if graph_db_backend in graph_db_backend_map: # Create MemCube config + default_cube_config = GeneralMemCubeConfig.model_validate( { "user_id": user_id, @@ -374,6 +420,7 @@ def create_user_config(user_name: str, user_id: str) -> tuple[MOSConfig, General "config": graph_db_backend_map[graph_db_backend], }, "embedder": APIConfig.get_embedder_config(), + "internet_retriever": internet_config, }, }, "act_mem": {} @@ -384,7 +431,10 @@ def create_user_config(user_name: str, user_id: str) -> tuple[MOSConfig, General ) else: raise ValueError(f"Invalid Neo4j backend: {graph_db_backend}") - + if os.getenv("ENABLE_INTERNET", "false").lower() == "true": + default_cube_config.text_mem.config["internet_retriever"] = ( + APIConfig.get_internet_config() + ) default_mem_cube = GeneralMemCube(default_cube_config) return default_config, default_mem_cube @@ -405,6 +455,11 @@ def get_default_cube_config() -> GeneralMemCubeConfig | None: "neo4j-community": neo4j_community_config, "neo4j": neo4j_config, } + internet_config = ( + APIConfig.get_internet_config() + if os.getenv("ENABLE_INTERNET", "false").lower() == "true" + else None + ) graph_db_backend = os.getenv("NEO4J_BACKEND", "neo4j-community").lower() if graph_db_backend in graph_db_backend_map: return GeneralMemCubeConfig.model_validate( @@ -423,6 +478,7 @@ def get_default_cube_config() -> GeneralMemCubeConfig | None: "embedder": APIConfig.get_embedder_config(), "reorganize": os.getenv("MOS_ENABLE_REORGANIZE", "false").lower() == "true", + "internet_retriever": internet_config, }, }, "act_mem": {} diff --git a/src/memos/api/product_models.py b/src/memos/api/product_models.py index d14a8e01a..f21ba2b6b 100644 --- a/src/memos/api/product_models.py +++ b/src/memos/api/product_models.py @@ -83,6 +83,7 @@ class ChatRequest(BaseRequest): query: str = Field(..., description="Chat query message") mem_cube_id: str | None = Field(None, description="Cube ID to use for chat") history: list[MessageDict] | None = Field(None, description="Chat history") + internet_search: bool = Field(True, description="Whether to use internet search") class UserCreate(BaseRequest): diff --git a/src/memos/api/routers/product_router.py b/src/memos/api/routers/product_router.py index dce15fcac..17125c031 100644 --- a/src/memos/api/routers/product_router.py +++ b/src/memos/api/routers/product_router.py @@ -234,6 +234,7 @@ async def generate_chat_response(): user_id=chat_req.user_id, cube_id=chat_req.mem_cube_id, history=chat_req.history, + internet_search=chat_req.internet_search, ): yield chunk await asyncio.sleep(0.00001) # 50ms delay between chunks diff --git a/src/memos/configs/mem_reader.py b/src/memos/configs/mem_reader.py index c867f30b5..1c62087a3 100644 --- a/src/memos/configs/mem_reader.py +++ b/src/memos/configs/mem_reader.py @@ -15,6 +15,15 @@ class BaseMemReaderConfig(BaseConfig): created_at: datetime = Field( default_factory=datetime.now, description="Creation timestamp for the MemReader" ) + + @field_validator("created_at", mode="before") + @classmethod + def parse_datetime(cls, value): + """Parse datetime from string if needed.""" + if isinstance(value, str): + return datetime.fromisoformat(value.replace("Z", "+00:00")) + return value + llm: LLMConfigFactory = Field(..., description="LLM configuration for the MemReader") embedder: EmbedderConfigFactory = Field( ..., description="Embedder configuration for the MemReader" diff --git a/src/memos/log.py b/src/memos/log.py index 12a92b5b5..3ba0fe485 100644 --- a/src/memos/log.py +++ b/src/memos/log.py @@ -48,7 +48,7 @@ def _setup_logfile() -> Path: "class": "logging.handlers.RotatingFileHandler", "filename": _setup_logfile(), "maxBytes": 1024**2 * 10, - "backupCount": 3, + "backupCount": 10, "formatter": "standard", }, }, diff --git a/src/memos/mem_cube/utils.py b/src/memos/mem_cube/utils.py index 0e7afaf39..7d5414b0f 100644 --- a/src/memos/mem_cube/utils.py +++ b/src/memos/mem_cube/utils.py @@ -71,10 +71,6 @@ def merge_config_with_default( # Define graph_db fields to preserve (user-specific) preserve_graph_fields = { - "uri", - "user", - "password", - "db_name", "auto_create", "user_name", "use_multi_db", diff --git a/src/memos/mem_os/core.py b/src/memos/mem_os/core.py index 015d3774c..7538db4ee 100644 --- a/src/memos/mem_os/core.py +++ b/src/memos/mem_os/core.py @@ -530,6 +530,8 @@ def search( user_id: str | None = None, install_cube_ids: list[str] | None = None, top_k: int | None = None, + mode: Literal["fast", "fine"] = "fast", + internet_search: bool = False, ) -> MOSSearchResult: """ Search for textual memories across all registered MemCubes. @@ -567,7 +569,11 @@ def search( and self.config.enable_textual_memory ): memories = mem_cube.text_mem.search( - query, top_k=top_k if top_k else self.config.top_k + query, + top_k=top_k if top_k else self.config.top_k, + mode=mode, + manual_close_internet=not internet_search, + info={"user_id": target_user_id, "session_id": str(uuid.uuid4())}, ) result["text_mem"].append({"cube_id": mem_cube_id, "memories": memories}) logger.info( diff --git a/src/memos/mem_os/product.py b/src/memos/mem_os/product.py index 95fd62a90..e6b14c7d9 100644 --- a/src/memos/mem_os/product.py +++ b/src/memos/mem_os/product.py @@ -363,7 +363,7 @@ def _build_system_prompt(self, user_id: str, memories_all: list[TextualMemoryIte for i, memory in enumerate(memories_all, 1): # Format: [memory_id]: memory_content memory_id = f"{memory.id.split('-')[0]}" if hasattr(memory, "id") else f"mem_{i}" - memory_content = memory.memory if hasattr(memory, "memory") else str(memory) + memory_content = memory.memory[:500] if hasattr(memory, "memory") else str(memory) memory_context += f"{memory_id}: {memory_content}\n" return base_prompt + memory_context @@ -694,7 +694,7 @@ def get_suggestion_query(self, user_id: str, language: str = "zh") -> list[str]: "text_mem" ] if text_mem_result: - memories = "\n".join([m.memory for m in text_mem_result[0]["memories"]]) + memories = "\n".join([m.memory[:200] for m in text_mem_result[0]["memories"]]) else: memories = "" message_list = [{"role": "system", "content": suggestion_prompt.format(memories=memories)}] @@ -710,6 +710,7 @@ def chat_with_references( cube_id: str | None = None, history: MessageList | None = None, top_k: int = 10, + internet_search: bool = False, ) -> Generator[str, None, None]: """ Chat with LLM with memory references and streaming output. @@ -729,7 +730,12 @@ def chat_with_references( memories_list = [] yield f"data: {json.dumps({'type': 'status', 'data': '0'})}\n\n" memories_result = super().search( - query, user_id, install_cube_ids=[cube_id] if cube_id else None, top_k=top_k + query, + user_id, + install_cube_ids=[cube_id] if cube_id else None, + top_k=top_k, + mode="fine", + internet_search=internet_search, )["text_mem"] yield f"data: {json.dumps({'type': 'status', 'data': '1'})}\n\n" self._send_message_to_scheduler( @@ -829,6 +835,7 @@ def chat_with_references( memories_json["metadata"]["embedding"] = [] memories_json["metadata"]["sources"] = [] memories_json["metadata"]["memory"] = memories.memory + memories_json["metadata"]["id"] = memories.id reference.append({"metadata": memories_json["metadata"]}) yield f"data: {json.dumps({'type': 'reference', 'data': reference})}\n\n" @@ -1008,13 +1015,18 @@ def get_subgraph( return reformat_memory_list def search( - self, query: str, user_id: str, install_cube_ids: list[str] | None = None, top_k: int = 10 + self, + query: str, + user_id: str, + install_cube_ids: list[str] | None = None, + top_k: int = 10, + mode: Literal["fast", "fine"] = "fast", ): """Search memories for a specific user.""" # Load user cubes if not already loaded self._load_user_cubes(user_id, self.default_cube_config) - search_result = super().search(query, user_id, install_cube_ids, top_k) + search_result = super().search(query, user_id, install_cube_ids, top_k, mode=mode) text_memory_list = search_result["text_mem"] reformat_memory_list = [] for memory in text_memory_list: diff --git a/src/memos/memories/textual/tree.py b/src/memos/memories/textual/tree.py index 3b50001d8..601597b19 100644 --- a/src/memos/memories/textual/tree.py +++ b/src/memos/memories/textual/tree.py @@ -117,13 +117,19 @@ def search( logger.warning( "Internet retriever is init by config , but this search set manual_close_internet is True and will close it" ) - self.internet_retriever = None - searcher = Searcher( - self.dispatcher_llm, - self.graph_store, - self.embedder, - internet_retriever=self.internet_retriever, - ) + searcher = Searcher( + self.dispatcher_llm, + self.graph_store, + self.embedder, + internet_retriever=None, + ) + else: + searcher = Searcher( + self.dispatcher_llm, + self.graph_store, + self.embedder, + internet_retriever=self.internet_retriever, + ) return searcher.search(query, top_k, info, mode, memory_type) def get_relevant_subgraph( diff --git a/src/memos/memories/textual/tree_text_memory/organize/reorganizer.py b/src/memos/memories/textual/tree_text_memory/organize/reorganizer.py index dfaa57563..301857076 100644 --- a/src/memos/memories/textual/tree_text_memory/organize/reorganizer.py +++ b/src/memos/memories/textual/tree_text_memory/organize/reorganizer.py @@ -125,8 +125,8 @@ def _run_structure_organizer_loop(self): """ import schedule - schedule.every(20).seconds.do(self.optimize_structure, scope="LongTermMemory") - schedule.every(20).seconds.do(self.optimize_structure, scope="UserMemory") + schedule.every(600).seconds.do(self.optimize_structure, scope="LongTermMemory") + schedule.every(600).seconds.do(self.optimize_structure, scope="UserMemory") logger.info("Structure optimizer schedule started.") while not getattr(self, "_stop_scheduler", False): 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 c198fa231..6cfd37903 100644 --- a/src/memos/memories/textual/tree_text_memory/retrieve/searcher.py +++ b/src/memos/memories/textual/tree_text_memory/retrieve/searcher.py @@ -149,7 +149,7 @@ def retrieve_from_internet(): query=query, query_embedding=query_embedding[0], graph_results=internet_items, - top_k=max(top_k, 10), + top_k=min(top_k, 5), parsed_goal=parsed_goal, ) return ranked_memories diff --git a/src/memos/settings.py b/src/memos/settings.py index d55a52e8a..3b3f05ebd 100644 --- a/src/memos/settings.py +++ b/src/memos/settings.py @@ -1,7 +1,9 @@ +import os + from pathlib import Path -MEMOS_DIR = Path.cwd() / ".memos" +MEMOS_DIR = Path(os.getenv("MEMOS_BASE_PATH", Path.cwd())) / ".memos" DEBUG = False # "memos" or "memos.submodules" ... to filter logs from specific packages diff --git a/tests/mem_os/test_memos_core.py b/tests/mem_os/test_memos_core.py index 88923fc94..0ebe49ff5 100644 --- a/tests/mem_os/test_memos_core.py +++ b/tests/mem_os/test_memos_core.py @@ -326,7 +326,16 @@ def test_search_memories( assert "para_mem" in result assert len(result["text_mem"]) == 1 assert result["text_mem"][0]["cube_id"] == "test_cube_1" - mock_mem_cube.text_mem.search.assert_called_once_with("football", top_k=5) + # Verify the search was called with the correct parameters + mock_mem_cube.text_mem.search.assert_called_once() + call_args = mock_mem_cube.text_mem.search.call_args + assert call_args[0] == ("football",) # positional args + assert call_args[1]["top_k"] == 5 + assert call_args[1]["mode"] == "fast" + assert call_args[1]["manual_close_internet"] + assert "info" in call_args[1] + assert call_args[1]["info"]["user_id"] == "test_user" + assert "session_id" in call_args[1]["info"] @patch("memos.mem_os.core.UserManager") @patch("memos.mem_os.core.MemReaderFactory")