MemoryOS 2.0.3__py3-none-any.whl

memos/context/context.py

@@ -0,0 +1,355 @@
+"""
+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 functools
+import os
+import threading
+
+from collections.abc import Callable
+from concurrent.futures import ThreadPoolExecutor
+from contextvars import ContextVar
+from typing import Any, TypeVar
+
+
+T = TypeVar("T")
+
+# 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,
+        api_path: str | None = None,
+        env: str | None = None,
+        user_type: str | None = None,
+        user_name: str | None = None,
+        source: str | None = None,
+    ):
+        self.trace_id = trace_id or "trace-id"
+        self.api_path = api_path
+        self.env = env
+        self.user_type = user_type
+        self.user_name = user_name
+        self.source = source
+        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 in (
+            "trace_id",
+            "api_path",
+            "env",
+            "user_type",
+            "user_name",
+            "source",
+        ):
+            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,
+            "api_path": self.api_path,
+            "env": self.env,
+            "user_type": self.user_type,
+            "user_name": self.user_name,
+            "source": self.source,
+            "data": self._data.copy(),
+        }
+
+
+def set_request_context(context: RequestContext | None) -> None:
+    """
+    Set the current request context.
+
+    This is typically called by the API dependency injection system.
+    """
+    if context:
+        _request_context.set(context.to_dict())
+    else:
+        _request_context.set(None)
+
+
+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_api_path() -> str | None:
+    """
+    Get the current request's api path.
+    """
+    context = _request_context.get()
+    if context:
+        return context.get("api_path")
+    return None
+
+
+def get_current_env() -> str | None:
+    """
+    Get the current request's env.
+    """
+    context = _request_context.get()
+    if context:
+        return context.get("env")
+    return "prod"
+
+
+def get_current_user_type() -> str | None:
+    """
+    Get the current request's user type.
+    """
+    context = _request_context.get()
+    if context:
+        return context.get("user_type")
+    return "opensource"
+
+
+def get_current_user_name() -> str | None:
+    """
+    Get the current request's user name.
+    """
+    context = _request_context.get()
+    if context:
+        return context.get("user_name")
+    return "memos"
+
+
+def get_current_source() -> str | None:
+    """
+    Get the current request's source (e.g., 'product_api' or 'server_api').
+    """
+    context = _request_context.get()
+    if context:
+        return context.get("source")
+    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"),
+            api_path=context_dict.get("api_path"),
+            env=context_dict.get("env"),
+            user_type=context_dict.get("user_type"),
+            user_name=context_dict.get("user_name"),
+            source=context_dict.get("source"),
+        )
+        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
+
+
+class ContextThread(threading.Thread):
+    """
+    Thread class that automatically propagates the main thread's trace_id to child threads.
+    """
+
+    def __init__(self, target, args=(), kwargs=None, **thread_kwargs):
+        super().__init__(**thread_kwargs)
+        self.target = target
+        self.args = args
+        self.kwargs = kwargs or {}
+
+        self.main_trace_id = get_current_trace_id()
+        self.main_api_path = get_current_api_path()
+        self.main_env = get_current_env()
+        self.main_user_type = get_current_user_type()
+        self.main_user_name = get_current_user_name()
+        self.main_context = get_current_context()
+
+    def run(self):
+        # Create a new RequestContext with the main thread's trace_id
+        if self.main_context:
+            # Copy the context data
+            child_context = RequestContext(
+                trace_id=self.main_trace_id,
+                api_path=self.main_api_path,
+                env=self.main_env,
+                user_type=self.main_user_type,
+                user_name=self.main_user_name,
+            )
+            child_context._data = self.main_context._data.copy()
+
+            # Set the context in the child thread
+            set_request_context(child_context)
+
+        # Run the target function
+        self.target(*self.args, **self.kwargs)
+
+
+class ContextThreadPoolExecutor(ThreadPoolExecutor):
+    """
+    ThreadPoolExecutor that automatically propagates the main thread's trace_id to worker threads.
+    """
+
+    def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Any:
+        """
+        Submit a callable to be executed with the given arguments.
+        Automatically propagates the current thread's context to the worker thread.
+        """
+        main_trace_id = get_current_trace_id()
+        main_api_path = get_current_api_path()
+        main_env = get_current_env()
+        main_user_type = get_current_user_type()
+        main_user_name = get_current_user_name()
+        main_context = get_current_context()
+
+        @functools.wraps(fn)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            if main_context:
+                # Create and set new context in worker thread
+                child_context = RequestContext(
+                    trace_id=main_trace_id,
+                    api_path=main_api_path,
+                    env=main_env,
+                    user_type=main_user_type,
+                    user_name=main_user_name,
+                )
+                child_context._data = main_context._data.copy()
+                set_request_context(child_context)
+
+            return fn(*args, **kwargs)
+
+        return super().submit(wrapper, *args, **kwargs)
+
+    def map(
+        self,
+        fn: Callable[..., T],
+        *iterables: Any,
+        timeout: float | None = None,
+        chunksize: int = 1,
+    ) -> Any:
+        """
+        Returns an iterator equivalent to map(fn, iter).
+        Automatically propagates the current thread's context to worker threads.
+        """
+        main_trace_id = get_current_trace_id()
+        main_api_path = get_current_api_path()
+        main_env = get_current_env()
+        main_user_type = get_current_user_type()
+        main_user_name = get_current_user_name()
+        main_context = get_current_context()
+
+        @functools.wraps(fn)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            if main_context:
+                # Create and set new context in worker thread
+                child_context = RequestContext(
+                    trace_id=main_trace_id,
+                    api_path=main_api_path,
+                    env=main_env,
+                    user_type=main_user_type,
+                    user_name=main_user_name,
+                )
+                child_context._data = main_context._data.copy()
+                set_request_context(child_context)
+
+            return fn(*args, **kwargs)
+
+        return super().map(wrapper, *iterables, timeout=timeout, chunksize=chunksize)
+
+
+# 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 generate_trace_id() -> str:
+    """Generate a random trace_id."""
+    return os.urandom(16).hex()
+
+
+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 general_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/dependency.py

@@ -0,0 +1,52 @@
+"""
+This utility provides tools for managing dependencies in MemOS.
+"""
+
+import functools
+import importlib
+
+
+def require_python_package(
+    import_name: str, install_command: str | None = None, install_link: str | None = None
+):
+    """Check if a package is available and provide installation hints on import failure.
+
+    Args:
+        import_name (str): The top-level importable module name a package provides.
+        install_command (str, optional): Installation command.
+        install_link (str, optional): URL link to installation guide.
+
+    Returns:
+        Callable: A decorator function that wraps the target function with package availability check.
+
+    Raises:
+        ImportError: When the specified package is not available, with installation
+            instructions included in the error message.
+
+    Example:
+        >>> @require_python_package(
+        ...     import_name='faiss',
+        ...     install_command='pip install faiss-cpu',
+        ...     install_link='https://github.com/facebookresearch/faiss/blob/main/INSTALL.md'
+        ... )
+        ... def create_faiss_index():
+        ...     from faiss import IndexFlatL2  # Actual import in function
+        ...     return IndexFlatL2(128)
+    """
+
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            try:
+                importlib.import_module(import_name)
+            except ImportError:
+                error_msg = f"Missing required module - '{import_name}'\n"
+                error_msg += f"💡 Install command: {install_command}\n" if install_command else ""
+                error_msg += f"💡 Install guide:   {install_link}\n" if install_link else ""
+
+                raise ImportError(error_msg) from None
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator

memos/deprecation.py

@@ -0,0 +1,262 @@
+"""
+This module provides utilities for marking functions, classes, and parameters
+as deprecated. It includes decorators for deprecation, a function to issue
+warnings, and utilities to check deprecation status.
+"""
+
+import functools
+import warnings
+
+from collections.abc import Callable
+from typing import Any, TypeVar
+
+
+warnings.simplefilter("default", DeprecationWarning)
+
+
+F = TypeVar("F", bound=Callable[..., Any])
+C = TypeVar("C", bound=type)
+
+
+def deprecated(
+    reason: str | None = None,
+    version: str | None = None,
+    alternative: str | None = None,
+    category: type[Warning] = DeprecationWarning,
+    stacklevel: int = 2,
+) -> Callable[[F], F]:
+    """
+    Decorator to mark functions as deprecated.
+
+    Args:
+        reason: Optional reason for deprecation
+        version: Version when the function was deprecated
+        alternative: Suggested alternative function/method
+        category: Warning category to use
+        stacklevel: Stack level for the warning
+
+    Example:
+        @deprecated(reason="Use new_function instead", version="1.2.0")
+        def old_function():
+            pass
+    """
+
+    def decorator(func: F) -> F:
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            # Build deprecation message
+            msg_parts = [f"Function '{func.__name__}' is deprecated"]
+
+            if version:
+                msg_parts.append(f"since version {version}")
+
+            if reason:
+                msg_parts.append(f"- {reason}")
+
+            if alternative:
+                msg_parts.append(f"Use '{alternative}' instead")
+
+            message = ". ".join(msg_parts) + "."
+
+            warnings.warn(message, category=category, stacklevel=stacklevel)
+            return func(*args, **kwargs)
+
+        # Mark the wrapper as deprecated for introspection
+        wrapper.__deprecated__ = True
+        wrapper.__deprecation_info__ = {
+            "reason": reason,
+            "version": version,
+            "alternative": alternative,
+            "category": category,
+        }
+
+        return wrapper
+
+    return decorator
+
+
+def deprecated_class(
+    reason: str | None = None,
+    version: str | None = None,
+    alternative: str | None = None,
+    category: type[Warning] = DeprecationWarning,
+    stacklevel: int = 2,
+) -> Callable[[C], C]:
+    """
+    Decorator to mark classes as deprecated.
+
+    Args:
+        reason: Optional reason for deprecation
+        version: Version when the class was deprecated
+        alternative: Suggested alternative class
+        category: Warning category to use
+        stacklevel: Stack level for the warning
+
+    Example:
+        @deprecated_class(reason="Use NewClass instead", version="1.2.0")
+        class OldClass:
+            pass
+    """
+
+    def decorator(cls: C) -> C:
+        # Store original __init__
+        original_init = cls.__init__
+
+        @functools.wraps(original_init)
+        def new_init(self, *args, **kwargs):
+            # Build deprecation message
+            msg_parts = [f"Class '{cls.__name__}' is deprecated"]
+
+            if version:
+                msg_parts.append(f"since version {version}")
+
+            if reason:
+                msg_parts.append(f"- {reason}")
+
+            if alternative:
+                msg_parts.append(f"Use '{alternative}' instead")
+
+            message = ". ".join(msg_parts) + "."
+
+            warnings.warn(message, category=category, stacklevel=stacklevel)
+            original_init(self, *args, **kwargs)
+
+        # Replace __init__
+        cls.__init__ = new_init
+
+        # Mark the class as deprecated for introspection
+        cls.__deprecated__ = True
+        cls.__deprecation_info__ = {
+            "reason": reason,
+            "version": version,
+            "alternative": alternative,
+            "category": category,
+        }
+
+        return cls
+
+    return decorator
+
+
+def deprecated_parameter(
+    parameter_name: str,
+    reason: str | None = None,
+    version: str | None = None,
+    alternative: str | None = None,
+    category: type[Warning] = DeprecationWarning,
+    stacklevel: int = 2,
+) -> Callable[[F], F]:
+    """
+    Decorator to mark specific parameters as deprecated.
+
+    Args:
+        parameter_name: Name of the deprecated parameter
+        reason: Optional reason for deprecation
+        version: Version when the parameter was deprecated
+        alternative: Suggested alternative parameter
+        category: Warning category to use
+        stacklevel: Stack level for the warning
+
+    Example:
+        @deprecated_parameter("old_param", alternative="new_param", version="1.2.0")
+        def my_function(new_param=None, old_param=None):
+            pass
+    """
+
+    def decorator(func: F) -> F:
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            # Check if deprecated parameter is used
+            if parameter_name in kwargs:
+                # Build deprecation message
+                msg_parts = [
+                    f"Parameter '{parameter_name}' in function '{func.__name__}' is deprecated"
+                ]
+
+                if version:
+                    msg_parts.append(f"since version {version}")
+
+                if reason:
+                    msg_parts.append(f"- {reason}")
+
+                if alternative:
+                    msg_parts.append(f"Use parameter '{alternative}' instead")
+
+                message = ". ".join(msg_parts) + "."
+
+                warnings.warn(message, category=category, stacklevel=stacklevel)
+
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+def warn_deprecated(
+    item_name: str,
+    item_type: str = "feature",
+    reason: str | None = None,
+    version: str | None = None,
+    alternative: str | None = None,
+    category: type[Warning] = DeprecationWarning,
+    stacklevel: int = 2,
+) -> None:
+    """
+    Issue a deprecation warning for any item.
+
+    Args:
+        item_name: Name of the deprecated item
+        item_type: Type of item (e.g., "function", "class", "parameter", "feature")
+        reason: Optional reason for deprecation
+        version: Version when the item was deprecated
+        alternative: Suggested alternative
+        category: Warning category to use
+        stacklevel: Stack level for the warning
+
+    Example:
+        warn_deprecated("old_method", "method", version="1.2.0", alternative="new_method")
+    """
+    # Build deprecation message
+    msg_parts = [f"{item_type.capitalize()} '{item_name}' is deprecated"]
+
+    if version:
+        msg_parts.append(f"since version {version}")
+
+    if reason:
+        msg_parts.append(f"- {reason}")
+
+    if alternative:
+        msg_parts.append(f"Use '{alternative}' instead")
+
+    message = ". ".join(msg_parts) + "."
+
+    warnings.warn(message, category=category, stacklevel=stacklevel)
+
+
+def is_deprecated(obj: Any) -> bool:
+    """
+    Check if an object is marked as deprecated.
+
+    Args:
+        obj: Object to check
+
+    Returns:
+        True if the object is deprecated, False otherwise
+    """
+    return getattr(obj, "__deprecated__", False)
+
+
+def get_deprecation_info(obj: Any) -> dict | None:
+    """
+    Get deprecation information for an object.
+
+    Args:
+        obj: Object to get deprecation info for
+
+    Returns:
+        Dictionary with deprecation info or None if not deprecated
+    """
+    if is_deprecated(obj):
+        return getattr(obj, "__deprecation_info__", None)
+    return None