MemoryOS 0.2.1__py3-none-any.whl → 0.2.2__py3-none-any.whl

memos/api/context/context.py

@@ -0,0 +1,147 @@
+"""
+Global request context management for trace_id and request-scoped data.
+
+This module provides optional trace_id functionality that can be enabled
+when using the API components. It uses ContextVar to ensure thread safety
+and request isolation.
+"""
+
+import uuid
+
+from collections.abc import Callable
+from contextvars import ContextVar
+from typing import Any
+
+
+# Global context variable for request-scoped data
+_request_context: ContextVar[dict[str, Any] | None] = ContextVar("request_context", default=None)
+
+
+class RequestContext:
+    """
+    Request-scoped context object that holds trace_id and other request data.
+
+    This provides a Flask g-like object for FastAPI applications.
+    """
+
+    def __init__(self, trace_id: str | None = None):
+        self.trace_id = trace_id or str(uuid.uuid4())
+        self._data: dict[str, Any] = {}
+
+    def set(self, key: str, value: Any) -> None:
+        """Set a value in the context."""
+        self._data[key] = value
+
+    def get(self, key: str, default: Any | None = None) -> Any:
+        """Get a value from the context."""
+        return self._data.get(key, default)
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        if name.startswith("_") or name == "trace_id":
+            super().__setattr__(name, value)
+        else:
+            if not hasattr(self, "_data"):
+                super().__setattr__(name, value)
+            else:
+                self._data[name] = value
+
+    def __getattr__(self, name: str) -> Any:
+        if hasattr(self, "_data") and name in self._data:
+            return self._data[name]
+        raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{name}'")
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert context to dictionary."""
+        return {"trace_id": self.trace_id, "data": self._data.copy()}
+
+
+def set_request_context(context: RequestContext) -> None:
+    """
+    Set the current request context.
+
+    This is typically called by the API dependency injection system.
+    """
+    _request_context.set(context.to_dict())
+
+
+def get_current_trace_id() -> str | None:
+    """
+    Get the current request's trace_id.
+
+    Returns:
+        The trace_id if available, None otherwise.
+    """
+    context = _request_context.get()
+    if context:
+        return context.get("trace_id")
+    return None
+
+
+def get_current_context() -> RequestContext | None:
+    """
+    Get the current request context.
+
+    Returns:
+        The current RequestContext if available, None otherwise.
+    """
+    context_dict = _request_context.get()
+    if context_dict:
+        ctx = RequestContext(trace_id=context_dict.get("trace_id"))
+        ctx._data = context_dict.get("data", {}).copy()
+        return ctx
+    return None
+
+
+def require_context() -> RequestContext:
+    """
+    Get the current request context, raising an error if not available.
+
+    Returns:
+        The current RequestContext.
+
+    Raises:
+        RuntimeError: If called outside of a request context.
+    """
+    context = get_current_context()
+    if context is None:
+        raise RuntimeError(
+            "No request context available. This function must be called within a request handler."
+        )
+    return context
+
+
+# Type for trace_id getter function
+TraceIdGetter = Callable[[], str | None]
+
+# Global variable to hold the trace_id getter function
+_trace_id_getter: TraceIdGetter | None = None
+
+
+def set_trace_id_getter(getter: TraceIdGetter) -> None:
+    """
+    Set a custom trace_id getter function.
+
+    This allows the logging system to retrieve trace_id without importing
+    API-specific modules.
+    """
+    global _trace_id_getter
+    _trace_id_getter = getter
+
+
+def get_trace_id_for_logging() -> str | None:
+    """
+    Get trace_id for logging purposes.
+
+    This function is used by the logging system and will use either
+    the custom getter function or fall back to the default context.
+    """
+    if _trace_id_getter:
+        try:
+            return _trace_id_getter()
+        except Exception:
+            pass
+    return get_current_trace_id()
+
+
+# Initialize the default trace_id getter
+set_trace_id_getter(get_current_trace_id)

memos/api/context/dependencies.py

@@ -0,0 +1,90 @@
+import logging
+
+from fastapi import Depends, Header, Request
+
+from memos.api.context.context import RequestContext, set_request_context
+
+
+logger = logging.getLogger(__name__)
+
+# Type alias for the RequestContext from context module
+G = RequestContext
+
+
+def get_trace_id_from_header(
+    trace_id: str | None = Header(None, alias="trace-id"),
+    x_trace_id: str | None = Header(None, alias="x-trace-id"),
+    g_trace_id: str | None = Header(None, alias="g-trace-id"),
+) -> str | None:
+    """
+    Extract trace_id from various possible headers.
+
+    Priority: g-trace-id > x-trace-id > trace-id
+    """
+    return g_trace_id or x_trace_id or trace_id
+
+
+def get_request_context(
+    request: Request, trace_id: str | None = Depends(get_trace_id_from_header)
+) -> RequestContext:
+    """
+    Get request context object with trace_id and request metadata.
+
+    This function creates a RequestContext and automatically sets it
+    in the global context for use throughout the request lifecycle.
+    """
+    # Create context object
+    ctx = RequestContext(trace_id=trace_id)
+
+    # Set the context globally for this request
+    set_request_context(ctx)
+
+    # Log request start
+    logger.info(f"Request started with trace_id: {ctx.trace_id}")
+
+    # Add request metadata to context
+    ctx.set("method", request.method)
+    ctx.set("path", request.url.path)
+    ctx.set("client_ip", request.client.host if request.client else None)
+
+    return ctx
+
+
+def get_g_object(trace_id: str | None = Depends(get_trace_id_from_header)) -> G:
+    """
+    Get Flask g-like object for the current request.
+
+    This creates a RequestContext and sets it globally for access
+    throughout the request lifecycle.
+    """
+    g = RequestContext(trace_id=trace_id)
+    set_request_context(g)
+    logger.info(f"Request g object created with trace_id: {g.trace_id}")
+    return g
+
+
+def get_current_g() -> G | None:
+    """
+    Get the current request's g object from anywhere in the application.
+
+    Returns:
+        The current request's g object if available, None otherwise.
+    """
+    from memos.context import get_current_context
+
+    return get_current_context()
+
+
+def require_g() -> G:
+    """
+    Get the current request's g object, raising an error if not available.
+
+    Returns:
+        The current request's g object.
+
+    Raises:
+        RuntimeError: If called outside of a request context.
+    """
+    from memos.context import require_context
+
+    return require_context()

memos/api/product_models.py

@@ -83,11 +83,12 @@ class ChatRequest(BaseRequest):
     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(True, description="Whether to use internet search")
 
 
 class UserCreate(BaseRequest):
     user_name: str | None = Field(None, description="Name of the user")
-    role: str = Field("user", description="Role of the user")
+    role: str = Field("USER", description="Role of the user")
     user_id: str = Field(..., description="User ID")
 
 
@@ -142,6 +143,8 @@ class MemoryCreateRequest(BaseRequest):
     memory_content: str | None = Field(None, description="Memory content to store")
     doc_path: str | None = Field(None, description="Path to document to store")
     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")
 
 
 class SearchRequest(BaseRequest):
@@ -150,6 +153,7 @@ class SearchRequest(BaseRequest):
     user_id: str = Field(..., description="User ID")
     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")
 
 
 class SuggestionRequest(BaseRequest):

memos/api/routers/product_router.py

@@ -2,10 +2,14 @@ import json
 import logging
 import traceback
 
-from fastapi import APIRouter, HTTPException
+from datetime import datetime
+from typing import Annotated
+
+from fastapi import APIRouter, Depends, 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,
     ChatRequest,
@@ -22,6 +26,7 @@ from memos.api.product_models import (
 )
 from memos.configs.mem_os import MOSConfig
 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__)
@@ -37,16 +42,25 @@ def get_mos_product_instance():
     global MOS_PRODUCT_INSTANCE
     if MOS_PRODUCT_INSTANCE is None:
         default_config = APIConfig.get_product_default_config()
-        print(default_config)
+        logger.info(f"*********init_default_mos_config********* {default_config}")
         from memos.configs.mem_os import MOSConfig
 
         mos_config = MOSConfig(**default_config)
 
         # Get default cube config from APIConfig (may be None if disabled)
         default_cube_config = APIConfig.get_default_cube_config()
-        print("*********default_cube_config*********", default_cube_config)
+        logger.info(f"*********initdefault_cube_config******** {default_cube_config}")
+
+        # Get DingDing bot functions
+        dingding_enabled = APIConfig.is_dingding_bot_enabled()
+        online_bot = get_online_bot_function() if dingding_enabled else None
+        error_bot = get_error_bot_function() if dingding_enabled else None
+
         MOS_PRODUCT_INSTANCE = MOSProduct(
-            default_config=mos_config, default_cube_config=default_cube_config
+            default_config=mos_config,
+            default_cube_config=default_cube_config,
+            online_bot=online_bot,
+            error_bot=error_bot,
         )
         logger.info("MOSProduct instance created successfully with inheritance architecture")
     return MOS_PRODUCT_INSTANCE
@@ -56,7 +70,7 @@ get_mos_product_instance()
 
 
 @router.post("/configure", summary="Configure MOSProduct", response_model=SimpleResponse)
-async def set_config(config):
+def set_config(config):
     """Set MOSProduct configuration."""
     global MOS_PRODUCT_INSTANCE
     MOS_PRODUCT_INSTANCE = MOSProduct(default_config=config)
@@ -64,9 +78,18 @@ async def set_config(config):
 
 
 @router.post("/users/register", summary="Register a new user", response_model=UserRegisterResponse)
-async def register_user(user_req: UserRegisterRequest):
+def register_user(user_req: UserRegisterRequest, g: Annotated[G, Depends(get_g_object)]):
     """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
         user_config, default_mem_cube = APIConfig.create_user_config(
             user_name=user_req.user_id, user_id=user_req.user_id
@@ -100,7 +123,7 @@ async def register_user(user_req: UserRegisterRequest):
 @router.get(
     "/suggestions/{user_id}", summary="Get suggestion queries", response_model=SuggestionResponse
 )
-async def get_suggestion_queries(user_id: str):
+def get_suggestion_queries(user_id: str):
     """Get suggestion queries for a specific user."""
     try:
         mos_product = get_mos_product_instance()
@@ -120,7 +143,7 @@ async def get_suggestion_queries(user_id: str):
     summary="Get suggestion queries with language",
     response_model=SuggestionResponse,
 )
