diff --git a/sentry_sdk/_types.py b/sentry_sdk/_types.py index d3c70ddfa4..59e501d7da 100644 --- a/sentry_sdk/_types.py +++ b/sentry_sdk/_types.py @@ -472,6 +472,15 @@ class TextPart(TypedDict): type: Literal["text"] content: str + class ReasoningPart(TypedDict): + type: Literal["reasoning"] + content: str + + class ToolCallPart(TypedDict): + type: Literal["tool_call"] + name: NotRequired[str] + arguments: NotRequired[Any] + class ToolDefinition(TypedDict): type: str name: NotRequired[str] diff --git a/sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py b/sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py index c2a033537c..3a7d2fa635 100644 --- a/sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py +++ b/sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py @@ -16,7 +16,9 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from typing import Any, Callable + from typing import Any, Callable, Optional + + from pydantic_ai.messages import ModelResponse def _extract_span_data(node: "Any", ctx: "Any") -> "tuple[list[Any], Any, Any]": @@ -63,7 +65,7 @@ async def wrapped_model_request_run(self: "Any", ctx: "Any") -> "Any": result = await original_model_request_run(self, ctx) # Extract response from result if available - model_response = None + model_response: "Optional[ModelResponse]" = None if hasattr(result, "model_response"): model_response = result.model_response @@ -93,7 +95,7 @@ async def wrapped_model_request_stream(self: "Any", ctx: "Any") -> "Any": # After streaming completes, update span with response data # The ModelRequestNode stores the final response in _result - model_response = None + model_response: "Optional[ModelResponse]" = None if hasattr(self, "_result") and self._result is not None: # _result is a NextNode containing the model_response if hasattr(self._result, "model_response"): diff --git a/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py b/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py index 263e280e11..c89b8609d8 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py @@ -32,11 +32,11 @@ ) if TYPE_CHECKING: - from typing import Any, Dict, List, Union + from typing import Any, Dict, List, Optional, Union - from pydantic_ai.messages import ModelMessage, SystemPromptPart + from pydantic_ai.messages import ModelMessage, ModelResponse, SystemPromptPart - from sentry_sdk._types import TextPart as SentryTextPart + from sentry_sdk import _types try: from pydantic_ai.messages import ( @@ -59,13 +59,14 @@ 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[SentryTextPart]": - text_parts: "list[SentryTextPart]" = [ +) -> "list[_types.TextPart]": + text_parts: "list[_types.TextPart]" = [ { "type": "text", "content": instruction.content, @@ -231,7 +232,8 @@ def _set_input_messages( def _set_output_data( - span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", response: "Any" + span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", + response: "Optional[ModelResponse]", ) -> None: """Set output data on a span.""" if not _should_send_prompts(): @@ -243,13 +245,11 @@ def _set_output_data( set_on_span = ( span.set_attribute if isinstance(span, StreamedSpan) else span.set_data ) - set_on_span(SPANDATA.GEN_AI_RESPONSE_MODEL, response.model_name) + set_on_span(SPANDATA.GEN_AI_RESPONSE_MODEL, response.model_name) # type: ignore[arg-type] try: - # Extract text from ModelResponse if hasattr(response, "parts"): - texts = [] - tool_calls = [] + parts: "list[Union[_types.TextPart, _types.ReasoningPart, _types.ToolCallPart]]" = [] for part in response.parts: if ( @@ -257,25 +257,30 @@ def _set_output_data( and isinstance(part, TextPart) and hasattr(part, "content") ): - texts.append(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_call_data = { - "type": "function", - } + tool_part: "_types.ToolCallPart" = {"type": "tool_call"} if hasattr(part, "tool_name"): - tool_call_data["name"] = part.tool_name + tool_part["name"] = part.tool_name if hasattr(part, "args"): - tool_call_data["arguments"] = safe_serialize(part.args) - tool_calls.append(tool_call_data) - - if texts: - set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, texts) + tool_part["arguments"] = safe_serialize(part.args) + parts.append(tool_part) - if tool_calls: + if parts: set_on_span( - SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, safe_serialize(tool_calls) + SPANDATA.GEN_AI_OUTPUT_MESSAGES, + json.dumps([{"role": "assistant", "parts": parts}]), ) except Exception: @@ -338,7 +343,8 @@ def ai_client_span( def update_ai_client_span( - span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", model_response: "Any" + span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", + model_response: "Optional[ModelResponse]", ) -> None: """Update the AI client span with response data.""" if not span: diff --git a/tests/integrations/pydantic_ai/test_pydantic_ai.py b/tests/integrations/pydantic_ai/test_pydantic_ai.py index 6bf6248192..075c0d2cc7 100644 --- a/tests/integrations/pydantic_ai/test_pydantic_ai.py +++ b/tests/integrations/pydantic_ai/test_pydantic_ai.py @@ -12,6 +12,7 @@ ImageUrl, ModelResponse, TextPart, + ThinkingPart, ToolCallPart, ToolReturnPart, UserPromptPart, @@ -3067,7 +3068,12 @@ async def test_output_data_transformations( def response_model(messages, info): if isinstance(messages[-1].parts[-1], ToolReturnPart): - return ModelResponse(parts=[TextPart(content="The answer is 15.")]) + return ModelResponse( + parts=[ + ThinkingPart(content="5 times 3 is 15."), + TextPart(content="The answer is 15."), + ] + ) return ModelResponse( parts=[ToolCallPart(tool_name="multiply", args={"a": 5, "b": 3})] @@ -3103,18 +3109,30 @@ def multiply(a: int, b: int) -> int: if span["attributes"].get("sentry.op") == "gen_ai.chat" ] assert json.loads( - chat_spans[0]["attributes"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS] + chat_spans[0]["attributes"][SPANDATA.GEN_AI_OUTPUT_MESSAGES] ) == [ { - "type": "function", - "name": "multiply", - "arguments": '{"a": 5, "b": 3}', + "role": "assistant", + "parts": [ + { + "type": "tool_call", + "name": "multiply", + "arguments": '{"a": 5, "b": 3}', + } + ], + } + ] + assert json.loads( + chat_spans[1]["attributes"][SPANDATA.GEN_AI_OUTPUT_MESSAGES] + ) == [ + { + "role": "assistant", + "parts": [ + {"type": "reasoning", "content": "5 times 3 is 15."}, + {"type": "text", "content": "The answer is 15."}, + ], } ] - assert ( - chat_spans[1]["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] - == "The answer is 15." - ) elif stream_gen_ai_spans: items = capture_items("transaction", "span") @@ -3133,18 +3151,30 @@ def multiply(a: int, b: int) -> int: if span["attributes"].get("sentry.op") == "gen_ai.chat" ] assert json.loads( - chat_spans[0]["attributes"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS] + chat_spans[0]["attributes"][SPANDATA.GEN_AI_OUTPUT_MESSAGES] ) == [ { - "type": "function", - "name": "multiply", - "arguments": '{"a": 5, "b": 3}', + "role": "assistant", + "parts": [ + { + "type": "tool_call", + "name": "multiply", + "arguments": '{"a": 5, "b": 3}', + } + ], + } + ] + assert json.loads( + chat_spans[1]["attributes"][SPANDATA.GEN_AI_OUTPUT_MESSAGES] + ) == [ + { + "role": "assistant", + "parts": [ + {"type": "reasoning", "content": "5 times 3 is 15."}, + {"type": "text", "content": "The answer is 15."}, + ], } ] - assert ( - chat_spans[1]["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] - == "The answer is 15." - ) else: events = capture_events() @@ -3159,18 +3189,27 @@ def multiply(a: int, b: int) -> int: chat_spans = [ span for span in transaction["spans"] if span["op"] == "gen_ai.chat" ] - assert json.loads( - chat_spans[0]["data"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS] - ) == [ + assert json.loads(chat_spans[0]["data"][SPANDATA.GEN_AI_OUTPUT_MESSAGES]) == [ { - "type": "function", - "name": "multiply", - "arguments": '{"a": 5, "b": 3}', + "role": "assistant", + "parts": [ + { + "type": "tool_call", + "name": "multiply", + "arguments": '{"a": 5, "b": 3}', + } + ], + } + ] + assert json.loads(chat_spans[1]["data"][SPANDATA.GEN_AI_OUTPUT_MESSAGES]) == [ + { + "role": "assistant", + "parts": [ + {"type": "reasoning", "content": "5 times 3 is 15."}, + {"type": "text", "content": "The answer is 15."}, + ], } ] - assert ( - chat_spans[1]["data"][SPANDATA.GEN_AI_RESPONSE_TEXT] == "The answer is 15." - ) @pytest.mark.asyncio