merge api refactor to dev

Author: yuan.wang

src/memos/api/handlers/chat_handler.py

@@ -11,6 +11,7 @@
 
 from memos.api.config import APIConfig
 from memos.api.handlers.config_builders import (
+    build_chat_llm_config,
     build_embedder_config,
     build_graph_db_config,
     build_internet_retriever_config,
@@ -77,6 +78,38 @@ def _get_default_memory_size(cube_config: Any) -> dict[str, int]:
     }
 
 
+def _init_chat_llms(chat_llm_configs: list[dict]) -> dict[str, Any]:
+    """
+    Initialize chat language models from configuration.
+
+    Args:
+        chat_llm_configs: List of chat LLM configuration dictionaries
+
+    Returns:
+        Dictionary mapping model names to initialized LLM instances
+    """
+
+    def _list_models(client):
+        try:
+            models = (
+                [model.id for model in client.models.list().data]
+                if client.models.list().data
+                else client.models.list().models
+            )
+        except Exception as e:
+            logger.error(f"Error listing models: {e}")
+            models = []
+        return models
+
+    model_name_instrance_maping = {}
+    for cfg in chat_llm_configs:
+        llm = LLMFactory.from_config(cfg["config_class"])
+        if cfg["support_models"]:
+            for model_name in cfg["support_models"]:
+                model_name_instrance_maping[model_name] = llm
+    return model_name_instrance_maping
+
+
 def init_server() -> dict[str, Any]:
     """
     Initialize all server components and configurations.
@@ -104,6 +137,7 @@ def init_server() -> dict[str, Any]:
     # Build component configurations
     graph_db_config = build_graph_db_config()
     llm_config = build_llm_config()
+    chat_llm_config = build_chat_llm_config()
     embedder_config = build_embedder_config()
     mem_reader_config = build_mem_reader_config()
     reranker_config = build_reranker_config()
@@ -123,13 +157,16 @@ def init_server() -> dict[str, Any]:
         else None
     )
     llm = LLMFactory.from_config(llm_config)
+    chat_llms = _init_chat_llms(chat_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 chat llms
+
     logger.debug("Core components instantiated")
 
     # Initialize memory manager
@@ -276,6 +313,7 @@ def init_server() -> dict[str, Any]:
         "graph_db": graph_db,
         "mem_reader": mem_reader,
         "llm": llm,
+        "chat_llms": chat_llms,
         "embedder": embedder,
         "reranker": reranker,
         "internet_retriever": internet_retriever,

src/memos/api/handlers/component_init.py

@@ -6,6 +6,7 @@
 a configuration dictionary using the appropriate ConfigFactory.
 """
 
+import json
 import os
 
 from typing import Any
@@ -81,6 +82,32 @@ def build_llm_config() -> dict[str, Any]:
     )
 
 
