MemoryOS 2.0.3__py3-none-any.whl

memos/api/routers/server_router.py

@@ -0,0 +1,394 @@
+"""
+Server API Router for MemOS (Class-based handlers version).
+
+This router demonstrates the improved architecture using class-based handlers
+with dependency injection, providing better modularity and maintainability.
+
+Comparison with function-based approach:
+- Cleaner code: No need to pass dependencies in every endpoint
+- Better testability: Easy to mock handler dependencies
+- Improved extensibility: Add new handlers or modify existing ones easily
+- Clear separation of concerns: Router focuses on routing, handlers handle business logic
+"""
+
+import os
+import random as _random
+import socket
+
+from fastapi import APIRouter, HTTPException, Query
+
+from memos.api import handlers
+from memos.api.handlers.add_handler import AddHandler
+from memos.api.handlers.base_handler import HandlerDependencies
+from memos.api.handlers.chat_handler import ChatHandler
+from memos.api.handlers.feedback_handler import FeedbackHandler
+from memos.api.handlers.search_handler import SearchHandler
+from memos.api.product_models import (
+    AllStatusResponse,
+    APIADDRequest,
+    APIChatCompleteRequest,
+    APIFeedbackRequest,
+    APISearchRequest,
+    ChatPlaygroundRequest,
+    ChatRequest,
+    DeleteMemoryRequest,
+    DeleteMemoryResponse,
+    ExistMemCubeIdRequest,
+    ExistMemCubeIdResponse,
+    GetMemoryPlaygroundRequest,
+    GetMemoryRequest,
+    GetMemoryResponse,
+    GetUserNamesByMemoryIdsRequest,
+    GetUserNamesByMemoryIdsResponse,
+    MemoryResponse,
+    SearchResponse,
+    StatusResponse,
+    SuggestionRequest,
+    SuggestionResponse,
+    TaskQueueResponse,
+)
+from memos.graph_dbs.polardb import PolarDBGraphDB
+from memos.log import get_logger
+from memos.mem_scheduler.base_scheduler import BaseScheduler
+from memos.mem_scheduler.utils.status_tracker import TaskStatusTracker
+
+
+logger = get_logger(__name__)
+
+router = APIRouter(prefix="/product", tags=["Server API"])
+
+# Instance ID for identifying this server instance in logs and responses
+INSTANCE_ID = f"{socket.gethostname()}:{os.getpid()}:{_random.randint(1000, 9999)}"
+
+# Initialize all server components
+components = handlers.init_server()
+
+# Create dependency container
+dependencies = HandlerDependencies.from_init_server(components)
+
+# Initialize all handlers with dependency injection
+search_handler = SearchHandler(dependencies)
+add_handler = AddHandler(dependencies)
+chat_handler = (
+    ChatHandler(
+        dependencies,
+        components["chat_llms"],
+        search_handler,
+        add_handler,
+        online_bot=components.get("online_bot"),
+    )
+    if os.getenv("ENABLE_CHAT_API", "false") == "true"
+    else None
+)
+feedback_handler = FeedbackHandler(dependencies)
+# Extract commonly used components for function-based handlers
+# (These can be accessed from the components dict without unpacking all of them)
+mem_scheduler: BaseScheduler = components["mem_scheduler"]
+llm = components["llm"]
+naive_mem_cube = components["naive_mem_cube"]
+redis_client = components["redis_client"]
+status_tracker = TaskStatusTracker(redis_client=redis_client)
+graph_db = components["graph_db"]
+vector_db = components["vector_db"]
+
+
+# =============================================================================
+# Search API Endpoints
+# =============================================================================
+
+
+@router.post("/search", summary="Search memories", response_model=SearchResponse)
+def search_memories(search_req: APISearchRequest):
+    """
+    Search memories for a specific user.
+
+    This endpoint uses the class-based SearchHandler for better code organization.
+    """
+    search_results = search_handler.handle_search_memories(search_req)
+    return search_results
+
+
+# =============================================================================
+# Add API Endpoints
+# =============================================================================
+
+
+@router.post("/add", summary="Add memories", response_model=MemoryResponse)
+def add_memories(add_req: APIADDRequest):
+    """
+    Add memories for a specific user.
+
+    This endpoint uses the class-based AddHandler for better code organization.
+    """
+    return add_handler.handle_add_memories(add_req)
+
+
+# =============================================================================
+# Scheduler API Endpoints
+# =============================================================================
+
+
+@router.get(  # Changed from post to get
+    "/scheduler/allstatus",
+    summary="Get detailed scheduler status",
+    response_model=AllStatusResponse,
+)
+def scheduler_allstatus():
+    """Get detailed scheduler status including running tasks and queue metrics."""
+    return handlers.scheduler_handler.handle_scheduler_allstatus(
+        mem_scheduler=mem_scheduler, status_tracker=status_tracker
+    )
+
+
+@router.get(  # Changed from post to get
+    "/scheduler/status", summary="Get scheduler running status", response_model=StatusResponse
+)
+def scheduler_status(
+    user_id: str = Query(..., description="User ID"),
+    task_id: str | None = Query(None, description="Optional Task ID to query a specific task"),
+):
+    """Get scheduler running status."""
+    return handlers.scheduler_handler.handle_scheduler_status(
+        user_id=user_id,
+        task_id=task_id,
+        status_tracker=status_tracker,
+    )
+
+
+@router.get(  # Changed from post to get
+    "/scheduler/task_queue_status",
+    summary="Get scheduler task queue status",
+    response_model=TaskQueueResponse,
+)
+def scheduler_task_queue_status(
+    user_id: str = Query(..., description="User ID whose queue status is requested"),
+):
+    """Get scheduler task queue backlog/pending status for a user."""
+    return handlers.scheduler_handler.handle_task_queue_status(
+        user_id=user_id, mem_scheduler=mem_scheduler
+    )
+
+
+@router.post("/scheduler/wait", summary="Wait until scheduler is idle for a specific user")
+def scheduler_wait(
+    user_name: str,
+    timeout_seconds: float = 120.0,
+    poll_interval: float = 0.5,
+):
+    """Wait until scheduler is idle for a specific user."""
+    return handlers.scheduler_handler.handle_scheduler_wait(
+        user_name=user_name,
+        status_tracker=status_tracker,
+        timeout_seconds=timeout_seconds,
+        poll_interval=poll_interval,
+    )
+
+
+@router.get("/scheduler/wait/stream", summary="Stream scheduler progress for a user")
+def scheduler_wait_stream(
+    user_name: str,
+    timeout_seconds: float = 120.0,
+    poll_interval: float = 0.5,
+):
+    """Stream scheduler progress via Server-Sent Events (SSE)."""
+    return handlers.scheduler_handler.handle_scheduler_wait_stream(
+        user_name=user_name,
+        status_tracker=status_tracker,
+        timeout_seconds=timeout_seconds,
+        poll_interval=poll_interval,
+        instance_id=INSTANCE_ID,
+    )
+
+
+# =============================================================================
+# Chat API Endpoints
+# =============================================================================
+
+
+@router.post("/chat/complete", summary="Chat with MemOS (Complete Response)")
+def chat_complete(chat_req: APIChatCompleteRequest):
+    """
+    Chat with MemOS for a specific user. Returns complete response (non-streaming).
+
+    This endpoint uses the class-based ChatHandler.
+    """
+    if chat_handler is None:
+        raise HTTPException(
+            status_code=503, detail="Chat service is not available. Chat handler not initialized."
+        )
+    return chat_handler.handle_chat_complete(chat_req)
+
+
+@router.post("/chat/stream", summary="Chat with MemOS")
+def chat_stream(chat_req: ChatRequest):
+    """
+    Chat with MemOS for a specific user. Returns SSE stream.
+
+    This endpoint uses the class-based ChatHandler which internally
+    composes SearchHandler and AddHandler for a clean architecture.
+    """
+    if chat_handler is None:
+        raise HTTPException(
+            status_code=503, detail="Chat service is not available. Chat handler not initialized."
+        )
+    return chat_handler.handle_chat_stream(chat_req)
+
+
+@router.post("/chat/stream/playground", summary="Chat with MemOS playground")
+def chat_stream_playground(chat_req: ChatPlaygroundRequest):
+    """
+    Chat with MemOS for a specific user. Returns SSE stream.
+
+    This endpoint uses the class-based ChatHandler which internally
+    composes SearchHandler and AddHandler for a clean architecture.
+    """
+    if chat_handler is None:
+        raise HTTPException(
+            status_code=503, detail="Chat service is not available. Chat handler not initialized."
+        )
+    return chat_handler.handle_chat_stream_playground(chat_req)
+
+
+# =============================================================================
+# Suggestion API Endpoints
+# =============================================================================
+
+
+@router.post(
+    "/suggestions",
+    summary="Get suggestion queries",
+    response_model=SuggestionResponse,
+)
+def get_suggestion_queries(suggestion_req: SuggestionRequest):
+    """Get suggestion queries for a specific user with language preference."""
+    return handlers.suggestion_handler.handle_get_suggestion_queries(
+        user_id=suggestion_req.mem_cube_id,
+        language=suggestion_req.language,
+        message=suggestion_req.message,
+        llm=llm,
+        naive_mem_cube=naive_mem_cube,
+    )
+
+
+# =============================================================================
+# Memory Retrieval Delete API Endpoints
+# =============================================================================
+
+
+@router.post("/get_all", summary="Get all memories for user", response_model=MemoryResponse)
+def get_all_memories(memory_req: GetMemoryPlaygroundRequest):
+    """
+    Get all memories or subgraph for a specific user.
+
+    If search_query is provided, returns a subgraph based on the query.
+    Otherwise, returns all memories of the specified type.
+    """
+    if memory_req.search_query:
+        return handlers.memory_handler.handle_get_subgraph(
+            user_id=memory_req.user_id,
+            mem_cube_id=(
+                memory_req.mem_cube_ids[0] if memory_req.mem_cube_ids else memory_req.user_id
+            ),
+            query=memory_req.search_query,
+            top_k=20,
+            naive_mem_cube=naive_mem_cube,
+        )
+    else:
+        return handlers.memory_handler.handle_get_all_memories(
+            user_id=memory_req.user_id,
+            mem_cube_id=(
+                memory_req.mem_cube_ids[0] if memory_req.mem_cube_ids else memory_req.user_id
+            ),
+            memory_type=memory_req.memory_type or "text_mem",
+            naive_mem_cube=naive_mem_cube,
+        )
+
+
+@router.post("/get_memory", summary="Get memories for user", response_model=GetMemoryResponse)
+def get_memories(memory_req: GetMemoryRequest):
+    return handlers.memory_handler.handle_get_memories(
+        get_mem_req=memory_req,
+        naive_mem_cube=naive_mem_cube,
+    )
+
+
+@router.get("/get_memory/{memory_id}", summary="Get memory by id", response_model=GetMemoryResponse)
+def get_memory_by_id(memory_id: str):
+    return handlers.memory_handler.handle_get_memory(
+        memory_id=memory_id,
+        naive_mem_cube=naive_mem_cube,
+    )
+
+
+@router.post(
+    "/delete_memory", summary="Delete memories for user", response_model=DeleteMemoryResponse
+)
+def delete_memories(memory_req: DeleteMemoryRequest):
+    return handlers.memory_handler.handle_delete_memories(
+        delete_mem_req=memory_req, naive_mem_cube=naive_mem_cube
+    )
+
+
+# =============================================================================
+# Feedback API Endpoints
+# =============================================================================
+
+
+@router.post("/feedback", summary="Feedback memories", response_model=MemoryResponse)
+def feedback_memories(feedback_req: APIFeedbackRequest):
+    """
+    Feedback memories for a specific user.
+
+    This endpoint uses the class-based FeedbackHandler for better code organization.
+    """
+    return feedback_handler.handle_feedback_memories(feedback_req)
+
+
+# =============================================================================
+# Other API Endpoints (for internal use)
+# =============================================================================
+
+
+@router.post(
+    "/get_user_names_by_memory_ids",
+    summary="Get user names by memory ids",
+    response_model=GetUserNamesByMemoryIdsResponse,
+)
+def get_user_names_by_memory_ids(request: GetUserNamesByMemoryIdsRequest):
+    """Get user names by memory ids."""
+    if not isinstance(graph_db, PolarDBGraphDB):
+        raise HTTPException(
+            status_code=400,
+            detail=(
+                "graph_db must be an instance of PolarDBGraphDB to use "
+                "get_user_names_by_memory_ids"
+                f"current graph_db is: {graph_db.__class__.__name__}"
+            ),
+        )
+    result = graph_db.get_user_names_by_memory_ids(memory_ids=request.memory_ids)
+    if vector_db:
+        prefs = []
+        for collection_name in ["explicit_preference", "implicit_preference"]:
+            prefs.extend(
+                vector_db.get_by_ids(collection_name=collection_name, ids=request.memory_ids)
+            )
+        result.update({pref.id: pref.payload.get("mem_cube_id", None) for pref in prefs})
+    return GetUserNamesByMemoryIdsResponse(
+        code=200,
+        message="Successfully",
+        data=result,
+    )
+
+
+@router.post(
+    "/exist_mem_cube_id",
+    summary="Check if mem cube id exists",
+    response_model=ExistMemCubeIdResponse,
+)
+def exist_mem_cube_id(request: ExistMemCubeIdRequest):
+    """Check if mem cube id exists."""
+    return ExistMemCubeIdResponse(
+        code=200,
+        message="Successfully",
+        data=graph_db.exist_user_name(user_name=request.mem_cube_id),
+    )

