feat(pydantic-ai): Record gen_ai.output.messages including ThinkingPart

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]

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"):

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,39 +245,42 @@ 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 (
                     TextPart is not None
                     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:

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