ai-pipeline-core 0.1.12__py3-none-any.whl → 0.4.1__py3-none-any.whl

ai_pipeline_core/tracing.py

@@ -1,475 +0,0 @@
-"""Tracing utilities that integrate Laminar (``lmnr``) with our code-base.
-
-@public
-
-This module centralizes:
-- ``TraceInfo`` - a small helper object for propagating contextual metadata.
-- ``trace`` decorator - augments a callable with Laminar tracing, automatic
-  ``observe`` instrumentation, and optional support for test runs.
-"""
-
-from __future__ import annotations
-
-import inspect
-import os
-from functools import wraps
-from typing import Any, Callable, Literal, ParamSpec, TypeVar, cast, overload
-
-from lmnr import Instruments, Laminar, observe
-from pydantic import BaseModel
-
-from ai_pipeline_core.settings import settings
-
-# ---------------------------------------------------------------------------
-# Typing helpers
-# ---------------------------------------------------------------------------
-P = ParamSpec("P")
-R = TypeVar("R")
-
-TraceLevel = Literal["always", "debug", "off"]
-"""Control level for tracing activation.
-
-@public
-
-Values:
-- "always": Always trace (default, production mode)
-- "debug": Only trace when LMNR_DEBUG == "true"
-- "off": Disable tracing completely
-"""
-
-
-# ---------------------------------------------------------------------------
-# ``TraceInfo`` – metadata container
-# ---------------------------------------------------------------------------
-class TraceInfo(BaseModel):
-    """Container for propagating trace context through the pipeline.
-
-    TraceInfo provides a structured way to pass tracing metadata through
-    function calls, ensuring consistent observability across the entire
-    execution flow. It integrates with Laminar (LMNR) for distributed
-    tracing and debugging.
-
-    Attributes:
-        session_id: Unique identifier for the current session/conversation.
-                   Falls back to LMNR_SESSION_ID environment variable.
-        user_id: Identifier for the user triggering the operation.
-                Falls back to LMNR_USER_ID environment variable.
-        metadata: Key-value pairs for additional trace context.
-                 Useful for filtering and searching in LMNR dashboard.
-        tags: List of tags for categorizing traces (e.g., ["production", "v2"]).
-
-    Environment fallbacks:
-        - LMNR_SESSION_ID: Default session_id if not explicitly set
-        - LMNR_USER_ID: Default user_id if not explicitly set
-        - LMNR_DEBUG: Controls debug-level tracing when set to "true"
-
-        Note: These variables are read directly by the tracing layer and are
-        not part of the Settings configuration.
-
-    Example:
-        >>> # Create trace context
-        >>> trace_info = TraceInfo(
-        ...     session_id="sess_123",
-        ...     user_id="user_456",
-        ...     metadata={"flow": "document_analysis", "version": "1.2"},
-        ...     tags=["production", "high_priority"]
-        ... )
-        >>>
-        >>> # Pass through function calls
-        >>> @trace
-        >>> async def process(data, trace_info: TraceInfo):
-        ...     # TraceInfo automatically propagates to nested calls
-        ...     result = await analyze(data, trace_info=trace_info)
-        ...     return result
-
-    Note:
-        TraceInfo is typically created at the entry point of a flow
-        and passed through all subsequent function calls for
-        consistent tracing context.
-    """
-
-    session_id: str | None = None
-    user_id: str | None = None
-    metadata: dict[str, str] = {}
-    tags: list[str] = []
-
-    def get_observe_kwargs(self) -> dict[str, Any]:
-        """Convert TraceInfo to kwargs for Laminar's observe decorator.
-
-        Transforms the TraceInfo fields into the format expected by
-        the lmnr.observe() decorator, applying environment variable
-        fallbacks for session_id and user_id.
-
-        Returns:
-            Dictionary with keys:
-            - session_id: From field or LMNR_SESSION_ID env var
-            - user_id: From field or LMNR_USER_ID env var
-            - metadata: Dictionary of custom metadata (if set)
-            - tags: List of tags (if set)
-
-            Only non-empty values are included in the output.
-
-        Example:
-            >>> trace_info = TraceInfo(session_id="sess_123", tags=["test"])
-            >>> kwargs = trace_info.get_observe_kwargs()
-            >>> # Returns: {"session_id": "sess_123", "tags": ["test"]}
-
-        Note:
-            This method is called internally by the trace decorator
-            to configure Laminar observation parameters.
-        """
-        kwargs: dict[str, Any] = {}
-
-        # Use environment variable fallback for session_id
-        session_id = self.session_id or os.getenv("LMNR_SESSION_ID")
-        if session_id:
-            kwargs["session_id"] = session_id
-
-        # Use environment variable fallback for user_id
-        user_id = self.user_id or os.getenv("LMNR_USER_ID")
-        if user_id:
-            kwargs["user_id"] = user_id
-
-        if self.metadata:
-            kwargs["metadata"] = self.metadata
-        if self.tags:
-            kwargs["tags"] = self.tags
-        return kwargs
-
-
-# ---------------------------------------------------------------------------
-# ``trace`` decorator
-# ---------------------------------------------------------------------------
-
-
-def _initialise_laminar() -> None:
-    """Initialize Laminar SDK with project configuration.
-
-    Sets up the Laminar observability client with the project API key
-    from settings. Disables automatic OpenAI instrumentation to avoid
-    conflicts with our custom tracing.
-
-    Configuration:
-        - Uses settings.lmnr_project_api_key for authentication
-        - Disables OPENAI instrument to prevent double-tracing
-        - Called automatically by trace decorator on first use
-
-    Note:
-        This is an internal function called once per process.
-        Multiple calls are safe (Laminar handles idempotency).
-    """
-    if settings.lmnr_project_api_key:
-        Laminar.initialize(
-            project_api_key=settings.lmnr_project_api_key,
-            disabled_instruments=[Instruments.OPENAI] if Instruments.OPENAI else [],
-        )
-
-
-# Overload for calls like @trace(name="...", level="debug")
-@overload
-def trace(
-    *,
-    level: TraceLevel = "always",
-    name: str | None = None,
-    session_id: str | None = None,
-    user_id: str | None = None,
-    metadata: dict[str, Any] | None = None,
-    tags: list[str] | None = None,
-    span_type: str | None = None,
-    ignore_input: bool = False,
-    ignore_output: bool = False,
-    ignore_inputs: list[str] | None = None,
-    input_formatter: Callable[..., str] | None = None,
-    output_formatter: Callable[..., str] | None = None,
-    ignore_exceptions: bool = False,
-    preserve_global_context: bool = True,
-) -> Callable[[Callable[P, R]], Callable[P, R]]: ...
-
-
-# Overload for the bare @trace call
-@overload
-def trace(func: Callable[P, R]) -> Callable[P, R]: ...
-
-
-# Actual implementation
-def trace(
-    func: Callable[P, R] | None = None,
-    *,
-    level: TraceLevel = "always",
-    name: str | None = None,
-    session_id: str | None = None,
-    user_id: str | None = None,
-    metadata: dict[str, Any] | None = None,
-    tags: list[str] | None = None,
-    span_type: str | None = None,
-    ignore_input: bool = False,
-    ignore_output: bool = False,
-    ignore_inputs: list[str] | None = None,
-    input_formatter: Callable[..., str] | None = None,
-    output_formatter: Callable[..., str] | None = None,
-    ignore_exceptions: bool = False,
-    preserve_global_context: bool = True,
-) -> Callable[[Callable[P, R]], Callable[P, R]] | Callable[P, R]:
-    """Add Laminar observability tracing to any function.
-
-    @public
-
-    The trace decorator integrates functions with Laminar (LMNR) for
-    distributed tracing, performance monitoring, and debugging. It
-    automatically handles both sync and async functions, propagates
-    trace context, and provides fine-grained control over what gets traced.
-
-    USAGE GUIDELINE - Defaults First:
-        In 90% of cases, use WITHOUT any parameters.
-        The defaults are optimized for most use cases.
-
-    Args:
-        func: Function to trace (when used without parentheses: @trace).
-
-        level: Controls when tracing is active:
-            - "always": Always trace (default, production mode)
-            - "debug": Only trace when LMNR_DEBUG == "true"
-            - "off": Disable tracing completely
-
-        name: Custom span name in traces (defaults to function.__name__).
-             Use descriptive names for better trace readability.
-
-        session_id: Override session ID for this function's traces.
-                   Typically propagated via TraceInfo instead.
-
-        user_id: Override user ID for this function's traces.
-                Typically propagated via TraceInfo instead.
-
-        metadata: Additional key-value metadata attached to spans.
-                 Searchable in LMNR dashboard. Merged with TraceInfo metadata.
-
-        tags: List of tags for categorizing spans (e.g., ["api", "critical"]).
-             Merged with TraceInfo tags.
-
-        span_type: Semantic type of the span (e.g., "LLM", "CHAIN", "TOOL").
-                  Affects visualization in LMNR dashboard.
-
-        ignore_input: Don't record function inputs in trace (privacy/size).
-
-        ignore_output: Don't record function output in trace (privacy/size).
-
-        ignore_inputs: List of parameter names to exclude from trace.
-                      Useful for sensitive data like API keys.
-
-        input_formatter: Custom function to format inputs for tracing.
-                        Receives all function args, returns display string.
-
-        output_formatter: Custom function to format output for tracing.
-                         Receives function result, returns display string.
-
-        ignore_exceptions: Don't record exceptions in traces (default False).
-
-        preserve_global_context: Maintain Laminar's global context across
-                                calls (default True). Set False for isolated traces.
-
-    Returns:
-        Decorated function with same signature but added tracing.
-
-    TraceInfo propagation:
-        If the decorated function has a 'trace_info' parameter, the decorator
-        automatically creates or propagates a TraceInfo instance, ensuring
-        consistent session/user tracking across the call chain.
-
-    Example:
-        >>> # RECOMMENDED - No parameters needed for most cases!
-        >>> @trace
-        >>> async def process_document(doc):
-        ...     return await analyze(doc)
-        >>>
-        >>> # With parameters (RARE - only when specifically needed):
-        >>> @trace(level="debug")  # Only for debug-specific tracing
-        >>> async def debug_operation():
-        ...     pass
-
-        >>> @trace(ignore_inputs=["api_key"])  # Only for sensitive data
-        >>> async def api_call(data, api_key):
-        ...     return await external_api(data, api_key)
-        >>>
-        >>> # AVOID unnecessary configuration - defaults handle:
-        >>> # - Automatic naming from function name
-        >>> # - Standard trace level ("always")
-        >>> # - Full input/output capture
-        >>> # - Proper span type inference
-        >>>
-        >>> # Custom formatting
-        >>> @trace(
-        ...     input_formatter=lambda doc: f"Document: {doc.id}",
-        ...     output_formatter=lambda res: f"Results: {len(res)} items"
-        >>> )
-        >>> def analyze(doc):
-        ...     return results
-
-    Environment variables:
-        - LMNR_DEBUG: Set to "true" to enable debug-level traces
-        - LMNR_SESSION_ID: Default session ID if not in TraceInfo
-        - LMNR_USER_ID: Default user ID if not in TraceInfo
-        - LMNR_PROJECT_API_KEY: Required for trace submission
-
-    Performance:
-        - Tracing overhead is minimal (~1-2ms per call)
-        - When level="off", decorator returns original function unchanged
-        - Large inputs/outputs can be excluded with ignore_* parameters
-
-    Note:
-        - Automatically initializes Laminar on first use
-        - Works with both sync and async functions
-        - Preserves function signature and metadata
-        - Thread-safe and async-safe
-
-    See Also:
-        - TraceInfo: Container for trace metadata
-        - pipeline_task: Task decorator with built-in tracing
-        - pipeline_flow: Flow decorator with built-in tracing
-    """
-    if level == "off":
-        if func:
-            return func
-        return lambda f: f
-
-    def decorator(f: Callable[P, R]) -> Callable[P, R]:
-        """Apply tracing to the target function.
-
-        Returns:
-            Wrapped function with LMNR observability.
-
-        Raises:
-            TypeError: If function is already decorated with @pipeline_task or @pipeline_flow.
-        """
-        # Check if this is already a traced pipeline_task or pipeline_flow
-        # This happens when @trace is applied after @pipeline_task/@pipeline_flow
-        if hasattr(f, "__is_traced__") and f.__is_traced__:  # type: ignore[attr-defined]
-            # Check if it's a Prefect Task or Flow object (they have specific attributes)
-            # Prefect objects have certain attributes that regular functions don't
-            is_prefect_task = hasattr(f, "fn") and hasattr(f, "submit") and hasattr(f, "map")
-            is_prefect_flow = hasattr(f, "fn") and hasattr(f, "serve")
-            if is_prefect_task or is_prefect_flow:
-                fname = getattr(f, "__name__", "function")
-                raise TypeError(
-                    f"Function '{fname}' is already decorated with @pipeline_task or "
-                    f"@pipeline_flow. Remove the @trace decorator - pipeline decorators "
-                    f"include tracing automatically."
-                )
-
-        # Handle 'debug' level logic - only trace when LMNR_DEBUG is "true"
-        if level == "debug" and os.getenv("LMNR_DEBUG", "").lower() != "true":
-            return f
-
-        # --- Pre-computation (done once when the function is decorated) ---
-        _initialise_laminar()
-        sig = inspect.signature(f)
-        is_coroutine = inspect.iscoroutinefunction(f)
-        observe_name = name or f.__name__
-        _observe = observe
-
-        # Store the new parameters
-        _session_id = session_id
-        _user_id = user_id
-        _metadata = metadata
-        _tags = tags or []
-        _span_type = span_type
-        _ignore_input = ignore_input
-        _ignore_output = ignore_output
-        _ignore_inputs = ignore_inputs
-        _input_formatter = input_formatter
-        _output_formatter = output_formatter
-        _ignore_exceptions = ignore_exceptions
-        _preserve_global_context = preserve_global_context
-
-        # --- Helper function for runtime logic ---
-        def _prepare_and_get_observe_params(runtime_kwargs: dict[str, Any]) -> dict[str, Any]:
-            """Inspects runtime args, manages TraceInfo, and returns params for lmnr.observe.
-
-            Modifies runtime_kwargs in place to inject TraceInfo if the function expects it.
-
-            Returns:
-                Dictionary of parameters for lmnr.observe decorator.
-            """
-            trace_info = runtime_kwargs.get("trace_info")
-            if not isinstance(trace_info, TraceInfo):
-                trace_info = TraceInfo()
-                if "trace_info" in sig.parameters:
-                    runtime_kwargs["trace_info"] = trace_info
-
-            observe_params = trace_info.get_observe_kwargs()
-            observe_params["name"] = observe_name
-
-            # Override with decorator-level session_id and user_id if provided
-            if _session_id:
-                observe_params["session_id"] = _session_id
-            if _user_id:
-                observe_params["user_id"] = _user_id
-
-            # Merge decorator-level metadata and tags
-            if _metadata:
-                observe_params["metadata"] = {**observe_params.get("metadata", {}), **_metadata}
-            if _tags:
-                observe_params["tags"] = observe_params.get("tags", []) + _tags
-            if _span_type:
-                observe_params["span_type"] = _span_type
-
-            # Add the new Laminar parameters
-            if _ignore_input:
-                observe_params["ignore_input"] = _ignore_input
-            if _ignore_output:
-                observe_params["ignore_output"] = _ignore_output
-            if _ignore_inputs is not None:
-                observe_params["ignore_inputs"] = _ignore_inputs
-            if _input_formatter is not None:
-                observe_params["input_formatter"] = _input_formatter
-            if _output_formatter is not None:
-                observe_params["output_formatter"] = _output_formatter
-            if _ignore_exceptions:
-                observe_params["ignore_exceptions"] = _ignore_exceptions
-            if _preserve_global_context:
-                observe_params["preserve_global_context"] = _preserve_global_context
-
-            return observe_params
-
-        # --- The actual wrappers ---
-        @wraps(f)
-        def sync_wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
-            """Synchronous wrapper for traced function.
-
-            Returns:
-                The result of the wrapped function.
-            """
-            observe_params = _prepare_and_get_observe_params(kwargs)
-            observed_func = _observe(**observe_params)(f)
-            return observed_func(*args, **kwargs)
-
-        @wraps(f)
-        async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
-            """Asynchronous wrapper for traced function.
-
-            Returns:
-                The result of the wrapped function.
-            """
-            observe_params = _prepare_and_get_observe_params(kwargs)
-            observed_func = _observe(**observe_params)(f)
-            return await observed_func(*args, **kwargs)
-
-        wrapper = async_wrapper if is_coroutine else sync_wrapper
-
-        # Mark function as traced for detection by pipeline decorators
-        wrapper.__is_traced__ = True  # type: ignore[attr-defined]
-
-        # Preserve the original function signature
-        try:
-            wrapper.__signature__ = sig  # type: ignore[attr-defined]
-        except (AttributeError, ValueError):
-            pass
-
-        return cast(Callable[P, R], wrapper)
-
-    if func:
-        return decorator(func)  # Called as @trace
-    else:
-        return decorator  # Called as @trace(...)
-
-
-__all__ = ["trace", "TraceLevel", "TraceInfo"]