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

mirascope/ops/_internal/configuration.py

@@ -0,0 +1,126 @@
+"""Configuration utilities for Mirascope ops module initialization and setup."""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from contextlib import contextmanager
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from opentelemetry.sdk.trace import TracerProvider
+    from opentelemetry.trace import Tracer
+else:
+    Tracer = None
+
+DEFAULT_TRACER_NAME = "mirascope.llm"
+
+try:
+    from opentelemetry import trace as otel_trace
+except ImportError:  # pragma: no cover
+    otel_trace = None
+
+_tracer_provider: TracerProvider | None = None
+_tracer_name: str = DEFAULT_TRACER_NAME
+_tracer_version: str | None = None
+_tracer: Tracer | None = None
+
+
+def configure(
+    *,
+    tracer_provider: TracerProvider | None = None,
+    tracer_name: str = DEFAULT_TRACER_NAME,
+    tracer_version: str | None = None,
+) -> None:
+    """Configure the ops module defaults for tracing.
+
+    Sets up default tracer settings for the ops module. If a tracer_provider
+    is supplied, it will be installed as the global OpenTelemetry tracer provider.
+
+    Args:
+        tracer_provider: Optional OpenTelemetry TracerProvider to use as default
+            and to install globally.
+        tracer_name: Tracer name to use when creating a tracer.
+            Defaults to "mirascope.llm".
+        tracer_version: Optional tracer version.
+
+    Example:
+
+        Configure custom tracer settings:
+        ```python
+        from mirascope import ops
+        from opentelemetry.sdk.trace import TracerProvider
+
+        provider = TracerProvider()
+        ops.configure(tracer_provider=provider)
+        ops.instrument_llm()
+        ```
+    """
+    # TODO: refactor alongside other import error handling improvements
+    if otel_trace is None:  # pragma: no cover
+        raise ImportError(
+            "OpenTelemetry is not installed. Run `pip install mirascope[otel]` "
+            "before calling `ops.configure(tracer_provider=...)`."
+        )
+
+    global _tracer_provider, _tracer_name, _tracer_version, _tracer
+
+    if tracer_provider is not None:
+        _tracer_provider = tracer_provider
+        otel_trace.set_tracer_provider(tracer_provider)
+
+    _tracer_name = tracer_name
+    _tracer_version = tracer_version
+
+    if otel_trace is not None:
+        provider = (
+            otel_trace.get_tracer_provider()
+            if _tracer_provider is None
+            else _tracer_provider
+        )
+        _tracer = provider.get_tracer(_tracer_name, _tracer_version)
+
+
+def set_tracer(tracer: Tracer | None) -> None:
+    """Set the configured tracer instance."""
+    global _tracer
+    _tracer = tracer
+
+
+def get_tracer() -> Tracer | None:
+    """Return the configured tracer instance."""
+    return _tracer
+
+
+@contextmanager
+def tracer_context(tracer: Tracer | None) -> Iterator[Tracer | None]:
+    """Context manager for temporarily setting a tracer.
+
+    Temporarily sets the tracer for the duration of the context and restores
+    the previous tracer when the context exits.
+
+    Args:
+        tracer: The tracer to use within the context.
+
+    Yields:
+        The tracer that was set.
+
+    Example:
+        ```python
+        from mirascope import ops
+        from opentelemetry.sdk.trace import TracerProvider
+
+        provider = TracerProvider()
+        tracer = provider.get_tracer("my-tracer")
+
+        with ops.tracer_context(tracer):
+            # Use the tracer within this context
+            ...
+        # Previous tracer is restored here
+        ```
+    """
+    previous_tracer = get_tracer()
+    set_tracer(tracer)
+    try:
+        yield tracer
+    finally:
+        set_tracer(previous_tracer)

mirascope/ops/_internal/context.py

