MemoryOS 1.0.0__py3-none-any.whl → 1.1.1__py3-none-any.whl

memos/api/middleware/request_context.py

@@ -0,0 +1,63 @@
+"""
+Request context middleware for automatic trace_id injection.
+"""
+
+from collections.abc import Callable
+
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.requests import Request
+from starlette.responses import Response
+
+import memos.log
+
+from memos.context.context import RequestContext, generate_trace_id, set_request_context
+
+
+logger = memos.log.get_logger(__name__)
+
+
+def extract_trace_id_from_headers(request: Request) -> str | None:
+    """Extract trace_id from various possible headers with priority: g-trace-id > x-trace-id > trace-id."""
+    for header in ["g-trace-id", "x-trace-id", "trace-id"]:
+        if trace_id := request.headers.get(header):
+            return trace_id
+    return None
+
+
+class RequestContextMiddleware(BaseHTTPMiddleware):
+    """
+    Middleware to automatically inject request context for every HTTP request.
+
+    This middleware:
+    1. Extracts trace_id from headers or generates a new one
+    2. Creates a RequestContext and sets it globally
+    3. Ensures the context is available throughout the request lifecycle
+    """
+
+    async def dispatch(self, request: Request, call_next: Callable) -> Response:
+        # Extract or generate trace_id
+        trace_id = extract_trace_id_from_headers(request) or generate_trace_id()
+
+        # Create and set request context
+        context = RequestContext(trace_id=trace_id, api_path=request.url.path)
+        set_request_context(context)
+
+        # Log request start with parameters
+        params_log = {}
+
+        # Get query parameters
+        if request.query_params:
+            params_log["query_params"] = dict(request.query_params)
+
+        logger.info(f"Request started: {request.method} {request.url.path}, {params_log}")
+
+        # Process the request
+        response = await call_next(request)
+
+        # Log request completion with output
+        logger.info(f"Request completed: {request.url.path}, status: {response.status_code}")
+
+        # Add trace_id to response headers for debugging
+        response.headers["x-trace-id"] = trace_id
+
+        return response

memos/api/product_api.py

@@ -3,6 +3,7 @@ import logging
 from fastapi import FastAPI
 
 from memos.api.exceptions import APIExceptionHandler
+from memos.api.middleware.request_context import RequestContextMiddleware
 from memos.api.routers.product_router import router as product_router
 
 
@@ -13,9 +14,10 @@ logger = logging.getLogger(__name__)
 app = FastAPI(
     title="MemOS Product REST APIs",
     description="A REST API for managing multiple users with MemOS Product.",
-    version="1.0.0",
+    version="1.0.1",
 )
 
+app.add_middleware(RequestContextMiddleware)
 # Include routers
 app.include_router(product_router)
 
@@ -31,5 +33,6 @@ if __name__ == "__main__":
 
     parser = argparse.ArgumentParser()
     parser.add_argument("--port", type=int, default=8001)
+    parser.add_argument("--workers", type=int, default=32)
     args = parser.parse_args()
-    uvicorn.run(app, host="0.0.0.0", port=args.port)
+    uvicorn.run("memos.api.product_api:app", host="0.0.0.0", port=args.port, workers=args.workers)

memos/api/product_models.py

@@ -1,26 +1,14 @@
 import uuid
 
-from typing import Generic, Literal, TypeAlias, TypeVar
+from typing import Generic, Literal, TypeVar
 
 from pydantic import BaseModel, Field
-from typing_extensions import TypedDict
 
+# Import message types from core types module
+from memos.types import MessageDict
 
-T = TypeVar("T")
-
-
-# ─── Message Types ──────────────────────────────────────────────────────────────
-
-# Chat message roles
-MessageRole: TypeAlias = Literal["user", "assistant", "system"]
 
-
-# Message structure
-class MessageDict(TypedDict):
-    """Typed dictionary for chat message dictionaries."""
-
-    role: MessageRole
-    content: str
+T = TypeVar("T")
 
 
 class BaseRequest(BaseModel):
@@ -42,6 +30,7 @@ class UserRegisterRequest(BaseRequest):
     user_id: str = Field(
         default_factory=lambda: str(uuid.uuid4()), description="User ID for registration"
     )
+    mem_cube_id: str | None = Field(None, description="Cube ID for registration")
     user_name: str | None = Field(None, description="User name for registration")
     interests: str | None = Field(None, description="User interests")
 