+def build_chat_llm_config() -> list[dict[str, Any]]:
+    """
+    Build chat LLM configuration.
+
+    Returns:
+        Validated chat LLM configuration dictionary
+    """
+    configs = json.loads(os.getenv("CHAT_MODEL_LIST"))
+    return [
+        {
+            "config_class": LLMConfigFactory.model_validate(
+                {
+                    "backend": cfg.get("backend", "openai"),
+                    "config": (
+                        {k: v for k, v in cfg.items() if k not in ["backend", "support_models"]}
+                    )
+                    if cfg
+                    else APIConfig.get_openai_config(),
+                }
+            ),
+            "support_models": cfg.get("support_models", None),
+        }
+        for cfg in configs
+    ]
+
+
 def build_embedder_config() -> dict[str, Any]:
     """
     Build embedder configuration.

src/memos/api/handlers/config_builders.py

@@ -6,7 +6,14 @@
 
 from typing import Any, Literal
 
-from memos.api.product_models import MemoryResponse
+from memos.api.handlers.formatters_handler import format_memory_item
+from memos.api.product_models import (
+    DeleteMemoryRequest,
+    DeleteMemoryResponse,
+    GetMemoryRequest,
+    GetMemoryResponse,
+    MemoryResponse,
+)
 from memos.log import get_logger
 from memos.mem_os.utils.format_utils import (
     convert_graph_to_tree_forworkmem,
@@ -149,3 +156,37 @@ def handle_get_subgraph(
     except Exception as e:
         logger.error(f"Failed to get subgraph: {e}", exc_info=True)
         raise
+
+
+def handle_get_memories(get_mem_req: GetMemoryRequest, naive_mem_cube: Any) -> GetMemoryResponse:
+    # TODO: Implement get memory with filter
+    memories = naive_mem_cube.text_mem.get_all(user_name=get_mem_req.mem_cube_id)["nodes"]
+    filter_params: dict[str, Any] = {}
+    if get_mem_req.user_id is not None:
+        filter_params["user_id"] = get_mem_req.user_id
+    if get_mem_req.mem_cube_id is not None:
+        filter_params["mem_cube_id"] = get_mem_req.mem_cube_id
+    preferences = naive_mem_cube.pref_mem.get_memory_by_filter(filter_params)
+    return GetMemoryResponse(
+        message="Memories retrieved successfully",
+        data={
+            "text_mem": memories,
+            "pref_mem": [format_memory_item(mem) for mem in preferences],
+        },
+    )
+
+
+def handle_delete_memories(delete_mem_req: DeleteMemoryRequest, naive_mem_cube: Any):
+    try:
+        naive_mem_cube.text_mem.delete(delete_mem_req.memory_ids)
+        naive_mem_cube.pref_mem.delete(delete_mem_req.memory_ids)
+    except Exception as e:
+        logger.error(f"Failed to delete memories: {e}", exc_info=True)
+        return DeleteMemoryResponse(
+            message="Failed to delete memories",
+            data="failure",
+        )
+    return DeleteMemoryResponse(
+        message="Memories deleted successfully",
+        data={"status": "success"},
+    )

src/memos/api/handlers/memory_handler.py

@@ -22,7 +22,7 @@
 
 
 def handle_scheduler_status(
-    user_name: str | None = None,
+    mem_cube_id: str | None = None,
     mem_scheduler: Any | None = None,
     instance_id: str = "",
 ) -> dict[str, Any]:
@@ -43,17 +43,17 @@ def handle_scheduler_status(
         HTTPException: If status retrieval fails
     """
     try:
