feat: update memcube for api

Author: fridayL

src/memos/api/product_models.py

@@ -5,7 +5,7 @@
 from pydantic import BaseModel, Field
 
 # Import message types from core types module
-from memos.types import MessageDict
+from memos.types import MessageDict, PermissionDict
 
 
 T = TypeVar("T")
@@ -168,26 +168,28 @@ class APISearchRequest(BaseRequest):
     """Request model for searching memories."""
 
     query: str = Field(..., description="Search query")
+    user_id: str = Field(None, description="User ID")
     mem_cube_id: str | None = Field(None, description="Cube ID to search in")
+    mode: str = Field("fast", description="search mode fast or fine")
     internet_search: bool = Field(False, description="Whether to use internet search")
     moscube: bool = Field(False, description="Whether to use MemOSCube")
     top_k: int = Field(10, description="Number of results to return")
+    chat_history: list[MessageDict] | None = Field(None, description="Chat history")
     session_id: str | None = Field(None, description="Session ID for soft-filtering memories")
+    operation: list[PermissionDict] | None = Field(None, description="operation ids for multi cubes")
 
 
-class APIMemoryADDRequest(BaseRequest):
+class APIADDRequest(BaseRequest):
     """Request model for creating memories."""
-
+    user_id: str = Field(None, description="User ID")
     mem_cube_id: str = Field(..., description="Cube ID")
-    mode: str = Field("fast", description="search mode fast or fine")
     messages: list[MessageDict] | None = Field(None, description="List of messages to store.")
     memory_content: str | None = Field(None, description="Memory content to store")
     doc_path: str | None = Field(None, description="Path to document to store")
     source: str | None = Field(None, description="Source of the memory")
     chat_history: list[MessageDict] | None = Field(None, description="Chat history")
-    session_id: str | None = Field(
-        default_factory=lambda: str(uuid.uuid4()), description="Session id"
-    )
+    session_id: str | None = Field(None, description="Session id")
+    operation: list[PermissionDict] | None = Field(None, description="operation ids for multi cubes")
 
 
 class SuggestionRequest(BaseRequest):

src/memos/api/routers/server_router.py

@@ -1,91 +1,284 @@
+import os
+import time
+from typing import Any, Dict, List
+
+from fastapi import APIRouter
+
 from memos.api.config import APIConfig
 from memos.api.product_models import (
-    APIMemoryADDRequest,
+    APIADDRequest,
     MemoryResponse,
+    APISearchRequest,
     SearchResponse,
 )
+from memos.configs.embedder import EmbedderConfigFactory
+from memos.configs.graph_db import GraphDBConfigFactory
+from memos.configs.internet_retriever import InternetRetrieverConfigFactory
+from memos.configs.llm import LLMConfigFactory
+from memos.configs.mem_reader import MemReaderConfigFactory
+from memos.configs.reranker import RerankerConfigFactory
+from memos.embedders.factory import EmbedderFactory
+from memos.graph_dbs.factory import GraphStoreFactory
+from memos.llms.factory import LLMFactory
 from memos.log import get_logger
 from memos.mem_cube.navie import NaiveMemCube
 from memos.mem_reader.factory import MemReaderFactory
-from memos.types import MOSSearchResult
+from memos.memories.textual.tree_text_memory.organize.manager import MemoryManager
+from memos.memories.textual.tree_text_memory.retrieve.internet_retriever_factory import (
+    InternetRetrieverFactory,
+)
+from memos.reranker.factory import RerankerFactory
+from memos.types import MOSSearchResult, UserContext
 
 
 logger = get_logger(__name__)
 
+router = APIRouter(prefix="/server", tags=["Server API"])
+
+
+def _build_graph_db_config(user_id: str = "default") -> Dict[str, Any]:
+    """Build graph database configuration."""
+    graph_db_backend_map = {
+        "neo4j-community": APIConfig.get_neo4j_community_config(user_id=user_id),
+        "neo4j": APIConfig.get_neo4j_config(user_id=user_id),
+        "nebular": APIConfig.get_nebular_config(user_id=user_id),
+    }
+    
+    graph_db_backend = os.getenv("NEO4J_BACKEND", "nebular").lower()
+    return GraphDBConfigFactory.model_validate({
+        "backend": graph_db_backend,
+        "config": graph_db_backend_map[graph_db_backend],
+    })
+
+
+def _build_llm_config() -> Dict[str, Any]:
+    """Build LLM configuration."""
+    return LLMConfigFactory.model_validate({
+        "backend": "openai",
+        "config": APIConfig.get_openai_config(),
+    })
+
+
+def _build_embedder_config() -> Dict[str, Any]:
+    """Build embedder configuration."""
+    return EmbedderConfigFactory.model_validate(
+        APIConfig.get_embedder_config()
+    )
+
+
+def _build_mem_reader_config() -> Dict[str, Any]:
+    """Build memory reader configuration."""
+    return MemReaderConfigFactory.model_validate(
+        APIConfig.get_product_default_config()["mem_reader"]
+    )
+
+
+def _build_reranker_config() -> Dict[str, Any]:
+    """Build reranker configuration."""
+    return RerankerConfigFactory.model_validate(
+        APIConfig.get_reranker_config()
+    )
+
+
+def _build_internet_retriever_config() -> Dict[str, Any]:
+    """Build internet retriever configuration."""
+    return InternetRetrieverConfigFactory.model_validate(
+        APIConfig.get_internet_config()
+    )
 
