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..286650d872 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) @@ -2293,33 +2561,10 @@ def __init__(self): @pytest.mark.parametrize( - "data_collection,send_default_pii,expected_present,expected_absent", + "data_collection,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, @@ -2331,7 +2576,6 @@ def __init__(self): ), pytest.param( {"gen_ai": {"inputs": False, "outputs": True}}, - False, [ SPANDATA.GEN_AI_RESPONSE_TEXT, ], @@ -2341,39 +2585,6 @@ def __init__(self): ], 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( @@ -2381,14 +2592,12 @@ def test_generate_content_data_collection( 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 @@ -2423,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 @@ -2434,28 +2642,9 @@ def test_generate_content_data_collection( 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], @@ -2468,24 +2657,6 @@ def test_generate_content_data_collection( [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( @@ -2493,14 +2664,12 @@ def test_generate_content_data_collection_tools( 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 @@ -2578,33 +2747,10 @@ def test_generate_content_data_collection_tools( @pytest.mark.parametrize( - "data_collection,send_default_pii,expected_present,expected_absent", + "data_collection,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, @@ -2616,7 +2762,6 @@ def test_generate_content_data_collection_tools( ), pytest.param( {"gen_ai": {"inputs": False, "outputs": True}}, - False, [ SPANDATA.GEN_AI_RESPONSE_TEXT, ], @@ -2626,39 +2771,6 @@ def test_generate_content_data_collection_tools( ], 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( @@ -2666,14 +2778,12 @@ def test_streaming_generate_content_data_collection( 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 @@ -2750,28 +2860,9 @@ def test_streaming_generate_content_data_collection( 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], @@ -2784,24 +2875,6 @@ def test_streaming_generate_content_data_collection( [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( @@ -2809,14 +2882,12 @@ def test_streaming_generate_content_data_collection_tools( 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 @@ -2899,29 +2970,10 @@ def test_streaming_generate_content_data_collection_tools( @pytest.mark.parametrize( - "data_collection,send_default_pii,expected_present,expected_absent", + "data_collection,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, ], @@ -2930,40 +2982,12 @@ def test_streaming_generate_content_data_collection_tools( ), 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( @@ -2971,14 +2995,12 @@ def test_embed_content_data_collection( 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 @@ -3016,33 +3038,10 @@ def test_embed_content_data_collection( @pytest.mark.asyncio @pytest.mark.parametrize( - "data_collection,send_default_pii,expected_present,expected_absent", + "data_collection,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, @@ -3054,7 +3053,6 @@ def test_embed_content_data_collection( ), pytest.param( {"gen_ai": {"inputs": False, "outputs": True}}, - False, [ SPANDATA.GEN_AI_RESPONSE_TEXT, ], @@ -3064,39 +3062,6 @@ def test_embed_content_data_collection( ], 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( @@ -3104,14 +3069,12 @@ async def test_async_generate_content_data_collection( 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 @@ -3150,29 +3113,10 @@ async def test_async_generate_content_data_collection( @pytest.mark.asyncio @pytest.mark.parametrize( - "data_collection,send_default_pii,expected_present,expected_absent", + "data_collection,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, ], @@ -3181,40 +3125,12 @@ async def test_async_generate_content_data_collection( ), 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( @@ -3222,14 +3138,12 @@ async def test_async_embed_content_data_collection( 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