-async def get_suggestion_queries_post(suggestion_req: SuggestionRequest):
+def get_suggestion_queries_post(suggestion_req: SuggestionRequest):
     """Get suggestion queries for a specific user with language preference."""
     try:
         mos_product = get_mos_product_instance()
@@ -138,7 +161,7 @@ async def get_suggestion_queries_post(suggestion_req: SuggestionRequest):
 
 
 @router.post("/get_all", summary="Get all memories for user", response_model=MemoryResponse)
-async def get_all_memories(memory_req: GetMemoryRequest):
+def get_all_memories(memory_req: GetMemoryRequest):
     """Get all memories for a specific user."""
     try:
         mos_product = get_mos_product_instance()
@@ -165,7 +188,7 @@ async def get_all_memories(memory_req: GetMemoryRequest):
 
 
 @router.post("/add", summary="add a new memory", response_model=SimpleResponse)
-async def create_memory(memory_req: MemoryCreateRequest):
+def create_memory(memory_req: MemoryCreateRequest):
     """Create a new memory for a specific user."""
     try:
         mos_product = get_mos_product_instance()
@@ -175,6 +198,8 @@ async def create_memory(memory_req: MemoryCreateRequest):
             messages=memory_req.messages,
             doc_path=memory_req.doc_path,
             mem_cube_id=memory_req.mem_cube_id,
