diff --git a/tests/integrations/pydantic_ai/test_pydantic_ai.py b/tests/integrations/pydantic_ai/test_pydantic_ai.py index f0ddc6c4ed..15627a705a 100644 --- a/tests/integrations/pydantic_ai/test_pydantic_ai.py +++ b/tests/integrations/pydantic_ai/test_pydantic_ai.py @@ -2794,8 +2794,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 +2812,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 +2842,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 ):