-        if user_name:
+        if mem_cube_id:
             running = mem_scheduler.dispatcher.get_running_tasks(
-                lambda task: getattr(task, "mem_cube_id", None) == user_name
+                lambda task: getattr(task, "mem_cube_id", None) == mem_cube_id
             )
             tasks_iter = to_iter(running)
             running_count = len(tasks_iter)
             return {
                 "message": "ok",
                 "data": {
                     "scope": "user",
-                    "user_name": user_name,
+                    "mem_cube_id": mem_cube_id,
                     "running_tasks": running_count,
                     "timestamp": time.time(),
                     "instance_id": instance_id,

src/memos/api/handlers/scheduler_handler.py

@@ -1,7 +1,6 @@
-import os
 import uuid
 
-from typing import Generic, Literal, TypeVar
+from typing import Any, Generic, Literal, TypeVar
 
 from pydantic import BaseModel, Field
 
@@ -37,7 +36,7 @@ class UserRegisterRequest(BaseRequest):
     interests: str | None = Field(None, description="User interests")
 
 
-class GetMemoryRequest(BaseRequest):
+class GetMemoryPlaygroundRequest(BaseRequest):
     """Request model for getting memories."""
 
     user_id: str = Field(..., description="User ID")
@@ -80,9 +79,20 @@ class ChatRequest(BaseRequest):
         None, description="List of cube IDs user can write for multi-cube chat"
     )
     history: list[MessageDict] | None = Field(None, description="Chat history")
+    mode: SearchMode = Field(SearchMode.FAST, description="search mode: fast, fine, or mixture")
     internet_search: bool = Field(True, description="Whether to use internet search")
-    moscube: bool = Field(False, description="Whether to use MemOSCube")
+    system_prompt: str | None = Field(None, description="Base system prompt to use for chat")
+    top_k: int = Field(10, description="Number of results to return")
+    threshold: float = Field(0.5, description="Threshold for filtering references")
     session_id: str | None = Field(None, description="Session ID for soft-filtering memories")
+    include_preference: bool = Field(True, description="Whether to handle preference memory")
+    pref_top_k: int = Field(6, description="Number of preference results to return")
+    filter: dict[str, Any] | None = Field(None, description="Filter for the memory")
+    model_name_or_path: str | None = Field(None, description="Model name to use for chat")
+    max_tokens: int | None = Field(None, description="Max tokens to generate")
+    temperature: float | None = Field(None, description="Temperature for sampling")
+    top_p: float | None = Field(None, description="Top-p (nucleus) sampling parameter")
+    add_message_on_answer: bool = Field(True, description="Add dialogs to memory after chat")
 
 
 class ChatCompleteRequest(BaseRequest):
@@ -93,11 +103,18 @@ class ChatCompleteRequest(BaseRequest):
     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(False, description="Whether to use internet search")
-    moscube: bool = Field(False, description="Whether to use MemOSCube")
-    base_prompt: str | None = Field(None, description="Base prompt to use for chat")
+    system_prompt: str | None = Field(None, description="Base prompt to use for chat")
     top_k: int = Field(10, description="Number of results to return")
     threshold: float = Field(0.5, description="Threshold for filtering references")
     session_id: str | None = Field(None, description="Session ID for soft-filtering memories")
+    include_preference: bool = Field(True, description="Whether to handle preference memory")
+    pref_top_k: int = Field(6, description="Number of preference results to return")
+    filter: dict[str, Any] | None = Field(None, description="Filter for the memory")
+    model_name_or_path: str | None = Field(None, description="Model name to use for chat")
+    max_tokens: int | None = Field(None, description="Max tokens to generate")
+    temperature: float | None = Field(None, description="Temperature for sampling")
+    top_p: float | None = Field(None, description="Top-p (nucleus) sampling parameter")
+    add_message_on_answer: bool = Field(True, description="Add dialogs to memory after chat")
 
 
 class UserCreate(BaseRequest):
@@ -129,6 +146,10 @@ class SuggestionResponse(BaseResponse[list]):
     data: dict[str, list[str]] | None = Field(None, description="Response data")
 
 
+class AddStatusResponse(BaseResponse[dict]):
+    """Response model for add status operations."""
+
+
 class ConfigResponse(BaseResponse[None]):
     """Response model for configuration endpoint."""
 
@@ -141,6 +162,14 @@ class ChatResponse(BaseResponse[str]):
     """Response model for chat operations."""
 
 
+class GetMemoryResponse(BaseResponse[dict]):
+    """Response model for getting memories."""
+
+
+class DeleteMemoryResponse(BaseResponse[dict]):
+    """Response model for deleting memories."""
+
+
 class UserResponse(BaseResponse[dict]):
     """Response model for user operations."""
 
@@ -181,11 +210,8 @@ class APISearchRequest(BaseRequest):
     readable_cube_ids: list[str] | None = Field(
         None, description="List of cube IDs user can read for multi-cube search"
     )
-    mode: SearchMode = Field(
-        os.getenv("SEARCH_MODE", SearchMode.FAST), description="search mode: fast, fine, or mixture"
-    )
+    mode: SearchMode = Field(SearchMode.FAST, description="search mode: fast, fine, or mixture")
     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")
@@ -194,6 +220,7 @@ class APISearchRequest(BaseRequest):
     )
     include_preference: bool = Field(True, description="Whether to handle preference memory")
     pref_top_k: int = Field(6, description="Number of preference results to return")