+            source=memory_req.source,
+            user_profile=memory_req.user_profile,
         )
         return SimpleResponse(message="Memory created successfully")
 
@@ -186,7 +211,7 @@ async def create_memory(memory_req: MemoryCreateRequest):
 
 
 @router.post("/search", summary="Search memories", response_model=SearchResponse)
-async def search_memories(search_req: SearchRequest):
+def search_memories(search_req: SearchRequest):
     """Search memories for a specific user."""
     try:
         mos_product = get_mos_product_instance()
@@ -194,6 +219,7 @@ async def search_memories(search_req: SearchRequest):
             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,
         )
         return SearchResponse(message="Search completed successfully", data=result)
 
@@ -205,24 +231,23 @@ async def search_memories(search_req: SearchRequest):
 
 
 @router.post("/chat", summary="Chat with MemOS")
-async def chat(chat_req: ChatRequest):
+def chat(chat_req: ChatRequest):
     """Chat with MemOS for a specific user. Returns SSE stream."""
     try:
         mos_product = get_mos_product_instance()
 
-        async def generate_chat_response():
+        def generate_chat_response():
             """Generate chat response as SSE stream."""
             try:
-                import asyncio
-
-                for chunk in mos_product.chat_with_references(
+                # Directly yield from the generator without async wrapper
+                yield from mos_product.chat_with_references(
                     query=chat_req.query,
                     user_id=chat_req.user_id,
                     cube_id=chat_req.mem_cube_id,
                     history=chat_req.history,
-                ):
-                    yield chunk
-                    await asyncio.sleep(0.00001)  # 50ms delay between chunks
+                    internet_search=chat_req.internet_search,
+                )
+
             except Exception as e:
                 logger.error(f"Error in chat stream: {e}")
                 error_data = f"data: {json.dumps({'type': 'error', 'content': str(traceback.format_exc())})}\n\n"
