Skip to content
Open
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 24 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,16 @@ 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):
breakpoint()
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 +141,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
Loading