From 92a6ff390f08f249b34b52a9a0098d245396119a Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Tue, 29 Sep 2026 09:29:46 +0200 Subject: [PATCH 1/3] ref(google-genai): Drop send_default_pii support --- .../integrations/google_genai/streaming.py | 20 +- sentry_sdk/integrations/google_genai/utils.py | 46 +- .../google_genai/test_google_genai.py | 1417 ++++------------- 3 files changed, 355 insertions(+), 1128 deletions(-) diff --git a/sentry_sdk/integrations/google_genai/streaming.py b/sentry_sdk/integrations/google_genai/streaming.py index cde266780c..55bc70ac84 100644 --- a/sentry_sdk/integrations/google_genai/streaming.py +++ b/sentry_sdk/integrations/google_genai/streaming.py @@ -3,10 +3,8 @@ import sentry_sdk from sentry_sdk.ai.utils import set_data_normalized from sentry_sdk.consts import SPANDATA -from sentry_sdk.scope import should_send_default_pii from sentry_sdk.traces import Span from sentry_sdk.utils import ( - has_data_collection_enabled, safe_serialize, ) @@ -154,28 +152,14 @@ def set_span_data_for_streaming_response( ) if accumulated_response.get("tool_calls"): - if has_data_collection_enabled(client.options): - if client.options["data_collection"]["gen_ai"]["outputs"]: - span.set_attribute( - SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, - safe_serialize(accumulated_response["tool_calls"]), - ) - else: - # Before data collection was introduced this was unconditionally set + if client.options["data_collection"]["gen_ai"]["outputs"]: span.set_attribute( SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, safe_serialize(accumulated_response["tool_calls"]), ) if accumulated_response.get("text"): - if has_data_collection_enabled(client.options): - if client.options["data_collection"]["gen_ai"]["outputs"]: - span.set_attribute( - SPANDATA.GEN_AI_RESPONSE_TEXT, - safe_serialize([accumulated_response["text"]]), - ) - - elif should_send_default_pii(): + if client.options["data_collection"]["gen_ai"]["outputs"]: span.set_attribute( SPANDATA.GEN_AI_RESPONSE_TEXT, safe_serialize([accumulated_response["text"]]), diff --git a/sentry_sdk/integrations/google_genai/utils.py b/sentry_sdk/integrations/google_genai/utils.py index 03081d908b..01d7f80dd0 100644 --- a/sentry_sdk/integrations/google_genai/utils.py +++ b/sentry_sdk/integrations/google_genai/utils.py @@ -27,11 +27,9 @@ transform_google_content_part, ) from sentry_sdk.consts import OP, SPANDATA -from sentry_sdk.scope import should_send_default_pii from sentry_sdk.utils import ( capture_internal_exceptions, event_from_exception, - has_data_collection_enabled, safe_serialize, ) @@ -919,16 +917,7 @@ def set_span_data_for_request( if tools: formatted_tools = _format_tools_for_span(tools) if formatted_tools: - if has_data_collection_enabled(client.options): - if client.options["data_collection"]["gen_ai"]["inputs"]: - set_data_normalized( - span, - SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, - formatted_tools, - unpack=False, - ) - else: - # To remove once data collection has been fully rolled out + if client.options["data_collection"]["gen_ai"]["inputs"]: set_data_normalized( span, SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, @@ -937,10 +926,7 @@ def set_span_data_for_request( ) record_inputs = False - if has_data_collection_enabled(client.options): - if client.options["data_collection"]["gen_ai"]["inputs"]: - record_inputs = True - elif should_send_default_pii(): + if client.options["data_collection"]["gen_ai"]["inputs"]: record_inputs = True if record_inputs: @@ -1043,29 +1029,14 @@ def set_span_data_for_response( tool_calls = extract_tool_calls(response) if tool_calls: - if has_data_collection_enabled(client.options): - if client.options["data_collection"]["gen_ai"]["outputs"]: - span.set_attribute( - SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, safe_serialize(tool_calls) - ) - else: - # Before data collection was introduced, this was set unconditionally + if client.options["data_collection"]["gen_ai"]["outputs"]: span.set_attribute( SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, safe_serialize(tool_calls) ) - if has_data_collection_enabled(client.options): - if client.options["data_collection"]["gen_ai"]["outputs"]: - response_texts = _extract_response_text(response) - if response_texts: - span.set_attribute( - SPANDATA.GEN_AI_RESPONSE_TEXT, safe_serialize(response_texts) - ) - elif should_send_default_pii(): - # TODO: Delete this block once data collection has been completely rolled out + if client.options["data_collection"]["gen_ai"]["outputs"]: response_texts = _extract_response_text(response) if response_texts: - # Format as JSON string array as per documentation span.set_attribute( SPANDATA.GEN_AI_RESPONSE_TEXT, safe_serialize(response_texts) ) @@ -1111,14 +1082,7 @@ def set_span_data_for_embed_request( """Set span data for embedding request.""" client = sentry_sdk.get_client() - record_inputs = False - if has_data_collection_enabled(client.options): - if client.options["data_collection"]["gen_ai"]["inputs"]: - record_inputs = True - elif should_send_default_pii(): - record_inputs = True - - if record_inputs: + if client.options["data_collection"]["gen_ai"]["inputs"]: if contents: # For embeddings, contents is typically a list of strings/texts input_texts = [] diff --git a/tests/integrations/google_genai/test_google_genai.py b/tests/integrations/google_genai/test_google_genai.py index f2d7d82479..61d9ed1e7c 100644 --- a/tests/integrations/google_genai/test_google_genai.py +++ b/tests/integrations/google_genai/test_google_genai.py @@ -114,20 +114,20 @@ def create_test_config( return genai_types.GenerateContentConfig(**config_dict) -@pytest.mark.parametrize( - "send_default_pii", - [True, False], -) def test_nonstreaming_generate_content( sentry_init, capture_items, - send_default_pii, mock_genai_client, ): sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - send_default_pii=send_default_pii, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) # Mock the HTTP response at the _api_client.request() level @@ -161,33 +161,86 @@ def test_nonstreaming_generate_content( assert chat_span["attributes"][SPANDATA.GEN_AI_PROVIDER_NAME] == "gcp.gemini" assert chat_span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] == "gemini-1.5-flash" - if send_default_pii: - assert json.loads( - chat_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - ) == [ - { - "role": "user", - "content": [ - { - "type": "text", - "text": "Message demonstrating the absence of truncation.", - }, - { - "type": "text", - "text": "Tell me a joke", - }, - ], + assert json.loads(chat_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) == [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Message demonstrating the absence of truncation.", + }, + { + "type": "text", + "text": "Tell me a joke", + }, + ], + } + ] + + # Response text is stored as a JSON array + response_text = chat_span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] + + # Parse the JSON array + response_texts = json.loads(response_text) + assert response_texts == ["Hello! How can I help you today?"] + + # Check token usage + assert chat_span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 + # Output tokens now include reasoning tokens: candidates_token_count (20) + thoughts_token_count (3) = 23 + assert chat_span["attributes"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 23 + assert chat_span["attributes"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 30 + assert chat_span["attributes"][SPANDATA.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS] == 5 + assert chat_span["attributes"][SPANDATA.GEN_AI_USAGE_REASONING_OUTPUT_TOKENS] == 3 + + +def test_nonstreaming_generate_content_no_sensitive_data( + sentry_init, + capture_items, + mock_genai_client, +): + sentry_init( + integrations=[GoogleGenAIIntegration()], + traces_sample_rate=1.0, + data_collection={ + "gen_ai": { + "inputs": False, + "outputs": False, } - ] + }, + ) - # Response text is stored as a JSON array - response_text = chat_span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] + # Mock the HTTP response at the _api_client.request() level + mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) + items = capture_items("span") - # Parse the JSON array - response_texts = json.loads(response_text) - assert response_texts == ["Hello! How can I help you today?"] - else: - assert SPANDATA.GEN_AI_RESPONSE_TEXT not in chat_span["attributes"] + with mock.patch.object( + mock_genai_client._api_client, + "request", + return_value=mock_http_response, + ): + config = create_test_config(temperature=0.7, max_output_tokens=100) + mock_genai_client.models.generate_content( + model="gemini-1.5-flash", + contents=[ + "Message demonstrating the absence of truncation.", + "Tell me a joke", + ], + config=config, + ) + + sentry_sdk.flush() + spans = [item.payload for item in items] + assert len(spans) == 1 + chat_span = next(item.payload for item in items if item.type == "span") + + # Check chat span + assert chat_span["attributes"]["sentry.op"] == OP.GEN_AI_CHAT + assert chat_span["name"] == "chat gemini-1.5-flash" + assert chat_span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" + assert chat_span["attributes"][SPANDATA.GEN_AI_PROVIDER_NAME] == "gcp.gemini" + assert chat_span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] == "gemini-1.5-flash" + + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in chat_span["attributes"] # Check token usage assert chat_span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 @@ -237,7 +290,12 @@ def test_generate_content_with_system_instruction( sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - send_default_pii=True, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) @@ -284,6 +342,11 @@ def test_generate_content_with_tools( sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, + data_collection={ + "gen_ai": { + "inputs": True, + } + }, ) # Create a mock tool function @@ -368,7 +431,12 @@ def test_tool_execution( sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - send_default_pii=True, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) # Create a mock tool function @@ -410,6 +478,7 @@ def test_error_handling( sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, + data_collection={}, ) items = capture_items("event") @@ -440,7 +509,12 @@ def test_streaming_generate_content( sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - send_default_pii=True, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) # Create streaming chunks - simulating a multi-chunk response @@ -584,6 +658,7 @@ def test_span_origin( sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, + data_collection={}, ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) @@ -612,6 +687,7 @@ def test_response_without_usage_metadata( sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, + data_collection={}, ) # Response without usage metadata @@ -656,7 +732,12 @@ def test_multiple_candidates( sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - send_default_pii=True, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) # Response with multiple candidates @@ -726,6 +807,7 @@ def test_all_configuration_parameters( sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, + data_collection={}, ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) @@ -769,6 +851,7 @@ def test_empty_response( sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, + data_collection={}, ) # Minimal response with empty candidates array @@ -802,6 +885,7 @@ def test_response_with_different_id_fields( sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, + data_collection={}, ) # Response with response_id and model_version @@ -844,6 +928,7 @@ def test_tool_with_async_function(sentry_init): sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, + data_collection={}, ) # Create an async tool function @@ -869,7 +954,12 @@ def test_contents_as_none( sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - send_default_pii=True, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) @@ -901,6 +991,11 @@ def test_tool_calls_extraction( sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, + data_collection={ + "gen_ai": { + "outputs": True, + } + }, ) # Response with function calls @@ -1003,20 +1098,20 @@ def test_tool_calls_extraction( } -@pytest.mark.parametrize( - "send_default_pii", - [True, False], -) def test_embed_content( sentry_init, capture_items, - send_default_pii, mock_genai_client, ): sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - send_default_pii=send_default_pii, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) # Mock the HTTP response at the _api_client.request() level @@ -1051,17 +1146,67 @@ def test_embed_content( embed_span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] == "text-embedding-004" ) - # Check input texts if PII is allowed - if send_default_pii: - input_texts = json.loads( - embed_span["attributes"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT] + input_texts = json.loads(embed_span["attributes"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT]) + assert input_texts == [ + "What is your name?", + "What is your favorite color?", + ] + + # Check usage data (sum of token counts from statistics: 10 + 15 = 25) + # Note: Only available in newer versions with ContentEmbeddingStatistics + if SPANDATA.GEN_AI_USAGE_INPUT_TOKENS in embed_span["attributes"]: + assert embed_span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 25 + + +def test_embed_content_no_sensitive_data( + sentry_init, + capture_items, + mock_genai_client, +): + sentry_init( + integrations=[GoogleGenAIIntegration()], + traces_sample_rate=1.0, + data_collection={ + "gen_ai": { + "inputs": False, + "outputs": False, + } + }, + ) + + # Mock the HTTP response at the _api_client.request() level + mock_http_response = create_mock_http_response(EXAMPLE_EMBED_RESPONSE_JSON) + items = capture_items("span") + + with mock.patch.object( + mock_genai_client._api_client, + "request", + return_value=mock_http_response, + ): + mock_genai_client.models.embed_content( + model="text-embedding-004", + contents=[ + "What is your name?", + "What is your favorite color?", + ], ) - assert input_texts == [ - "What is your name?", - "What is your favorite color?", - ] - else: - assert SPANDATA.GEN_AI_EMBEDDINGS_INPUT not in embed_span["attributes"] + + # Should have 1 span for embeddings + sentry_sdk.flush() + spans = [item.payload for item in items] + assert len(spans) == 1 + (embed_span,) = spans + + # Check embeddings span + assert embed_span["attributes"]["sentry.op"] == OP.GEN_AI_EMBEDDINGS + assert embed_span["name"] == "embeddings text-embedding-004" + assert embed_span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "embeddings" + assert embed_span["attributes"][SPANDATA.GEN_AI_PROVIDER_NAME] == "gcp.gemini" + assert ( + embed_span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] == "text-embedding-004" + ) + + assert SPANDATA.GEN_AI_EMBEDDINGS_INPUT not in embed_span["attributes"] # Check usage data (sum of token counts from statistics: 10 + 15 = 25) # Note: Only available in newer versions with ContentEmbeddingStatistics @@ -1078,7 +1223,12 @@ def test_embed_content_string_input( sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - send_default_pii=True, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) # Mock response with single embedding @@ -1130,6 +1280,7 @@ def test_embed_content_error_handling( sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, + data_collection={}, ) items = capture_items("event") @@ -1161,6 +1312,7 @@ def test_embed_content_without_statistics( sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, + data_collection={}, ) # Response without statistics (typical for older google-genai versions) @@ -1204,6 +1356,7 @@ def test_embed_content_span_origin( sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, + data_collection={}, ) mock_http_response = create_mock_http_response(EXAMPLE_EMBED_RESPONSE_JSON) @@ -1223,21 +1376,21 @@ def test_embed_content_span_origin( @pytest.mark.asyncio -@pytest.mark.parametrize( - "send_default_pii", - [True, False], -) async def test_async_embed_content( sentry_init, capture_items, - send_default_pii, mock_genai_client, ): """Test async embed_content method.""" sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - send_default_pii=send_default_pii, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) # Mock the async HTTP response @@ -1272,17 +1425,69 @@ async def test_async_embed_content( embed_span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] == "text-embedding-004" ) - # Check input texts if PII is allowed - if send_default_pii: - input_texts = json.loads( - embed_span["attributes"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT] + input_texts = json.loads(embed_span["attributes"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT]) + assert input_texts == [ + "What is your name?", + "What is your favorite color?", + ] + + # Check usage data (sum of token counts from statistics: 10 + 15 = 25) + # Note: Only available in newer versions with ContentEmbeddingStatistics + if SPANDATA.GEN_AI_USAGE_INPUT_TOKENS in embed_span["attributes"]: + assert embed_span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 25 + + +@pytest.mark.asyncio +async def test_async_embed_content_no_sensitive_data( + sentry_init, + capture_items, + mock_genai_client, +): + """Test async embed_content method.""" + sentry_init( + integrations=[GoogleGenAIIntegration()], + traces_sample_rate=1.0, + data_collection={ + "gen_ai": { + "inputs": False, + "outputs": False, + } + }, + ) + + # Mock the async HTTP response + mock_http_response = create_mock_http_response(EXAMPLE_EMBED_RESPONSE_JSON) + items = capture_items("span") + + with mock.patch.object( + mock_genai_client._api_client, + "async_request", + return_value=mock_http_response, + ): + await mock_genai_client.aio.models.embed_content( + model="text-embedding-004", + contents=[ + "What is your name?", + "What is your favorite color?", + ], ) - assert input_texts == [ - "What is your name?", - "What is your favorite color?", - ] - else: - assert SPANDATA.GEN_AI_EMBEDDINGS_INPUT not in embed_span["attributes"] + + # Should have 1 span for embeddings + sentry_sdk.flush() + spans = [item.payload for item in items] + assert len(spans) == 1 + (embed_span,) = spans + + # Check embeddings span + assert embed_span["attributes"]["sentry.op"] == OP.GEN_AI_EMBEDDINGS + assert embed_span["name"] == "embeddings text-embedding-004" + assert embed_span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "embeddings" + assert embed_span["attributes"][SPANDATA.GEN_AI_PROVIDER_NAME] == "gcp.gemini" + assert ( + embed_span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] == "text-embedding-004" + ) + + assert SPANDATA.GEN_AI_EMBEDDINGS_INPUT not in embed_span["attributes"] # Check usage data (sum of token counts from statistics: 10 + 15 = 25) # Note: Only available in newer versions with ContentEmbeddingStatistics @@ -1300,7 +1505,12 @@ async def test_async_embed_content_string_input( sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - send_default_pii=True, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) # Mock response with single embedding @@ -1354,6 +1564,7 @@ async def test_async_embed_content_error_handling( sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, + data_collection={}, ) items = capture_items("event") @@ -1386,6 +1597,7 @@ async def test_async_embed_content_without_statistics( sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, + data_collection={}, ) # Response without statistics (typical for older google-genai versions) @@ -1432,6 +1644,7 @@ async def test_async_embed_content_span_origin( sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, + data_collection={}, ) mock_http_response = create_mock_http_response(EXAMPLE_EMBED_RESPONSE_JSON) @@ -1463,7 +1676,12 @@ def test_generate_content_with_content_object( sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - send_default_pii=True, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) @@ -1502,7 +1720,12 @@ def test_generate_content_with_dict_format( sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - send_default_pii=True, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) @@ -1539,7 +1762,12 @@ def test_generate_content_with_file_data( sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - send_default_pii=True, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) @@ -1591,7 +1819,12 @@ def test_generate_content_with_inline_data( sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - send_default_pii=True, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) @@ -1637,7 +1870,12 @@ def test_generate_content_with_function_response( sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - send_default_pii=True, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) @@ -1695,7 +1933,12 @@ def test_generate_content_with_mixed_string_and_content( sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - send_default_pii=True, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) @@ -1742,7 +1985,12 @@ def test_generate_content_with_part_object_directly( sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - send_default_pii=True, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) @@ -1781,7 +2029,12 @@ def test_generate_content_with_list_of_dicts( sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - send_default_pii=True, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) @@ -1821,7 +2074,12 @@ def test_generate_content_with_dict_inline_data( sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - send_default_pii=True, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) @@ -1868,7 +2126,12 @@ def test_generate_content_without_parts_property_inline_data( sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - send_default_pii=True, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) @@ -1913,7 +2176,12 @@ def test_generate_content_without_parts_property_inline_data_and_binary_data_wit sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - send_default_pii=True, + data_collection={ + "gen_ai": { + "inputs": True, + "outputs": True, + } + }, ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) @@ -2274,992 +2542,3 @@ def __init__(self): assert len(result) == 1 assert result[0]["role"] == "user" assert result[0]["content"] == [{"text": "Object text", "type": "text"}] - - -DATA_COLLECTION_CHAT_EXPECTED_VALUES = { - SPANDATA.GEN_AI_REQUEST_MESSAGES: [{"role": "user", "content": "Tell me a joke"}], - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS: [ - {"type": "text", "content": "You are a helpful assistant."} - ], - SPANDATA.GEN_AI_RESPONSE_TEXT: ["Hello! How can I help you today?"], -} - -DATA_COLLECTION_EMBED_EXPECTED_VALUES = { - SPANDATA.GEN_AI_EMBEDDINGS_INPUT: [ - "What is your name?", - "What is your favorite color?", - ], -} - - -@pytest.mark.parametrize( - "data_collection,send_default_pii,expected_present,expected_absent", - [ - pytest.param( - {"gen_ai": {"inputs": True, "outputs": True}}, - False, - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, - SPANDATA.GEN_AI_RESPONSE_TEXT, - ], - [], - id="gen-ai-inputs-and-outputs-enabled-override-legacy-off", - ), - pytest.param( - {"gen_ai": {"inputs": False, "outputs": False}}, - True, - [], - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, - SPANDATA.GEN_AI_RESPONSE_TEXT, - ], - id="gen-ai-inputs-and-outputs-disabled-override-legacy-on", - ), - pytest.param( - {"gen_ai": {"inputs": True, "outputs": False}}, - False, - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, - ], - [ - SPANDATA.GEN_AI_RESPONSE_TEXT, - ], - id="gen-ai-inputs-enabled-outputs-disabled", - ), - pytest.param( - {"gen_ai": {"inputs": False, "outputs": True}}, - False, - [ - SPANDATA.GEN_AI_RESPONSE_TEXT, - ], - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, - ], - id="gen-ai-outputs-enabled-inputs-disabled", - ), - pytest.param( - {"gen_ai": {}}, - False, - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, - SPANDATA.GEN_AI_RESPONSE_TEXT, - ], - [], - id="gen-ai-inputs-and-outputs-omitted-default-to-enabled", - ), - pytest.param( - None, - True, - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, - SPANDATA.GEN_AI_RESPONSE_TEXT, - ], - [], - id="no-gen-ai-config-legacy-pii", - ), - pytest.param( - None, - False, - [], - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, - SPANDATA.GEN_AI_RESPONSE_TEXT, - ], - id="no-gen-ai-config-legacy-pii-disabled", - ), - ], -) -def test_generate_content_data_collection( - sentry_init, - capture_items, - mock_genai_client, - data_collection, - send_default_pii, - expected_present, - expected_absent, -): - sentry_init_kwargs = dict( - integrations=[GoogleGenAIIntegration()], - traces_sample_rate=1.0, - send_default_pii=send_default_pii, - ) - if data_collection is not None: - sentry_init_kwargs["data_collection"] = data_collection - - sentry_init(**sentry_init_kwargs) - - mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) - captured = capture_items("span") - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ): - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", - contents="Tell me a joke", - config=create_test_config( - temperature=0.7, - max_output_tokens=100, - system_instruction="You are a helpful assistant.", - ), - ) - - sentry_sdk.flush() - spans = [item.payload for item in captured if item.type == "span"] - (span,) = [s for s in spans if s["attributes"].get("sentry.op") == OP.GEN_AI_CHAT] - span_data = span["attributes"] - - for key in expected_present: - assert key in span_data, f"{key} should have been collected" - assert json.loads(span_data[key]) == DATA_COLLECTION_CHAT_EXPECTED_VALUES[key] - - for key in expected_absent: - assert key not in span_data, f"{key} should not have been collected" - - # Data collection never gates non-PII attributes - assert span_data[SPANDATA.GEN_AI_PROVIDER_NAME] == "gcp.gemini" - assert span_data[SPANDATA.GEN_AI_REQUEST_MODEL] == "gemini-1.5-flash" - assert span_data[SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7 - assert span_data[SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100 - assert span_data[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 - assert span_data[SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 23 - assert span_data[SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 30 - assert span_data[SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS] == "STOP" - - -@pytest.mark.parametrize("send_default_pii", [True, False]) -@pytest.mark.parametrize( - "data_collection,expected_present,expected_absent", - [ - pytest.param( - {"gen_ai": {"inputs": True, "outputs": True}}, - [ - SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, - SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, - ], - [], - id="gen-ai-inputs-and-outputs-enabled-tools-collected", - ), - pytest.param( - {"gen_ai": {"inputs": False, "outputs": False}}, - [], - [ - SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, - SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, - ], - id="gen-ai-inputs-and-outputs-disabled-tools-not-collected", - ), - pytest.param( - {"gen_ai": {"inputs": True, "outputs": False}}, - [SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS], - [SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS], - id="gen-ai-inputs-enabled-outputs-disabled-drops-tool-calls-only", - ), - pytest.param( - {"gen_ai": {"inputs": False, "outputs": True}}, - [SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS], - [SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS], - id="gen-ai-inputs-disabled-outputs-enabled-drops-available-tools-only", - ), - pytest.param( - {"gen_ai": {}}, - [ - SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, - SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, - ], - [], - id="gen-ai-inputs-and-outputs-omitted-tools-collected", - ), - pytest.param( - None, - [ - SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, - SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, - ], - [], - id="no-gen-ai-config-tools-collected-regardless-of-pii", - ), - ], -) -def test_generate_content_data_collection_tools( - sentry_init, - capture_items, - mock_genai_client, - data_collection, - send_default_pii, - expected_present, - expected_absent, -): - sentry_init_kwargs = dict( - integrations=[GoogleGenAIIntegration()], - traces_sample_rate=1.0, - send_default_pii=send_default_pii, - ) - if data_collection is not None: - sentry_init_kwargs["data_collection"] = data_collection - - sentry_init(**sentry_init_kwargs) - - # A declaration rather than a Python callable, so that the client does not - # automatically call the function and issue a second request - weather_tool = genai_types.Tool( - function_declarations=[ - genai_types.FunctionDeclaration( - name="get_weather", - description="Get the weather for a location", - ) - ] - ) - - mock_http_response = create_mock_http_response( - { - "candidates": [ - { - "content": { - "role": "model", - "parts": [ - {"text": "I'll help you with that."}, - { - "functionCall": { - "name": "get_weather", - "args": {"location": "San Francisco"}, - } - }, - ], - }, - "finishReason": "STOP", - } - ], - "usageMetadata": { - "promptTokenCount": 20, - "candidatesTokenCount": 30, - "totalTokenCount": 50, - }, - } - ) - - captured = capture_items("span") - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ): - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", - contents="What's the weather?", - config=create_test_config(tools=[weather_tool]), - ) - - sentry_sdk.flush() - spans = [item.payload for item in captured if item.type == "span"] - (span,) = [s for s in spans if s["attributes"].get("sentry.op") == OP.GEN_AI_CHAT] - span_data = span["attributes"] - - for key in expected_present: - assert key in span_data, f"{key} should have been collected" - - for key in expected_absent: - assert key not in span_data, f"{key} should not have been collected" - - if SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS in expected_present: - available_tools = json.loads(span_data[SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS]) - assert [tool["name"] for tool in available_tools] == ["get_weather"] - - if SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS in expected_present: - tool_calls = json.loads(span_data[SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS]) - assert [tool_call["name"] for tool_call in tool_calls] == ["get_weather"] - assert json.loads(tool_calls[0]["arguments"]) == {"location": "San Francisco"} - - -@pytest.mark.parametrize( - "data_collection,send_default_pii,expected_present,expected_absent", - [ - pytest.param( - {"gen_ai": {"inputs": True, "outputs": True}}, - False, - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, - SPANDATA.GEN_AI_RESPONSE_TEXT, - ], - [], - id="gen-ai-inputs-and-outputs-enabled-override-legacy-off", - ), - pytest.param( - {"gen_ai": {"inputs": False, "outputs": False}}, - True, - [], - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, - SPANDATA.GEN_AI_RESPONSE_TEXT, - ], - id="gen-ai-inputs-and-outputs-disabled-override-legacy-on", - ), - pytest.param( - {"gen_ai": {"inputs": True, "outputs": False}}, - False, - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, - ], - [ - SPANDATA.GEN_AI_RESPONSE_TEXT, - ], - id="gen-ai-inputs-enabled-outputs-disabled", - ), - pytest.param( - {"gen_ai": {"inputs": False, "outputs": True}}, - False, - [ - SPANDATA.GEN_AI_RESPONSE_TEXT, - ], - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, - ], - id="gen-ai-outputs-enabled-inputs-disabled", - ), - pytest.param( - {"gen_ai": {}}, - False, - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, - SPANDATA.GEN_AI_RESPONSE_TEXT, - ], - [], - id="gen-ai-inputs-and-outputs-omitted-default-to-enabled", - ), - pytest.param( - None, - True, - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, - SPANDATA.GEN_AI_RESPONSE_TEXT, - ], - [], - id="no-gen-ai-config-legacy-pii", - ), - pytest.param( - None, - False, - [], - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, - SPANDATA.GEN_AI_RESPONSE_TEXT, - ], - id="no-gen-ai-config-legacy-pii-disabled", - ), - ], -) -def test_streaming_generate_content_data_collection( - sentry_init, - capture_items, - mock_genai_client, - data_collection, - send_default_pii, - expected_present, - expected_absent, -): - sentry_init_kwargs = dict( - integrations=[GoogleGenAIIntegration()], - traces_sample_rate=1.0, - send_default_pii=send_default_pii, - ) - if data_collection is not None: - sentry_init_kwargs["data_collection"] = data_collection - - sentry_init(**sentry_init_kwargs) - - mock_stream = create_mock_streaming_responses( - [ - { - "candidates": [ - {"content": {"role": "model", "parts": [{"text": "Hello! "}]}} - ], - "usageMetadata": { - "promptTokenCount": 10, - "candidatesTokenCount": 2, - "totalTokenCount": 12, - }, - "modelVersion": "gemini-1.5-flash", - }, - { - "candidates": [ - {"content": {"role": "model", "parts": [{"text": "How can I "}]}} - ], - }, - { - "candidates": [ - { - "content": { - "role": "model", - "parts": [{"text": "help you today?"}], - }, - "finishReason": "STOP", - } - ], - "usageMetadata": { - "promptTokenCount": 10, - "candidatesTokenCount": 7, - "totalTokenCount": 25, - }, - }, - ] - ) - - captured = capture_items("span") - - with mock.patch.object( - mock_genai_client._api_client, "request_streamed", return_value=mock_stream - ): - stream = mock_genai_client.models.generate_content_stream( - model="gemini-1.5-flash", - contents="Tell me a joke", - config=create_test_config( - system_instruction="You are a helpful assistant." - ), - ) - list(stream) - - sentry_sdk.flush() - spans = [item.payload for item in captured if item.type == "span"] - (span,) = [s for s in spans if s["attributes"].get("sentry.op") == OP.GEN_AI_CHAT] - span_data = span["attributes"] - - for key in expected_present: - assert key in span_data, f"{key} should have been collected" - assert json.loads(span_data[key]) == DATA_COLLECTION_CHAT_EXPECTED_VALUES[key] - - for key in expected_absent: - assert key not in span_data, f"{key} should not have been collected" - - assert span_data[SPANDATA.GEN_AI_RESPONSE_STREAMING] is True - assert span_data[SPANDATA.GEN_AI_REQUEST_MODEL] == "gemini-1.5-flash" - assert span_data[SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS] == "STOP" - assert span_data[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 - assert span_data[SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 25 - - -@pytest.mark.parametrize("send_default_pii", [True, False]) -@pytest.mark.parametrize( - "data_collection,expected_present,expected_absent", - [ - pytest.param( - {"gen_ai": {"inputs": True, "outputs": True}}, - [ - SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, - SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, - ], - [], - id="gen-ai-inputs-and-outputs-enabled-tools-collected", - ), - pytest.param( - {"gen_ai": {"inputs": False, "outputs": False}}, - [], - [ - SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, - SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, - ], - id="gen-ai-inputs-and-outputs-disabled-tools-not-collected", - ), - pytest.param( - {"gen_ai": {"inputs": True, "outputs": False}}, - [SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS], - [SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS], - id="gen-ai-inputs-enabled-outputs-disabled-drops-tool-calls-only", - ), - pytest.param( - {"gen_ai": {"inputs": False, "outputs": True}}, - [SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS], - [SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS], - id="gen-ai-inputs-disabled-outputs-enabled-drops-available-tools-only", - ), - pytest.param( - {"gen_ai": {}}, - [ - SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, - SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, - ], - [], - id="gen-ai-inputs-and-outputs-omitted-tools-collected", - ), - pytest.param( - None, - [ - SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, - SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, - ], - [], - id="no-gen-ai-config-tools-collected-regardless-of-pii", - ), - ], -) -def test_streaming_generate_content_data_collection_tools( - sentry_init, - capture_items, - mock_genai_client, - data_collection, - send_default_pii, - expected_present, - expected_absent, -): - sentry_init_kwargs = dict( - integrations=[GoogleGenAIIntegration()], - traces_sample_rate=1.0, - send_default_pii=send_default_pii, - ) - if data_collection is not None: - sentry_init_kwargs["data_collection"] = data_collection - - sentry_init(**sentry_init_kwargs) - - weather_tool = genai_types.Tool( - function_declarations=[ - genai_types.FunctionDeclaration( - name="get_weather", - description="Get the weather for a location", - ) - ] - ) - - mock_stream = create_mock_streaming_responses( - [ - { - "candidates": [ - { - "content": { - "role": "model", - "parts": [{"text": "I'll help you with that."}], - } - } - ], - "modelVersion": "gemini-1.5-flash", - }, - { - "candidates": [ - { - "content": { - "role": "model", - "parts": [ - { - "functionCall": { - "name": "get_weather", - "args": {"location": "San Francisco"}, - } - } - ], - }, - "finishReason": "STOP", - } - ], - }, - ] - ) - - captured = capture_items("span") - - with mock.patch.object( - mock_genai_client._api_client, "request_streamed", return_value=mock_stream - ): - stream = mock_genai_client.models.generate_content_stream( - model="gemini-1.5-flash", - contents="What's the weather?", - config=create_test_config(tools=[weather_tool]), - ) - list(stream) - - sentry_sdk.flush() - spans = [item.payload for item in captured if item.type == "span"] - (span,) = [s for s in spans if s["attributes"].get("sentry.op") == OP.GEN_AI_CHAT] - span_data = span["attributes"] - - for key in expected_present: - assert key in span_data, f"{key} should have been collected" - - for key in expected_absent: - assert key not in span_data, f"{key} should not have been collected" - - if SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS in expected_present: - available_tools = json.loads(span_data[SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS]) - assert [tool["name"] for tool in available_tools] == ["get_weather"] - - if SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS in expected_present: - tool_calls = json.loads(span_data[SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS]) - assert [tool_call["name"] for tool_call in tool_calls] == ["get_weather"] - - -@pytest.mark.parametrize( - "data_collection,send_default_pii,expected_present,expected_absent", - [ - pytest.param( - {"gen_ai": {"inputs": True, "outputs": True}}, - False, - [ - SPANDATA.GEN_AI_EMBEDDINGS_INPUT, - ], - [], - id="gen-ai-inputs-and-outputs-enabled-override-legacy-off", - ), - pytest.param( - {"gen_ai": {"inputs": False, "outputs": False}}, - True, - [], - [ - SPANDATA.GEN_AI_EMBEDDINGS_INPUT, - ], - id="gen-ai-inputs-and-outputs-disabled-override-legacy-on", - ), - pytest.param( - {"gen_ai": {"inputs": True, "outputs": False}}, - False, - [ - SPANDATA.GEN_AI_EMBEDDINGS_INPUT, - ], - [], - id="gen-ai-inputs-enabled-outputs-disabled", - ), - pytest.param( - {"gen_ai": {"inputs": False, "outputs": True}}, - False, - [], - [ - SPANDATA.GEN_AI_EMBEDDINGS_INPUT, - ], - id="gen-ai-outputs-enabled-inputs-disabled", - ), - pytest.param( - {"gen_ai": {}}, - False, - [ - SPANDATA.GEN_AI_EMBEDDINGS_INPUT, - ], - [], - id="gen-ai-inputs-and-outputs-omitted-default-to-enabled", - ), - pytest.param( - None, - True, - [ - SPANDATA.GEN_AI_EMBEDDINGS_INPUT, - ], - [], - id="no-gen-ai-config-legacy-pii", - ), - pytest.param( - None, - False, - [], - [ - SPANDATA.GEN_AI_EMBEDDINGS_INPUT, - ], - id="no-gen-ai-config-legacy-pii-disabled", - ), - ], -) -def test_embed_content_data_collection( - sentry_init, - capture_items, - mock_genai_client, - data_collection, - send_default_pii, - expected_present, - expected_absent, -): - sentry_init_kwargs = dict( - integrations=[GoogleGenAIIntegration()], - traces_sample_rate=1.0, - send_default_pii=send_default_pii, - ) - if data_collection is not None: - sentry_init_kwargs["data_collection"] = data_collection - - sentry_init(**sentry_init_kwargs) - - mock_http_response = create_mock_http_response(EXAMPLE_EMBED_RESPONSE_JSON) - captured = capture_items("span") - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ): - mock_genai_client.models.embed_content( - model="text-embedding-004", - contents=["What is your name?", "What is your favorite color?"], - ) - - sentry_sdk.flush() - spans = [item.payload for item in captured if item.type == "span"] - (span,) = [ - s for s in spans if s["attributes"].get("sentry.op") == OP.GEN_AI_EMBEDDINGS - ] - span_data = span["attributes"] - - for key in expected_present: - assert key in span_data, f"{key} should have been collected" - assert json.loads(span_data[key]) == DATA_COLLECTION_EMBED_EXPECTED_VALUES[key] - - for key in expected_absent: - assert key not in span_data, f"{key} should not have been collected" - - assert span_data[SPANDATA.GEN_AI_OPERATION_NAME] == "embeddings" - assert span_data[SPANDATA.GEN_AI_REQUEST_MODEL] == "text-embedding-004" - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "data_collection,send_default_pii,expected_present,expected_absent", - [ - pytest.param( - {"gen_ai": {"inputs": True, "outputs": True}}, - False, - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, - SPANDATA.GEN_AI_RESPONSE_TEXT, - ], - [], - id="gen-ai-inputs-and-outputs-enabled-override-legacy-off", - ), - pytest.param( - {"gen_ai": {"inputs": False, "outputs": False}}, - True, - [], - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, - SPANDATA.GEN_AI_RESPONSE_TEXT, - ], - id="gen-ai-inputs-and-outputs-disabled-override-legacy-on", - ), - pytest.param( - {"gen_ai": {"inputs": True, "outputs": False}}, - False, - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, - ], - [ - SPANDATA.GEN_AI_RESPONSE_TEXT, - ], - id="gen-ai-inputs-enabled-outputs-disabled", - ), - pytest.param( - {"gen_ai": {"inputs": False, "outputs": True}}, - False, - [ - SPANDATA.GEN_AI_RESPONSE_TEXT, - ], - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, - ], - id="gen-ai-outputs-enabled-inputs-disabled", - ), - pytest.param( - {"gen_ai": {}}, - False, - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, - SPANDATA.GEN_AI_RESPONSE_TEXT, - ], - [], - id="gen-ai-inputs-and-outputs-omitted-default-to-enabled", - ), - pytest.param( - None, - True, - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, - SPANDATA.GEN_AI_RESPONSE_TEXT, - ], - [], - id="no-gen-ai-config-legacy-pii", - ), - pytest.param( - None, - False, - [], - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, - SPANDATA.GEN_AI_RESPONSE_TEXT, - ], - id="no-gen-ai-config-legacy-pii-disabled", - ), - ], -) -async def test_async_generate_content_data_collection( - sentry_init, - capture_items, - mock_genai_client, - data_collection, - send_default_pii, - expected_present, - expected_absent, -): - sentry_init_kwargs = dict( - integrations=[GoogleGenAIIntegration()], - traces_sample_rate=1.0, - send_default_pii=send_default_pii, - ) - if data_collection is not None: - sentry_init_kwargs["data_collection"] = data_collection - - sentry_init(**sentry_init_kwargs) - - mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) - captured = capture_items("span") - - with mock.patch.object( - mock_genai_client._api_client, "async_request", return_value=mock_http_response - ): - await mock_genai_client.aio.models.generate_content( - model="gemini-1.5-flash", - contents="Tell me a joke", - config=create_test_config( - system_instruction="You are a helpful assistant." - ), - ) - - sentry_sdk.flush() - spans = [item.payload for item in captured if item.type == "span"] - (span,) = [s for s in spans if s["attributes"].get("sentry.op") == OP.GEN_AI_CHAT] - span_data = span["attributes"] - - for key in expected_present: - assert key in span_data, f"{key} should have been collected" - assert json.loads(span_data[key]) == DATA_COLLECTION_CHAT_EXPECTED_VALUES[key] - - for key in expected_absent: - assert key not in span_data, f"{key} should not have been collected" - - assert span_data[SPANDATA.GEN_AI_REQUEST_MODEL] == "gemini-1.5-flash" - assert span_data[SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 30 - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "data_collection,send_default_pii,expected_present,expected_absent", - [ - pytest.param( - {"gen_ai": {"inputs": True, "outputs": True}}, - False, - [ - SPANDATA.GEN_AI_EMBEDDINGS_INPUT, - ], - [], - id="gen-ai-inputs-and-outputs-enabled-override-legacy-off", - ), - pytest.param( - {"gen_ai": {"inputs": False, "outputs": False}}, - True, - [], - [ - SPANDATA.GEN_AI_EMBEDDINGS_INPUT, - ], - id="gen-ai-inputs-and-outputs-disabled-override-legacy-on", - ), - pytest.param( - {"gen_ai": {"inputs": True, "outputs": False}}, - False, - [ - SPANDATA.GEN_AI_EMBEDDINGS_INPUT, - ], - [], - id="gen-ai-inputs-enabled-outputs-disabled", - ), - pytest.param( - {"gen_ai": {"inputs": False, "outputs": True}}, - False, - [], - [ - SPANDATA.GEN_AI_EMBEDDINGS_INPUT, - ], - id="gen-ai-outputs-enabled-inputs-disabled", - ), - pytest.param( - {"gen_ai": {}}, - False, - [ - SPANDATA.GEN_AI_EMBEDDINGS_INPUT, - ], - [], - id="gen-ai-inputs-and-outputs-omitted-default-to-enabled", - ), - pytest.param( - None, - True, - [ - SPANDATA.GEN_AI_EMBEDDINGS_INPUT, - ], - [], - id="no-gen-ai-config-legacy-pii", - ), - pytest.param( - None, - False, - [], - [ - SPANDATA.GEN_AI_EMBEDDINGS_INPUT, - ], - id="no-gen-ai-config-legacy-pii-disabled", - ), - ], -) -async def test_async_embed_content_data_collection( - sentry_init, - capture_items, - mock_genai_client, - data_collection, - send_default_pii, - expected_present, - expected_absent, -): - sentry_init_kwargs = dict( - integrations=[GoogleGenAIIntegration()], - traces_sample_rate=1.0, - send_default_pii=send_default_pii, - ) - if data_collection is not None: - sentry_init_kwargs["data_collection"] = data_collection - - sentry_init(**sentry_init_kwargs) - - mock_http_response = create_mock_http_response(EXAMPLE_EMBED_RESPONSE_JSON) - captured = capture_items("span") - - with mock.patch.object( - mock_genai_client._api_client, "async_request", return_value=mock_http_response - ): - await mock_genai_client.aio.models.embed_content( - model="text-embedding-004", - contents=["What is your name?", "What is your favorite color?"], - ) - - sentry_sdk.flush() - spans = [item.payload for item in captured if item.type == "span"] - (span,) = [ - s for s in spans if s["attributes"].get("sentry.op") == OP.GEN_AI_EMBEDDINGS - ] - span_data = span["attributes"] - - for key in expected_present: - assert key in span_data, f"{key} should have been collected" - assert json.loads(span_data[key]) == DATA_COLLECTION_EMBED_EXPECTED_VALUES[key] - - for key in expected_absent: - assert key not in span_data, f"{key} should not have been collected" - - assert span_data[SPANDATA.GEN_AI_OPERATION_NAME] == "embeddings" - assert span_data[SPANDATA.GEN_AI_REQUEST_MODEL] == "text-embedding-004" From 6d428db2c394ab8838074b2582997191cd675f21 Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Tue, 29 Sep 2026 10:38:08 +0200 Subject: [PATCH 2/3] restore tests with mixed data collection settings --- .../google_genai/test_google_genai.py | 636 ++++++++++++++++++ 1 file changed, 636 insertions(+) diff --git a/tests/integrations/google_genai/test_google_genai.py b/tests/integrations/google_genai/test_google_genai.py index 61d9ed1e7c..88469ae725 100644 --- a/tests/integrations/google_genai/test_google_genai.py +++ b/tests/integrations/google_genai/test_google_genai.py @@ -2542,3 +2542,639 @@ def __init__(self): assert len(result) == 1 assert result[0]["role"] == "user" assert result[0]["content"] == [{"text": "Object text", "type": "text"}] + + +DATA_COLLECTION_CHAT_EXPECTED_VALUES = { + SPANDATA.GEN_AI_REQUEST_MESSAGES: [{"role": "user", "content": "Tell me a joke"}], + SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS: [ + {"type": "text", "content": "You are a helpful assistant."} + ], + SPANDATA.GEN_AI_RESPONSE_TEXT: ["Hello! How can I help you today?"], +} + +DATA_COLLECTION_EMBED_EXPECTED_VALUES = { + SPANDATA.GEN_AI_EMBEDDINGS_INPUT: [ + "What is your name?", + "What is your favorite color?", + ], +} + + +@pytest.mark.parametrize( + "data_collection,expected_present,expected_absent", + [ + pytest.param( + {"gen_ai": {"inputs": True, "outputs": False}}, + [ + SPANDATA.GEN_AI_REQUEST_MESSAGES, + SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, + ], + [ + SPANDATA.GEN_AI_RESPONSE_TEXT, + ], + id="gen-ai-inputs-enabled-outputs-disabled", + ), + pytest.param( + {"gen_ai": {"inputs": False, "outputs": True}}, + [ + SPANDATA.GEN_AI_RESPONSE_TEXT, + ], + [ + SPANDATA.GEN_AI_REQUEST_MESSAGES, + SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, + ], + id="gen-ai-outputs-enabled-inputs-disabled", + ), + ], +) +def test_generate_content_data_collection( + sentry_init, + capture_items, + mock_genai_client, + data_collection, + expected_present, + expected_absent, +): + sentry_init_kwargs = dict( + integrations=[GoogleGenAIIntegration()], + traces_sample_rate=1.0, + ) + if data_collection is not None: + sentry_init_kwargs["data_collection"] = data_collection + + sentry_init(**sentry_init_kwargs) + + mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) + captured = capture_items("span") + + with mock.patch.object( + mock_genai_client._api_client, "request", return_value=mock_http_response + ): + mock_genai_client.models.generate_content( + model="gemini-1.5-flash", + contents="Tell me a joke", + config=create_test_config( + temperature=0.7, + max_output_tokens=100, + system_instruction="You are a helpful assistant.", + ), + ) + + sentry_sdk.flush() + spans = [item.payload for item in captured if item.type == "span"] + (span,) = [s for s in spans if s["attributes"].get("sentry.op") == OP.GEN_AI_CHAT] + span_data = span["attributes"] + + for key in expected_present: + assert key in span_data, f"{key} should have been collected" + assert json.loads(span_data[key]) == DATA_COLLECTION_CHAT_EXPECTED_VALUES[key] + + for key in expected_absent: + assert key not in span_data, f"{key} should not have been collected" + + # Data collection never gates non-PII attributes + assert span_data[SPANDATA.GEN_AI_PROVIDER_NAME] == "gcp.gemini" + assert span_data[SPANDATA.GEN_AI_REQUEST_MODEL] == "gemini-1.5-flash" + assert span_data[SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7 + assert span_data[SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100 + assert span_data[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 + assert span_data[SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 23 + assert span_data[SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 30 + assert span_data[SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS] == "STOP" + + +@pytest.mark.parametrize( + "data_collection,expected_present,expected_absent", + [ + pytest.param( + {"gen_ai": {"inputs": True, "outputs": False}}, + [SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS], + [SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS], + id="gen-ai-inputs-enabled-outputs-disabled-drops-tool-calls-only", + ), + pytest.param( + {"gen_ai": {"inputs": False, "outputs": True}}, + [SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS], + [SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS], + id="gen-ai-inputs-disabled-outputs-enabled-drops-available-tools-only", + ), + ], +) +def test_generate_content_data_collection_tools( + sentry_init, + capture_items, + mock_genai_client, + data_collection, + expected_present, + expected_absent, +): + sentry_init_kwargs = dict( + integrations=[GoogleGenAIIntegration()], + traces_sample_rate=1.0, + ) + if data_collection is not None: + sentry_init_kwargs["data_collection"] = data_collection + + sentry_init(**sentry_init_kwargs) + + # A declaration rather than a Python callable, so that the client does not + # automatically call the function and issue a second request + weather_tool = genai_types.Tool( + function_declarations=[ + genai_types.FunctionDeclaration( + name="get_weather", + description="Get the weather for a location", + ) + ] + ) + + mock_http_response = create_mock_http_response( + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + {"text": "I'll help you with that."}, + { + "functionCall": { + "name": "get_weather", + "args": {"location": "San Francisco"}, + } + }, + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 20, + "candidatesTokenCount": 30, + "totalTokenCount": 50, + }, + } + ) + + captured = capture_items("span") + + with mock.patch.object( + mock_genai_client._api_client, "request", return_value=mock_http_response + ): + mock_genai_client.models.generate_content( + model="gemini-1.5-flash", + contents="What's the weather?", + config=create_test_config(tools=[weather_tool]), + ) + + sentry_sdk.flush() + spans = [item.payload for item in captured if item.type == "span"] + (span,) = [s for s in spans if s["attributes"].get("sentry.op") == OP.GEN_AI_CHAT] + span_data = span["attributes"] + + for key in expected_present: + assert key in span_data, f"{key} should have been collected" + + for key in expected_absent: + assert key not in span_data, f"{key} should not have been collected" + + if SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS in expected_present: + available_tools = json.loads(span_data[SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS]) + assert [tool["name"] for tool in available_tools] == ["get_weather"] + + if SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS in expected_present: + tool_calls = json.loads(span_data[SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS]) + assert [tool_call["name"] for tool_call in tool_calls] == ["get_weather"] + assert json.loads(tool_calls[0]["arguments"]) == {"location": "San Francisco"} + + +@pytest.mark.parametrize( + "data_collection,expected_present,expected_absent", + [ + pytest.param( + {"gen_ai": {"inputs": True, "outputs": False}}, + [ + SPANDATA.GEN_AI_REQUEST_MESSAGES, + SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, + ], + [ + SPANDATA.GEN_AI_RESPONSE_TEXT, + ], + id="gen-ai-inputs-enabled-outputs-disabled", + ), + pytest.param( + {"gen_ai": {"inputs": False, "outputs": True}}, + [ + SPANDATA.GEN_AI_RESPONSE_TEXT, + ], + [ + SPANDATA.GEN_AI_REQUEST_MESSAGES, + SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, + ], + id="gen-ai-outputs-enabled-inputs-disabled", + ), + ], +) +def test_streaming_generate_content_data_collection( + sentry_init, + capture_items, + mock_genai_client, + data_collection, + expected_present, + expected_absent, +): + sentry_init_kwargs = dict( + integrations=[GoogleGenAIIntegration()], + traces_sample_rate=1.0, + ) + if data_collection is not None: + sentry_init_kwargs["data_collection"] = data_collection + + sentry_init(**sentry_init_kwargs) + + mock_stream = create_mock_streaming_responses( + [ + { + "candidates": [ + {"content": {"role": "model", "parts": [{"text": "Hello! "}]}} + ], + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 2, + "totalTokenCount": 12, + }, + "modelVersion": "gemini-1.5-flash", + }, + { + "candidates": [ + {"content": {"role": "model", "parts": [{"text": "How can I "}]}} + ], + }, + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "help you today?"}], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 7, + "totalTokenCount": 25, + }, + }, + ] + ) + + captured = capture_items("span") + + with mock.patch.object( + mock_genai_client._api_client, "request_streamed", return_value=mock_stream + ): + stream = mock_genai_client.models.generate_content_stream( + model="gemini-1.5-flash", + contents="Tell me a joke", + config=create_test_config( + system_instruction="You are a helpful assistant." + ), + ) + list(stream) + + sentry_sdk.flush() + spans = [item.payload for item in captured if item.type == "span"] + (span,) = [s for s in spans if s["attributes"].get("sentry.op") == OP.GEN_AI_CHAT] + span_data = span["attributes"] + + for key in expected_present: + assert key in span_data, f"{key} should have been collected" + assert json.loads(span_data[key]) == DATA_COLLECTION_CHAT_EXPECTED_VALUES[key] + + for key in expected_absent: + assert key not in span_data, f"{key} should not have been collected" + + assert span_data[SPANDATA.GEN_AI_RESPONSE_STREAMING] is True + assert span_data[SPANDATA.GEN_AI_REQUEST_MODEL] == "gemini-1.5-flash" + assert span_data[SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS] == "STOP" + assert span_data[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 + assert span_data[SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 25 + + +@pytest.mark.parametrize( + "data_collection,expected_present,expected_absent", + [ + pytest.param( + {"gen_ai": {"inputs": True, "outputs": False}}, + [SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS], + [SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS], + id="gen-ai-inputs-enabled-outputs-disabled-drops-tool-calls-only", + ), + pytest.param( + {"gen_ai": {"inputs": False, "outputs": True}}, + [SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS], + [SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS], + id="gen-ai-inputs-disabled-outputs-enabled-drops-available-tools-only", + ), + ], +) +def test_streaming_generate_content_data_collection_tools( + sentry_init, + capture_items, + mock_genai_client, + data_collection, + expected_present, + expected_absent, +): + sentry_init_kwargs = dict( + integrations=[GoogleGenAIIntegration()], + traces_sample_rate=1.0, + ) + if data_collection is not None: + sentry_init_kwargs["data_collection"] = data_collection + + sentry_init(**sentry_init_kwargs) + + weather_tool = genai_types.Tool( + function_declarations=[ + genai_types.FunctionDeclaration( + name="get_weather", + description="Get the weather for a location", + ) + ] + ) + + mock_stream = create_mock_streaming_responses( + [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "I'll help you with that."}], + } + } + ], + "modelVersion": "gemini-1.5-flash", + }, + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "functionCall": { + "name": "get_weather", + "args": {"location": "San Francisco"}, + } + } + ], + }, + "finishReason": "STOP", + } + ], + }, + ] + ) + + captured = capture_items("span") + + with mock.patch.object( + mock_genai_client._api_client, "request_streamed", return_value=mock_stream + ): + stream = mock_genai_client.models.generate_content_stream( + model="gemini-1.5-flash", + contents="What's the weather?", + config=create_test_config(tools=[weather_tool]), + ) + list(stream) + + sentry_sdk.flush() + spans = [item.payload for item in captured if item.type == "span"] + (span,) = [s for s in spans if s["attributes"].get("sentry.op") == OP.GEN_AI_CHAT] + span_data = span["attributes"] + + for key in expected_present: + assert key in span_data, f"{key} should have been collected" + + for key in expected_absent: + assert key not in span_data, f"{key} should not have been collected" + + if SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS in expected_present: + available_tools = json.loads(span_data[SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS]) + assert [tool["name"] for tool in available_tools] == ["get_weather"] + + if SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS in expected_present: + tool_calls = json.loads(span_data[SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS]) + assert [tool_call["name"] for tool_call in tool_calls] == ["get_weather"] + + +@pytest.mark.parametrize( + "data_collection,expected_present,expected_absent", + [ + pytest.param( + {"gen_ai": {"inputs": True, "outputs": False}}, + [ + SPANDATA.GEN_AI_EMBEDDINGS_INPUT, + ], + [], + id="gen-ai-inputs-enabled-outputs-disabled", + ), + pytest.param( + {"gen_ai": {"inputs": False, "outputs": True}}, + [], + [ + SPANDATA.GEN_AI_EMBEDDINGS_INPUT, + ], + id="gen-ai-outputs-enabled-inputs-disabled", + ), + ], +) +def test_embed_content_data_collection( + sentry_init, + capture_items, + mock_genai_client, + data_collection, + expected_present, + expected_absent, +): + sentry_init_kwargs = dict( + integrations=[GoogleGenAIIntegration()], + traces_sample_rate=1.0, + ) + if data_collection is not None: + sentry_init_kwargs["data_collection"] = data_collection + + sentry_init(**sentry_init_kwargs) + + mock_http_response = create_mock_http_response(EXAMPLE_EMBED_RESPONSE_JSON) + captured = capture_items("span") + + with mock.patch.object( + mock_genai_client._api_client, "request", return_value=mock_http_response + ): + mock_genai_client.models.embed_content( + model="text-embedding-004", + contents=["What is your name?", "What is your favorite color?"], + ) + + sentry_sdk.flush() + spans = [item.payload for item in captured if item.type == "span"] + (span,) = [ + s for s in spans if s["attributes"].get("sentry.op") == OP.GEN_AI_EMBEDDINGS + ] + span_data = span["attributes"] + + for key in expected_present: + assert key in span_data, f"{key} should have been collected" + assert json.loads(span_data[key]) == DATA_COLLECTION_EMBED_EXPECTED_VALUES[key] + + for key in expected_absent: + assert key not in span_data, f"{key} should not have been collected" + + assert span_data[SPANDATA.GEN_AI_OPERATION_NAME] == "embeddings" + assert span_data[SPANDATA.GEN_AI_REQUEST_MODEL] == "text-embedding-004" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "data_collection,expected_present,expected_absent", + [ + pytest.param( + {"gen_ai": {"inputs": True, "outputs": False}}, + [ + SPANDATA.GEN_AI_REQUEST_MESSAGES, + SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, + ], + [ + SPANDATA.GEN_AI_RESPONSE_TEXT, + ], + id="gen-ai-inputs-enabled-outputs-disabled", + ), + pytest.param( + {"gen_ai": {"inputs": False, "outputs": True}}, + [ + SPANDATA.GEN_AI_RESPONSE_TEXT, + ], + [ + SPANDATA.GEN_AI_REQUEST_MESSAGES, + SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, + ], + id="gen-ai-outputs-enabled-inputs-disabled", + ), + ], +) +async def test_async_generate_content_data_collection( + sentry_init, + capture_items, + mock_genai_client, + data_collection, + expected_present, + expected_absent, +): + sentry_init_kwargs = dict( + integrations=[GoogleGenAIIntegration()], + traces_sample_rate=1.0, + ) + if data_collection is not None: + sentry_init_kwargs["data_collection"] = data_collection + + sentry_init(**sentry_init_kwargs) + + mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) + captured = capture_items("span") + + with mock.patch.object( + mock_genai_client._api_client, "async_request", return_value=mock_http_response + ): + await mock_genai_client.aio.models.generate_content( + model="gemini-1.5-flash", + contents="Tell me a joke", + config=create_test_config( + system_instruction="You are a helpful assistant." + ), + ) + + sentry_sdk.flush() + spans = [item.payload for item in captured if item.type == "span"] + (span,) = [s for s in spans if s["attributes"].get("sentry.op") == OP.GEN_AI_CHAT] + span_data = span["attributes"] + + for key in expected_present: + assert key in span_data, f"{key} should have been collected" + assert json.loads(span_data[key]) == DATA_COLLECTION_CHAT_EXPECTED_VALUES[key] + + for key in expected_absent: + assert key not in span_data, f"{key} should not have been collected" + + assert span_data[SPANDATA.GEN_AI_REQUEST_MODEL] == "gemini-1.5-flash" + assert span_data[SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 30 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "data_collection,expected_present,expected_absent", + [ + pytest.param( + {"gen_ai": {"inputs": True, "outputs": False}}, + [ + SPANDATA.GEN_AI_EMBEDDINGS_INPUT, + ], + [], + id="gen-ai-inputs-enabled-outputs-disabled", + ), + pytest.param( + {"gen_ai": {"inputs": False, "outputs": True}}, + [], + [ + SPANDATA.GEN_AI_EMBEDDINGS_INPUT, + ], + id="gen-ai-outputs-enabled-inputs-disabled", + ), + ], +) +async def test_async_embed_content_data_collection( + sentry_init, + capture_items, + mock_genai_client, + data_collection, + expected_present, + expected_absent, +): + sentry_init_kwargs = dict( + integrations=[GoogleGenAIIntegration()], + traces_sample_rate=1.0, + ) + if data_collection is not None: + sentry_init_kwargs["data_collection"] = data_collection + + sentry_init(**sentry_init_kwargs) + + mock_http_response = create_mock_http_response(EXAMPLE_EMBED_RESPONSE_JSON) + captured = capture_items("span") + + with mock.patch.object( + mock_genai_client._api_client, "async_request", return_value=mock_http_response + ): + await mock_genai_client.aio.models.embed_content( + model="text-embedding-004", + contents=["What is your name?", "What is your favorite color?"], + ) + + sentry_sdk.flush() + spans = [item.payload for item in captured if item.type == "span"] + (span,) = [ + s for s in spans if s["attributes"].get("sentry.op") == OP.GEN_AI_EMBEDDINGS + ] + span_data = span["attributes"] + + for key in expected_present: + assert key in span_data, f"{key} should have been collected" + assert json.loads(span_data[key]) == DATA_COLLECTION_EMBED_EXPECTED_VALUES[key] + + for key in expected_absent: + assert key not in span_data, f"{key} should not have been collected" + + assert span_data[SPANDATA.GEN_AI_OPERATION_NAME] == "embeddings" + assert span_data[SPANDATA.GEN_AI_REQUEST_MODEL] == "text-embedding-004" From e7ffa8443a47be98541e8763b31d1e991977c34d Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Tue, 29 Sep 2026 13:17:06 +0200 Subject: [PATCH 3/3] remove stale comment --- tests/integrations/google_genai/test_google_genai.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/integrations/google_genai/test_google_genai.py b/tests/integrations/google_genai/test_google_genai.py index 88469ae725..286650d872 100644 --- a/tests/integrations/google_genai/test_google_genai.py +++ b/tests/integrations/google_genai/test_google_genai.py @@ -2632,7 +2632,6 @@ def test_generate_content_data_collection( for key in expected_absent: assert key not in span_data, f"{key} should not have been collected" - # Data collection never gates non-PII attributes assert span_data[SPANDATA.GEN_AI_PROVIDER_NAME] == "gcp.gemini" assert span_data[SPANDATA.GEN_AI_REQUEST_MODEL] == "gemini-1.5-flash" assert span_data[SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7