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

memos/api/middleware/request_context.py

@@ -2,40 +2,25 @@
 Request context middleware for automatic trace_id injection.
 """
 
-import logging
-import os
-
 from collections.abc import Callable
 
 from starlette.middleware.base import BaseHTTPMiddleware
 from starlette.requests import Request
 from starlette.responses import Response
 
-from memos.api.context.context import RequestContext, set_request_context
-
+import memos.log
 
-logger = logging.getLogger(__name__)
+from memos.context.context import RequestContext, generate_trace_id, set_request_context
 
 
-def generate_trace_id() -> str:
-    """Generate a random trace_id."""
-    return os.urandom(16).hex()
+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."""
-    trace_id = request.headers.get("g-trace-id")
-    if trace_id:
-        return trace_id
-
-    trace_id = request.headers.get("x-trace-id")
-    if trace_id:
-        return trace_id
-
-    trace_id = request.headers.get("trace-id")
-    if trace_id:
-        return 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
 
 
@@ -51,19 +36,12 @@ class RequestContextMiddleware(BaseHTTPMiddleware):
 
     async def dispatch(self, request: Request, call_next: Callable) -> Response:
         # Extract or generate trace_id
-        trace_id = extract_trace_id_from_headers(request)
-        if not trace_id:
-            trace_id = 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)
+        context = RequestContext(trace_id=trace_id, api_path=request.url.path)
         set_request_context(context)
 
-        # Add request metadata to context
-        context.set("method", request.method)
-        context.set("path", request.url.path)
-        context.set("client_ip", request.client.host if request.client else None)
-
         # Log request start with parameters
         params_log = {}
 
@@ -71,16 +49,7 @@ class RequestContextMiddleware(BaseHTTPMiddleware):
         if request.query_params:
             params_log["query_params"] = dict(request.query_params)
 
-        # Get request body if it's available
-        try:
-            params_log = await request.json()
-        except Exception as e:
-            logger.error(f"Error getting request body: {e}")
-            # If body is not JSON or empty, ignore it
-
-        logger.info(
-            f"Request started: {request.method} {request.url.path} - Parameters: {params_log}"
-        )
+        logger.info(f"Request started: {request.method} {request.url.path}, {params_log}")
 
         # Process the request
         response = await call_next(request)

memos/api/product_api.py

@@ -17,9 +17,7 @@ app = FastAPI(
     version="1.0.1",
 )
 
-# Add request context middleware (must be added first)
 app.add_middleware(RequestContextMiddleware)
-
 # Include routers
 app.include_router(product_router)
 
@@ -35,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")
 
@@ -85,6 +74,7 @@ class ChatRequest(BaseRequest):
     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):
@@ -99,6 +89,7 @@ class ChatCompleteRequest(BaseRequest):
     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):
@@ -160,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):
@@ -169,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):
@@ -177,3 +170,85 @@ 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,14 +1,11 @@
 import json
+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,
@@ -79,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
@@ -106,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",
@@ -194,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,
@@ -203,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")
 
@@ -217,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)
 
@@ -250,6 +255,7 @@ def chat(chat_req: ChatRequest):
                     history=chat_req.history,
                     internet_search=chat_req.internet_search,
                     moscube=chat_req.moscube,
+                    session_id=chat_req.session_id,
                 )
 
             except Exception as e:
@@ -294,6 +300,7 @@ def chat_complete(chat_req: ChatCompleteRequest):
             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

memos/api/start_api.py

@@ -421,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