|
12 | 12 | from sentry_sdk.integrations.pydantic_ai import PydanticAIIntegration |
13 | 13 | from sentry_sdk.integrations.pydantic_ai.spans.ai_client import _set_input_messages |
14 | 14 | from sentry_sdk.integrations.pydantic_ai.spans.utils import _set_usage_data |
15 | | - |
16 | 15 | from pydantic_ai import Agent |
17 | | -from pydantic_ai.messages import BinaryContent, UserPromptPart |
| 16 | +from pydantic_ai.messages import BinaryContent, ImageUrl, UserPromptPart |
18 | 17 | from pydantic_ai.usage import RequestUsage |
19 | 18 | from pydantic_ai.exceptions import ModelRetry, UnexpectedModelBehavior |
20 | 19 |
|
@@ -2797,6 +2796,150 @@ async def test_set_usage_data_with_cache_tokens(sentry_init, capture_events): |
2797 | 2796 | assert span_data["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE] == 20 |
2798 | 2797 |
|
2799 | 2798 |
|
| 2799 | +@pytest.mark.parametrize( |
| 2800 | + "url,image_url_kwargs,expected_content", |
| 2801 | + [ |
| 2802 | + pytest.param( |
| 2803 | + "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAYAAACNMs", |
| 2804 | + {}, |
| 2805 | + BLOB_DATA_SUBSTITUTE, |
| 2806 | + id="base64_data_url", |
| 2807 | + ), |
| 2808 | + pytest.param( |
| 2809 | + "https://example.com/image.png", |
| 2810 | + {}, |
| 2811 | + "https://example.com/image.png", |
| 2812 | + id="http_url_no_redaction", |
| 2813 | + ), |
| 2814 | + pytest.param( |
| 2815 | + "https://example.com/api?data=iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAYAAACNMs", |
| 2816 | + {"media_type": "image/png"}, |
| 2817 | + "https://example.com/api?data=iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAYAAACNMs", |
| 2818 | + id="http_url_with_base64_query_param", |
| 2819 | + ), |
| 2820 | + pytest.param( |
| 2821 | + "data:image/svg+xml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciLz4=", |
| 2822 | + {}, |
| 2823 | + BLOB_DATA_SUBSTITUTE, |
| 2824 | + id="complex_mime_type", |
| 2825 | + ), |
| 2826 | + pytest.param( |
| 2827 | + "data:image/png;name=file.png;base64,iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAYAAACNMs", |
| 2828 | + {}, |
| 2829 | + BLOB_DATA_SUBSTITUTE, |
| 2830 | + id="optional_parameters", |
| 2831 | + ), |
| 2832 | + pytest.param( |
| 2833 | + "data:text/plain;charset=utf-8;name=hello.txt;base64,SGVsbG8sIFdvcmxkIQ==", |
| 2834 | + {}, |
| 2835 | + BLOB_DATA_SUBSTITUTE, |
| 2836 | + id="multiple_optional_parameters", |
| 2837 | + ), |
| 2838 | + ], |
| 2839 | +) |
| 2840 | +def test_image_url_base64_content_in_span( |
| 2841 | + sentry_init, capture_events, url, image_url_kwargs, expected_content |
| 2842 | +): |
| 2843 | + from sentry_sdk.integrations.pydantic_ai.spans.ai_client import ai_client_span |
| 2844 | + |
| 2845 | + sentry_init( |
| 2846 | + integrations=[PydanticAIIntegration()], |
| 2847 | + traces_sample_rate=1.0, |
| 2848 | + send_default_pii=True, |
| 2849 | + ) |
| 2850 | + |
| 2851 | + events = capture_events() |
| 2852 | + |
| 2853 | + with sentry_sdk.start_transaction(op="test", name="test"): |
| 2854 | + image_url = ImageUrl(url=url, **image_url_kwargs) |
| 2855 | + user_part = UserPromptPart(content=["Look at this image:", image_url]) |
| 2856 | + mock_msg = MagicMock() |
| 2857 | + mock_msg.parts = [user_part] |
| 2858 | + mock_msg.instructions = None |
| 2859 | + |
| 2860 | + span = ai_client_span([mock_msg], None, None, None) |
| 2861 | + span.finish() |
| 2862 | + |
| 2863 | + (event,) = events |
| 2864 | + chat_spans = [s for s in event["spans"] if s["op"] == "gen_ai.chat"] |
| 2865 | + assert len(chat_spans) >= 1 |
| 2866 | + messages_data = _get_messages_from_span(chat_spans[0]["data"]) |
| 2867 | + |
| 2868 | + found_image = False |
| 2869 | + for msg in messages_data: |
| 2870 | + if "content" not in msg: |
| 2871 | + continue |
| 2872 | + for content_item in msg["content"]: |
| 2873 | + if content_item.get("type") == "image": |
| 2874 | + found_image = True |
| 2875 | + assert content_item["content"] == expected_content |
| 2876 | + |
| 2877 | + assert found_image, "Image content item should be found in messages data" |
| 2878 | + |
| 2879 | + |
| 2880 | +@pytest.mark.asyncio |
| 2881 | +@pytest.mark.parametrize( |
| 2882 | + "url, image_url_kwargs, expected_content", |
| 2883 | + [ |
| 2884 | + pytest.param( |
| 2885 | + "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAYAAACNMs", |
| 2886 | + {}, |
| 2887 | + BLOB_DATA_SUBSTITUTE, |
| 2888 | + id="base64_data_url_redacted", |
| 2889 | + ), |
| 2890 | + pytest.param( |
| 2891 | + "https://example.com/image.png", |
| 2892 | + {}, |
| 2893 | + "https://example.com/image.png", |
| 2894 | + id="http_url_no_redaction", |
| 2895 | + ), |
| 2896 | + pytest.param( |
| 2897 | + "https://example.com/api?data=iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAYAAACNMs", |
| 2898 | + {}, |
| 2899 | + "https://example.com/api?data=iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAYAAACNMs", |
| 2900 | + id="http_url_with_base64_query_param", |
| 2901 | + ), |
| 2902 | + pytest.param( |
| 2903 | + "https://example.com/api?data=iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAYAAACNMs", |
| 2904 | + {"media_type": "image/png"}, |
| 2905 | + "https://example.com/api?data=iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAYAAACNMs", |
| 2906 | + id="http_url_with_base64_query_param_and_media_type", |
| 2907 | + ), |
| 2908 | + ], |
| 2909 | +) |
| 2910 | +async def test_invoke_agent_image_url( |
| 2911 | + sentry_init, capture_events, url, image_url_kwargs, expected_content |
| 2912 | +): |
| 2913 | + sentry_init( |
| 2914 | + integrations=[PydanticAIIntegration()], |
| 2915 | + traces_sample_rate=1.0, |
| 2916 | + send_default_pii=True, |
| 2917 | + ) |
| 2918 | + |
| 2919 | + agent = Agent("test", name="test_image_url_agent") |
| 2920 | + |
| 2921 | + events = capture_events() |
| 2922 | + image_url = ImageUrl(url=url, **image_url_kwargs) |
| 2923 | + await agent.run([image_url, "Describe this image"]) |
| 2924 | + |
| 2925 | + (transaction,) = events |
| 2926 | + |
| 2927 | + found_image = False |
| 2928 | + |
| 2929 | + chat_spans = [s for s in transaction["spans"] if s["op"] == "gen_ai.chat"] |
| 2930 | + for chat_span in chat_spans: |
| 2931 | + messages_data = _get_messages_from_span(chat_span["data"]) |
| 2932 | + for msg in messages_data: |
| 2933 | + if "content" not in msg: |
| 2934 | + continue |
| 2935 | + for content_item in msg["content"]: |
| 2936 | + if content_item.get("type") == "image": |
| 2937 | + assert content_item["content"] == expected_content |
| 2938 | + found_image = True |
| 2939 | + |
| 2940 | + assert found_image, "Image content item should be found in messages data" |
| 2941 | + |
| 2942 | + |
2800 | 2943 | @pytest.mark.asyncio |
2801 | 2944 | async def test_tool_description_in_execute_tool_span(sentry_init, capture_events): |
2802 | 2945 | """ |
|
0 commit comments