MemoryOS 0.2.0__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/exceptions.py

@@ -0,0 +1,28 @@
+import logging
+
+from fastapi.requests import Request
+from fastapi.responses import JSONResponse
+
+
+logger = logging.getLogger(__name__)
+
+
+class APIExceptionHandler:
+    """Centralized exception handling for MemOS APIs."""
+
+    @staticmethod
+    async def value_error_handler(request: Request, exc: ValueError):
+        """Handle ValueError exceptions globally."""
+        return JSONResponse(
+            status_code=400,
+            content={"code": 400, "message": str(exc), "data": None},
+        )
+
+    @staticmethod
+    async def global_exception_handler(request: Request, exc: Exception):
+        """Handle all unhandled exceptions globally."""
+        logger.exception("Unhandled error:")
+        return JSONResponse(
+            status_code=500,
+            content={"code": 500, "message": str(exc), "data": None},
+        )