fix: fix code format

Author: fridayL

examples/mem_api/pipeline_test.py

@@ -5,8 +5,6 @@
 uvicorn memos.api.server_api:app --host 0.0.0.0 --port 8002 --workers 4
 """
 
-import time
-
 from typing import Any
 
 from dotenv import load_dotenv
@@ -45,8 +43,6 @@ def test_add_memories(
     """
     logger.info(f"Testing add memories for user: {user_id}, mem_cube: {mem_cube_id}")
 
-    time_start = time.time()
-
     # Create NaiveMemCube using server_router function
     naive_mem_cube = _create_naive_mem_cube()
 
@@ -63,8 +59,6 @@ def test_add_memories(
     # Flatten memory list
     flattened_memories = [mm for m in memories for mm in m]
 
-    elapsed_time = time.time() - time_start
-
     # Add memories to the system
     mem_id_list: list[str] = naive_mem_cube.text_mem.add(
         flattened_memories,
@@ -117,8 +111,6 @@ def test_search_memories(
     # Prepare search filter
     search_filter = {"session_id": session_id} if session_id != "default_session" else None
 
-    # Perform search
-    time_start = time.time()
     search_results = naive_mem_cube.text_mem.search(
         query=query,
         user_name=mem_cube_id,

src/memos/api/routers/server_router.py

@@ -1,5 +1,4 @@
 import os
-import time
 
 from typing import Any
 
@@ -189,9 +188,11 @@ def search_memories(search_req: APISearchRequest):
     """Search memories for a specific user."""
     # Create UserContext object - how to assign values
     user_context = UserContext(
-        user_id=search_req.user_id, session_id=search_req.session_id or "default_session"
+        user_id=search_req.user_id,
+        mem_cube_id=search_req.mem_cube_id,
+        session_id=search_req.session_id or "default_session",
     )
-
+    logger.info(f"Search user_id is: {user_context.mem_cube_id}")
     memories_result: MOSSearchResult = {
         "text_mem": [],
         "act_mem": [],
@@ -206,7 +207,7 @@ def search_memories(search_req: APISearchRequest):
     naive_mem_cube = _create_naive_mem_cube()
     search_results = naive_mem_cube.text_mem.search(
         query=search_req.query,
-        user_name=search_req.mem_cube_id,
+        user_name=user_context.mem_cube_id,
         top_k=search_req.top_k,
         mode=search_req.mode,
         manual_close_internet=not search_req.internet_search,
@@ -237,9 +238,11 @@ def search_memories(search_req: APISearchRequest):
 def add_memories(add_req: APIADDRequest):
     """Add memories for a specific user."""
     # Create UserContext object - how to assign values
-    user_context = UserContext(user_id=add_req.user_id, session_id=add_req.session_id)
-
-    time_start = time.time()
+    user_context = UserContext(
+        user_id=add_req.user_id,
+        mem_cube_id=add_req.mem_cube_id,
+        session_id=add_req.session_id or "default_session",
+    )
     naive_mem_cube = _create_naive_mem_cube()
     target_session_id = add_req.session_id
     if not target_session_id:
@@ -255,12 +258,10 @@ def add_memories(add_req: APIADDRequest):
 
     # Flatten memory list
     flattened_memories = [mm for m in memories for mm in m]
-
-    elapsed_time = time.time() - time_start
-    logger.info(f"Memory extraction completed for user {add_req.user_id} in {elapsed_time:.2f}s")
+    logger.info(f"Memory extraction completed for user {add_req.user_id}")
     mem_id_list: list[str] = naive_mem_cube.text_mem.add(
         flattened_memories,
-        user_name=add_req.mem_cube_id,
+        user_name=user_context.mem_cube_id,
     )
 
     logger.info(

src/memos/graph_dbs/nebular.py

@@ -1330,7 +1330,6 @@ def get_all_memory_items(
 
         where_clause = f"WHERE n.memory_type = '{scope}'"
         where_clause += f" AND n.user_name = '{user_name}'"
-        # where_clause = f"WHERE n.user_name = '{user_name}'"
 
         return_fields = self._build_return_fields(include_embedding)
 

src/memos/mem_user/redis_persistent_user_manager.py

@@ -5,6 +5,7 @@
 
 import json
 
+from memos.configs.mem_os import MOSConfig
 from memos.dependency import require_python_package
 from memos.log import get_logger
 
@@ -103,7 +104,7 @@ def save_user_config(self, user_id: str, config: MOSConfig) -> bool:
             logger.error(f"Error saving configuration for user {user_id}: {e}")
             return False
 
-    def get_user_config(self, user_id: str) -> MOSConfig | None:
+    def get_user_config(self, user_id: str) -> dict | None:
         """Get user configuration from Redis (search interface).
 
         Args:
@@ -176,15 +177,15 @@ def exists_user_config(self, user_id: str) -> bool:
 
     def list_user_configs(
         self, pattern: str = "user_config:*", count: int = 100
-    ) -> dict[str, MOSConfig]:
+    ) -> dict[str, dict]:
         """List all user configurations.
 
         Args:
             pattern (str): Redis key matching pattern. Defaults to "user_config:*".
             count (int): Number of keys to return per scan. Defaults to 100.
 
         Returns:
-            dict[str, MOSConfig]: Dictionary mapping user_id to MOSConfig objects.
+            dict[str, dict]: Dictionary mapping user_id to dict objects.
         """
         result = {}
         try:

src/memos/memories/textual/simple_tree.py

@@ -1,15 +1,12 @@
 import time
 
 from datetime import datetime
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from memos.configs.memory import TreeTextMemoryConfig
 from memos.embedders.base import BaseEmbedder
-from memos.embedders.factory import OllamaEmbedder
 from memos.graph_dbs.base import BaseGraphDB
-from memos.graph_dbs.factory import Neo4jGraphDB
 from memos.llms.base import BaseLLM
-from memos.llms.factory import AzureLLM, OllamaLLM, OpenAILLM
 from memos.log import get_logger
 from memos.mem_reader.base import BaseMemReader
 from memos.memories.textual.item import TextualMemoryItem, TreeNodeTextualMemoryMetadata
@@ -20,6 +17,12 @@
 from memos.types import MessageList
 
 
+if TYPE_CHECKING:
+    from memos.embedders.factory import OllamaEmbedder
+    from memos.graph_dbs.factory import Neo4jGraphDB
+    from memos.llms.factory import AzureLLM, OllamaLLM, OpenAILLM
+
+
 logger = get_logger(__name__)
 
 

src/memos/types.py

@@ -74,6 +74,7 @@ class PermissionDict(TypedDict, total=False):
 class UserContext(BaseModel):
     """Model to represent user context."""
 
-    user_id: str
+    user_id: str | None = None
+    mem_cube_id: str | None = None
     session_id: str | None = None
     operation: list[PermissionDict] | None = None