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

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

@@ -1,300 +0,0 @@
-"""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

@@ -1,198 +0,0 @@
-"""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

@@ -1,133 +0,0 @@
-"""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