mirascope 1.22.4__py3-none-any.whl → 2.1.1__py3-none-any.whl

mirascope/ops/_internal/instrumentation/providers/openai.py

@@ -0,0 +1,82 @@
+"""OpenTelemetry instrumentation for OpenAI SDK."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from opentelemetry.instrumentation.openai_v2 import OpenAIInstrumentor
+
+from .base import BaseInstrumentation, ContentCaptureMode
+
+if TYPE_CHECKING:
+    from opentelemetry.trace import TracerProvider
+
+
+class OpenAIInstrumentation(BaseInstrumentation[OpenAIInstrumentor]):
+    """Manages OpenTelemetry instrumentation lifecycle for the OpenAI SDK."""
+
+    def _create_instrumentor(self) -> OpenAIInstrumentor:
+        """Create a new OpenAI instrumentor instance."""
+        return OpenAIInstrumentor()
+
+    def _configure_capture_content(self, capture_content: ContentCaptureMode) -> None:
+        """Configure environment variables for OpenAI content capture."""
+        if capture_content == "enabled":
+            self._set_env_var(
+                "OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT", "true"
+            )
+        elif capture_content == "disabled":
+            self._set_env_var(
+                "OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT", "false"
+            )
+
+
+def instrument_openai(
+    *,
+    tracer_provider: TracerProvider | None = None,
+    capture_content: ContentCaptureMode = "default",
+) -> None:
+    """Enable OpenTelemetry instrumentation for the OpenAI Python SDK.
+
+    Uses the provided tracer_provider or the global OpenTelemetry tracer provider.
+
+    Args:
+        tracer_provider: Optional tracer provider to use. If not provided,
+            uses the global OpenTelemetry tracer provider.
+        capture_content: Controls whether to capture message content in spans.
+            - "enabled": Capture message content
+            - "disabled": Do not capture message content
+            - "default": Use the library's default behavior
+
+    Example:
+
+        Enable instrumentation with Mirascope Cloud:
+        ```python
+        from mirascope import ops
+
+        ops.configure()
+        ops.instrument_openai()
+        ```
+
+        Enable instrumentation with content capture:
+        ```python
+        from mirascope import ops
+
+        ops.configure()
+        ops.instrument_openai(capture_content="enabled")
+        ```
+    """
+    OpenAIInstrumentation().instrument(
+        tracer_provider=tracer_provider,
+        capture_content=capture_content,
+    )
+
+
+def uninstrument_openai() -> None:
+    """Disable previously configured OpenAI instrumentation."""
+    OpenAIInstrumentation().uninstrument()
+
+
+def is_openai_instrumented() -> bool:
+    """Return whether OpenAI instrumentation is currently active."""
+    return OpenAIInstrumentation().is_instrumented

mirascope/ops/_internal/propagation.py

@@ -0,0 +1,198 @@
+"""Context propagation utilities for distributed tracing."""
+
+from __future__ import annotations
+
+import logging
+import os
+from collections.abc import Mapping, MutableMapping
+from typing import Literal, TypeAlias
+
+from opentelemetry import propagate
+from opentelemetry.context import Context
+from opentelemetry.propagators.b3 import B3MultiFormat, B3SingleFormat
+from opentelemetry.propagators.composite import CompositePropagator
+from opentelemetry.propagators.jaeger import JaegerPropagator
+from opentelemetry.propagators.textmap import Setter, TextMapPropagator
+from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
+
+from ..exceptions import ConfigurationError
+from .session import SESSION_HEADER_NAME, current_session
+
+logger = logging.getLogger(__name__)
+
+PropagatorFormat: TypeAlias = Literal[
+    "tracecontext", "b3", "b3multi", "jaeger", "composite"
+]
+CarrierValue: TypeAlias = str | list[str]
+
+# Environment variable names
+ENV_PROPAGATOR_FORMAT = "MIRASCOPE_PROPAGATOR"
+ENV_PROPAGATOR_SET_GLOBAL = "_MIRASCOPE_PROPAGATOR_SET_GLOBAL"
+
+_PROPAGATOR: ContextPropagator | None = None
+
+
+class _StrSetter(Setter[MutableMapping[str, str]]):
+    """Setter that writes string header values into the carrier."""
+
+    def set(self, carrier: MutableMapping[str, str], key: str, value: str) -> None:
+        """Set a single header value on the carrier."""
+        carrier[key] = value
+
+
+_STR_SETTER = _StrSetter()
+
+
+class ContextPropagator:
+    """Manages OpenTelemetry context propagation across service boundaries."""
+
+    _propagator: TextMapPropagator
+
+    def __init__(self, *, set_global: bool = True) -> None:
+        """Initialize propagator and optionally set as global.
+
+        Propagator format is determined by MIRASCOPE_PROPAGATOR environment variable.
+        Defaults to "tracecontext" if not set.
+
+        Args:
+            set_global: Whether to set this propagator as the global textmap propagator.
+
+        Raises:
+            ConfigurationError: If the propagator format is invalid.
+        """
+        env_format = os.environ.get(ENV_PROPAGATOR_FORMAT, "tracecontext")
+        propagator_name: PropagatorFormat
+        if env_format in ("tracecontext", "b3", "b3multi", "jaeger", "composite"):
+            propagator_name = env_format
+        else:
+            error_message = (
+                f"Invalid propagator format: {env_format}. "
+                f"Valid options: tracecontext, b3, b3multi, jaeger, composite"
+            )
+            logger.error(error_message)
+            raise ConfigurationError(error_message)
+        logger.debug(f"Initializing ContextPropagator with format: {propagator_name}")
+
+        match propagator_name:
+            case "tracecontext":
+                self._propagator = TraceContextTextMapPropagator()
+            case "b3":
+                self._propagator = B3SingleFormat()
+            case "b3multi":
+                self._propagator = B3MultiFormat()
+            case "jaeger":
+                self._propagator = JaegerPropagator()
+            case "composite":
+                propagators: list[TextMapPropagator] = [
+                    TraceContextTextMapPropagator(),
+                    B3SingleFormat(),
+                    B3MultiFormat(),
+                    JaegerPropagator(),
+                ]
+                self._propagator = CompositePropagator(propagators)
+
+        if set_global:
+            should_set_global = os.environ.get(ENV_PROPAGATOR_SET_GLOBAL, "true")
+            if should_set_global.lower() == "true":
+                propagate.set_global_textmap(self._propagator)
+                logger.debug(f"Set {propagator_name} as global textmap propagator")
+
+    def extract_context(self, carrier: Mapping[str, CarrierValue]) -> Context:
+        """Extract OTEL context from carrier headers.
+
+        Args:
+            carrier: Dictionary containing HTTP headers or similar carrier data.
+
+        Returns:
+            Extracted OpenTelemetry context. Returns empty context if extraction fails.
+        """
+        try:
+            context = self._propagator.extract(carrier=carrier)
+            logger.debug("Successfully extracted context from carrier")
+            return context
+        except Exception as exception:
+            logger.debug(
+                f"Failed to extract context from carrier: {type(exception).__name__}: {exception}"
+            )
+            return Context()
+
+    def inject_context(
+        self,
+        carrier: MutableMapping[str, str],
+        context: Context | None = None,
+    ) -> None:
+        """Inject current OTEL context into carrier headers.
+
+        This method also injects session context if one is active, adding the
+        SESSION_HEADER_NAME header to the carrier.
+
+        Args:
+            carrier: Mutable mapping (e.g., HTTP headers dict) to inject context into.
+            context: Optional specific context to inject. If None, uses current context.
+        """
+        try:
+            self._propagator.inject(
+                carrier=carrier, context=context, setter=_STR_SETTER
+            )
+            logger.debug("Successfully injected context into carrier")
+        except Exception as exception:
+            logger.debug(
+                f"Failed to inject context into carrier: {type(exception).__name__}: {exception}"
+            )
+
+        session_ctx = current_session()
+        if session_ctx is not None:
+            carrier[SESSION_HEADER_NAME] = session_ctx.id
+
+
+def get_propagator() -> ContextPropagator:
+    """Get or create the singleton ContextPropagator instance.
+
+    Reads propagator format from MIRASCOPE_PROPAGATOR environment variable.
+
+    Returns:
+        The global ContextPropagator instance.
+    """
+    global _PROPAGATOR
+    if _PROPAGATOR is None:
+        _PROPAGATOR = ContextPropagator()
+    return _PROPAGATOR
+
+
+def reset_propagator() -> None:
+    """Reset the singleton ContextPropagator instance.
+
+    This is primarily useful for testing to ensure a clean state between tests.
+    The next call to get_propagator() will create a new instance.
+    """
+    global _PROPAGATOR
+    _PROPAGATOR = None
+
+
+def extract_context(carrier: Mapping[str, CarrierValue]) -> Context:
+    """Extract OTEL context from carrier headers using the global propagator.
+
+    Args:
+        carrier: Dictionary containing HTTP headers or similar carrier data.
+
+    Returns:
+        Extracted OpenTelemetry context.
+    """
+    return get_propagator().extract_context(carrier)
+
+
+def inject_context(
+    carrier: MutableMapping[str, str],
+    *,
+    context: Context | None = None,
+) -> None:
+    """Inject current OTEL context into carrier headers using the global propagator.
+
+    This function also injects session context if one is active, adding the
+    SESSION_HEADER_NAME header to the carrier.
+
+    Args:
+        carrier: Mutable mapping (e.g., HTTP headers dict) to inject context into.
+        context: Optional specific context to inject. If None, uses current context.
+    """
+    get_propagator().inject_context(carrier, context)

mirascope/ops/_internal/protocols.py

@@ -0,0 +1,133 @@
+"""Call protocol helpers for ops.tracing."""
+
+from __future__ import annotations
+
+import inspect
+from typing import TYPE_CHECKING, Protocol
+from typing_extensions import TypeIs
+
+from ...llm.context import Context, DepsT
+from .types import P, R
+
+if TYPE_CHECKING:
+    from .spans import Span
+
+
+class SyncFunction(Protocol[P, R]):
+    """Protocol for synchronous callable functions."""
+
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R:
+        """The function required a synchronous call method."""
+        ...  # pragma: no cover
+
+
+class AsyncFunction(Protocol[P, R]):
+    """Protocol for asynchronous callable functions."""
+
+    async def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R:
+        """The function's required asynchronous call method."""
+        ...  # pragma: no cover
+
+
+class SyncSpanFunction(Protocol[P, R]):
+    """Protocol for synchronous functions that receive injected Span.
+
+    Functions matching this protocol have `trace_ctx: Span` as their first
+    parameter. The trace decorator will inject the span automatically.
+    """
+
+    def __call__(self, trace_ctx: Span, *args: P.args, **kwargs: P.kwargs) -> R:
+        """The function receives a Span as first parameter."""
+        ...  # pragma: no cover
+
+
+class AsyncSpanFunction(Protocol[P, R]):
+    """Protocol for asynchronous functions that receive injected Span.
+
+    Functions matching this protocol have `trace_ctx: Span` as their first
+    parameter. The trace decorator will inject the span automatically.
+    """
+
+    async def __call__(self, trace_ctx: Span, *args: P.args, **kwargs: P.kwargs) -> R:
+        """The function receives a Span as first parameter."""
+        ...  # pragma: no cover
+
+
+class SyncContextFunction(Protocol[P, DepsT, R]):
+    """Protocol for synchronous callable functions with Context parameter."""
+
+    def __call__(self, ctx: Context[DepsT], *args: P.args, **kwargs: P.kwargs) -> R:
+        """The function required a synchronous call method with context."""
+        ...  # pragma: no cover
+
+
+class AsyncContextFunction(Protocol[P, DepsT, R]):
+    """Protocol for asynchronous callable functions with Context parameter."""
+
+    async def __call__(
+        self, ctx: Context[DepsT], *args: P.args, **kwargs: P.kwargs
+    ) -> R:
+        """The function's required asynchronous call method with context."""
+        ...  # pragma: no cover
+
+
+def fn_is_async(
+    fn: SyncFunction[P, R] | AsyncFunction[P, R],
+) -> TypeIs[AsyncFunction[P, R]]:
+    """Type check to determine if a given function is asynchronous."""
+    return inspect.iscoroutinefunction(fn) or inspect.isasyncgenfunction(fn)
+
+
+def fn_wants_span(
+    fn: (
+        SyncFunction[P, R]
+        | AsyncFunction[P, R]
+        | SyncSpanFunction[P, R]
+        | AsyncSpanFunction[P, R]
+    ),
+) -> TypeIs[SyncSpanFunction[P, R] | AsyncSpanFunction[P, R]]:
+    """Check if function wants Span injection as first parameter.
+
+    Returns True if the function has a first parameter named `trace_ctx`
+    with type annotation `Span`.
+    """
+    # Import here to avoid circular imports
+    from .spans import Span
+
+    try:
+        sig = inspect.signature(fn)
+    except (ValueError, TypeError):
+        return False
+
+    params = list(sig.parameters.values())
+    if not params:
+        return False
+
+    first_param = params[0]
+    if first_param.name != "trace_ctx":
+        return False
+
+    # Check annotation
+    annotation = first_param.annotation
+    if annotation is inspect.Parameter.empty:
+        return False
+
+    # Handle string annotations (forward references)
+    # The annotation could be "Span", "ops.Span", "mirascope.ops.Span", etc.
+    if isinstance(annotation, str):
+        return annotation == "Span" or annotation.endswith(".Span")
+
+    # Check by identity first
+    if annotation is Span:
+        return True
+
+    # Fallback: check by class name and module for robustness
+    # This handles cases where the same class might have different identities
+    # due to module reloading or import issues in test environments
+    if isinstance(annotation, type):
+        return (
+            annotation.__name__ == "Span"
+            and annotation.__module__ == "mirascope.ops._internal.spans"
+        )
+
+    return False

