Skip to content
25 changes: 23 additions & 2 deletions sentry_sdk/integrations/pydantic_ai/spans/ai_client.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,10 @@
import base64
import sentry_sdk
from sentry_sdk.ai.utils import set_data_normalized
from sentry_sdk.ai.utils import (
normalize_message_roles,
set_data_normalized,
truncate_and_annotate_messages,
)
from sentry_sdk.consts import OP, SPANDATA
from sentry_sdk.utils import safe_serialize

Expand Down Expand Up @@ -29,6 +34,7 @@
UserPromptPart,
TextPart,
ThinkingPart,
BinaryContent,
)
except ImportError:
# Fallback if these classes are not available
Expand All @@ -38,6 +44,7 @@
UserPromptPart = None
TextPart = None
ThinkingPart = None
BinaryContent = None


def _set_input_messages(span: "sentry_sdk.tracing.Span", messages: "Any") -> None:
Expand Down Expand Up @@ -107,6 +114,15 @@ def _set_input_messages(span: "sentry_sdk.tracing.Span", messages: "Any") -> Non
for item in part.content:
if isinstance(item, str):
content.append({"type": "text", "text": item})
elif BinaryContent and isinstance(item, BinaryContent):
content.append(
{
"type": "blob",
"modality": item.media_type.split("/")[0],
"mime_type": item.media_type,
"content": f"data:{item.media_type};base64,{base64.b64encode(item.data).decode('utf-8')}",
}
)
else:
content.append(safe_serialize(item))
else:
Expand All @@ -124,8 +140,13 @@ def _set_input_messages(span: "sentry_sdk.tracing.Span", messages: "Any") -> Non
formatted_messages.append(message)

if formatted_messages:
normalized_messages = normalize_message_roles(formatted_messages)
scope = sentry_sdk.get_current_scope()
messages_data = truncate_and_annotate_messages(
normalized_messages, span, scope
)
set_data_normalized(
span, SPANDATA.GEN_AI_REQUEST_MESSAGES, formatted_messages, unpack=False
span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False
)
except Exception:
# If we fail to format messages, just skip it
Expand Down
14 changes: 12 additions & 2 deletions sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,10 @@
import sentry_sdk
from sentry_sdk.ai.utils import get_start_span_function, set_data_normalized
from sentry_sdk.ai.utils import (
get_start_span_function,
normalize_message_roles,
set_data_normalized,
truncate_and_annotate_messages,
)
from sentry_sdk.consts import OP, SPANDATA

from ..consts import SPAN_ORIGIN
Expand Down Expand Up @@ -102,8 +107,13 @@ def invoke_agent_span(
)

if messages:
normalized_messages = normalize_message_roles(messages)
scope = sentry_sdk.get_current_scope()
messages_data = truncate_and_annotate_messages(
normalized_messages, span, scope
)
set_data_normalized(
span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages, unpack=False
span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False
)

return span
Expand Down
2 changes: 1 addition & 1 deletion sentry_sdk/integrations/pydantic_ai/spans/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
from typing import TYPE_CHECKING

if TYPE_CHECKING:
from typing import Union
from typing import Union, Dict, Any, List
from pydantic_ai.usage import RequestUsage, RunUsage # type: ignore


Expand Down
131 changes: 131 additions & 0 deletions tests/integrations/pydantic_ai/test_pydantic_ai.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,18 @@
import asyncio
import json
import pytest
from unittest.mock import MagicMock

from typing import Annotated
from pydantic import Field

import sentry_sdk
from sentry_sdk._types import BLOB_DATA_SUBSTITUTE
from sentry_sdk.integrations.pydantic_ai import PydanticAIIntegration
from sentry_sdk.integrations.pydantic_ai.spans.ai_client import _set_input_messages

from pydantic_ai import Agent
from pydantic_ai.messages import BinaryContent, UserPromptPart
from pydantic_ai.models.test import TestModel
from pydantic_ai.exceptions import ModelRetry, UnexpectedModelBehavior

