mirascope 1.25.6__py3-none-any.whl → 2.0.0__py3-none-any.whl

mirascope/ops/_internal/instrumentation/llm/serialize.py

@@ -0,0 +1,300 @@
+"""Mirascope-specific serialization for span attributes."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any, Protocol
+
+from .....llm.content import (
+    Audio,
+    Base64ImageSource,
+    Document,
+    Image,
+    Text,
+    Thought,
+    ToolCall,
+    ToolOutput,
+)
+from .....llm.content.document import Base64DocumentSource, TextDocumentSource
+from .....llm.messages import AssistantMessage, Message, SystemMessage, UserMessage
+from .....llm.responses.usage import Usage
+from ...utils import json_dumps
+from .cost import calculate_cost_async, calculate_cost_sync
+
+if TYPE_CHECKING:
+    from opentelemetry.util.types import AttributeValue
+
+    from .....llm.responses.root_response import RootResponse
+    from .....llm.types import Jsonable
+
+logger = logging.getLogger(__name__)
+
+
+class SpanProtocol(Protocol):
+    """Protocol for span objects that support setting attributes."""
+
+    def set(self, **attributes: AttributeValue) -> None:
+        """Set attributes on the span."""
+        ...
+
+
+def _serialize_content_part(
+    part: Text | ToolCall | Thought | Image | Audio | Document | ToolOutput[Jsonable],
+) -> dict[str, Any]:
+    """Serialize a single content part to a dict matching the Mirascope dataclass structure."""
+    if isinstance(part, Text):
+        return {"type": "text", "text": part.text}
+    elif isinstance(part, ToolCall):
+        return {
+            "type": "tool_call",
+            "id": part.id,
+            "name": part.name,
+            "args": part.args,
+        }
+    elif isinstance(part, Thought):
+        return {"type": "thought", "thought": part.thought}
+    elif isinstance(part, ToolOutput):
+        return {
+            "type": "tool_output",
+            "id": part.id,
+            "name": part.name,
+            "result": part.result,
+        }
+    elif isinstance(part, Image):
+        if isinstance(part.source, Base64ImageSource):
+            return {
+                "type": "image",
+                "source": {
+                    "type": "base64_image_source",
+                    "mime_type": part.source.mime_type,
+                    "data": part.source.data,
+                },
+            }
+        else:  # URLImageSource
+            return {
+                "type": "image",
+                "source": {"type": "url_image_source", "url": part.source.url},
+            }
+    elif isinstance(part, Audio):
+        return {
+            "type": "audio",
+            "source": {
+                "type": "base64_audio_source",
+                "mime_type": part.source.mime_type,
+                "data": part.source.data,
+            },
+        }
+    elif isinstance(part, Document):
+        # Document has multiple source types - serialize based on actual type
+        if isinstance(part.source, Base64DocumentSource):
+            return {
+                "type": "document",
+                "source": {
+                    "type": "base64_document_source",
+                    "data": part.source.data,
+                    "media_type": part.source.media_type,
+                },
+            }
+        elif isinstance(part.source, TextDocumentSource):
+            return {
+                "type": "document",
+                "source": {
+                    "type": "text_document_source",
+                    "data": part.source.data,
+                    "media_type": part.source.media_type,
+                },
+            }
+        else:  # URLDocumentSource
+            return {
+                "type": "document",
+                "source": {
+                    "type": "url_document_source",
+                    "url": part.source.url,
+                },
+            }
+    return {"type": "unknown"}  # pragma: no cover
+
+
+def _serialize_message(message: Message) -> dict[str, Any]:
+    """Serialize a Message to a dict matching the Mirascope dataclass structure."""
+    if isinstance(message, SystemMessage):
+        return {
+            "role": "system",
+            "content": _serialize_content_part(message.content),
+        }
+    elif isinstance(message, UserMessage):
+        return {
+            "role": "user",
+            "content": [_serialize_content_part(p) for p in message.content],
+            "name": message.name,
+        }
+    elif isinstance(message, AssistantMessage):
+        return {
+            "role": "assistant",
+            "content": [_serialize_content_part(p) for p in message.content],
+            "name": message.name,
+        }
+    return {"role": "unknown"}  # pragma: no cover
+
+
+def serialize_mirascope_messages(messages: Sequence[Message]) -> str:
+    """Serialize input messages to JSON for span attributes."""
+    return json_dumps([_serialize_message(m) for m in messages])
+
+
+def serialize_mirascope_content(
+    content: Sequence[Text | ToolCall | Thought],
+) -> str:
+    """Serialize response content to JSON for span attributes."""
+    return json_dumps([_serialize_content_part(p) for p in content])
+
+
+def serialize_mirascope_usage(usage: Usage | None) -> AttributeValue | None:
+    """Serialize response usage to JSON for span attributes. Returns None if usage is None."""
+    if usage is None:
+        return None
+    return json_dumps(
+        {
+            "input_tokens": usage.input_tokens,
+            "output_tokens": usage.output_tokens,
+            "cache_read_tokens": usage.cache_read_tokens,
+            "cache_write_tokens": usage.cache_write_tokens,
+            "reasoning_tokens": usage.reasoning_tokens,
+            "total_tokens": usage.total_tokens,
+        }
+    )
+
+
+def serialize_mirascope_cost(
+    input_cost: float,
+    output_cost: float,
+    total_cost: float,
+    cache_read_cost: float | None = None,
+    cache_write_cost: float | None = None,
+) -> str:
+    """Serialize cost to JSON for span attributes.
+
+    All costs are in centicents (1 centicent = $0.0001).
+    Consumers can divide by 10,000 to get dollar amounts.
+    """
+    return json_dumps(
+        {
+            "input_cost": input_cost,
+            "output_cost": output_cost,
+            "cache_read_cost": cache_read_cost,
+            "cache_write_cost": cache_write_cost,
+            "total_cost": total_cost,
+        }
+    )
+
+
+def attach_mirascope_response(
+    span: SpanProtocol, response: RootResponse[Any, Any]
+) -> None:
+    """Attach Mirascope-specific response attributes to a span.
+
+    Sets the following attributes:
+    - mirascope.trace.output: Pretty-printed response
+    - mirascope.response.messages: Serialized input messages (excluding final assistant message)
+    - mirascope.response.content: Serialized response content
+    - mirascope.response.usage: Serialized usage (if available)
+    - mirascope.response.cost: Serialized cost (if available)
+    """
+    span.set(
+        **{
+            "mirascope.response.provider_id": response.provider_id,
+            "mirascope.response.model_id": response.model_id,
+            "mirascope.trace.output": response.pretty(),
+            "mirascope.response.messages": serialize_mirascope_messages(
+                response.messages[:-1]
+            ),
+            "mirascope.response.content": serialize_mirascope_content(response.content),
+        }
+    )
+    if (usage_json := serialize_mirascope_usage(response.usage)) is not None:
+        span.set(**{"mirascope.response.usage": usage_json})
+        logger.debug("Attached usage to span")
+    else:
+        logger.debug("No usage available, skipping cost calculation")
+
+    # Calculate and attach cost if usage is available
+    if response.usage is not None:
+        logger.debug("Attempting cost calculation (sync)")
+        cost = calculate_cost_sync(
+            response.provider_id, response.model_id, response.usage
+        )
+        if cost is not None:
+            span.set(
+                **{
+                    "mirascope.response.cost": serialize_mirascope_cost(
+                        input_cost=cost.input_cost_centicents,
+                        output_cost=cost.output_cost_centicents,
+                        total_cost=cost.total_cost_centicents,
+                        cache_read_cost=cost.cache_read_cost_centicents,
+                        cache_write_cost=cost.cache_write_cost_centicents,
+                    )
+                }
+            )
+            logger.debug(
+                "Attached cost to span: total=%s centicents", cost.total_cost_centicents
+            )
+        else:
+            logger.debug("Cost calculation returned None, not attaching cost to span")
+
+
+async def attach_mirascope_response_async(
+    span: SpanProtocol, response: RootResponse[Any, Any]
+) -> None:
+    """Attach Mirascope-specific response attributes to a span (async version).
+
+    Sets the following attributes:
+    - mirascope.trace.output: Pretty-printed response
+    - mirascope.response.messages: Serialized input messages (excluding final assistant message)
+    - mirascope.response.content: Serialized response content
+    - mirascope.response.usage: Serialized usage (if available)
+    - mirascope.response.cost: Serialized cost (if available)
+    """
+    span.set(
+        **{
+            "mirascope.response.provider_id": response.provider_id,
+            "mirascope.response.model_id": response.model_id,
+            "mirascope.trace.output": response.pretty(),
+            "mirascope.response.messages": serialize_mirascope_messages(
+                response.messages[:-1]
+            ),
+            "mirascope.response.content": serialize_mirascope_content(response.content),
+        }
+    )
+    if (usage_json := serialize_mirascope_usage(response.usage)) is not None:
+        span.set(**{"mirascope.response.usage": usage_json})
+        logger.debug("Attached usage to span (async)")
+    else:
+        logger.debug("No usage available, skipping cost calculation (async)")
+
+    # Calculate and attach cost if usage is available (async)
+    if response.usage is not None:
+        logger.debug("Attempting cost calculation (async)")
+        cost = await calculate_cost_async(
+            response.provider_id, response.model_id, response.usage
+        )
+        if cost is not None:
+            span.set(
+                **{
+                    "mirascope.response.cost": serialize_mirascope_cost(
+                        input_cost=cost.input_cost_centicents,
+                        output_cost=cost.output_cost_centicents,
+                        total_cost=cost.total_cost_centicents,
+                        cache_read_cost=cost.cache_read_cost_centicents,
+                        cache_write_cost=cost.cache_write_cost_centicents,
+                    )
+                }
+            )
+            logger.debug(
+                "Attached cost to span (async): total=%s centicents",
+                cost.total_cost_centicents,
+            )
+        else:
+            logger.debug(
+                "Cost calculation returned None, not attaching cost to span (async)"
+            )

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