-def init_mem_cube():
+
+def _get_default_memory_size(cube_config) -> Dict[str, int]:
+    """Get default memory size configuration."""
+    return getattr(cube_config.text_mem.config, "memory_size", None) or {
+        "WorkingMemory": 20,
+        "LongTermMemory": 1500,
+        "UserMemory": 480,
+    }
+
+
+def init_server():
+    """Initialize server components and configurations."""
+    # Get default cube configuration
     default_cube_config = APIConfig.get_default_cube_config()
-    mos_config = APIConfig.get_product_default_config()
-    mem_reader = MemReaderFactory.from_config(mos_config["mem_reader"])
-    naive_mem_cube = NaiveMemCube(default_cube_config)
-    return naive_mem_cube, mem_reader
+    
+    # Build component configurations
+    graph_db_config = _build_graph_db_config()
+    print(graph_db_config)
+    llm_config = _build_llm_config()
+    embedder_config = _build_embedder_config()
+    mem_reader_config = _build_mem_reader_config()
+    reranker_config = _build_reranker_config()
+    internet_retriever_config = _build_internet_retriever_config()
+    
+    # Create component instances
+    graph_db = GraphStoreFactory.from_config(graph_db_config)
+    llm = LLMFactory.from_config(llm_config)
+    embedder = EmbedderFactory.from_config(embedder_config)
+    mem_reader = MemReaderFactory.from_config(mem_reader_config)
+    reranker = RerankerFactory.from_config(reranker_config)
+    internet_retriever = InternetRetrieverFactory.from_config(internet_retriever_config, embedder=embedder)
+    
+    # Initialize memory manager
+    memory_manager = MemoryManager(
+        graph_db,
+        embedder,
+        llm,
+        memory_size=_get_default_memory_size(default_cube_config),
+        is_reorganize=getattr(default_cube_config.text_mem.config, "reorganize", False),
+    )
+    
+    return (
+        graph_db,
+        mem_reader,
+        llm,
+        embedder,
+        reranker,
+        internet_retriever,
+        memory_manager,
+        default_cube_config,
+    )
+
+
+# Initialize global components
+(
+    graph_db,
+    mem_reader,
+    llm,
+    embedder,
+    reranker,
+    internet_retriever,
+    memory_manager,
+    default_cube_config,
+) = init_server()
 
 
-naive_mem_cube, mem_reader = init_mem_cube()
+def _create_naive_mem_cube() -> NaiveMemCube:
+    """Create a NaiveMemCube instance with initialized components."""
+    naive_mem_cube = NaiveMemCube(
+        llm=llm,
+        embedder=embedder,
+        mem_reader=mem_reader,
+        graph_db=graph_db,
+        reranker=reranker,
+        internet_retriever=internet_retriever,
+        memory_manager=memory_manager,
+        default_cube_config=default_cube_config,
+    )
+    return naive_mem_cube
+
+
+def _format_memory_item(memory_data: Any) -> Dict[str, Any]:
+    """Format a single memory item for API response."""
+    memory = memory_data.model_dump()
+    memory_id = memory["id"]
+    ref_id = f"[{memory_id.split('-')[0]}]"
+    
+    memory["ref_id"] = ref_id
+    memory["metadata"]["embedding"] = []
+    memory["metadata"]["sources"] = []
+    memory["metadata"]["ref_id"] = ref_id
+    memory["metadata"]["id"] = memory_id
+    memory["metadata"]["memory"] = memory["memory"]
+    
+    return memory
 
 
 @router.post("/search", summary="Search memories", response_model=SearchResponse)
-def search_memories(search_req: APIMemoryADDRequest):
+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"
+    )
+    
     memories_result: MOSSearchResult = {
         "text_mem": [],
         "act_mem": [],
         "para_mem": [],
     }