@@ -230,11 +255,14 @@ async def chat(chat_req: ChatRequest):
 
         return StreamingResponse(
             generate_chat_response(),
-            media_type="text/plain",
+            media_type="text/event-stream",
             headers={
                 "Cache-Control": "no-cache",
                 "Connection": "keep-alive",
                 "Content-Type": "text/event-stream",
+                "Access-Control-Allow-Origin": "*",
+                "Access-Control-Allow-Headers": "*",
+                "Access-Control-Allow-Methods": "*",
             },
         )
 
@@ -246,7 +274,7 @@ async def chat(chat_req: ChatRequest):
 
 
 @router.get("/users", summary="List all users", response_model=BaseResponse[list])
-async def list_users():
+def list_users():
     """List all registered users."""
     try:
         mos_product = get_mos_product_instance()
@@ -274,7 +302,7 @@ async def get_user_info(user_id: str):
 @router.get(
     "/configure/{user_id}", summary="Get MOSProduct configuration", response_model=SimpleResponse
 )
-async def get_config(user_id: str):
+def get_config(user_id: str):
     """Get MOSProduct configuration."""
     global MOS_PRODUCT_INSTANCE
     config = MOS_PRODUCT_INSTANCE.default_config
@@ -284,7 +312,7 @@ async def get_config(user_id: str):
 @router.get(
     "/users/{user_id}/config", summary="Get user configuration", response_model=BaseResponse[dict]
 )
-async def get_user_config(user_id: str):
+def get_user_config(user_id: str):
     """Get user-specific configuration."""
     try:
         mos_product = get_mos_product_instance()
@@ -308,7 +336,7 @@ async def get_user_config(user_id: str):
 @router.put(
     "/users/{user_id}/config", summary="Update user configuration", response_model=SimpleResponse
 )
-async def update_user_config(user_id: str, config_data: dict):
+def update_user_config(user_id: str, config_data: dict):
     """Update user-specific configuration."""
     try:
         mos_product = get_mos_product_instance()
@@ -333,7 +361,7 @@ async def update_user_config(user_id: str, config_data: dict):
 @router.get(
     "/instances/status", summary="Get user configuration status", response_model=BaseResponse[dict]
 )