@@ -84,6 +73,23 @@ class ChatRequest(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(True, description="Whether to use internet search")
+    moscube: bool = Field(False, description="Whether to use MemOSCube")
+    session_id: str | None = Field(None, description="Session ID for soft-filtering memories")
+
+
+class ChatCompleteRequest(BaseRequest):
+    """Request model for chat operations."""
+
+    user_id: str = Field(..., description="User ID")
+    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(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")
+    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")
 
 
 class UserCreate(BaseRequest):
@@ -145,6 +151,7 @@ class MemoryCreateRequest(BaseRequest):
     mem_cube_id: str | None = Field(None, description="Cube ID")
     source: str | None = Field(None, description="Source of the memory")
     user_profile: bool = Field(False, description="User profile memory")
+    session_id: str | None = Field(None, description="Session id")
 
 
 class SearchRequest(BaseRequest):
@@ -154,6 +161,7 @@ class SearchRequest(BaseRequest):
     query: str = Field(..., description="Search query")
     mem_cube_id: str | None = Field(None, description="Cube ID to search in")
     top_k: int = Field(10, description="Number of results to return")
+    session_id: str | None = Field(None, description="Session ID for soft-filtering memories")
 
 
 class SuggestionRequest(BaseRequest):
@@ -161,3 +169,86 @@ class SuggestionRequest(BaseRequest):
 
     user_id: str = Field(..., description="User ID")
     language: Literal["zh", "en"] = Field("zh", description="Language for suggestions")
+    message: list[MessageDict] | None = Field(None, description="List of messages to store.")
+
+
+# ─── MemOS Client Response Models ──────────────────────────────────────────────
+
+
+class MessageDetail(BaseModel):
+    """Individual message detail model based on actual API response."""
+
+    model_config = {"extra": "allow"}
+
+
+class MemoryDetail(BaseModel):
+    """Individual memory detail model based on actual API response."""
+
+    model_config = {"extra": "allow"}
+
+
+class GetMessagesData(BaseModel):
+    """Data model for get messages response based on actual API."""
+
+    message_detail_list: list[MessageDetail] = Field(
+        default_factory=list, alias="memory_detail_list", description="List of message details"
+    )
+
+
+class SearchMemoryData(BaseModel):
+    """Data model for search memory response based on actual API."""
+
+    memory_detail_list: list[MemoryDetail] = Field(
+        default_factory=list, alias="memory_detail_list", description="List of memory details"
+    )
+    message_detail_list: list[MessageDetail] | None = Field(
+        None, alias="message_detail_list", description="List of message details (usually None)"
+    )
+
+
+class AddMessageData(BaseModel):
+    """Data model for add message response based on actual API."""
+
+    success: bool = Field(..., description="Operation success status")
+
+
+# ─── MemOS Response Models (Similar to OpenAI ChatCompletion) ──────────────────
+
+
+class MemOSGetMessagesResponse(BaseModel):
+    """Response model for get messages operation based on actual API."""
+
+    code: int = Field(..., description="Response status code")
+    message: str = Field(..., description="Response message")
+    data: GetMessagesData = Field(..., description="Messages data")
+
+    @property
+    def messages(self) -> list[MessageDetail]:
+        """Convenient access to message list."""
+        return self.data.message_detail_list
+
+
+class MemOSSearchResponse(BaseModel):
+    """Response model for search memory operation based on actual API."""
+
+    code: int = Field(..., description="Response status code")
+    message: str = Field(..., description="Response message")
+    data: SearchMemoryData = Field(..., description="Search results data")
+
+    @property
+    def memories(self) -> list[MemoryDetail]:
+        """Convenient access to memory list."""
+        return self.data.memory_detail_list
+
+
+class MemOSAddResponse(BaseModel):
+    """Response model for add message operation based on actual API."""
+
+    code: int = Field(..., description="Response status code")
+    message: str = Field(..., description="Response message")
+    data: AddMessageData = Field(..., description="Add operation data")
+
+    @property
+    def success(self) -> bool:
+        """Convenient access to success status."""
+        return self.data.success

memos/api/routers/product_router.py

@@ -1,17 +1,14 @@
 import json
-import logging
+import time
 import traceback
 
-from datetime import datetime
-from typing import Annotated
-
-from fastapi import APIRouter, Depends, HTTPException
+from fastapi import APIRouter, HTTPException
 from fastapi.responses import StreamingResponse
 
 from memos.api.config import APIConfig
-from memos.api.context.dependencies import G, get_g_object
 from memos.api.product_models import (
     BaseResponse,
+    ChatCompleteRequest,
     ChatRequest,
     GetMemoryRequest,
     MemoryCreateRequest,
@@ -25,11 +22,12 @@ from memos.api.product_models import (
     UserRegisterResponse,
 )
 from memos.configs.mem_os import MOSConfig
+from memos.log import get_logger
 from memos.mem_os.product import MOSProduct
 from memos.memos_tools.notification_service import get_error_bot_function, get_online_bot_function
 
 
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 router = APIRouter(prefix="/product", tags=["Product API"])
 
@@ -78,24 +76,19 @@ def set_config(config):
 
 
 @router.post("/users/register", summary="Register a new user", response_model=UserRegisterResponse)
-def register_user(user_req: UserRegisterRequest, g: Annotated[G, Depends(get_g_object)]):
+def register_user(user_req: UserRegisterRequest):
     """Register a new user with configuration and default cube."""
     try:
-        # Set request-related information in g object
-        g.user_id = user_req.user_id
-        g.action = "user_register"
-        g.timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-
-        logger.info(f"Starting user registration for user_id: {user_req.user_id}")
-        logger.info(f"Request trace_id: {g.trace_id}")
-        logger.info(f"Request timestamp: {g.timestamp}")
-
         # Get configuration for the user
+        time_start_register = time.time()
         user_config, default_mem_cube = APIConfig.create_user_config(
             user_name=user_req.user_id, user_id=user_req.user_id
         )
         logger.info(f"user_config: {user_config.model_dump(mode='json')}")
         logger.info(f"default_mem_cube: {default_mem_cube.config.model_dump(mode='json')}")
+        logger.info(
+            f"time register api : create user config time user_id: {user_req.user_id} time is: {time.time() - time_start_register}"
+        )
         mos_product = get_mos_product_instance()
 
         # Register user with default config and mem cube
@@ -105,8 +98,11 @@ def register_user(user_req: UserRegisterRequest, g: Annotated[G, Depends(get_g_o
             interests=user_req.interests,
             config=user_config,
             default_mem_cube=default_mem_cube,
+            mem_cube_id=user_req.mem_cube_id,
+        )
+        logger.info(
+            f"time register api : register time user_id: {user_req.user_id} time is: {time.time() - time_start_register}"
         )
-
         if result["status"] == "success":
             return UserRegisterResponse(
                 message="User registered successfully",
@@ -148,7 +144,9 @@ def get_suggestion_queries_post(suggestion_req: SuggestionRequest):
     try:
         mos_product = get_mos_product_instance()
         suggestions = mos_product.get_suggestion_query(
-            user_id=suggestion_req.user_id, language=suggestion_req.language
+            user_id=suggestion_req.user_id,
+            language=suggestion_req.language,
+            message=suggestion_req.message,
         )
         return SuggestionResponse(
             message="Suggestions retrieved successfully", data={"query": suggestions}
@@ -191,6 +189,7 @@ def get_all_memories(memory_req: GetMemoryRequest):
 def create_memory(memory_req: MemoryCreateRequest):
     """Create a new memory for a specific user."""
     try:
+        time_start_add = time.time()
         mos_product = get_mos_product_instance()
         mos_product.add(
             user_id=memory_req.user_id,
@@ -200,6 +199,10 @@ def create_memory(memory_req: MemoryCreateRequest):
             mem_cube_id=memory_req.mem_cube_id,
             source=memory_req.source,
             user_profile=memory_req.user_profile,
+            session_id=memory_req.session_id,
+        )
+        logger.info(
+            f"time add api : add time user_id: {memory_req.user_id} time is: {time.time() - time_start_add}"
         )
         return SimpleResponse(message="Memory created successfully")
 
@@ -214,12 +217,17 @@ def create_memory(memory_req: MemoryCreateRequest):
 def search_memories(search_req: SearchRequest):
     """Search memories for a specific user."""
     try:
+        time_start_search = time.time()
         mos_product = get_mos_product_instance()
         result = mos_product.search(
             query=search_req.query,
             user_id=search_req.user_id,
             install_cube_ids=[search_req.mem_cube_id] if search_req.mem_cube_id else None,
             top_k=search_req.top_k,
+            session_id=search_req.session_id,
+        )
+        logger.info(
+            f"time search api : add time user_id: {search_req.user_id} time is: {time.time() - time_start_search}"
         )
         return SearchResponse(message="Search completed successfully", data=result)
 
@@ -246,6 +254,8 @@ def chat(chat_req: ChatRequest):
                     cube_id=chat_req.mem_cube_id,
                     history=chat_req.history,
                     internet_search=chat_req.internet_search,
+                    moscube=chat_req.moscube,
+                    session_id=chat_req.session_id,
                 )
 
             except Exception as e:
@@ -273,6 +283,39 @@ def chat(chat_req: ChatRequest):
         raise HTTPException(status_code=500, detail=str(traceback.format_exc())) from err
 
 
+@router.post("/chat/complete", summary="Chat with MemOS (Complete Response)")
+def chat_complete(chat_req: ChatCompleteRequest):
+    """Chat with MemOS for a specific user. Returns complete response (non-streaming)."""
+    try:
+        mos_product = get_mos_product_instance()
+
+        # Collect all responses from the generator
+        content, references = mos_product.chat(
+            query=chat_req.query,
+            user_id=chat_req.user_id,
+            cube_id=chat_req.mem_cube_id,
+            history=chat_req.history,
+            internet_search=chat_req.internet_search,
+            moscube=chat_req.moscube,
+            base_prompt=chat_req.base_prompt,
+            top_k=chat_req.top_k,
+            threshold=chat_req.threshold,
+            session_id=chat_req.session_id,
+        )
+
+        # Return the complete response
+        return {
+            "message": "Chat completed successfully",
+            "data": {"response": content, "references": references},
+        }
+
+    except ValueError as err:
+        raise HTTPException(status_code=404, detail=str(traceback.format_exc())) from err
+    except Exception as err:
+        logger.error(f"Failed to start chat: {traceback.format_exc()}")
+        raise HTTPException(status_code=500, detail=str(traceback.format_exc())) from err
+
+
 @router.get("/users", summary="List all users", response_model=BaseResponse[list])
 def list_users():
     """List all registered users."""

memos/api/start_api.py

@@ -9,6 +9,7 @@ from fastapi.requests import Request
 from fastapi.responses import JSONResponse, RedirectResponse
 from pydantic import BaseModel, Field
 
+from memos.api.middleware.request_context import RequestContextMiddleware
 from memos.configs.mem_os import MOSConfig
 from memos.mem_os.main import MOS
 from memos.mem_user.user_manager import UserManager, UserRole
@@ -78,6 +79,8 @@ app = FastAPI(
     version="1.0.0",
 )
 
+app.add_middleware(RequestContextMiddleware)
+
 
 class BaseRequest(BaseModel):
     """Base model for all requests."""
@@ -418,3 +421,13 @@ async def global_exception_handler(request: Request, exc: Exception):
         status_code=500,
         content={"code": 500, "message": str(exc), "data": None},
     )
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--port", type=int, default=8000, help="Port to run the server on")
+    parser.add_argument("--host", type=str, default="0.0.0.0", help="Host to run the server on")
+    parser.add_argument("--reload", action="store_true", help="Enable auto-reload for development")
+    args = parser.parse_args()

memos/configs/graph_db.py

@@ -140,6 +140,10 @@ class NebulaGraphDBConfig(BaseGraphDBConfig):
             "If False: use a single shared database with logical isolation by user_name."
         ),
     )
+    max_client: int = Field(
+        default=1000,
+        description=("max_client"),
+    )
     embedding_dimension: int = Field(default=3072, description="Dimension of vector embedding")
 
     @model_validator(mode="after")

memos/configs/mem_scheduler.py

@@ -6,7 +6,7 @@ from typing import Any, ClassVar
 from pydantic import ConfigDict, Field, field_validator, model_validator
 
 from memos.configs.base import BaseConfig
-from memos.mem_scheduler.general_modules.misc import DictConversionMixin
+from memos.mem_scheduler.general_modules.misc import DictConversionMixin, EnvConfigMixin
 from memos.mem_scheduler.schemas.general_schemas import (
     BASE_DIR,
     DEFAULT_ACT_MEM_DUMP_PATH,
@@ -64,6 +64,19 @@ class GeneralSchedulerConfig(BaseSchedulerConfig):
         default=20, description="Capacity of the activation memory monitor"
     )
 
+    # Database configuration for ORM persistence
+    db_path: str | None = Field(
+        default=None,
+        description="Path to SQLite database file for ORM persistence. If None, uses default scheduler_orm.db",
+    )
+    db_url: str | None = Field(
+        default=None,
+        description="Database URL for ORM persistence (e.g., mysql://user:pass@host/db). Takes precedence over db_path",
+    )
+    enable_orm_persistence: bool = Field(
+        default=True, description="Whether to enable ORM-based persistence for monitors"
+    )
+
 
 class SchedulerConfigFactory(BaseConfig):
     """Factory class for creating scheduler configurations."""
@@ -74,6 +87,7 @@ class SchedulerConfigFactory(BaseConfig):
     model_config = ConfigDict(extra="forbid", strict=True)
     backend_to_class: ClassVar[dict[str, Any]] = {
         "general_scheduler": GeneralSchedulerConfig,
+        "optimized_scheduler": GeneralSchedulerConfig,  # optimized_scheduler uses same config as general_scheduler
     }
 
     @field_validator("backend")
@@ -94,6 +108,8 @@ class SchedulerConfigFactory(BaseConfig):
 # ************************* Auth *************************
 class RabbitMQConfig(
     BaseConfig,
+    DictConversionMixin,
+    EnvConfigMixin,
 ):
     host_name: str = Field(default="", description="Endpoint for RabbitMQ instance access")
     user_name: str = Field(default="", description="Static username for RabbitMQ instance")
@@ -110,7 +126,7 @@ class RabbitMQConfig(
     )
 
 
-class GraphDBAuthConfig(BaseConfig):
+class GraphDBAuthConfig(BaseConfig, DictConversionMixin, EnvConfigMixin):
     uri: str = Field(
         default="bolt://localhost:7687",
         description="URI for graph database access (e.g., bolt://host:port)",
@@ -127,7 +143,7 @@ class GraphDBAuthConfig(BaseConfig):
     )
 
 
-class OpenAIConfig(BaseConfig):
+class OpenAIConfig(BaseConfig, DictConversionMixin, EnvConfigMixin):
     api_key: str = Field(default="", description="API key for OpenAI service")
     base_url: str = Field(default="", description="Base URL for API endpoint")
     default_model: str = Field(default="", description="Default model to use")
@@ -183,6 +199,25 @@ class AuthConfig(BaseConfig, DictConversionMixin):
                 "Please use YAML (.yaml, .yml) or JSON (.json) files."
             )
 
+    @classmethod
+    def from_local_env(cls) -> "AuthConfig":
+        """Creates an AuthConfig instance by loading configuration from environment variables.
+
+        This method loads configuration for all nested components (RabbitMQ, OpenAI, GraphDB)
+        from their respective environment variables using each component's specific prefix.
+
+        Returns:
+            AuthConfig: Configured instance with values from environment variables
+
+        Raises:
+            ValueError: If any required environment variables are missing
+        """
+        return cls(
+            rabbitmq=RabbitMQConfig.from_env(),
+            openai=OpenAIConfig.from_env(),
+            graph_db=GraphDBAuthConfig.from_env(),
+        )
+
     def set_openai_config_to_environment(self):
         # Set environment variables
         os.environ["OPENAI_API_KEY"] = self.openai.api_key

memos/configs/memory.py

@@ -7,6 +7,7 @@ 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.reranker import RerankerConfigFactory
 from memos.configs.vec_db import VectorDBConfigFactory
 from memos.exceptions import ConfigurationError
 
@@ -151,6 +152,10 @@ class TreeTextMemoryConfig(BaseTextMemoryConfig):
         default_factory=EmbedderConfigFactory,
         description="Embedder configuration for the memory embedding",
     )
+    reranker: RerankerConfigFactory | None = Field(
+        None,
+        description="Reranker configuration (optional, defaults to cosine_local).",
+    )
     graph_db: GraphDBConfigFactory = Field(
         ...,
         default_factory=GraphDBConfigFactory,
@@ -166,6 +171,14 @@ class TreeTextMemoryConfig(BaseTextMemoryConfig):
         description="Optional description for this memory configuration.",
     )
 
+    memory_size: dict[str, Any] | None = Field(
+        default=None,
+        description=(
+            "Maximum item counts per memory bucket, e.g.: "
+            '{"WorkingMemory": 20, "LongTermMemory": 10000, "UserMemory": 10000}'
+        ),
+    )
+
 
 # ─── 3. Global Memory Config Factory ──────────────────────────────────────────
 

memos/configs/reranker.py

@@ -0,0 +1,18 @@
+# memos/configs/reranker.py
+from __future__ import annotations
+
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class RerankerConfigFactory(BaseModel):
+    """
+    {
+      "backend": "http_bge" | "cosine_local" | "noop",
+      "config": { ... backend-specific ... }
+    }
+    """
+
+    backend: str = Field(..., description="Reranker backend id")
+    config: dict[str, Any] = Field(default_factory=dict, description="Backend-specific options")