diff --git a/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py b/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py index 03cb707141..e9fdc90fe0 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py @@ -19,7 +19,6 @@ from .utils import ( _serialize_binary_content_item, _serialize_image_url_item, - _set_usage_data, ) if TYPE_CHECKING: @@ -150,16 +149,6 @@ def update_invoke_agent_span(span: "sentry_sdk.tracing.Span", result: "Any") -> span, SPANDATA.GEN_AI_RESPONSE_TEXT, str(output), unpack=False ) - # Set token usage data if available - if hasattr(result, "usage") and callable(result.usage): - try: - usage = result.usage() - if usage: - _set_usage_data(span, usage) - except Exception: - # If usage() call fails, continue without setting usage data - pass - # Set model name from response if available if hasattr(result, "response"): try: diff --git a/tests/integrations/pydantic_ai/test_pydantic_ai.py b/tests/integrations/pydantic_ai/test_pydantic_ai.py index 9a1fe79167..7aad94d852 100644 --- a/tests/integrations/pydantic_ai/test_pydantic_ai.py +++ b/tests/integrations/pydantic_ai/test_pydantic_ai.py @@ -205,71 +205,6 @@ def failing_model(messages, info): assert spans[0]["status"] == "internal_error" -@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) -@pytest.mark.asyncio -async def test_agent_run_async_usage_data( - sentry_init, - capture_events, - capture_items, - get_test_agent, - stream_gen_ai_spans, -): - """ - Test that the invoke_agent span includes token usage and model data. - """ - sentry_init( - integrations=[PydanticAIIntegration()], - traces_sample_rate=1.0, - send_default_pii=True, - stream_gen_ai_spans=stream_gen_ai_spans, - ) - - test_agent = get_test_agent() - - if stream_gen_ai_spans: - items = capture_items("transaction", "span") - - result = await test_agent.run("Test input") - - assert result is not None - assert result.output is not None - - (transaction,) = (item.payload for item in items if item.type == "transaction") - else: - events = capture_events() - - result = await test_agent.run("Test input") - - assert result is not None - assert result.output is not None - - (transaction,) = events - - # Verify transaction (the transaction IS the invoke_agent span) - assert transaction["transaction"] == "invoke_agent test_agent" - - # The invoke_agent span should have token usage data - trace_data = transaction["contexts"]["trace"].get("data", {}) - assert "gen_ai.usage.input_tokens" in trace_data, ( - "Missing input_tokens on invoke_agent span" - ) - assert "gen_ai.usage.output_tokens" in trace_data, ( - "Missing output_tokens on invoke_agent span" - ) - assert "gen_ai.usage.total_tokens" in trace_data, ( - "Missing total_tokens on invoke_agent span" - ) - assert "gen_ai.response.model" in trace_data, ( - "Missing response.model on invoke_agent span" - ) - - # Verify the values are reasonable - assert trace_data["gen_ai.usage.input_tokens"] > 0 - assert trace_data["gen_ai.usage.output_tokens"] > 0 - assert trace_data["gen_ai.usage.total_tokens"] > 0 - assert trace_data["gen_ai.response.model"] == "test" # Test model name - - @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_agent_run_sync( sentry_init,