-    search_filter = None
-    target_session_id = (
-        search_req.session_id if search_req.session_id is not None else "default_session"
-    )
-    if search_req.session_id is not None:
-        search_filter = {"session_id": search_req.session_id}
-    memories_list = naive_mem_cube.search(
+    target_session_id = search_req.session_id
+    if not target_session_id:
+        target_session_id = "default_session"
+    search_filter = {"session_id": search_req.session_id} if search_req.session_id else None
+    
+    # Create MemCube and perform search
+    naive_mem_cube = _create_naive_mem_cube()
+    search_results = naive_mem_cube.text_mem.search(
         query=search_req.query,
-        user_id=search_req.mem_cube_id,
+        user_name=search_req.mem_cube_id,
         top_k=search_req.top_k,
         mode=search_req.mode,
         internet_search=not search_req.internet_search,
         moscube=search_req.moscube,
         search_filter=search_filter,
         info={
-            "user_id": search_req.mem_cube_id,
+            "user_id": search_req.user_id,
             "session_id": target_session_id,
             "chat_history": search_req.chat_history,
         },
     )
-    memories_list = []
-    for data in memories_list:
-        memories = data.model_dump()
-        memories["ref_id"] = f"[{memories['id'].split('-')[0]}]"
-        memories["metadata"]["embedding"] = []
-        memories["metadata"]["sources"] = []
-        memories["metadata"]["ref_id"] = f"[{memories['id'].split('-')[0]}]"
-        memories["metadata"]["id"] = memories["id"]
-        memories["metadata"]["memory"] = memories["memory"]
-        memories_list.append(memories)
-    memories_result["text_mem"].append(
-        {"cube_id": search_req.mem_cube_id, "memories": memories_list}
+    formatted_memories = [_format_memory_item(data) for data in search_results]
+    
+    memories_result["text_mem"].append({
+        "cube_id": search_req.mem_cube_id,
+        "memories": formatted_memories,
+    })
+    
+    return SearchResponse(
+        message="Search completed successfully",
+        data=memories_result,
     )
-    return SearchResponse(message="Search completed successfully", data=memories_result)
 
 
-@router.post("/add", summary="add memories", response_model=MemoryResponse)
-def add_memories(add_req: APIMemoryADDRequest):
+@router.post("/add", summary="Add memories", response_model=MemoryResponse)
+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()
+    naive_mem_cube = _create_naive_mem_cube()
+    target_session_id = add_req.session_id
+    if not target_session_id:
+        target_session_id = "default_session"
     memories = mem_reader.get_memory(
         [add_req.messages],
         type="chat",
-        info={"user_id": add_req.user_id, "session_id": add_req.session_id},
+        info={
+            "user_id": add_req.user_id,
+            "session_id": target_session_id,
+        },
     )
-    memories = [mm for m in memories for mm in m]
+    
+    # Flatten memory list
+    flattened_memories = [mm for m in memories for mm in m]
+    
+    elapsed_time = time.time() - time_start
     logger.info(
-        f"time add: get mem_reader time user_id: {add_req.user_id} time is: {time.time() - time_start:.2f}s"
+        f"Memory extraction completed for user {add_req.user_id} in {elapsed_time:.2f}s"
     )
-    mem_id_list: list[str] = memory_add_obj.add(memories, user_name=add_req.user_id)
+    mem_id_list: List[str] = naive_mem_cube.text_mem.add(
+        flattened_memories,
+        user_name=add_req.mem_cube_id,
+    )
+    
     logger.info(
-        f"Added memory for user {add_req.user_id} in session {add_req.session_id}: {mem_id_list}"
+        f"Added {len(mem_id_list)} memories for user {add_req.user_id} "
+        f"in session {add_req.session_id}: {mem_id_list}"
+    )
+    response_data = [
+        {
+            "memory": memory.memory,
+            "memory_id": memory_id,
+            "memory_type": memory.metadata.memory_type,
+        }
+        for memory_id, memory in zip(mem_id_list, flattened_memories, strict=False)
+    ]
+    return MemoryResponse(
+        message="Memory added successfully",
+        data=response_data,
     )
-    data = []
-    for m_id, m in zip(mem_id_list, memories, strict=False):
-        data.append({"memory": m.memory, "memory_id": m_id, "memory_type": m.metadata.memory_type})
-    return MemoryResponse(message="Memory added successfully", data=data)

src/memos/configs/mem_user.py

@@ -30,6 +30,16 @@ class MySQLUserManagerConfig(BaseUserManagerConfig):
     database: str = Field(default="memos_users", description="MySQL database name")
     charset: str = Field(default="utf8mb4", description="MySQL charset")
 
+class RedisUserManagerConfig(BaseUserManagerConfig):
+    """Redis user manager configuration."""
+
+    host: str = Field(default="localhost", description="Redis server host")
+    port: int = Field(default=6379, description="Redis server port")
+    username: str = Field(default="root", description="Redis username")
+    password: str = Field(default="", description="Redis password")
+    database: str = Field(default="memos_users", description="Redis database name")
+    charset: str = Field(default="utf8mb4", description="Redis charset")
+
 
 class UserManagerConfigFactory(BaseModel):
     """Factory for user manager configurations."""
@@ -42,6 +52,7 @@ class UserManagerConfigFactory(BaseModel):
     backend_to_class: ClassVar[dict[str, Any]] = {
         "sqlite": SQLiteUserManagerConfig,
         "mysql": MySQLUserManagerConfig,
+        "redis": RedisUserManagerConfig,
     }
 
     @field_validator("backend")