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

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

@@ -0,0 +1,324 @@
+"""Mirascope-specific serialization for span attributes."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Sequence
+from typing import Any, Protocol
+
+from opentelemetry.util.types import AttributeValue
+
+from .....llm import (
+    AnyToolSchema,
+    AssistantMessage,
+    Audio,
+    Base64ImageSource,
+    Document,
+    Image,
+    Jsonable,
+    Message,
+    ProviderTool,
+    RootResponse,
+    SystemMessage,
+    Text,
+    Thought,
+    ToolCall,
+    ToolOutput,
+    Usage,
+    UserMessage,
+)
+from .....llm.content.document import Base64DocumentSource, TextDocumentSource
+from ...utils import json_dumps
+from .cost import calculate_cost_async, calculate_cost_sync
+
+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, Jsonable]:
+    """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, Jsonable]:
+    """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_tool(tool: AnyToolSchema | ProviderTool) -> dict[str, Jsonable]:
+    if isinstance(tool, ProviderTool):
+        return {"name": tool.name, "type": "extension"}
+    result: dict[str, Jsonable] = {
+        "name": tool.name,
+        "description": tool.description,
+        "type": "function",
+        "parameters": tool.parameters.model_dump(by_alias=True, mode="json"),
+    }
+    if tool.strict is not None:
+        result["strict"] = tool.strict
+    return result
+
+
+def serialize_tools(tools: Sequence[AnyToolSchema | ProviderTool]) -> str | None:
+    """Serialize a sequence of Mirascope tools"""
+    if not tools:
+        return None
+    return json_dumps([_serialize_tool(t) for t in tools])
+
+
+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/instrumentation/providers/__init__.py

@@ -0,0 +1,29 @@
+"""Provider instrumentation modules."""
+
+from .anthropic import (
+    instrument_anthropic,
+    is_anthropic_instrumented,
+    uninstrument_anthropic,
+)
+from .google_genai import (
+    instrument_google_genai,
+    is_google_genai_instrumented,
+    uninstrument_google_genai,
+)
+from .openai import (
+    instrument_openai,
+    is_openai_instrumented,
+    uninstrument_openai,
+)
+
+__all__ = [
+    "instrument_anthropic",
+    "instrument_google_genai",
+    "instrument_openai",
+    "is_anthropic_instrumented",
+    "is_google_genai_instrumented",
+    "is_openai_instrumented",
+    "uninstrument_anthropic",
+    "uninstrument_google_genai",
+    "uninstrument_openai",
+]

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

@@ -0,0 +1,78 @@
+"""OpenTelemetry instrumentation for Anthropic SDK."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from opentelemetry.instrumentation.anthropic import AnthropicInstrumentor
+
+from .base import BaseInstrumentation, ContentCaptureMode
+
+if TYPE_CHECKING:
+    from opentelemetry.trace import TracerProvider
+
+
+class AnthropicInstrumentation(BaseInstrumentation[AnthropicInstrumentor]):
+    """Manages OpenTelemetry instrumentation lifecycle for the Anthropic SDK."""
+
+    def _create_instrumentor(self) -> AnthropicInstrumentor:
+        """Create a new Anthropic instrumentor instance."""
+        return AnthropicInstrumentor()
+
+    def _configure_capture_content(self, capture_content: ContentCaptureMode) -> None:
+        """Configure environment variables for Anthropic content capture."""
+        if capture_content == "enabled":
+            self._set_env_var("TRACELOOP_TRACE_CONTENT", "true")
+        elif capture_content == "disabled":
+            self._set_env_var("TRACELOOP_TRACE_CONTENT", "false")
+
+
+def instrument_anthropic(
+    *,
+    tracer_provider: TracerProvider | None = None,
+    capture_content: ContentCaptureMode = "default",
+) -> None:
+    """Enable OpenTelemetry instrumentation for the Anthropic 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_anthropic()
+        ```
+
+        Enable instrumentation with content capture:
+        ```python
+        from mirascope import ops
+
+        ops.configure()
+        ops.instrument_anthropic(capture_content="enabled")
+        ```
+    """
+    AnthropicInstrumentation().instrument(
+        tracer_provider=tracer_provider,
+        capture_content=capture_content,
+    )
+
+
+def uninstrument_anthropic() -> None:
+    """Disable previously configured Anthropic instrumentation."""
+    AnthropicInstrumentation().uninstrument()
+
+
+def is_anthropic_instrumented() -> bool:
+    """Return whether Anthropic instrumentation is currently active."""
+    return AnthropicInstrumentation().is_instrumented

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