-async def get_instance_status():
+def get_instance_status():
     """Get information about active user configurations in memory."""
     try:
         mos_product = get_mos_product_instance()
@@ -347,7 +375,7 @@ async def get_instance_status():
 
 
 @router.get("/instances/count", summary="Get active user count", response_model=BaseResponse[int])
-async def get_active_user_count():
+def get_active_user_count():
     """Get the number of active user configurations in memory."""
     try:
         mos_product = get_mos_product_instance()

memos/configs/graph_db.py

@@ -9,7 +9,7 @@ from memos.configs.vec_db import VectorDBConfigFactory
 class BaseGraphDBConfig(BaseConfig):
     """Base class for all graph database configurations."""
 
-    uri: str
+    uri: str | list
     user: str
     password: str
 
@@ -103,6 +103,53 @@ class Neo4jCommunityGraphDBConfig(Neo4jGraphDBConfig):
         return self
 
 
+class NebulaGraphDBConfig(BaseGraphDBConfig):
+    """
+    NebulaGraph-specific configuration.
+
+    Key concepts:
+    - `space`: Equivalent to a database or namespace. All tag/edge/schema live within a space.
+    - `user_name`: Used for logical tenant isolation if needed.
+    - `auto_create`: Whether to automatically create the target space if it does not exist.
+
+    Example:
+    ---
+    hosts = ["127.0.0.1:9669"]
+    user = "root"
+    password = "nebula"
+    space = "shared_graph"
+    user_name = "alice"
+    """
+
+    space: str = Field(
+        ..., description="The name of the target NebulaGraph space (like a database)"
+    )
+    user_name: str | None = Field(
+        default=None,
+        description="Logical user or tenant ID for data isolation (optional, used in metadata tagging)",
+    )
+    auto_create: bool = Field(
+        default=False,
+        description="Whether to auto-create the space if it does not exist",
+    )
+    use_multi_db: bool = Field(
+        default=True,
+        description=(
+            "If True: use Neo4j's multi-database feature for physical isolation; "
+            "each user typically gets a separate database. "
+            "If False: use a single shared database with logical isolation by user_name."
+        ),
+    )
+    embedding_dimension: int = Field(default=3072, description="Dimension of vector embedding")
+
+    @model_validator(mode="after")
+    def validate_config(self):
+        """Validate config."""
+        if not self.space:
+            raise ValueError("`space` must be provided")
+        return self
+
+
 class GraphDBConfigFactory(BaseModel):
     backend: str = Field(..., description="Backend for graph database")
     config: dict[str, Any] = Field(..., description="Configuration for the graph database backend")
@@ -110,6 +157,7 @@ class GraphDBConfigFactory(BaseModel):
     backend_to_class: ClassVar[dict[str, Any]] = {
         "neo4j": Neo4jGraphDBConfig,
         "neo4j-community": Neo4jCommunityGraphDBConfig,
+        "nebular": NebulaGraphDBConfig,
     }
 
     @field_validator("backend")

memos/configs/internet_retriever.py

@@ -6,6 +6,7 @@ from pydantic import Field, field_validator, model_validator
 
 from memos.configs.base import BaseConfig
 from memos.exceptions import ConfigurationError
+from memos.mem_reader.factory import MemReaderConfigFactory
 
 
 class BaseInternetRetrieverConfig(BaseConfig):
@@ -47,6 +48,11 @@ class XinyuSearchConfig(BaseInternetRetrieverConfig):
     num_per_request: int = Field(
         default=10, description="Number of results per API request (not used for Xinyu)"
     )
+    reader: MemReaderConfigFactory = Field(
+        ...,
+        default_factory=MemReaderConfigFactory,
+        description="Reader configuration",
+    )
 
 
 class InternetRetrieverConfigFactory(BaseConfig):

memos/configs/mem_os.py