+    filter: dict[str, Any] | None = Field(None, description="Filter for the memory")
 
 
 class APIADDRequest(BaseRequest):
@@ -213,8 +240,13 @@ class APIADDRequest(BaseRequest):
     operation: list[PermissionDict] | None = Field(
         None, description="operation ids for multi cubes"
     )
-    async_mode: Literal["async", "sync"] | None = Field(
-        None, description="Whether to add memory in async mode"
+    async_mode: Literal["async", "sync"] = Field(
+        "async", description="Whether to add memory in async mode"
+    )
+    custom_tags: list[str] | None = Field(None, description="Custom tags for the memory")
+    info: dict[str, str] | None = Field(None, description="Additional information for the memory")
+    is_feedback: bool = Field(
+        False, description="Whether the user feedback in knowladge base service"
     )
 
 
@@ -232,13 +264,43 @@ class APIChatCompleteRequest(BaseRequest):
     )
     history: list[MessageDict] | None = Field(None, description="Chat history")
     internet_search: bool = Field(False, description="Whether to use internet search")
-    moscube: bool = Field(True, description="Whether to use MemOSCube")
-    base_prompt: str | None = Field(None, description="Base prompt to use for chat")
+    system_prompt: str | None = Field(None, description="Base system prompt to use for chat")
+    mode: SearchMode = Field(SearchMode.FAST, description="search mode: fast, fine, or mixture")
     top_k: int = Field(10, description="Number of results to return")
     threshold: float = Field(0.5, description="Threshold for filtering references")
     session_id: str | None = Field(
         "default_session", description="Session ID for soft-filtering memories"
     )
+    include_preference: bool = Field(True, description="Whether to handle preference memory")
+    pref_top_k: int = Field(6, description="Number of preference results to return")
+    filter: dict[str, Any] | None = Field(None, description="Filter for the memory")
+    model_name_or_path: str | None = Field(None, description="Model name to use for chat")
+    max_tokens: int | None = Field(None, description="Max tokens to generate")
+    temperature: float | None = Field(None, description="Temperature for sampling")
+    top_p: float | None = Field(None, description="Top-p (nucleus) sampling parameter")
+    add_message_on_answer: bool = Field(True, description="Add dialogs to memory after chat")
+
+
+class AddStatusRequest(BaseRequest):
+    """Request model for checking add status."""
+
+    mem_cube_id: str = Field(..., description="Cube ID")
+    user_id: str | None = Field(None, description="User ID")
+    session_id: str | None = Field(None, description="Session ID")
+
+
+class GetMemoryRequest(BaseRequest):
+    """Request model for getting memories."""
+
+    mem_cube_id: str = Field(..., description="Cube ID")
+    user_id: str | None = Field(None, description="User ID")
+    include_preference: bool = Field(True, description="Whether to handle preference memory")
+
+
+class DeleteMemoryRequest(BaseRequest):
+    """Request model for deleting memories."""
+
+    memory_ids: list[str] = Field(..., description="Memory IDs")
 
 
 class SuggestionRequest(BaseRequest):

src/memos/api/product_models.py

@@ -10,7 +10,7 @@
     BaseResponse,
     ChatCompleteRequest,
     ChatRequest,
-    GetMemoryRequest,
+    GetMemoryPlaygroundRequest,
     MemoryCreateRequest,
     MemoryResponse,
     SearchRequest,
@@ -159,7 +159,7 @@ def get_suggestion_queries_post(suggestion_req: SuggestionRequest):
 
 
 @router.post("/get_all", summary="Get all memories for user", response_model=MemoryResponse)
-def get_all_memories(memory_req: GetMemoryRequest):
+def get_all_memories(memory_req: GetMemoryPlaygroundRequest):
     """Get all memories for a specific user."""
     try:
         mos_product = get_mos_product_instance()