From d29dcf42cf6ce16b0bd2858fcb39ac22cd84a521 Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Wed, 5 Aug 2026 09:11:01 +0200 Subject: [PATCH 1/2] test(pydantic-ai): Use FunctionModel for testing output data transformations --- .../pydantic_ai/test_pydantic_ai.py | 177 +++++++++--------- 1 file changed, 90 insertions(+), 87 deletions(-) diff --git a/tests/integrations/pydantic_ai/test_pydantic_ai.py b/tests/integrations/pydantic_ai/test_pydantic_ai.py index 51a1072690..8cea9e859e 100644 --- a/tests/integrations/pydantic_ai/test_pydantic_ai.py +++ b/tests/integrations/pydantic_ai/test_pydantic_ai.py @@ -7,7 +7,15 @@ from pydantic import Field from pydantic_ai import Agent from pydantic_ai.exceptions import ModelRetry, UnexpectedModelBehavior -from pydantic_ai.messages import BinaryContent, ImageUrl, UserPromptPart +from pydantic_ai.messages import ( + BinaryContent, + ImageUrl, + ModelResponse, + TextPart, + ToolCallPart, + ToolReturnPart, + UserPromptPart, +) from pydantic_ai.models.function import FunctionModel from pydantic_ai.usage import RequestUsage @@ -2827,39 +2835,6 @@ async def test_agent_without_name( assert "invoke_agent" in transaction["transaction"] -@pytest.mark.asyncio -async def test_model_response_without_parts(sentry_init, capture_items): - """ - Test handling of model response without parts attribute. - """ - from unittest.mock import MagicMock - - 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") as transaction: - span = sentry_sdk.start_span(op="test_span") - - # Create mock response without parts - mock_response = MagicMock() - mock_response.model_name = "test-model" - del mock_response.parts # Remove parts attribute - - # Should not raise, just skip formatting - _set_output_data(span, mock_response) - - span.finish() - - # Should not crash - assert transaction is not None - - @pytest.mark.asyncio async def test_input_messages_error_handling(sentry_init, capture_items): """ @@ -3073,44 +3048,102 @@ async def test_message_parts_with_list_content(sentry_init, capture_items): assert transaction is not None +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio -async def test_output_data_with_text_and_tool_calls(sentry_init, capture_items): +async def test_output_data_transformations( + sentry_init, capture_items, capture_events, span_streaming +): """ - Test that _set_output_data handles both text and tool calls in response. + Test transformation of outputs from `Model.get_response()` and `Model.stream_response()`. """ - from unittest.mock import MagicMock - - 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, + stream_gen_ai_spans=False, + trace_lifecycle="stream" if span_streaming else "static", ) - with sentry_sdk.start_transaction(op="test", name="test") as transaction: - span = sentry_sdk.start_span(op="test_span") + def response_model(messages, info): + if isinstance(messages[-1].parts[-1], ToolReturnPart): + return ModelResponse(parts=[TextPart(content="The answer is 15.")]) - # Create mock response with both TextPart and ToolCallPart - from pydantic_ai import messages + return ModelResponse( + parts=[ToolCallPart(tool_name="multiply", args={"a": 5, "b": 3})] + ) - text_part = messages.TextPart(content="Here's the result") - tool_call_part = MagicMock() - tool_call_part.tool_name = "test_tool" - tool_call_part.args = {"x": 5} + agent = Agent(FunctionModel(response_model), name="test_agent") - mock_response = MagicMock() - mock_response.model_name = "test-model" - mock_response.parts = [text_part, tool_call_part] + @agent.tool_plain + def multiply(a: int, b: int) -> int: + """Multiply two numbers.""" + return a * b - # Should handle both text and tool calls - _set_output_data(span, mock_response) + if span_streaming: + items = capture_items("span") - span.finish() + await agent.run( + "What is 5 times 3?", + ) + sentry_sdk.flush() - # Should not crash - assert transaction is not None + spans = [item.payload for item in items] + + invoke_agent_span = next( + span + for span in spans + if span["attributes"].get("sentry.op") == "gen_ai.invoke_agent" + ) + assert invoke_agent_span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] == ( + "The answer is 15." + ) + + chat_spans = [ + span + for span in spans + if span["attributes"].get("sentry.op") == "gen_ai.chat" + ] + assert json.loads( + chat_spans[0]["attributes"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS] + ) == [ + { + "type": "function", + "name": "multiply", + "arguments": '{"a": 5, "b": 3}', + } + ] + assert ( + chat_spans[1]["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] + == "The answer is 15." + ) + else: + events = capture_events() + + await agent.run( + "What is 5 times 3?", + ) + + (transaction,) = events + assert transaction["contexts"]["trace"]["op"] == "gen_ai.invoke_agent" + assert transaction["contexts"]["trace"]["data"][ + SPANDATA.GEN_AI_RESPONSE_TEXT + ] == ("The answer is 15.") + + 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] + ) == [ + { + "type": "function", + "name": "multiply", + "arguments": '{"a": 5, "b": 3}', + } + ] + assert ( + chat_spans[1]["data"][SPANDATA.GEN_AI_RESPONSE_TEXT] == "The answer is 15." + ) @pytest.mark.asyncio @@ -3246,36 +3279,6 @@ async def test_set_input_messages_without_prompts(sentry_init, capture_items): assert transaction is not None -@pytest.mark.asyncio -async def test_set_output_data_without_prompts(sentry_init, capture_items): - """ - Test that _set_output_data respects _should_send_prompts(). - """ - from unittest.mock import MagicMock - - import sentry_sdk - from sentry_sdk.integrations.pydantic_ai.spans.ai_client import _set_output_data - - sentry_init( - integrations=[PydanticAIIntegration(include_prompts=False)], - traces_sample_rate=1.0, - send_default_pii=True, - ) - - with sentry_sdk.start_transaction(op="test", name="test") as transaction: - span = sentry_sdk.start_span(op="test_span") - - # Even with response, should not set output data - mock_response = MagicMock() - mock_response.model_name = "test" - _set_output_data(span, mock_response) - - span.finish() - - # Should not crash and should not set output - assert transaction is not None - - @pytest.mark.asyncio async def test_get_model_name_with_exception_in_callable(sentry_init, capture_items): """ From bff16f306fd92ad0b1ff507b22c4d45724ae3d5c Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Wed, 5 Aug 2026 10:10:54 +0200 Subject: [PATCH 2/2] test cleanup --- .../pydantic_ai/test_pydantic_ai.py | 45 +++++++++++++++---- 1 file changed, 36 insertions(+), 9 deletions(-) diff --git a/tests/integrations/pydantic_ai/test_pydantic_ai.py b/tests/integrations/pydantic_ai/test_pydantic_ai.py index 8cea9e859e..6bf6248192 100644 --- a/tests/integrations/pydantic_ai/test_pydantic_ai.py +++ b/tests/integrations/pydantic_ai/test_pydantic_ai.py @@ -3049,18 +3049,19 @@ async def test_message_parts_with_list_content(sentry_init, capture_items): @pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio async def test_output_data_transformations( - sentry_init, capture_items, capture_events, span_streaming + sentry_init, capture_items, capture_events, span_streaming, stream_gen_ai_spans ): """ - Test transformation of outputs from `Model.get_response()` and `Model.stream_response()`. + Test transformation of the model response from `Hooks.on.after_model_request`. """ sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, send_default_pii=True, - stream_gen_ai_spans=False, + stream_gen_ai_spans=stream_gen_ai_spans, trace_lifecycle="stream" if span_streaming else "static", ) @@ -3082,9 +3083,7 @@ def multiply(a: int, b: int) -> int: if span_streaming: items = capture_items("span") - await agent.run( - "What is 5 times 3?", - ) + await agent.run("What is 5 times 3?") sentry_sdk.flush() spans = [item.payload for item in items] @@ -3098,6 +3097,36 @@ def multiply(a: int, b: int) -> int: "The answer is 15." ) + chat_spans = [ + span + for span in spans + if span["attributes"].get("sentry.op") == "gen_ai.chat" + ] + assert json.loads( + chat_spans[0]["attributes"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS] + ) == [ + { + "type": "function", + "name": "multiply", + "arguments": '{"a": 5, "b": 3}', + } + ] + assert ( + chat_spans[1]["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] + == "The answer is 15." + ) + elif stream_gen_ai_spans: + items = capture_items("transaction", "span") + + await agent.run("What is 5 times 3?") + + (transaction,) = (item.payload for item in items if item.type == "transaction") + assert transaction["contexts"]["trace"]["op"] == "gen_ai.invoke_agent" + assert transaction["contexts"]["trace"]["data"][ + SPANDATA.GEN_AI_RESPONSE_TEXT + ] == ("The answer is 15.") + + spans = [item.payload for item in items if item.type == "span"] chat_spans = [ span for span in spans @@ -3119,9 +3148,7 @@ def multiply(a: int, b: int) -> int: else: events = capture_events() - await agent.run( - "What is 5 times 3?", - ) + await agent.run("What is 5 times 3?") (transaction,) = events assert transaction["contexts"]["trace"]["op"] == "gen_ai.invoke_agent"