Skip to content

ref(pydantic-ai): Move pydantic-ai object reads behind an extraction layer #7194

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
ericapisani wants to merge 1 commit into from
Draft

ref(pydantic-ai): Move pydantic-ai object reads behind an extraction layer #7194

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
488 changes: 488 additions & 0 deletions sentry_sdk/integrations/pydantic_ai/_extract.py
View file
Open in desktop

Large diffs are not rendered by default.

30 changes: 3 additions & 27 deletions sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@

from sentry_sdk.integrations import DidNotEnable

from .._extract import extract_graph_request_data
from ..spans import (
ai_client_span,
update_ai_client_span,
Expand All @@ -21,31 +22,6 @@
from pydantic_ai.messages import ModelResponse


def _extract_span_data(node: "Any", ctx: "Any") -> "tuple[list[Any], Any, Any]":
"""Extract common data needed for creating chat spans.

Returns:
Tuple of (messages, model, model_settings)
"""
# Extract model and settings from context
model = None
model_settings = None
if hasattr(ctx, "deps"):
model = getattr(ctx.deps, "model", None)
model_settings = getattr(ctx.deps, "model_settings", None)

# Build full message list: history + current request
messages = []
if hasattr(ctx, "state") and hasattr(ctx.state, "message_history"):
messages.extend(ctx.state.message_history)

current_request = getattr(node, "request", None)
if current_request:
messages.append(current_request)

return messages, model, model_settings


def _patch_graph_nodes() -> None:
"""
Patches the graph node execution to create appropriate spans.
Expand All @@ -67,7 +43,7 @@ async def wrapped_model_request_run(self: "Any", ctx: "Any") -> "Any":
if did_stream or cached_result is not None:
return await original_model_request_run(self, ctx)

messages, model, model_settings = _extract_span_data(self, ctx)
messages, model, model_settings = extract_graph_request_data(self, ctx)

with ai_client_span(messages, None, model, model_settings) as span:
result = await original_model_request_run(self, ctx)
Expand Down Expand Up @@ -101,7 +77,7 @@ async def wrapped_model_request_stream(self: "Any", ctx: "Any") -> "Any":
yield stream
return

messages, model, model_settings = _extract_span_data(self, ctx)
messages, model, model_settings = extract_graph_request_data(self, ctx)

# Create chat span for streaming request
with ai_client_span(messages, None, model, model_settings) as span:
Expand Down
11 changes: 3 additions & 8 deletions sentry_sdk/integrations/pydantic_ai/patches/tools.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
from sentry_sdk.integrations import DidNotEnable
from sentry_sdk.utils import capture_internal_exceptions, reraise

from .._extract import extract_tool_call_args
from ..spans import execute_tool_span, update_execute_tool_span
from ..utils import _capture_exception, get_current_agent

Expand Down Expand Up @@ -52,10 +53,7 @@ async def wrapped_execute_tool_call(
agent = get_current_agent()

if agent and tool:
try:
args_dict = call.args_as_dict()
except Exception:
args_dict = call.args if isinstance(call.args, dict) else {}
args_dict = extract_tool_call_args(call)

# Create execute_tool span
# Nesting is handled by isolation_scope() to ensure proper parent-child relationships
Expand Down Expand Up @@ -125,10 +123,7 @@ async def wrapped_call_tool(
agent = get_current_agent()

if agent and tool:
try:
args_dict = call.args_as_dict()
except Exception:
args_dict = call.args if isinstance(call.args, dict) else {}
args_dict = extract_tool_call_args(call)

# Create execute_tool span
# Nesting is handled by isolation_scope() to ensure proper parent-child relationships
Expand Down
226 changes: 25 additions & 201 deletions sentry_sdk/integrations/pydantic_ai/spans/ai_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,96 +13,31 @@
has_span_streaming_enabled,
should_truncate_gen_ai_input,
)
from sentry_sdk.utils import safe_serialize

from .._extract import (
extract_model_info,
extract_request_messages,
extract_response_parts,
extract_system_instructions,
)
from ..consts import SPAN_ORIGIN
from ..utils import (
_get_model_name,
_set_agent_data,
_set_available_tools,
_set_model_data,
_should_send_prompts,
get_current_agent,
get_is_streaming,
)
from .utils import (
_serialize_binary_content_item,
_serialize_image_url_item,
_set_usage_data,
)
from .utils import _set_usage_data

if TYPE_CHECKING:
from typing import Any, Dict, List, Optional, Union
from typing import Any, Optional, Union

from pydantic_ai.messages import ModelMessage, ModelResponse, SystemPromptPart
from pydantic_ai.messages import ModelResponse

from sentry_sdk import _types
from sentry_sdk.traces import StreamedSpan

try:
from pydantic_ai.messages import (
BaseToolCallPart,
BaseToolReturnPart,
BinaryContent,
ImageUrl,
SystemPromptPart,
TextPart,
ThinkingPart,
UserPromptPart,
)
except ImportError:
# Fallback if these classes are not available
BaseToolCallPart = None # type: ignore[misc,assignment]
BaseToolReturnPart = None # type: ignore[misc,assignment]
SystemPromptPart = None # type: ignore[misc,assignment]
UserPromptPart = None # type: ignore[misc,assignment]
TextPart = None # type: ignore[misc,assignment]
ThinkingPart = None # type: ignore[misc,assignment]
BinaryContent = None # type: ignore[misc,assignment]
ImageUrl = None # type: ignore[misc,assignment]
ThinkingPart = None # type: ignore[misc,assignment]


def _transform_system_instructions(
permanent_instructions: "list[SystemPromptPart]",
current_instructions: "list[str]",
) -> "list[_types.TextPart]":
text_parts: "list[_types.TextPart]" = [
{
"type": "text",
"content": instruction.content,
}
for instruction in permanent_instructions
]

text_parts.extend(
{
"type": "text",
"content": instruction,
}
for instruction in current_instructions
)

return text_parts


def _get_system_instructions(
messages: "list[ModelMessage]",
) -> "tuple[list[SystemPromptPart], list[str]]":
permanent_instructions = []
current_instructions = []

for msg in messages:
if hasattr(msg, "parts"):
for part in msg.parts:
if SystemPromptPart is not None and isinstance(part, SystemPromptPart):
permanent_instructions.append(part)

if hasattr(msg, "instructions") and msg.instructions is not None:
current_instructions.append(msg.instructions)

return permanent_instructions, current_instructions


def _set_input_messages(
span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", messages: "Any"
Expand All @@ -114,97 +49,16 @@ def _set_input_messages(
if not messages:
return

permanent_instructions, current_instructions = _get_system_instructions(messages)
if len(permanent_instructions) > 0 or len(current_instructions) > 0:
system_instructions = extract_system_instructions(messages)
if system_instructions:
_set_span_data_attribute(
span,
SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS,
json.dumps(
_transform_system_instructions(
permanent_instructions, current_instructions
)
),
json.dumps(system_instructions),
)

try:
formatted_messages = []

for msg in messages:
if hasattr(msg, "parts"):
for part in msg.parts:
role = "user"
# Use isinstance checks with proper base classes
if SystemPromptPart is not None and isinstance(
part, SystemPromptPart
):
continue
elif (
(TextPart is not None and isinstance(part, TextPart))
or (ThinkingPart is not None and isinstance(part, ThinkingPart))
or (
BaseToolCallPart is not None
and isinstance(part, BaseToolCallPart)
)
):
role = "assistant"
elif BaseToolReturnPart is not None and isinstance(
part, BaseToolReturnPart
):
role = "tool"

content: "List[Dict[str, Any] | str]" = []
tool_calls = None
tool_call_id = None

# Handle ToolCallPart (assistant requesting tool use)
if BaseToolCallPart is not None and isinstance(
part, BaseToolCallPart
):
tool_call_data = {}
if hasattr(part, "tool_name"):
tool_call_data["name"] = part.tool_name
if hasattr(part, "args"):
tool_call_data["arguments"] = safe_serialize(part.args)
if tool_call_data:
tool_calls = [tool_call_data]
# Handle ToolReturnPart (tool result)
elif BaseToolReturnPart is not None and isinstance(
part, BaseToolReturnPart
):
if hasattr(part, "tool_name"):
tool_call_id = part.tool_name
if hasattr(part, "content"):
content.append({"type": "text", "text": str(part.content)})
# Handle regular content
elif hasattr(part, "content"):
if isinstance(part.content, str):
content.append({"type": "text", "text": part.content})
elif isinstance(part.content, list):
for item in part.content:
if isinstance(item, str):
content.append({"type": "text", "text": item})
elif ImageUrl is not None and isinstance(
item, ImageUrl
):
content.append(_serialize_image_url_item(item))
elif BinaryContent is not None and isinstance(
item, BinaryContent
):
content.append(_serialize_binary_content_item(item))
else:
content.append(safe_serialize(item))
else:
content.append({"type": "text", "text": str(part.content)})
# Add message if we have content or tool calls
if content or tool_calls:
message: "Dict[str, Any]" = {"role": role}
if content:
message["content"] = content
if tool_calls:
message["tool_calls"] = tool_calls
if tool_call_id:
message["tool_call_id"] = tool_call_id
formatted_messages.append(message)
formatted_messages = extract_request_messages(messages)

if formatted_messages:
normalized_messages = normalize_message_roles(formatted_messages)
Expand Down Expand Up @@ -240,42 +94,13 @@ def _set_output_data(
)

try:
if hasattr(response, "parts"):
parts: "list[Union[_types.TextPart, _types.ReasoningPart, _types.ToolCallPart]]" = []

for part in response.parts:
if (
TextPart is not None
and isinstance(part, TextPart)
and hasattr(part, "content")
):
parts.append({"type": "text", "content": part.content})

elif ThinkingPart is not None and isinstance(part, ThinkingPart):
parts.append(
{
"type": "reasoning",
"content": part.content,
}
)

elif BaseToolCallPart is not None and isinstance(
part, BaseToolCallPart
):
tool_part: "_types.ToolCallPart" = {"type": "tool_call"}
if hasattr(part, "tool_name"):
tool_part["name"] = part.tool_name
if hasattr(part, "args"):
tool_part["arguments"] = safe_serialize(part.args)
parts.append(tool_part)

if parts:
_set_span_data_attribute(
span,
SPANDATA.GEN_AI_OUTPUT_MESSAGES,
json.dumps([{"role": "assistant", "parts": parts}]),
)

parts = extract_response_parts(response)
if parts:
_set_span_data_attribute(
span,
SPANDATA.GEN_AI_OUTPUT_MESSAGES,
json.dumps([{"role": "assistant", "parts": parts}]),
)
except Exception:
# If we fail to format output, just skip it
pass
Expand All @@ -292,12 +117,11 @@ def ai_client_span(
model: Model object
model_settings: Model settings
"""
# Determine model name for span name
model_obj = model
if agent and hasattr(agent, "model"):
model_obj = agent.model

model_name = _get_model_name(model_obj) or "unknown"
# Determine model name for span name, resolving the same way as
# _set_model_data so the span name and gen_ai.request.model agree
model_name = (
extract_model_info(model, None, agent or get_current_agent()).name or "unknown"
)

span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options)
if span_streaming:
Expand Down
Loading
Loading