memos/api/server_api.py

@@ -0,0 +1,44 @@
+import logging
+
+from fastapi import FastAPI, HTTPException
+from fastapi.exceptions import RequestValidationError
+
+from memos.api.exceptions import APIExceptionHandler
+from memos.api.middleware.request_context import RequestContextMiddleware
+from memos.api.routers.server_router import router as server_router
+
+
+# Configure logging
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+logger = logging.getLogger(__name__)
+
+app = FastAPI(
+    title="MemOS Server REST APIs",
+    description="A REST API for managing multiple users with MemOS Server.",
+    version="1.0.1",
+)
+
+app.add_middleware(RequestContextMiddleware, source="server_api")
+# Include routers
+app.include_router(server_router)
+
+# Request validation failed
+app.exception_handler(RequestValidationError)(APIExceptionHandler.validation_error_handler)
+# Invalid business code parameters
+app.exception_handler(ValueError)(APIExceptionHandler.value_error_handler)
+# Business layer manual exception
+app.exception_handler(HTTPException)(APIExceptionHandler.http_error_handler)
+# Fallback for unknown errors
+app.exception_handler(Exception)(APIExceptionHandler.global_exception_handler)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    import uvicorn
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--port", type=int, default=8001)
+    parser.add_argument("--workers", type=int, default=1)
+    args = parser.parse_args()
+    uvicorn.run("memos.api.server_api:app", host="0.0.0.0", port=args.port, workers=args.workers)