"""Utility functions for PydanticAI span instrumentation.""" from typing import TYPE_CHECKING import sentry_sdk from sentry_sdk._types import BLOB_DATA_SUBSTITUTE from sentry_sdk.ai.consts import DATA_URL_BASE64_REGEX from sentry_sdk.ai.utils import get_modality_from_mime_type from sentry_sdk.consts import SPANDATA from sentry_sdk.traces import StreamedSpan if TYPE_CHECKING: from typing import Any, Dict, Union from pydantic_ai.usage import RequestUsage, RunUsage # type: ignore def _serialize_image_url_item(item: "Any") -> "Dict[str, Any]": """Serialize an ImageUrl content item for span data. For data URLs containing base64-encoded images, the content is redacted. For regular HTTP URLs, the URL string is preserved. """ url = str(item.url) data_url_match = DATA_URL_BASE64_REGEX.match(url) if data_url_match: return { "type": "image", "content": BLOB_DATA_SUBSTITUTE, } return { "type": "image", "content": url, } def _serialize_binary_content_item(item: "Any") -> "Dict[str, Any]": """Serialize a BinaryContent item for span data, redacting the blob data.""" return { "type": "blob", "modality": get_modality_from_mime_type(item.media_type), "mime_type": item.media_type, "content": BLOB_DATA_SUBSTITUTE, } def _set_usage_data( span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", usage: "Union[RequestUsage, RunUsage]", ) -> None: """Set token usage data on a span. This function works with both RequestUsage (single request) and RunUsage (agent run) objects from pydantic_ai. Args: span: The Sentry span to set data on. usage: RequestUsage or RunUsage object containing token usage information. """ if usage is None: return set_on_span = ( span.set_attribute if isinstance(span, StreamedSpan) else span.set_data ) if hasattr(usage, "input_tokens") and usage.input_tokens is not None: set_on_span(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, usage.input_tokens) # Pydantic AI uses cache_read_tokens (not input_tokens_cached) if hasattr(usage, "cache_read_tokens") and usage.cache_read_tokens is not None: set_on_span(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED, usage.cache_read_tokens) # Pydantic AI uses cache_write_tokens (not input_tokens_cache_write) if hasattr(usage, "cache_write_tokens") and usage.cache_write_tokens is not None: set_on_span( SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE, usage.cache_write_tokens, ) if hasattr(usage, "output_tokens") and usage.output_tokens is not None: set_on_span(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, usage.output_tokens) if hasattr(usage, "total_tokens") and usage.total_tokens is not None: set_on_span(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, usage.total_tokens)