fix(pydantic-ai): Stop double reporting model requests

sentry_sdk/integrations/pydantic_ai/__init__.py

@@ -10,7 +10,6 @@
 from .patches import (
     _patch_agent_run,
     _patch_graph_nodes,
-    _patch_model_request,
     _patch_tool_execution,
 )
 
@@ -46,5 +45,4 @@ def setup_once() -> None:
         """
         _patch_agent_run()
         _patch_graph_nodes()
-        _patch_model_request()
         _patch_tool_execution()

sentry_sdk/integrations/pydantic_ai/patches/__init__.py

@@ -1,4 +1,3 @@
 from .agent_run import _patch_agent_run  # noqa: F401
 from .graph_nodes import _patch_graph_nodes  # noqa: F401
-from .model_request import _patch_model_request  # noqa: F401
 from .tools import _patch_tool_execution  # noqa: F401

sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py

@@ -57,6 +57,11 @@ def _patch_graph_nodes() -> None:
 
     @wraps(original_model_request_run)
     async def wrapped_model_request_run(self: "Any", ctx: "Any") -> "Any":
+        did_stream = getattr(self, "_did_stream", None)
+        cached_result = getattr(self, "_result", None)
+        if did_stream or cached_result is not None:
+            return await original_model_request_run(self, ctx)
+
         messages, model, model_settings = _extract_span_data(self, ctx)
 
         with ai_client_span(messages, None, model, model_settings) as span:
@@ -83,6 +88,11 @@ def create_wrapped_stream(
         @asynccontextmanager
         @wraps(original_stream_method)
         async def wrapped_model_request_stream(self: "Any", ctx: "Any") -> "Any":
+            did_stream = getattr(self, "_did_stream", None)
+            if did_stream:
+                async with original_stream_method(self, ctx) as stream:
+                    yield stream
+
             messages, model, model_settings = _extract_span_data(self, ctx)
 
             # Create chat span for streaming request

tests/integrations/pydantic_ai/test_pydantic_ai.py

@@ -75,7 +75,7 @@ async def test_agent_run_async(sentry_init, capture_events, test_agent):
 
     # Find child span types (invoke_agent is the transaction, not a child span)
     chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"]
-    assert len(chat_spans) >= 1
+    assert len(chat_spans) == 1
 
     # Check chat span
     chat_span = chat_spans[0]
@@ -158,7 +158,7 @@ def test_agent_run_sync(sentry_init, capture_events, test_agent):
 
     # Find span types
     chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"]
-    assert len(chat_spans) >= 1
+    assert len(chat_spans) == 1
 
     # Verify streaming flag is False for sync
     for chat_span in chat_spans:
@@ -192,7 +192,7 @@ async def test_agent_run_stream(sentry_init, capture_events, test_agent):
 
     # Find chat spans
     chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"]
-    assert len(chat_spans) >= 1
+    assert len(chat_spans) == 1
 
     # Verify streaming flag is True for streaming
     for chat_span in chat_spans:
@@ -231,9 +231,8 @@ async def test_agent_run_stream_events(sentry_init, capture_events, test_agent):
     # Find chat spans
     spans = transaction["spans"]
     chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"]
-    assert len(chat_spans) >= 1
+    assert len(chat_spans) == 1
 
-    # run_stream_events uses run() internally, so streaming should be False
     for chat_span in chat_spans:
         assert chat_span["data"]["gen_ai.response.streaming"] is False
 
@@ -269,7 +268,7 @@ def add_numbers(a: int, b: int) -> int:
     tool_spans = [s for s in spans if s["op"] == "gen_ai.execute_tool"]
 
     # Should have tool spans
-    assert len(tool_spans) >= 1
+    assert len(tool_spans) == 1
 
     # Check tool span
     tool_span = tool_spans[0]
@@ -502,7 +501,7 @@ async def test_model_settings(sentry_init, capture_events, test_agent_with_setti
 
     # Find chat span
     chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"]
-    assert len(chat_spans) >= 1
+    assert len(chat_spans) == 1
 
     chat_span = chat_spans[0]
     # Check that model settings are captured
@@ -548,7 +547,7 @@ async def test_system_prompt_attribute(
 
     # The transaction IS the invoke_agent span, check for messages in chat spans instead
     chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"]
-    assert len(chat_spans) >= 1
+    assert len(chat_spans) == 1
 
     chat_span = chat_spans[0]
 
@@ -587,7 +586,7 @@ async def test_error_handling(sentry_init, capture_events):
     await agent.run("Hello")
 
     # At minimum, we should have a transaction
-    assert len(events) >= 1
+    assert len(events) == 1
     transaction = [e for e in events if e.get("type") == "transaction"][0]
     assert transaction["transaction"] == "invoke_agent test_error"
     # Transaction should complete successfully (status key may not exist if no error)
@@ -681,7 +680,7 @@ async def run_agent(input_text):
         assert transaction["type"] == "transaction"
         assert transaction["transaction"] == "invoke_agent test_agent"
         # Each should have its own spans
-        assert len(transaction["spans"]) >= 1
+        assert len(transaction["spans"]) == 1
 
 
 @pytest.mark.asyncio
@@ -721,7 +720,7 @@ async def test_message_history(sentry_init, capture_events):
     await agent.run("What is my name?", message_history=history)
 
     # We should have 2 transactions
-    assert len(events) >= 2
+    assert len(events) == 2
 
     # Check the second transaction has the full history
     second_transaction = events[1]
@@ -755,7 +754,7 @@ async def test_gen_ai_system(sentry_init, capture_events, test_agent):
 
     # Find chat span
     chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"]
-    assert len(chat_spans) >= 1
+    assert len(chat_spans) == 1
 
     chat_span = chat_spans[0]
     # gen_ai.system should be set from the model (TestModel -> 'test')
@@ -812,7 +811,7 @@ async def test_include_prompts_true(sentry_init, capture_events, test_agent):
     chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"]
 
     # Verify that messages are captured in chat spans
-    assert len(chat_spans) >= 1
+    assert len(chat_spans) == 1
     for chat_span in chat_spans:
         assert "gen_ai.request.messages" in chat_span["data"]
 
@@ -1242,7 +1241,7 @@ async def test_invoke_agent_with_instructions(
 
     # The transaction IS the invoke_agent span, check for messages in chat spans instead
     chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"]
-    assert len(chat_spans) >= 1
+    assert len(chat_spans) == 1
 
     chat_span = chat_spans[0]
 
@@ -1366,7 +1365,7 @@ async def test_usage_data_partial(sentry_init, capture_events):
     spans = transaction["spans"]
 
     chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"]
-    assert len(chat_spans) >= 1
+    assert len(chat_spans) == 1
 
     # Check that usage data fields exist (they may or may not be set depending on TestModel)
     chat_span = chat_spans[0]
@@ -1461,7 +1460,7 @@ def calc_tool(value: int) -> int:
     chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"]
 
     # At least one chat span should exist
-    assert len(chat_spans) >= 1
+    assert len(chat_spans) == 2
 
     # Check if tool calls are captured in response
     for chat_span in chat_spans:
@@ -1509,7 +1508,7 @@ async def test_message_formatting_with_different_parts(sentry_init, capture_even
     chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"]
 
     # Should have chat spans
-    assert len(chat_spans) >= 1
+    assert len(chat_spans) == 1
 
     # Check that messages are captured
     chat_span = chat_spans[0]
@@ -1781,7 +1780,7 @@ def test_tool(x: int) -> int:
     chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"]
 
     # Should have chat spans
-    assert len(chat_spans) >= 1
+    assert len(chat_spans) == 2
 
 
 @pytest.mark.asyncio
@@ -2762,7 +2761,7 @@ async def test_binary_content_in_agent_run(sentry_init, capture_events):
 
     (transaction,) = events
     chat_spans = [s for s in transaction["spans"] if s["op"] == "gen_ai.chat"]
-    assert len(chat_spans) >= 1
+    assert len(chat_spans) == 1
 
     chat_span = chat_spans[0]
     if "gen_ai.request.messages" in chat_span["data"]:
@@ -2794,8 +2793,9 @@ async def test_set_usage_data_with_cache_tokens(sentry_init, capture_events):
     assert span_data["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE] == 20
 
 
+@pytest.mark.asyncio
 @pytest.mark.parametrize(
-    "url,image_url_kwargs,expected_content",
+    "url, image_url_kwargs, expected_content",
     [
         pytest.param(
             "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAYAAACNMs",
@@ -2811,10 +2811,16 @@ async def test_set_usage_data_with_cache_tokens(sentry_init, capture_events):
         ),
         pytest.param(
             "https://example.com/api?data=iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAYAAACNMs",
-            {"media_type": "image/png"},
+            {},
             "https://example.com/api?data=iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAYAAACNMs",
             id="http_url_with_base64_query_param",
         ),
+        pytest.param(
+            "https://example.com/api?data=iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAYAAACNMs",
+            {"media_type": "image/png"},
+            "https://example.com/api?data=iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAYAAACNMs",
+            id="http_url_with_base64_query_param_and_media_type",
+        ),
         pytest.param(
             "data:image/svg+xml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciLz4=",
             {},
@@ -2835,76 +2841,6 @@ async def test_set_usage_data_with_cache_tokens(sentry_init, capture_events):
         ),
     ],
 )
-def test_image_url_base64_content_in_span(
-    sentry_init, capture_events, url, image_url_kwargs, expected_content
-):
-    from sentry_sdk.integrations.pydantic_ai.spans.ai_client import ai_client_span
-
-    sentry_init(
-        integrations=[PydanticAIIntegration()],
-        traces_sample_rate=1.0,
-        send_default_pii=True,
-    )
-
-    events = capture_events()
-
-    with sentry_sdk.start_transaction(op="test", name="test"):
-        image_url = ImageUrl(url=url, **image_url_kwargs)
-        user_part = UserPromptPart(content=["Look at this image:", image_url])
-        mock_msg = MagicMock()
-        mock_msg.parts = [user_part]
-        mock_msg.instructions = None
-
-        span = ai_client_span([mock_msg], None, None, None)
-        span.finish()
-
-    (event,) = events
-    chat_spans = [s for s in event["spans"] if s["op"] == "gen_ai.chat"]
-    assert len(chat_spans) >= 1
-    messages_data = _get_messages_from_span(chat_spans[0]["data"])
-
-    found_image = False
-    for msg in messages_data:
-        if "content" not in msg:
-            continue
-        for content_item in msg["content"]:
-            if content_item.get("type") == "image":
-                found_image = True
-                assert content_item["content"] == expected_content
-
-    assert found_image, "Image content item should be found in messages data"
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize(
-    "url, image_url_kwargs, expected_content",
-    [
-        pytest.param(
-            "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAYAAACNMs",
-            {},
-            BLOB_DATA_SUBSTITUTE,
-            id="base64_data_url_redacted",
-        ),
-        pytest.param(
-            "https://example.com/image.png",
-            {},
-            "https://example.com/image.png",
-            id="http_url_no_redaction",
-        ),
-        pytest.param(
-            "https://example.com/api?data=iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAYAAACNMs",
-            {},
-            "https://example.com/api?data=iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAYAAACNMs",
-            id="http_url_with_base64_query_param",
-        ),
-        pytest.param(
-            "https://example.com/api?data=iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAYAAACNMs",
-            {"media_type": "image/png"},
-            "https://example.com/api?data=iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAYAAACNMs",
-            id="http_url_with_base64_query_param_and_media_type",
-        ),
-    ],
-)
 async def test_invoke_agent_image_url(
     sentry_init, capture_events, url, image_url_kwargs, expected_content
 ):