@@ -0,0 +1,76 @@
+"""Context management utilities for distributed tracing."""
+
+from __future__ import annotations
+
+from collections.abc import Iterator, Mapping
+from contextlib import ExitStack, contextmanager
+from typing import Any
+
+from opentelemetry import context as otel_context
+from opentelemetry.context import Context
+
+from .propagation import extract_context
+from .session import extract_session_id, session
+
+
+@contextmanager
+def propagated_context(
+    *,
+    parent: Context | None = None,
+    extract_from: Mapping[str, Any] | None = None,
+) -> Iterator[None]:
+    """Attach a parent context or extract context from carrier headers.
+
+    This context manager is used to establish trace context continuity,
+    typically on the server side when receiving requests. It either extracts
+    context from incoming headers or attaches a pre-existing context.
+
+    Args:
+        parent: Pre-existing OTEL context to attach. Mutually exclusive with extract_from.
+        extract_from: Dictionary of headers to extract context from (e.g., request.headers).
+            Mutually exclusive with parent.
+
+    Raises:
+        ValueError: If both parent and extract_from are provided, or if neither is provided.
+
+    Example:
+        Server-side context extraction from FastAPI request:
+
+        ```python
+        @app.post("/endpoint")
+        async def endpoint(request: Request):
+            with propagated_context(extract_from=dict(request.headers)):
+                result = process_request()
+            return result
+        ```
+
+        Using a pre-existing context:
+
+        ```python
+        with propagated_context(parent=existing_context):
+            do_work()
+        ```
+    """
+    if parent is not None and extract_from is not None:
+        raise ValueError("Cannot specify both 'parent' and 'extract_from' parameters")
+
+    if parent is None and extract_from is None:
+        raise ValueError("Must specify either 'parent' or 'extract_from' parameter")
+
+    if extract_from is not None:
+        with ExitStack() as stack:
+            session_id = extract_session_id(extract_from)
+            if session_id:
+                stack.enter_context(session(id=session_id))
+
+            extracted_context = extract_context(extract_from)
+            token = otel_context.attach(extracted_context)
+            stack.callback(otel_context.detach, token)
+
+            yield
+    elif parent is not None:
+        token = otel_context.attach(parent)
+        try:
+            yield
+        finally:
+            otel_context.detach(token)

mirascope/ops/_internal/exporters/__init__.py

@@ -0,0 +1,26 @@
+"""Mirascope OpenTelemetry exporters for two-phase telemetry.
+
+This package provides a two-phase export system for OpenTelemetry tracing:
+1. Immediate start event transmission for real-time visibility
+2. Batched end event transmission for efficiency
+"""
+
+from .exporters import MirascopeOTLPExporter
+from .processors import MirascopeSpanProcessor
+from .types import (
+    Link,
+    SpanContextDict,
+    SpanEvent,
+    SpanEventType,
+    Status,
+)
+
+__all__ = [
+    "Link",
+    "MirascopeOTLPExporter",
+    "MirascopeSpanProcessor",
+    "SpanContextDict",
+    "SpanEvent",
+    "SpanEventType",
+    "Status",
+]

mirascope/ops/_internal/exporters/exporters.py