@@ -8,6 +8,7 @@ from memos.configs.base import BaseConfig
 from memos.configs.llm import LLMConfigFactory
 from memos.configs.mem_reader import MemReaderConfigFactory
 from memos.configs.mem_scheduler import SchedulerConfigFactory
+from memos.configs.mem_user import UserManagerConfigFactory
 
 
 class MOSConfig(BaseConfig):
@@ -33,6 +34,10 @@ class MOSConfig(BaseConfig):
         default=None,
         description="Memory scheduler configuration for managing memory operations",
     )
+    user_manager: UserManagerConfigFactory = Field(
+        default_factory=lambda: UserManagerConfigFactory(backend="sqlite", config={}),
+        description="User manager configuration for database operations",
+    )
     max_turns_window: int = Field(
         default=15,
         description="Maximum number of turns to keep in the conversation history",

memos/configs/mem_reader.py

@@ -15,6 +15,15 @@ class BaseMemReaderConfig(BaseConfig):
     created_at: datetime = Field(
         default_factory=datetime.now, description="Creation timestamp for the MemReader"
     )
+
+    @field_validator("created_at", mode="before")
+    @classmethod
+    def parse_datetime(cls, value):
+        """Parse datetime from string if needed."""
+        if isinstance(value, str):
+            return datetime.fromisoformat(value.replace("Z", "+00:00"))
+        return value
+
     llm: LLMConfigFactory = Field(..., description="LLM configuration for the MemReader")
     embedder: EmbedderConfigFactory = Field(
         ..., description="Embedder configuration for the MemReader"

memos/configs/mem_scheduler.py

@@ -6,12 +6,12 @@ from typing import Any, ClassVar
 from pydantic import ConfigDict, Field, field_validator, model_validator
 
 from memos.configs.base import BaseConfig
-from memos.mem_scheduler.modules.schemas import (
+from memos.mem_scheduler.modules.misc import DictConversionMixin
+from memos.mem_scheduler.schemas.general_schemas import (
     BASE_DIR,
     DEFAULT_ACT_MEM_DUMP_PATH,
     DEFAULT_CONSUME_INTERVAL_SECONDS,
     DEFAULT_THREAD__POOL_MAX_WORKERS,
-    DictConversionMixin,
 )
 
 
@@ -21,6 +21,7 @@ class BaseSchedulerConfig(BaseConfig):
     top_k: int = Field(
         default=10, description="Number of top candidates to consider in initial retrieval"
     )
+    # TODO: The 'top_n' field is deprecated and will be removed in future versions.
     top_n: int = Field(default=5, description="Number of final results to return after processing")
     enable_parallel_dispatch: bool = Field(
         default=True, description="Whether to enable parallel message processing using thread pool"
@@ -48,7 +49,7 @@ class GeneralSchedulerConfig(BaseSchedulerConfig):
         default=300, description="Interval in seconds for updating activation memory"
     )
     context_window_size: int | None = Field(
-        default=5, description="Size of the context window for conversation history"
+        default=10, description="Size of the context window for conversation history"
     )
     act_mem_dump_path: str | None = Field(
         default=DEFAULT_ACT_MEM_DUMP_PATH,  # Replace with DEFAULT_ACT_MEM_DUMP_PATH
@@ -105,7 +106,20 @@ class RabbitMQConfig(
 
 
 class GraphDBAuthConfig(BaseConfig):
-    uri: str = Field(default="localhost", description="URI for graph database access")
+    uri: str = Field(
+        default="bolt://localhost:7687",
+        description="URI for graph database access (e.g., bolt://host:port)",
+    )
+    user: str = Field(default="neo4j", description="Username for graph database authentication")
+    password: str = Field(
+        default="",
+        description="Password for graph database authentication",
+        min_length=8,  # 建议密码最小长度
+    )
+    db_name: str = Field(default="neo4j", description="Database name to connect to")
+    auto_create: bool = Field(
+        default=True, description="Whether to automatically create the database if it doesn't exist"
+    )
 
 
 class OpenAIConfig(BaseConfig):