Expand Down Expand Up @@ -2604,3 +2610,128 @@ async def test_ai_client_span_gets_agent_from_scope(sentry_init, capture_events)

# Should not crash
assert transaction is not None


def _get_messages_from_span(span_data):
"""Helper to extract and parse messages from span data."""
messages_data = span_data["gen_ai.request.messages"]
return (
json.loads(messages_data) if isinstance(messages_data, str) else messages_data
)


def _find_binary_content(messages_data, expected_modality, expected_mime_type):
"""Helper to find and verify binary content in messages."""
for msg in messages_data:
if "content" not in msg:
continue
for content_item in msg["content"]:
if content_item.get("type") == "blob":
assert content_item["modality"] == expected_modality
assert content_item["mime_type"] == expected_mime_type
assert "content" in content_item
content_str = str(content_item["content"])
assert (
f"data:{expected_mime_type};base64," in content_str
or BLOB_DATA_SUBSTITUTE in content_str
)
return True
return False


@pytest.mark.asyncio
async def test_binary_content_encoding_image(sentry_init, capture_events):
"""Test that BinaryContent with image data is properly encoded in messages."""
sentry_init(
integrations=[PydanticAIIntegration()],
traces_sample_rate=1.0,
send_default_pii=True,
)

events = capture_events()

with sentry_sdk.start_transaction(op="test", name="test"):
span = sentry_sdk.start_span(op="test_span")
binary_content = BinaryContent(
data=b"fake_image_data_12345", media_type="image/png"
)
user_part = UserPromptPart(content=["Look at this image:", binary_content])
mock_msg = MagicMock()
mock_msg.parts = [user_part]
mock_msg.instructions = None

_set_input_messages(span, [mock_msg])
span.finish()

(event,) = events
span_data = event["spans"][0]["data"]
messages_data = _get_messages_from_span(span_data)
assert _find_binary_content(messages_data, "image", "image/png")


@pytest.mark.asyncio
async def test_binary_content_encoding_mixed_content(sentry_init, capture_events):
"""Test that BinaryContent mixed with text content is properly handled."""
sentry_init(
integrations=[PydanticAIIntegration()],
traces_sample_rate=1.0,
send_default_pii=True,
)

events = capture_events()

with sentry_sdk.start_transaction(op="test", name="test"):
span = sentry_sdk.start_span(op="test_span")
binary_content = BinaryContent(
data=b"fake_image_bytes", media_type="image/jpeg"
)
user_part = UserPromptPart(
content=["Here is an image:", binary_content, "What do you see?"]
)
mock_msg = MagicMock()
mock_msg.parts = [user_part]
mock_msg.instructions = None

_set_input_messages(span, [mock_msg])
span.finish()

(event,) = events
span_data = event["spans"][0]["data"]
messages_data = _get_messages_from_span(span_data)

# Verify both text and binary content are present
found_text = any(
content_item.get("type") == "text"
for msg in messages_data
if "content" in msg
for content_item in msg["content"]
)
assert found_text, "Text content should be found"
assert _find_binary_content(messages_data, "image", "image/jpeg")


@pytest.mark.asyncio
async def test_binary_content_in_agent_run(sentry_init, capture_events):
"""Test that BinaryContent in actual agent run is properly captured in spans."""
agent = Agent("test", name="test_binary_agent")

sentry_init(
integrations=[PydanticAIIntegration()],
traces_sample_rate=1.0,
send_default_pii=True,
)

events = capture_events()
binary_content = BinaryContent(
data=b"fake_image_data_for_testing", media_type="image/png"
)
await agent.run(["Analyze this image:", binary_content])

(transaction,) = events
chat_spans = [s for s in transaction["spans"] if s["op"] == "gen_ai.chat"]
assert len(chat_spans) >= 1

chat_span = chat_spans[0]
if "gen_ai.request.messages" in chat_span["data"]:
messages_str = str(chat_span["data"]["gen_ai.request.messages"])
assert any(keyword in messages_str for keyword in ["blob", "image", "base64"])
Loading