mirascope 2.0.0a2__py3-none-any.whl → 2.0.0a3__py3-none-any.whl

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,51 @@
+"""Call protocol helpers for ops.tracing."""
+
+from __future__ import annotations
+
+import inspect
+from typing import Protocol
+from typing_extensions import TypeIs
+
+from ...llm.context import Context, DepsT
+from .types import P, R
+
+
+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 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)

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",
+]

mirascope/ops/_internal/spans.py

@@ -0,0 +1,232 @@
+"""Explicit span management utilities for `mirascope.ops`."""
+
+from __future__ import annotations
+
+import logging
+from contextvars import Token
+from typing import TYPE_CHECKING, Any
+
+from opentelemetry import context as otel_context, trace as otel_trace
+from opentelemetry.trace import (
+    SpanContext,
+    Status,
+    StatusCode,
+    format_span_id,
+    format_trace_id,
+)
+from opentelemetry.util.types import AttributeValue
+
+from .session import current_session
+from .utils import json_dumps
+
+if TYPE_CHECKING:
+    from opentelemetry.context import Context
+
+logger = logging.getLogger("mirascope.ops")
+_warned_noop = False
+
+
+class Span:
+    """Context-managed span for explicit tracing.
+
+    Creates a child span within the current trace context. Acts as a no-op
+    if tracing is not configured.
+    """
+
+    def __init__(self, name: str, **attributes: AttributeValue) -> None:
+        """Initialize a new span with the given name.
+
+        Args:
+            name: Name for the span.
+            **attributes: Initial attributes to set on the span.
+        """
+        self._name = name
+        self._initial_attributes = attributes
+        self._span: otel_trace.Span | None = None
+        self._token: Token[Context] | None = None
+        self._is_noop = True
+        self._finished = False
+
+    def __enter__(self) -> Span:
+        """Enter the span context.
+
+        Returns:
+            This span instance for use within the context.
+        """
+        tracer = otel_trace.get_tracer("mirascope.ops")
+        self._span = tracer.start_span(self._name)
+
+        if self._span.__class__.__name__ == "NonRecordingSpan":
+            self._is_noop = True
+            self._span = None
+            global _warned_noop
+            if not _warned_noop:
+                logger.warning(
+                    f"mirascope tracing is not configured; Span('{self._name}') is a no-op."
+                )
+                _warned_noop = True
+        else:
+            self._is_noop = False
+            self._span.set_attribute("mirascope.type", "trace")
+
+            session_ctx = current_session()
+            if session_ctx is not None:
+                self._span.set_attribute("mirascope.ops.session.id", session_ctx.id)
+                if session_ctx.attributes is not None:
+                    self._span.set_attribute(
+                        "mirascope.ops.session.attributes",
+                        json_dumps(session_ctx.attributes),
+                    )
+
+            self._token = otel_context.attach(
+                otel_trace.set_span_in_context(self._span)
+            )
+
+            if self._initial_attributes:
+                self.set(**self._initial_attributes)
+
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: Any,  # noqa: ANN401
+    ) -> None:
+        """Exit the span context.
+
+        Args:
+            exc_type: Exception type if an exception was raised.
+            exc: Exception instance if an exception was raised.
+            tb: Traceback if an exception was raised.
+        """
+        if self._span and not self._finished:
+            if exc is not None:
+                self._span.record_exception(exc)
+                self._span.set_status(Status(StatusCode.ERROR))
+            self.finish()
+
+    def set(self, **attributes: AttributeValue) -> None:
+        """Set attributes on the current span.
+
+        Args:
+            **attributes: Key-value pairs to set as span attributes.
+        """
+        if self._span and not self._finished:
+            for key, value in attributes.items():
+                self._span.set_attribute(key, value)
+
+    def event(self, name: str, **attributes: AttributeValue) -> None:
+        """Record an event within the span.
+
+        Args:
+            name: Name of the event.
+            **attributes: Event attributes as key-value pairs.
+        """
+        if self._span and not self._finished:
+            self._span.add_event(name, attributes=attributes)
+
+    def debug(self, message: str, **additional_attributes: AttributeValue) -> None:
+        """Log a debug message within the span.
+
+        Args:
+            message: Debug message text.
+            **additional_attributes: Additional structured attributes for the log entry.
+        """
+        self.event("debug", level="debug", message=message, **additional_attributes)
+
+    def info(self, message: str, **additional_attributes: AttributeValue) -> None:
+        """Log an info message within the span.
+
+        Args:
+            message: Info message text.
+            **additional_attributes: Additional structured attributes for the log entry.
+        """
+        self.event("info", level="info", message=message, **additional_attributes)
+
+    def warning(self, message: str, **additional_attributes: AttributeValue) -> None:
+        """Log a warning message within the span.
+
+        Args:
+            message: Warning message text.
+            **additional_attributes: Additional structured attributes for the log entry.
+        """
+        self.event("warning", level="warning", message=message, **additional_attributes)
+
+    def error(self, message: str, **additional_attributes: AttributeValue) -> None:
+        """Log an error message within the span.
+
+        Args:
+            message: Error message text.
+            **additional_attributes: Additional structured attributes for the log entry.
+        """
+        self.event("error", level="error", message=message, **additional_attributes)
+        if self._span and not self._finished:
+            self._span.set_status(Status(StatusCode.ERROR))
+
+    def critical(self, message: str, **additional_attributes: AttributeValue) -> None:
+        """Log a critical message within the span.
+
+        Args:
+            message: Critical message text.
+            **additional_attributes: Additional structured attributes for the log entry.
+        """
+        self.event("critical", level="error", message=message, **additional_attributes)
+        if self._span and not self._finished:
+            self._span.set_status(Status(StatusCode.ERROR))
+
+    def finish(self) -> None:
+        """Explicitly finish the span."""
+        if not self._finished:
+            self._finished = True
+            if self._span:
+                self._span.end()
+            if self._token:
+                otel_context.detach(self._token)
+                self._token = None
+
+    @property
+    def span_id(self) -> str | None:
+        """Get the span ID if available.
+
+        Returns:
+            The span ID or None if not available.
+        """
+        if span_context := self.span_context:
+            return format_span_id(span_context.span_id)
+        return None
+
+    @property
+    def trace_id(self) -> str | None:
+        if span_context := self.span_context:
+            return format_trace_id(span_context.trace_id)
+        return None
+
+    @property
+    def is_noop(self) -> bool:
+        """Check if this is a no-op span.
+
+        Returns:
+            True if tracing is disabled, False otherwise.
+        """
+        return self._is_noop
+
+    @property
+    def span_context(self) -> SpanContext | None:
+        """Return the span context if available."""
+        if otel_span := self._span:
+            return otel_span.get_span_context()
+        return None
+
+
+def span(name: str, **attributes: AttributeValue) -> Span:
+    """Create a new span context manager.
+
+    Args:
+        name: Name for the new span.
+        **attributes: Initial attributes to set on the span.
+
+    Returns:
+        A Span context manager that creates a child span when entered.
+    """
+    return Span(name, **attributes)