From 300048835f448315cdbbcf5ed32ffff46edc22a7 Mon Sep 17 00:00:00 2001 From: Erica Pisani Date: Fri, 7 Aug 2026 14:10:35 -0400 Subject: [PATCH] feat(google-genai): Gate prompt/response collection on data_collection option Modify the GoogleGenAI integration to respect the data_collection config for controlling whether prompts, responses, and tool calls are captured in spans. When data collection is enabled, the new gen_ai.inputs and gen_ai.outputs flags control what data is collected. When data collection is not configured, falls back to legacy send_default_pii and include_prompts settings for compatibility. Refs PY-2588 --- .../integrations/google_genai/streaming.py | 47 +- sentry_sdk/integrations/google_genai/utils.py | 97 +- .../google_genai/test_google_genai.py | 1126 +++++++++++++++++ 3 files changed, 1226 insertions(+), 44 deletions(-) diff --git a/sentry_sdk/integrations/google_genai/streaming.py b/sentry_sdk/integrations/google_genai/streaming.py index 8414ea4f21..86cdcf29ba 100644 --- a/sentry_sdk/integrations/google_genai/streaming.py +++ b/sentry_sdk/integrations/google_genai/streaming.py @@ -1,10 +1,12 @@ from typing import TYPE_CHECKING, Any, List, Optional, TypedDict, Union +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 StreamedSpan from sentry_sdk.utils import ( + has_data_collection_enabled, safe_serialize, ) @@ -106,16 +108,7 @@ def set_span_data_for_streaming_response( set_on_span = ( span.set_attribute if isinstance(span, StreamedSpan) else span.set_data ) - - if ( - should_send_default_pii() - and integration.include_prompts - and accumulated_response.get("text") - ): - set_on_span( - SPANDATA.GEN_AI_RESPONSE_TEXT, - safe_serialize([accumulated_response["text"]]), - ) + client = sentry_sdk.get_client() if accumulated_response.get("finish_reasons"): set_data_normalized( @@ -124,12 +117,6 @@ def set_span_data_for_streaming_response( accumulated_response["finish_reasons"], ) - if accumulated_response.get("tool_calls"): - set_on_span( - SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, - safe_serialize(accumulated_response["tool_calls"]), - ) - response_id = accumulated_response.get("id") if response_id is not None: set_on_span(SPANDATA.GEN_AI_RESPONSE_ID, response_id) @@ -170,3 +157,31 @@ def set_span_data_for_streaming_response( SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, accumulated_response["usage_metadata"]["total_tokens"], ) + + if accumulated_response.get("tool_calls"): + if has_data_collection_enabled(client.options): + if client.options["data_collection"]["gen_ai"]["outputs"]: + set_on_span( + SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, + safe_serialize(accumulated_response["tool_calls"]), + ) + else: + # Before data collection was introduced this was unconditionally set + set_on_span( + 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"]: + set_on_span( + SPANDATA.GEN_AI_RESPONSE_TEXT, + safe_serialize([accumulated_response["text"]]), + ) + + elif should_send_default_pii() and integration.include_prompts: + set_on_span( + 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 464a812680..75ba30c199 100644 --- a/sentry_sdk/integrations/google_genai/utils.py +++ b/sentry_sdk/integrations/google_genai/utils.py @@ -36,6 +36,7 @@ from sentry_sdk.utils import ( capture_internal_exceptions, event_from_exception, + has_data_collection_enabled, safe_serialize, ) @@ -887,6 +888,7 @@ def set_span_data_for_request( kwargs: "dict[str, Any]", ) -> None: """Set span data for the request.""" + client = sentry_sdk.get_client() set_on_span = ( span.set_attribute if isinstance(span, StreamedSpan) else span.set_data ) @@ -898,8 +900,37 @@ def set_span_data_for_request( config: "Optional[GenerateContentConfig]" = kwargs.get("config") - # Set input messages/prompts if PII is allowed - if should_send_default_pii() and integration.include_prompts: + # Set tools if available + if config is not None and hasattr(config, "tools"): + tools = config.tools + 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 + set_data_normalized( + span, + SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, + formatted_tools, + unpack=False, + ) + + 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() and integration.include_prompts: + record_inputs = True + + if record_inputs: messages = [] # Add system instruction if present @@ -951,42 +982,19 @@ def set_span_data_for_request( if value is not None: set_on_span(span_key, value) - # Set tools if available - if config is not None and hasattr(config, "tools"): - tools = config.tools - if tools: - formatted_tools = _format_tools_for_span(tools) - if formatted_tools: - set_data_normalized( - span, - SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, - formatted_tools, - unpack=False, - ) - def set_span_data_for_response( span: "Union[Span, StreamedSpan]", integration: "Any", response: "GenerateContentResponse", ) -> None: - """Set span data for the response.""" if not response: return + client = sentry_sdk.get_client() set_on_span = ( span.set_attribute if isinstance(span, StreamedSpan) else span.set_data ) - if should_send_default_pii() and integration.include_prompts: - response_texts = _extract_response_text(response) - if response_texts: - # Format as JSON string array as per documentation - set_on_span(SPANDATA.GEN_AI_RESPONSE_TEXT, safe_serialize(response_texts)) - - tool_calls = extract_tool_calls(response) - if tool_calls: - # Tool calls should be JSON serialized - set_on_span(SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, safe_serialize(tool_calls)) finish_reasons = extract_finish_reasons(response) if finish_reasons: @@ -1023,6 +1031,31 @@ def set_span_data_for_response( if usage_data["total_tokens"]: set_on_span(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, usage_data["total_tokens"]) + tool_calls = extract_tool_calls(response) + if tool_calls: + if has_data_collection_enabled(client.options): + if client.options["data_collection"]["gen_ai"]["outputs"]: + set_on_span( + SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, safe_serialize(tool_calls) + ) + else: + # Before data collection was introduced, this was set unconditionally + set_on_span(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: + set_on_span( + SPANDATA.GEN_AI_RESPONSE_TEXT, safe_serialize(response_texts) + ) + elif should_send_default_pii() and integration.include_prompts: + # TODO: Delete this block once data collection has been completely rolled out + response_texts = _extract_response_text(response) + if response_texts: + # Format as JSON string array as per documentation + set_on_span(SPANDATA.GEN_AI_RESPONSE_TEXT, safe_serialize(response_texts)) + def prepare_generate_content_args( args: "tuple[Any, ...]", kwargs: "dict[str, Any]" @@ -1062,8 +1095,16 @@ def set_span_data_for_embed_request( kwargs: "dict[str, Any]", ) -> None: """Set span data for embedding request.""" - # Include input contents if PII is allowed - if should_send_default_pii() and integration.include_prompts: + 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() and integration.include_prompts: + record_inputs = True + + if record_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 494b86cf5a..756d83572a 100644 --- a/tests/integrations/google_genai/test_google_genai.py +++ b/tests/integrations/google_genai/test_google_genai.py @@ -3632,3 +3632,1129 @@ 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("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +@pytest.mark.parametrize( + "data_collection,send_default_pii,include_prompts,expected_present,expected_absent", + [ + pytest.param( + {"gen_ai": {"inputs": True, "outputs": True}}, + False, + 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, + 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, + 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, + 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, + 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, + True, + [ + SPANDATA.GEN_AI_REQUEST_MESSAGES, + SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, + SPANDATA.GEN_AI_RESPONSE_TEXT, + ], + [], + id="no-gen-ai-config-legacy-pii-and-include-prompts-enabled", + ), + pytest.param( + None, + False, + True, + [], + [ + 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_events, + capture_items, + mock_genai_client, + data_collection, + send_default_pii, + include_prompts, + expected_present, + expected_absent, + stream_gen_ai_spans, + span_streaming, +): + sentry_init_kwargs = dict( + integrations=[GoogleGenAIIntegration(include_prompts=include_prompts)], + traces_sample_rate=1.0, + send_default_pii=send_default_pii, + stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", + ) + if data_collection is not None: + sentry_init_kwargs["_experiments"] = {"data_collection": data_collection} + + sentry_init(**sentry_init_kwargs) + + mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) + streamed = span_streaming or stream_gen_ai_spans + captured = capture_items("span") if streamed else capture_events() + + with mock.patch.object( + mock_genai_client._api_client, "request", return_value=mock_http_response + ), start_transaction(name="google_genai"): + 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.", + ), + ) + + if streamed: + 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"] + else: + (event,) = captured + (span,) = [s for s in event["spans"] if s["op"] == OP.GEN_AI_CHAT] + span_data = span["data"] + + 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_SYSTEM] == "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("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +@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-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_events, + capture_items, + mock_genai_client, + data_collection, + send_default_pii, + expected_present, + expected_absent, + stream_gen_ai_spans, + span_streaming, +): + sentry_init_kwargs = dict( + integrations=[GoogleGenAIIntegration(include_prompts=False)], + traces_sample_rate=1.0, + send_default_pii=send_default_pii, + stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", + ) + if data_collection is not None: + sentry_init_kwargs["_experiments"] = {"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, + }, + } + ) + + streamed = span_streaming or stream_gen_ai_spans + captured = capture_items("span") if streamed else capture_events() + + with mock.patch.object( + mock_genai_client._api_client, "request", return_value=mock_http_response + ), start_transaction(name="google_genai"): + mock_genai_client.models.generate_content( + model="gemini-1.5-flash", + contents="What's the weather?", + config=create_test_config(tools=[weather_tool]), + ) + + if streamed: + 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"] + else: + (event,) = captured + (span,) = [s for s in event["spans"] if s["op"] == OP.GEN_AI_CHAT] + span_data = span["data"] + + 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("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +@pytest.mark.parametrize( + "data_collection,send_default_pii,include_prompts,expected_present,expected_absent", + [ + pytest.param( + {"gen_ai": {"inputs": True, "outputs": True}}, + False, + 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, + 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, + 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, + 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, + 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, + True, + [ + SPANDATA.GEN_AI_REQUEST_MESSAGES, + SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, + SPANDATA.GEN_AI_RESPONSE_TEXT, + ], + [], + id="no-gen-ai-config-legacy-pii-and-include-prompts-enabled", + ), + pytest.param( + None, + False, + True, + [], + [ + 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_events, + capture_items, + mock_genai_client, + data_collection, + send_default_pii, + include_prompts, + expected_present, + expected_absent, + stream_gen_ai_spans, + span_streaming, +): + sentry_init_kwargs = dict( + integrations=[GoogleGenAIIntegration(include_prompts=include_prompts)], + traces_sample_rate=1.0, + send_default_pii=send_default_pii, + stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", + ) + if data_collection is not None: + sentry_init_kwargs["_experiments"] = {"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, + }, + }, + ] + ) + + streamed = span_streaming or stream_gen_ai_spans + captured = capture_items("span") if streamed else capture_events() + + with mock.patch.object( + mock_genai_client._api_client, "request_streamed", return_value=mock_stream + ), start_transaction(name="google_genai"): + 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) + + if streamed: + 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"] + else: + (event,) = captured + (span,) = [s for s in event["spans"] if s["op"] == OP.GEN_AI_CHAT] + span_data = span["data"] + + 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("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +@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-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_events, + capture_items, + mock_genai_client, + data_collection, + send_default_pii, + expected_present, + expected_absent, + stream_gen_ai_spans, + span_streaming, +): + sentry_init_kwargs = dict( + integrations=[GoogleGenAIIntegration(include_prompts=False)], + traces_sample_rate=1.0, + send_default_pii=send_default_pii, + stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", + ) + if data_collection is not None: + sentry_init_kwargs["_experiments"] = {"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", + } + ], + }, + ] + ) + + streamed = span_streaming or stream_gen_ai_spans + captured = capture_items("span") if streamed else capture_events() + + with mock.patch.object( + mock_genai_client._api_client, "request_streamed", return_value=mock_stream + ), start_transaction(name="google_genai"): + 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) + + if streamed: + 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"] + else: + (event,) = captured + (span,) = [s for s in event["spans"] if s["op"] == OP.GEN_AI_CHAT] + span_data = span["data"] + + 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("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +@pytest.mark.parametrize( + "data_collection,send_default_pii,include_prompts,expected_present,expected_absent", + [ + pytest.param( + {"gen_ai": {"inputs": True, "outputs": True}}, + False, + 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, + 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, + False, + [ + SPANDATA.GEN_AI_EMBEDDINGS_INPUT, + ], + [], + id="gen-ai-inputs-enabled-outputs-disabled", + ), + pytest.param( + {"gen_ai": {"inputs": False, "outputs": True}}, + False, + False, + [], + [ + SPANDATA.GEN_AI_EMBEDDINGS_INPUT, + ], + id="gen-ai-outputs-enabled-inputs-disabled", + ), + pytest.param( + {"gen_ai": {}}, + False, + False, + [ + SPANDATA.GEN_AI_EMBEDDINGS_INPUT, + ], + [], + id="gen-ai-inputs-and-outputs-omitted-default-to-enabled", + ), + pytest.param( + None, + True, + True, + [ + SPANDATA.GEN_AI_EMBEDDINGS_INPUT, + ], + [], + id="no-gen-ai-config-legacy-pii-and-include-prompts-enabled", + ), + pytest.param( + None, + False, + True, + [], + [ + SPANDATA.GEN_AI_EMBEDDINGS_INPUT, + ], + id="no-gen-ai-config-legacy-pii-disabled", + ), + ], +) +def test_embed_content_data_collection( + sentry_init, + capture_events, + capture_items, + mock_genai_client, + data_collection, + send_default_pii, + include_prompts, + expected_present, + expected_absent, + stream_gen_ai_spans, + span_streaming, +): + sentry_init_kwargs = dict( + integrations=[GoogleGenAIIntegration(include_prompts=include_prompts)], + traces_sample_rate=1.0, + send_default_pii=send_default_pii, + stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", + ) + if data_collection is not None: + sentry_init_kwargs["_experiments"] = {"data_collection": data_collection} + + sentry_init(**sentry_init_kwargs) + + mock_http_response = create_mock_http_response(EXAMPLE_EMBED_RESPONSE_JSON) + streamed = span_streaming or stream_gen_ai_spans + captured = capture_items("span") if streamed else capture_events() + + with mock.patch.object( + mock_genai_client._api_client, "request", return_value=mock_http_response + ), start_transaction(name="google_genai_embeddings"): + mock_genai_client.models.embed_content( + model="text-embedding-004", + contents=["What is your name?", "What is your favorite color?"], + ) + + if streamed: + 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"] + else: + (event,) = captured + (span,) = [s for s in event["spans"] if s["op"] == OP.GEN_AI_EMBEDDINGS] + span_data = span["data"] + + 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("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +@pytest.mark.parametrize( + "data_collection,send_default_pii,include_prompts,expected_present,expected_absent", + [ + pytest.param( + {"gen_ai": {"inputs": True, "outputs": True}}, + False, + 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, + 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, + 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, + 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, + 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, + True, + [ + SPANDATA.GEN_AI_REQUEST_MESSAGES, + SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, + SPANDATA.GEN_AI_RESPONSE_TEXT, + ], + [], + id="no-gen-ai-config-legacy-pii-and-include-prompts-enabled", + ), + pytest.param( + None, + False, + True, + [], + [ + 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_events, + capture_items, + mock_genai_client, + data_collection, + send_default_pii, + include_prompts, + expected_present, + expected_absent, + stream_gen_ai_spans, + span_streaming, +): + sentry_init_kwargs = dict( + integrations=[GoogleGenAIIntegration(include_prompts=include_prompts)], + traces_sample_rate=1.0, + send_default_pii=send_default_pii, + stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", + ) + if data_collection is not None: + sentry_init_kwargs["_experiments"] = {"data_collection": data_collection} + + sentry_init(**sentry_init_kwargs) + + mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) + streamed = span_streaming or stream_gen_ai_spans + captured = capture_items("span") if streamed else capture_events() + + with mock.patch.object( + mock_genai_client._api_client, "async_request", return_value=mock_http_response + ), start_transaction(name="google_genai"): + 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." + ), + ) + + if streamed: + 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"] + else: + (event,) = captured + (span,) = [s for s in event["spans"] if s["op"] == OP.GEN_AI_CHAT] + span_data = span["data"] + + 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("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +@pytest.mark.parametrize( + "data_collection,send_default_pii,include_prompts,expected_present,expected_absent", + [ + pytest.param( + {"gen_ai": {"inputs": True, "outputs": True}}, + False, + 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, + 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, + False, + [ + SPANDATA.GEN_AI_EMBEDDINGS_INPUT, + ], + [], + id="gen-ai-inputs-enabled-outputs-disabled", + ), + pytest.param( + {"gen_ai": {"inputs": False, "outputs": True}}, + False, + False, + [], + [ + SPANDATA.GEN_AI_EMBEDDINGS_INPUT, + ], + id="gen-ai-outputs-enabled-inputs-disabled", + ), + pytest.param( + {"gen_ai": {}}, + False, + False, + [ + SPANDATA.GEN_AI_EMBEDDINGS_INPUT, + ], + [], + id="gen-ai-inputs-and-outputs-omitted-default-to-enabled", + ), + pytest.param( + None, + True, + True, + [ + SPANDATA.GEN_AI_EMBEDDINGS_INPUT, + ], + [], + id="no-gen-ai-config-legacy-pii-and-include-prompts-enabled", + ), + pytest.param( + None, + False, + True, + [], + [ + SPANDATA.GEN_AI_EMBEDDINGS_INPUT, + ], + id="no-gen-ai-config-legacy-pii-disabled", + ), + ], +) +async def test_async_embed_content_data_collection( + sentry_init, + capture_events, + capture_items, + mock_genai_client, + data_collection, + send_default_pii, + include_prompts, + expected_present, + expected_absent, + stream_gen_ai_spans, + span_streaming, +): + sentry_init_kwargs = dict( + integrations=[GoogleGenAIIntegration(include_prompts=include_prompts)], + traces_sample_rate=1.0, + send_default_pii=send_default_pii, + stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", + ) + if data_collection is not None: + sentry_init_kwargs["_experiments"] = {"data_collection": data_collection} + + sentry_init(**sentry_init_kwargs) + + mock_http_response = create_mock_http_response(EXAMPLE_EMBED_RESPONSE_JSON) + streamed = span_streaming or stream_gen_ai_spans + captured = capture_items("span") if streamed else capture_events() + + with mock.patch.object( + mock_genai_client._api_client, "async_request", return_value=mock_http_response + ), start_transaction(name="google_genai_embeddings"): + await mock_genai_client.aio.models.embed_content( + model="text-embedding-004", + contents=["What is your name?", "What is your favorite color?"], + ) + + if streamed: + 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"] + else: + (event,) = captured + (span,) = [s for s in event["spans"] if s["op"] == OP.GEN_AI_EMBEDDINGS] + span_data = span["data"] + + 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"