@@ -0,0 +1,179 @@
+"""Base class for OpenTelemetry SDK instrumentation."""
+
+from __future__ import annotations
+
+import os
+import threading
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Generic, Literal, Protocol, TypeVar
+
+if TYPE_CHECKING:
+    from opentelemetry.trace import TracerProvider
+
+ContentCaptureMode = Literal["enabled", "disabled", "default"]
+
+# OpenTelemetry semantic conventions environment variable
+OTEL_SEMCONV_STABILITY_OPT_IN = "OTEL_SEMCONV_STABILITY_OPT_IN"
+OTEL_SEMCONV_STABILITY_VALUE = "gen_ai_latest_experimental"
+
+
+class Instrumentor(Protocol):
+    """Protocol for OpenTelemetry instrumentors."""
+
+    def instrument(self, *, tracer_provider: TracerProvider | None = None) -> None:
+        """Enable instrumentation."""
+        ...
+
+    def uninstrument(self) -> None:
+        """Disable instrumentation."""
+        ...
+
+
+InstrumentorT = TypeVar("InstrumentorT", bound=Instrumentor)
+
+
+class BaseInstrumentation(ABC, Generic[InstrumentorT]):
+    """Base class for managing OpenTelemetry instrumentation lifecycle.
+
+    This class provides a thread-safe singleton pattern for SDK instrumentation.
+    Subclasses must implement `_create_instrumentor()` and `_configure_capture_content()`.
+
+    Each subclass gets its own `_instance` and `_instance_lock` via `__init_subclass__`,
+    ensuring that different providers can be initialized independently without blocking.
+    """
+
+    _instance: BaseInstrumentation[InstrumentorT] | None = None
+    _instance_lock: threading.Lock
+    _lock: threading.Lock
+    _instrumentor: InstrumentorT | None
+    _original_env: dict[str, str | None]
+
+    def __init_subclass__(cls, **kwargs: object) -> None:
+        """Initialize subclass with its own singleton instance and lock."""
+        super().__init_subclass__(**kwargs)
+        cls._instance = None
+        cls._instance_lock = threading.Lock()
+
+    def __new__(cls) -> BaseInstrumentation[InstrumentorT]:
+        """Create or return the singleton instance."""
+        if cls._instance is None:
+            with cls._instance_lock:
+                if cls._instance is None:
+                    instance = super().__new__(cls)
+                    instance._lock = threading.Lock()
+                    instance._instrumentor = None
+                    instance._original_env = {}
+                    cls._instance = instance
+        return cls._instance
+
+    @property
+    def is_instrumented(self) -> bool:
+        """Return whether instrumentation is currently active."""
+        return self._instrumentor is not None
+
+    @abstractmethod
+    def _create_instrumentor(self) -> InstrumentorT:
+        """Create and return a new instrumentor instance.
+
+        Returns:
+            A new instance of the SDK-specific instrumentor.
+        """
+        ...
+
+    @abstractmethod
+    def _configure_capture_content(self, capture_content: ContentCaptureMode) -> None:
+        """Configure environment variables for content capture.
+
+        Args:
+            capture_content: The capture content mode to configure.
+        """
+        ...
+
+    def instrument(
+        self,
+        *,
+        tracer_provider: TracerProvider | None = None,
+        capture_content: ContentCaptureMode = "default",
+    ) -> None:
+        """Enable OpenTelemetry instrumentation for the SDK.
+
+        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.
+        """
+        with self._lock:
+            if self.is_instrumented:
+                return
+
+            self._set_env_var(
+                OTEL_SEMCONV_STABILITY_OPT_IN,
+                OTEL_SEMCONV_STABILITY_VALUE,
+                use_setdefault=True,
+            )
+
+            self._configure_capture_content(capture_content)
+
+            instrumentor = self._create_instrumentor()
+            try:
+                if tracer_provider is not None:
+                    instrumentor.instrument(tracer_provider=tracer_provider)
+                else:
+                    instrumentor.instrument()
+            except Exception:
+                self._restore_env_vars()
+                raise
+
+            self._instrumentor = instrumentor
+
+    def uninstrument(self) -> None:
+        """Disable previously configured instrumentation."""
+        with self._lock:
+            if self._instrumentor is None:
+                return
+
+            self._instrumentor.uninstrument()
+            self._instrumentor = None
+            self._restore_env_vars()
+
+    def _set_env_var(
+        self, key: str, value: str, *, use_setdefault: bool = False
+    ) -> None:
+        """Set an environment variable and track the original value.
+
+        Args:
+            key: The environment variable name.
+            value: The value to set.
+            use_setdefault: If True, only set if not already present.
+        """
+        if key not in self._original_env:
+            self._original_env[key] = os.environ.get(key)
+
+        if use_setdefault:
+            os.environ.setdefault(key, value)
+        else:
+            os.environ[key] = value
+
+    def _restore_env_vars(self) -> None:
+        """Restore all environment variables to their original values."""
+        for key, value in self._original_env.items():
+            if value is None:
+                os.environ.pop(key, None)
+            else:
+                os.environ[key] = value
+        self._original_env.clear()
+
+    @classmethod
+    def reset_for_testing(cls) -> None:
+        """Reset singleton instance for testing.
+
+        This properly cleans up the instrumentor and restores environment variables
+        before resetting the instance.
+        """
+        with cls._instance_lock:
+            if cls._instance is not None:
+                if cls._instance._instrumentor is not None:
+                    cls._instance._instrumentor.uninstrument()
+                    cls._instance._instrumentor = None
+                cls._instance._restore_env_vars()
+            cls._instance = None

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

