diff --git a/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py b/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py index 27deb0c55c..cfbe49d82c 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py @@ -32,7 +32,7 @@ ) 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 # type: ignore @@ -61,6 +61,22 @@ ImageUrl = None +def _nonempty_content(value: "Any") -> "Optional[str]": + if value is None: + return None + return str(value) or None + + +def _tool_call_arguments(args: "Any") -> "Any": + """Prefer structured args for OTEL tool_call parts; keep raw string if not JSON.""" + if not isinstance(args, str) or not args.strip(): + return args + try: + return json.loads(args) + except ValueError: + return args + + def _transform_system_instructions( permanent_instructions: "list[SystemPromptPart]", current_instructions: "list[str]", @@ -171,7 +187,10 @@ def _set_input_messages( tool_call_id = part.tool_name if hasattr(part, "content"): content.append({"type": "text", "text": str(part.content)}) - # Handle regular content + elif ThinkingPart and isinstance(part, ThinkingPart): + reasoning = _nonempty_content(part.content) + if reasoning is not None: + content.append({"type": "reasoning", "text": reasoning}) elif hasattr(part, "content"): if isinstance(part.content, str): content.append({"type": "text", "text": part.content}) @@ -231,31 +250,56 @@ def _set_output_data( set_on_span(SPANDATA.GEN_AI_RESPONSE_MODEL, response.model_name) try: - # Extract text from ModelResponse - if hasattr(response, "parts"): - texts = [] - tool_calls = [] - - for part in response.parts: - if TextPart and isinstance(part, TextPart) and hasattr(part, "content"): - texts.append(part.content) - elif BaseToolCallPart and isinstance(part, BaseToolCallPart): - tool_call_data = { - "type": "function", - } - if hasattr(part, "tool_name"): - tool_call_data["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) - - if tool_calls: - set_on_span( - SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, safe_serialize(tool_calls) - ) + # OTEL gen_ai.output.messages (text + reasoning + tool_call). + if not hasattr(response, "parts"): + return + + message_parts = [] # type: List[Dict[str, Any]] + + for part in response.parts: + if ThinkingPart and isinstance(part, ThinkingPart): + reasoning = _nonempty_content(getattr(part, "content", None)) + if reasoning is not None: + message_parts.append({"type": "reasoning", "content": reasoning}) + continue + + if TextPart and isinstance(part, TextPart) and hasattr(part, "content"): + text = _nonempty_content(part.content) + if text is not None: + message_parts.append({"type": "text", "content": text}) + continue + + if not (BaseToolCallPart and isinstance(part, BaseToolCallPart)): + continue + + name = getattr(part, "tool_name", None) + if not name: + continue + + otel_tool_call = {"type": "tool_call", "name": name} # type: Dict[str, Any] + tool_call_id = getattr(part, "tool_call_id", None) or getattr( + part, "id", None + ) + if tool_call_id: + otel_tool_call["id"] = tool_call_id + if hasattr(part, "args"): + otel_tool_call["arguments"] = _tool_call_arguments(part.args) + message_parts.append(otel_tool_call) + + if message_parts: + output_message = { + "role": "assistant", + "parts": message_parts, + } # type: Dict[str, Any] + finish_reason = getattr(response, "finish_reason", None) + if finish_reason is not None: + output_message["finish_reason"] = str(finish_reason) + set_data_normalized( + span, + SPANDATA.GEN_AI_OUTPUT_MESSAGES, + [output_message], + unpack=False, + ) except Exception: # If we fail to format output, just skip it diff --git a/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py b/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py index f0c68e85ba..c2e756d754 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py @@ -162,10 +162,17 @@ def update_invoke_agent_span( # Extract output from result output = getattr(result, "output", None) - # Set response text if prompts are enabled if _should_send_prompts() and output: set_data_normalized( - span, SPANDATA.GEN_AI_RESPONSE_TEXT, str(output), unpack=False + span, + SPANDATA.GEN_AI_OUTPUT_MESSAGES, + [ + { + "role": "assistant", + "parts": [{"type": "text", "content": str(output)}], + } + ], + unpack=False, ) # Set model name from response if available diff --git a/sentry_sdk/integrations/pydantic_ai/spans/utils.py b/sentry_sdk/integrations/pydantic_ai/spans/utils.py index 330496c6b2..8e3262a783 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/utils.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/utils.py @@ -10,7 +10,7 @@ from sentry_sdk.traces import StreamedSpan if TYPE_CHECKING: - from typing import Any, Dict, Union + from typing import Any, Dict, Optional, Union from pydantic_ai.usage import RequestUsage, RunUsage # type: ignore @@ -85,3 +85,26 @@ def _set_usage_data( if hasattr(usage, "total_tokens") and usage.total_tokens is not None: set_on_span(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, usage.total_tokens) + + reasoning_tokens = _reasoning_token_count(usage) + if reasoning_tokens is not None: + set_on_span(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING, reasoning_tokens) + + +def _reasoning_token_count(usage: "Any") -> "Optional[int]": + """Provider-specific reasoning token counts live on usage.details or the object itself.""" + details = getattr(usage, "details", None) + if isinstance(details, dict): + for key in ( + "reasoning_tokens", + "thinking_tokens", + "thoughts_token_count", + "thoughts_tokens", + "output_tokens.reasoning", + ): + value = details.get(key) + if isinstance(value, int) and value > 0: + return value + + value = getattr(usage, "reasoning_tokens", None) + return value if isinstance(value, int) and value > 0 else None diff --git a/tests/integrations/pydantic_ai/test_pydantic_ai.py b/tests/integrations/pydantic_ai/test_pydantic_ai.py index 51a1072690..3645b2cfb1 100644 --- a/tests/integrations/pydantic_ai/test_pydantic_ai.py +++ b/tests/integrations/pydantic_ai/test_pydantic_ai.py @@ -523,7 +523,7 @@ async def test_agent_run_stream( assert "gen_ai.usage.input_tokens" in chat_span["attributes"] # Streaming responses should still have output data assert ( - "gen_ai.response.text" in chat_span["attributes"] + "gen_ai.output.messages" in chat_span["attributes"] or "gen_ai.response.model" in chat_span["attributes"] ) elif stream_gen_ai_spans: @@ -573,7 +573,7 @@ async def test_agent_run_stream( assert "gen_ai.usage.input_tokens" in chat_span["attributes"] # Streaming responses should still have output data assert ( - "gen_ai.response.text" in chat_span["attributes"] + "gen_ai.output.messages" in chat_span["attributes"] or "gen_ai.response.model" in chat_span["attributes"] ) else: @@ -602,7 +602,7 @@ async def test_agent_run_stream( assert "gen_ai.usage.input_tokens" in chat_span["data"] # Streaming responses should still have output data assert ( - "gen_ai.response.text" in chat_span["data"] + "gen_ai.output.messages" in chat_span["data"] or "gen_ai.response.model" in chat_span["data"] ) @@ -3113,6 +3113,143 @@ async def test_output_data_with_text_and_tool_calls(sentry_init, capture_items): assert transaction is not None +@pytest.mark.asyncio +async def test_output_data_with_thinking_part(sentry_init): + """ThinkingPart is captured as OTEL reasoning in gen_ai.output.messages.""" + from pydantic_ai import messages + + import sentry_sdk + from sentry_sdk.integrations.pydantic_ai.spans.ai_client import _set_output_data + + sentry_init( + integrations=[PydanticAIIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with sentry_sdk.start_transaction(op="test", name="test"): + span = sentry_sdk.start_span(op="test_span") + + mock_response = MagicMock() + mock_response.model_name = "test-model" + mock_response.finish_reason = "stop" + mock_response.parts = [ + messages.ThinkingPart(content="Let me reason step by step"), + messages.TextPart(content="Final answer"), + ] + + _set_output_data(span, mock_response) + span.finish() + + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span._data + output_messages = json.loads(span._data[SPANDATA.GEN_AI_OUTPUT_MESSAGES]) + assert output_messages == [ + { + "role": "assistant", + "parts": [ + {"type": "reasoning", "content": "Let me reason step by step"}, + {"type": "text", "content": "Final answer"}, + ], + "finish_reason": "stop", + } + ] + + +@pytest.mark.asyncio +async def test_input_messages_with_thinking_part(sentry_init): + """ThinkingPart in history is serialized as type=reasoning, not plain text.""" + from pydantic_ai import messages + + import sentry_sdk + + sentry_init( + integrations=[PydanticAIIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + history = [ + messages.ModelRequest( + parts=[messages.UserPromptPart(content="Explain gravity")] + ), + messages.ModelResponse( + parts=[ + messages.ThinkingPart(content="Newtonian approximation first"), + messages.TextPart(content="Gravity attracts mass."), + ], + model_name="test", + ), + ] + + with sentry_sdk.start_transaction(op="test", name="test"): + span = sentry_sdk.start_span(op="test_span") + _set_input_messages(span, history) + span.finish() + + request_messages = json.loads(span._data[SPANDATA.GEN_AI_REQUEST_MESSAGES]) + assert { + "role": "assistant", + "content": [ + { + "type": "reasoning", + "text": "Newtonian approximation first", + } + ], + } in request_messages + assert { + "role": "assistant", + "content": [{"type": "text", "text": "Gravity attracts mass."}], + } in request_messages + + +@pytest.mark.asyncio +async def test_output_data_with_otel_tool_call(sentry_init): + """ToolCallPart is captured as OTEL tool_call in gen_ai.output.messages.""" + from pydantic_ai import messages + + import sentry_sdk + from sentry_sdk.integrations.pydantic_ai.spans.ai_client import _set_output_data + + sentry_init( + integrations=[PydanticAIIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + ) + + with sentry_sdk.start_transaction(op="test", name="test"): + span = sentry_sdk.start_span(op="test_span") + + mock_response = MagicMock() + mock_response.model_name = "test-model" + mock_response.finish_reason = "tool_call" + mock_response.parts = [ + messages.ToolCallPart( + tool_name="get_weather", + args='{"city": "Berlin"}', + tool_call_id="call_1", + ), + ] + + _set_output_data(span, mock_response) + span.finish() + + output_messages = json.loads(span._data[SPANDATA.GEN_AI_OUTPUT_MESSAGES]) + assert output_messages == [ + { + "role": "assistant", + "parts": [ + { + "type": "tool_call", + "name": "get_weather", + "id": "call_1", + "arguments": {"city": "Berlin"}, + } + ], + "finish_reason": "tool_call", + } + ] + + @pytest.mark.asyncio async def test_output_data_error_handling(sentry_init, capture_items): """ @@ -4287,6 +4424,64 @@ async def test_set_usage_data_with_cache_tokens( assert span_data["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE] == 20 +@pytest.mark.parametrize( + "details_key", + [ + "reasoning_tokens", + "thinking_tokens", + "thoughts_token_count", + "thoughts_tokens", + "output_tokens.reasoning", + ], +) +def test_set_usage_data_reasoning_token_detail_keys( + sentry_init, capture_events, details_key +): + """Provider-specific reasoning counts live under usage.details with different keys.""" + sentry_init( + integrations=[PydanticAIIntegration()], + traces_sample_rate=1.0, + ) + events = capture_events() + + with sentry_sdk.start_transaction(op="test", name="test"): + span = sentry_sdk.start_span(op="test_span") + usage = RequestUsage( + input_tokens=100, + output_tokens=50, + details={details_key: 12}, + ) + _set_usage_data(span, usage) + span.finish() + + (event,) = events + (span_data,) = event["spans"] + assert span_data["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING] == 12 + + +def test_set_usage_data_omits_zero_reasoning_tokens(sentry_init, capture_events): + """Zero reasoning token counts must not be written (providers often default to 0).""" + sentry_init( + integrations=[PydanticAIIntegration()], + traces_sample_rate=1.0, + ) + events = capture_events() + + with sentry_sdk.start_transaction(op="test", name="test"): + span = sentry_sdk.start_span(op="test_span") + usage = RequestUsage( + input_tokens=100, + output_tokens=50, + details={"reasoning_tokens": 0}, + ) + _set_usage_data(span, usage) + span.finish() + + (event,) = events + (span_data,) = event["spans"] + assert SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING not in span_data["data"] + + @pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.parametrize(