mirascope/ops/_internal/session.py

@@ -0,0 +1,139 @@
+"""Session context helpers for grouping traces."""
+
+from __future__ import annotations
+
+import uuid
+from collections.abc import Iterator, Mapping
+from contextlib import contextmanager
+from contextvars import ContextVar, Token
+from dataclasses import dataclass, field
+
+from mirascope.ops._internal.types import Jsonable
+
+SESSION_HEADER_NAME = "Mirascope-Session-Id"
+
+
+SESSION_CONTEXT: ContextVar[SessionContext | None] = ContextVar(
+    "MIRASCOPE_SESSION_CONTEXT", default=None
+)
+
+
+@dataclass(slots=True)
+class SessionContext:
+    """Represents a session context for grouping related spans."""
+
+    id: str = field(default_factory=lambda: str(uuid.uuid4()))
+    """Unique identifier for the session. Auto-generated if not provided."""
+
+    attributes: Mapping[str, Jsonable] | None = None
+    """Optional JSON-serializable metadata associated with the session."""
+
+
+@contextmanager
+def session(
+    *, id: str | None = None, attributes: Mapping[str, Jsonable] | None = None
+) -> Iterator[SessionContext]:
+    """Context manager for setting session context.
+
+    Sessions are used to group related traces together. The session ID and
+    optional attributes are automatically propagated to all spans created
+    within the session context and are included in outgoing HTTP requests.
+
+    Args:
+        id: Unique identifier for the session. If not provided, a UUID will be
+            automatically generated.
+        attributes: Optional dictionary of JSON-serializable attributes to
+            attach to the session.
+
+    Example:
+        ```python
+        # With explicit ID
+        with mirascope.ops.session(id="user-123") as ctx:
+            print(ctx.id)  # "user-123"
+            result = requests.get("https://api.example.com")
+
+        # With auto-generated ID
+        with mirascope.ops.session() as ctx:
+            print(ctx.id)  # Auto-generated UUID
+            result = requests.get("https://api.example.com")
+
+        # Nested sessions override parent session
+        with mirascope.ops.session(id="1"):
+            # Session ID: 1
+            with mirascope.ops.session(id="2"):
+                # Session ID: 2
+            # Session ID: 1
+        ```
+    """
+    if id is not None:
+        session_ctx = SessionContext(id=id, attributes=attributes)
+    else:
+        session_ctx = SessionContext(attributes=attributes)
+    token: Token[SessionContext | None] = SESSION_CONTEXT.set(session_ctx)
+    try:
+        yield session_ctx
+    finally:
+        SESSION_CONTEXT.reset(token)
+
+
+def current_session() -> SessionContext | None:
+    """Get the current session context if one is active.
+
+    Returns:
+        The active SessionContext or None if no session is active.
+
+    Example:
+        ```python
+        with mirascope.ops.session(id="user-123"):
+            ctx = mirascope.ops.current_session()
+            print(ctx.id)  # "user-123"
+
+        ctx = mirascope.ops.current_session()
+        print(ctx)  # None
+        ```
+    """
+    return SESSION_CONTEXT.get()
+
+
+def extract_session_id(headers: Mapping[str, str | list[str]]) -> str | None:
+    """Extract session ID from carrier headers.
+
+    Performs case-insensitive header name matching and handles both string
+    and list[str] header values.
+
+    Args:
+        headers: Dictionary of HTTP headers or similar carrier data.
+
+    Returns:
+        The extracted session ID or None if not present.
+
+    Example:
+        ```python
+        headers = {"mirascope-session-id": "session-123"}
+        session_id = extract_session_id(headers)
+        print(session_id)  # "session-123"
+        ```
+    """
+    normalized_headers = {key.lower(): value for key, value in headers.items()}
+    header_name_lower = SESSION_HEADER_NAME.lower()
+
+    if header_name_lower not in normalized_headers:
+        return None
+
+    value = normalized_headers[header_name_lower]
+
+    if isinstance(value, list):
+        if len(value) > 0:
+            return value[0]
+        return None
+
+    return value if value else None
+
+
+__all__ = [
+    "SESSION_HEADER_NAME",
+    "SessionContext",
+    "current_session",
+    "extract_session_id",
+    "session",
+]