@@ -0,0 +1,85 @@
+"""OpenTelemetry instrumentation for Google GenAI SDK."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from opentelemetry.instrumentation.google_genai import GoogleGenAiSdkInstrumentor
+
+from .base import BaseInstrumentation, ContentCaptureMode
+
+if TYPE_CHECKING:
+    from opentelemetry.trace import TracerProvider
+
+
+class GoogleGenAIInstrumentation(BaseInstrumentation[GoogleGenAiSdkInstrumentor]):
+    """Manages OpenTelemetry instrumentation lifecycle for the Google GenAI SDK."""
+
+    def _create_instrumentor(self) -> GoogleGenAiSdkInstrumentor:
+        """Create a new Google GenAI instrumentor instance."""
+        return GoogleGenAiSdkInstrumentor()
+
+    def _configure_capture_content(self, capture_content: ContentCaptureMode) -> None:
+        """Configure environment variables for Google GenAI content capture."""
+        # Google GenAI uses ContentCapturingMode enum instead of true/false.
+        # Valid values: NO_CONTENT, SPAN_ONLY, EVENT_ONLY, SPAN_AND_EVENT
+        if capture_content == "enabled":
+            self._set_env_var(
+                "OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT",
+                "SPAN_AND_EVENT",
+            )
+        elif capture_content == "disabled":
+            self._set_env_var(
+                "OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT", "NO_CONTENT"
+            )
+
+
+def instrument_google_genai(
+    *,
+    tracer_provider: TracerProvider | None = None,
+    capture_content: ContentCaptureMode = "default",
+) -> None:
+    """Enable OpenTelemetry instrumentation for the Google GenAI 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_google_genai()
+        ```
+
+        Enable instrumentation with content capture:
+        ```python
+        from mirascope import ops
+
+        ops.configure()
+        ops.instrument_google_genai(capture_content="enabled")
+        ```
+    """
+    GoogleGenAIInstrumentation().instrument(
+        tracer_provider=tracer_provider,
+        capture_content=capture_content,
+    )
+
+
+def uninstrument_google_genai() -> None:
+    """Disable previously configured Google GenAI instrumentation."""
+    GoogleGenAIInstrumentation().uninstrument()
+
+
+def is_google_genai_instrumented() -> bool:
+    """Return whether Google GenAI instrumentation is currently active."""
+    return GoogleGenAIInstrumentation().is_instrumented