@@ -0,0 +1,342 @@
+"""Exporter implementation for OpenTelemetry exporters.
+
+This module provides the export layer for sending OpenTelemetry span
+events to the Mirascope ingestion endpoint. It wraps the Fern-generated
+Mirascope client to provide the interface needed by OpenTelemetry exporters.
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from collections.abc import Sequence
+
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
+from opentelemetry.util.types import AttributeValue
+
+from ....api._generated.traces.types import (
+    TracesCreateRequestResourceSpansItem,
+    TracesCreateRequestResourceSpansItemResource,
+    TracesCreateRequestResourceSpansItemResourceAttributesItem,
+    TracesCreateRequestResourceSpansItemResourceAttributesItemValue,
+    TracesCreateRequestResourceSpansItemScopeSpansItem,
+    TracesCreateRequestResourceSpansItemScopeSpansItemScope,
+    TracesCreateRequestResourceSpansItemScopeSpansItemSpansItem,
+    TracesCreateRequestResourceSpansItemScopeSpansItemSpansItemAttributesItem,
+    TracesCreateRequestResourceSpansItemScopeSpansItemSpansItemAttributesItemValue,
+    TracesCreateRequestResourceSpansItemScopeSpansItemSpansItemStatus,
+)
+from ....api.client import Mirascope
+
+logger = logging.getLogger(__name__)
+
+
+class MirascopeOTLPExporter(SpanExporter):
+    """OTLP/HTTP exporter for completed spans.
+
+    This exporter implements the OpenTelemetry SpanExporter interface
+    for exporting completed spans in OTLP format over HTTP. It's
+    designed to work with BatchSpanProcessor for efficient batching.
+
+    This uses the Fern auto-generated client for sending converted spans.
+
+    Attributes:
+        exporter: Export client for sending events.
+        timeout: Request timeout in seconds.
+    """
+
+    def __init__(
+        self,
+        client: Mirascope,
+        timeout: float = 30.0,
+        max_retry_attempts: int = 3,
+    ) -> None:
+        """Initialize the telemetry exporter.
+
+        Args:
+            client: The Fern-generated Mirascope client instance.
+                In the future, this will accept the enhanced client from
+                mirascope.api.client that provides error handling and caching
+                capabilities.
+            timeout: Request timeout in seconds for telemetry operations.
+            max_retry_attempts: Maximum number of retry attempts for failed exports.
+        """
+        self.client = client
+        self.timeout = timeout
+        self.max_retry_attempts = max_retry_attempts
+        self._shutdown = False
+
+    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        """Export a batch of spans to the telemetry endpoint.
+
+        This is the standard OpenTelemetry export interface.
+
+        Args:
+            spans: Sequence of ReadableSpan objects to export.
+
+        Returns:
+            SpanExportResult indicating success or failure.
+        """
+        if self._shutdown:
+            return SpanExportResult.FAILURE
+
+        if not spans:
+            return SpanExportResult.SUCCESS
+
+        exceptions: list[Exception] = []
+        delay = 0.1
+
+        for i in range(self.max_retry_attempts):
+            if i > 0:
+                time.sleep(delay)
+                delay = min(delay * 2, 5.0)
+
+            try:
+                otlp_data = self._convert_spans_to_otlp(spans)
+                response = self.client.traces.create(resource_spans=otlp_data)
+
+                if (
+                    response
+                    and hasattr(response, "partial_success")
+                    and response.partial_success
+                ):
+                    partial_success = response.partial_success
+                    if hasattr(partial_success, "rejected_spans"):
+                        rejected = partial_success.rejected_spans
+                        if rejected is not None and rejected > 0:
+                            return SpanExportResult.FAILURE
+
+                return SpanExportResult.SUCCESS
+
+            except Exception as e:
+                exceptions.append(e)
+                logger.warning(
+                    f"Export attempt {i + 1} failed, retrying in {delay}s: {e}"
+                )
+
+        logger.error(
+            f"Failed to export spans after {self.max_retry_attempts} attempts: {exceptions}"
+        )
+
+        return SpanExportResult.FAILURE
+
+    def _convert_spans_to_otlp(
+        self, spans: Sequence[ReadableSpan]
+    ) -> list[TracesCreateRequestResourceSpansItem]:
+        """Convert OpenTelemetry spans to OTLP format.
+
+        Args:
+            spans: Sequence of ReadableSpan objects.
+
+        Returns:
+            List of ResourceSpans in OTLP format.
+        """
+        resource_spans_map = {}
+
+        for span in spans:
+            try:
+                otlp_span = self._convert_span(span)
+            except ValueError as e:
+                logger.warning(f"Skipping span due to error: {e}")
+                continue
+
+            resource_key = id(span.resource) if span.resource else "default"
+
+            if resource_key not in resource_spans_map:
+                resource = None
+                if span.resource:
+                    resource_attrs = []
+                    for key, value in span.resource.attributes.items():
+                        attr_value = self._convert_resource_attribute_value(value)
+                        resource_attrs.append(
+                            TracesCreateRequestResourceSpansItemResourceAttributesItem(
+                                key=key,
+                                value=attr_value,
+                            )
+                        )
+                    resource = TracesCreateRequestResourceSpansItemResource(
+                        attributes=resource_attrs
+                    )
+
+                resource_spans_map[resource_key] = {
+                    "resource": resource,
+                    "scope_spans": {},
+                }
+
+            scope_key = (
+                span.instrumentation_scope.name
+                if span.instrumentation_scope
+                else "unknown"
+            )
+
+            if scope_key not in resource_spans_map[resource_key]["scope_spans"]:
+                scope = None
+                if span.instrumentation_scope:
+                    scope = TracesCreateRequestResourceSpansItemScopeSpansItemScope(
+                        name=span.instrumentation_scope.name,
+                        version=span.instrumentation_scope.version,
+                    )
+
+                resource_spans_map[resource_key]["scope_spans"][scope_key] = {
+                    "scope": scope,
+                    "spans": [],
+                }
+
+            resource_spans_map[resource_key]["scope_spans"][scope_key]["spans"].append(
+                otlp_span
+            )
+
+        result = []
+        for resource_data in resource_spans_map.values():
+            scope_spans = []
+            for scope_data in resource_data["scope_spans"].values():
+                scope_spans.append(
+                    TracesCreateRequestResourceSpansItemScopeSpansItem(
+                        scope=scope_data["scope"],
+                        spans=scope_data["spans"],
+                    )
+                )
+
+            result.append(
+                TracesCreateRequestResourceSpansItem(
+                    resource=resource_data["resource"],
+                    scope_spans=scope_spans,
+                )
+            )
+
+        return result
+
+    def _convert_span(
+        self, span: ReadableSpan
+    ) -> TracesCreateRequestResourceSpansItemScopeSpansItemSpansItem:
+        """Convert a single ReadableSpan to OTLP format."""
+        context = span.get_span_context()
+        if not context or not context.is_valid:
+            raise ValueError(f"Cannot export span without valid context: {span.name}")
+
+        attributes = []
+        if span.attributes:
+            for key, value in span.attributes.items():
+                attr_value = self._convert_attribute_value(value)
+                attributes.append(
+                    TracesCreateRequestResourceSpansItemScopeSpansItemSpansItemAttributesItem(
+                        key=key,
+                        value=attr_value,
+                    )
+                )
+
+        status = None
+        if span.status:
+            status = TracesCreateRequestResourceSpansItemScopeSpansItemSpansItemStatus(
+                code=span.status.status_code.value,
+                message=span.status.description or "",
+            )
+
+        trace_id = format(context.trace_id, "032x")
+        span_id = format(context.span_id, "016x")
+
+        return TracesCreateRequestResourceSpansItemScopeSpansItemSpansItem(
+            trace_id=trace_id,
+            span_id=span_id,
+            parent_span_id=(
+                format(span.parent.span_id, "016x")
+                if span.parent and span.parent.span_id
+                else None
+            ),
+            name=span.name,
+            kind=span.kind.value if span.kind else 0,
+            start_time_unix_nano=str(span.start_time) if span.start_time else "0",
+            end_time_unix_nano=str(span.end_time) if span.end_time else "0",
+            attributes=attributes or None,
+            status=status,
+        )
+
+    def _convert_attribute_value(
+        self, value: AttributeValue
+    ) -> TracesCreateRequestResourceSpansItemScopeSpansItemSpansItemAttributesItemValue:
+        """Convert OpenTelemetry AttributeValue to Mirascope API's KeyValueValue.
+
+        This conversion is necessary because the Fern-generated API client
+        expects KeyValueValue objects, not OpenTelemetry's AttributeValue types.
+
+        Args:
+            value: An OpenTelemetry AttributeValue (bool, int, float, str, or Sequence)
+
+        Returns:
+            A KeyValueValue object for the Mirascope API
+        """
+        match value:
+            case str():
+                return TracesCreateRequestResourceSpansItemScopeSpansItemSpansItemAttributesItemValue(
+                    string_value=value
+                )
+            case bool():
+                return TracesCreateRequestResourceSpansItemScopeSpansItemSpansItemAttributesItemValue(
+                    bool_value=value
+                )
+            case int():
+                return TracesCreateRequestResourceSpansItemScopeSpansItemSpansItemAttributesItemValue(
+                    int_value=str(value)
+                )
+            case float():
+                return TracesCreateRequestResourceSpansItemScopeSpansItemSpansItemAttributesItemValue(
+                    double_value=value
+                )
+            case _:
+                return TracesCreateRequestResourceSpansItemScopeSpansItemSpansItemAttributesItemValue(
+                    string_value=str(list(value))
+                )
+
+    def _convert_resource_attribute_value(
+        self, value: AttributeValue
+    ) -> TracesCreateRequestResourceSpansItemResourceAttributesItemValue:
+        """Convert OpenTelemetry AttributeValue to Mirascope API's resource KeyValueValue.
+
+        This conversion is necessary because the Fern-generated API client
+        expects KeyValueValue objects, not OpenTelemetry's AttributeValue types.
+
+        Args:
+            value: An OpenTelemetry AttributeValue (bool, int, float, str, or Sequence)
+
+        Returns:
+            A KeyValueValue object for the Mirascope API resource attributes
+        """
+        match value:
+            case str():
+                return TracesCreateRequestResourceSpansItemResourceAttributesItemValue(
+                    string_value=value
+                )
+            case bool():
+                return TracesCreateRequestResourceSpansItemResourceAttributesItemValue(
+                    bool_value=value
+                )
+            case int():
+                return TracesCreateRequestResourceSpansItemResourceAttributesItemValue(
+                    int_value=str(value)
+                )
+            case float():
+                return TracesCreateRequestResourceSpansItemResourceAttributesItemValue(
+                    double_value=value
+                )
+            case _:
+                return TracesCreateRequestResourceSpansItemResourceAttributesItemValue(
+                    string_value=str(list(value))
+                )
+
+    def shutdown(self) -> None:
+        """Shutdown the exporter. Subsequent exports will return FAILURE."""
+        self._shutdown = True
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        """Force flush any pending data.
+
+        No-op since this exporter does not buffer data internally.
+
+        Args:
+            timeout_millis: Maximum time to wait in milliseconds (unused).
+
+        Returns:
+            Always True since there is no internal buffer to flush.
+        """
+        return True

mirascope/ops/_internal/exporters/processors.py

@@ -0,0 +1,104 @@
+"""Span processors for two-phase export system.
+
+This module implements a custom SpanProcessor that sends immediate
+start events and batches end events for efficient export.
+"""
+
+from concurrent.futures import ThreadPoolExecutor
+
+from opentelemetry.context import Context
+from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor
+from opentelemetry.sdk.trace.export import BatchSpanProcessor
+
+from .exporters import MirascopeOTLPExporter
+
+
+class MirascopeSpanProcessor(SpanProcessor):
+    """Two-phase span processor for Mirascope telemetry.
+
+    This processor implements a two-phase export strategy:
+    1. Immediate transmission of minimal start events for real-time visibility
+    2. Batched transmission of complete events for efficiency
+
+    The processor uses a thread pool to ensure start events don't block
+    the application while maintaining compatibility with OpenTelemetry's
+    synchronous SDK.
+
+    Attributes:
+        start_exporter: Exporter for immediate start events.
+        batch_processor: Standard batch processor for completed spans.
+        executor: Thread pool for non-blocking start event transmission.
+    """
+
+    def __init__(
+        self,
+        otlp_exporter: MirascopeOTLPExporter,
+        batch_processor: BatchSpanProcessor | None = None,
+        executor: ThreadPoolExecutor | None = None,
+    ) -> None:
+        """Initialize the two-phase processor.
+
+        Args:
+            start_exporter: Exporter for immediate start events.
+            batch_processor: Optional batch processor for end events.
+            executor: Optional thread pool executor (creates default if None).
+        """
+        self.otlp_exporter = otlp_exporter
+        self.batch_processor = batch_processor
+        self.executor = executor or ThreadPoolExecutor(
+            max_workers=2, thread_name_prefix="mirascope-span-processor"
+        )
+        self._shutdown = False
+
+    def on_start(
+        self, span: ReadableSpan, parent_context: Context | None = None
+    ) -> None:
+        """Handle span start by sending immediate start event.
+
+        This method extracts minimal span data and sends it immediately
+        via the start exporter in a non-blocking manner.
+
+        Args:
+            span: The span that just started.
+            parent_context: Optional parent context for the span.
+        """
+        if self._shutdown:
+            return
+
+        self.executor.submit(self.otlp_exporter.export, [span])
+
+    def on_end(self, span: ReadableSpan) -> None:
+        """Handle span end by delegating to batch processor.
+
+        Args:
+            span: The span that just ended.
+        """
+        if self.batch_processor and not self._shutdown:
+            self.batch_processor.on_end(span)
+
+    def shutdown(self) -> None:
+        """Gracefully shutdown the processor.
+
+        This ensures all pending start events are sent and the
+        batch processor is properly shutdown.
+        """
+        self._shutdown = True
+
+        if self.batch_processor:
+            self.batch_processor.shutdown()
+
+        self.executor.shutdown(wait=True)
+        self.otlp_exporter.shutdown()
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        """Force flush all pending data.
+
+        Args:
+            timeout_millis: Maximum time to wait in milliseconds.
+
+        Returns:
+            True if flush completed successfully.
+        """
+        if self.batch_processor:
+            return self.batch_processor.force_flush(timeout_millis)
+        return True