From ba727b31b2b3c02ee2cb1278ccf43fb5e6f1af21 Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Tue, 10 Feb 2026 14:20:44 +0100 Subject: [PATCH] fix(google-genai): Remove agent spans for simple requests --- .../integrations/google_genai/__init__.py | 126 +++++------------- .../google_genai/test_google_genai.py | 57 ++------ 2 files changed, 46 insertions(+), 137 deletions(-) diff --git a/sentry_sdk/integrations/google_genai/__init__.py b/sentry_sdk/integrations/google_genai/__init__.py index 27a42f4f6a..7a8130252f 100644 --- a/sentry_sdk/integrations/google_genai/__init__.py +++ b/sentry_sdk/integrations/google_genai/__init__.py @@ -73,17 +73,6 @@ def new_generate_content_stream( _model, contents, model_name = prepare_generate_content_args(args, kwargs) - span = get_start_span_function()( - op=OP.GEN_AI_INVOKE_AGENT, - name="invoke_agent", - origin=ORIGIN, - ) - span.__enter__() - span.set_data(SPANDATA.GEN_AI_AGENT_NAME, model_name) - span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent") - set_span_data_for_request(span, integration, model_name, contents, kwargs) - span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, True) - chat_span = sentry_sdk.start_span( op=OP.GEN_AI_CHAT, name=f"chat {model_name}", @@ -118,18 +107,13 @@ def new_iterator() -> "Iterator[Any]": set_span_data_for_streaming_response( chat_span, integration, accumulated_response ) - set_span_data_for_streaming_response( - span, integration, accumulated_response - ) chat_span.__exit__(None, None, None) - span.__exit__(None, None, None) return new_iterator() except Exception as exc: _capture_exception(exc) chat_span.__exit__(None, None, None) - span.__exit__(None, None, None) raise return new_generate_content_stream @@ -148,17 +132,6 @@ async def new_async_generate_content_stream( _model, contents, model_name = prepare_generate_content_args(args, kwargs) - span = get_start_span_function()( - op=OP.GEN_AI_INVOKE_AGENT, - name="invoke_agent", - origin=ORIGIN, - ) - span.__enter__() - span.set_data(SPANDATA.GEN_AI_AGENT_NAME, model_name) - span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent") - set_span_data_for_request(span, integration, model_name, contents, kwargs) - span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, True) - chat_span = sentry_sdk.start_span( op=OP.GEN_AI_CHAT, name=f"chat {model_name}", @@ -193,18 +166,13 @@ async def new_async_iterator() -> "AsyncIterator[Any]": set_span_data_for_streaming_response( chat_span, integration, accumulated_response ) - set_span_data_for_streaming_response( - span, integration, accumulated_response - ) chat_span.__exit__(None, None, None) - span.__exit__(None, None, None) return new_async_iterator() except Exception as exc: _capture_exception(exc) chat_span.__exit__(None, None, None) - span.__exit__(None, None, None) raise return new_async_generate_content_stream @@ -219,39 +187,29 @@ def new_generate_content(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": model, contents, model_name = prepare_generate_content_args(args, kwargs) - with get_start_span_function()( - op=OP.GEN_AI_INVOKE_AGENT, - name="invoke_agent", + with sentry_sdk.start_span( + op=OP.GEN_AI_CHAT, + name=f"chat {model_name}", origin=ORIGIN, - ) as span: - span.set_data(SPANDATA.GEN_AI_AGENT_NAME, model_name) - span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent") - set_span_data_for_request(span, integration, model_name, contents, kwargs) - - with sentry_sdk.start_span( - op=OP.GEN_AI_CHAT, - name=f"chat {model_name}", - origin=ORIGIN, - ) as chat_span: - chat_span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat") - chat_span.set_data(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM) - chat_span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model_name) - chat_span.set_data(SPANDATA.GEN_AI_AGENT_NAME, model_name) - set_span_data_for_request( - chat_span, integration, model_name, contents, kwargs - ) + ) as chat_span: + chat_span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat") + chat_span.set_data(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM) + chat_span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model_name) + chat_span.set_data(SPANDATA.GEN_AI_AGENT_NAME, model_name) + set_span_data_for_request( + chat_span, integration, model_name, contents, kwargs + ) - try: - response = f(self, *args, **kwargs) - except Exception as exc: - _capture_exception(exc) - chat_span.set_status(SPANSTATUS.INTERNAL_ERROR) - raise + try: + response = f(self, *args, **kwargs) + except Exception as exc: + _capture_exception(exc) + chat_span.set_status(SPANSTATUS.INTERNAL_ERROR) + raise - set_span_data_for_response(chat_span, integration, response) - set_span_data_for_response(span, integration, response) + set_span_data_for_response(chat_span, integration, response) - return response + return response return new_generate_content @@ -267,37 +225,27 @@ async def new_async_generate_content( model, contents, model_name = prepare_generate_content_args(args, kwargs) - with get_start_span_function()( - op=OP.GEN_AI_INVOKE_AGENT, - name="invoke_agent", + with sentry_sdk.start_span( + op=OP.GEN_AI_CHAT, + name=f"chat {model_name}", origin=ORIGIN, - ) as span: - span.set_data(SPANDATA.GEN_AI_AGENT_NAME, model_name) - span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent") - set_span_data_for_request(span, integration, model_name, contents, kwargs) - - with sentry_sdk.start_span( - op=OP.GEN_AI_CHAT, - name=f"chat {model_name}", - origin=ORIGIN, - ) as chat_span: - chat_span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat") - chat_span.set_data(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM) - chat_span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model_name) - set_span_data_for_request( - chat_span, integration, model_name, contents, kwargs - ) - try: - response = await f(self, *args, **kwargs) - except Exception as exc: - _capture_exception(exc) - chat_span.set_status(SPANSTATUS.INTERNAL_ERROR) - raise + ) as chat_span: + chat_span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat") + chat_span.set_data(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM) + chat_span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model_name) + set_span_data_for_request( + chat_span, integration, model_name, contents, kwargs + ) + try: + response = await f(self, *args, **kwargs) + except Exception as exc: + _capture_exception(exc) + chat_span.set_status(SPANSTATUS.INTERNAL_ERROR) + raise - set_span_data_for_response(chat_span, integration, response) - set_span_data_for_response(span, integration, response) + set_span_data_for_response(chat_span, integration, response) - return response + return response return new_async_generate_content diff --git a/tests/integrations/google_genai/test_google_genai.py b/tests/integrations/google_genai/test_google_genai.py index 7448dd630a..fc21216be6 100644 --- a/tests/integrations/google_genai/test_google_genai.py +++ b/tests/integrations/google_genai/test_google_genai.py @@ -152,17 +152,8 @@ def test_nonstreaming_generate_content( assert event["type"] == "transaction" assert event["transaction"] == "google_genai" - # Should have 2 spans: invoke_agent and chat - assert len(event["spans"]) == 2 - invoke_span, chat_span = event["spans"] - - # Check invoke_agent span - assert invoke_span["op"] == OP.GEN_AI_INVOKE_AGENT - assert invoke_span["description"] == "invoke_agent" - assert invoke_span["data"][SPANDATA.GEN_AI_AGENT_NAME] == "gemini-1.5-flash" - assert invoke_span["data"][SPANDATA.GEN_AI_SYSTEM] == "gcp.gemini" - assert invoke_span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "gemini-1.5-flash" - assert invoke_span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "invoke_agent" + assert len(event["spans"]) == 1 + chat_span = event["spans"][0] # Check chat span assert chat_span["op"] == OP.GEN_AI_CHAT @@ -172,18 +163,12 @@ def test_nonstreaming_generate_content( assert chat_span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "gemini-1.5-flash" if send_default_pii and include_prompts: - # Messages are serialized as JSON strings - messages = json.loads(invoke_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) - assert messages == [{"role": "user", "content": "Tell me a joke"}] - # Response text is stored as a JSON array response_text = chat_span["data"][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?"] else: - assert SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS not in invoke_span["data"] - assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in invoke_span["data"] assert SPANDATA.GEN_AI_RESPONSE_TEXT not in chat_span["data"] # Check token usage @@ -194,10 +179,6 @@ def test_nonstreaming_generate_content( assert chat_span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 5 assert chat_span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING] == 3 - # Check configuration parameters - assert invoke_span["data"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7 - assert invoke_span["data"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100 - @pytest.mark.parametrize("generate_content_config", (False, True)) @pytest.mark.parametrize( @@ -519,50 +500,30 @@ def test_streaming_generate_content(sentry_init, capture_events, mock_genai_clie (event,) = events - # There should be 2 spans: invoke_agent and chat - assert len(event["spans"]) == 2 - invoke_span = event["spans"][0] - chat_span = event["spans"][1] + assert len(event["spans"]) == 1 + chat_span = event["spans"][0] # Check that streaming flag is set on both spans - assert invoke_span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True assert chat_span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True # Verify accumulated response text (all chunks combined) expected_full_text = "Hello! How can I help you today?" # Response text is stored as a JSON string chat_response_text = json.loads(chat_span["data"][SPANDATA.GEN_AI_RESPONSE_TEXT]) - invoke_response_text = json.loads( - invoke_span["data"][SPANDATA.GEN_AI_RESPONSE_TEXT] - ) assert chat_response_text == [expected_full_text] - assert invoke_response_text == [expected_full_text] # Verify finish reasons (only the final chunk has a finish reason) # When there's a single finish reason, it's stored as a plain string (not JSON) assert SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS in chat_span["data"] - assert SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS in invoke_span["data"] assert chat_span["data"][SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS] == "STOP" - assert invoke_span["data"][SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS] == "STOP" - assert chat_span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 - assert invoke_span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 - assert chat_span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 10 - assert invoke_span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 10 - assert chat_span["data"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 25 - assert invoke_span["data"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 25 - assert chat_span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 5 - assert invoke_span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 5 - assert chat_span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING] == 3 - assert invoke_span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING] == 3 # Verify model name assert chat_span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "gemini-1.5-flash" - assert invoke_span["data"][SPANDATA.GEN_AI_AGENT_NAME] == "gemini-1.5-flash" def test_span_origin(sentry_init, capture_events, mock_genai_client): @@ -625,7 +586,7 @@ def test_response_without_usage_metadata( ) (event,) = events - chat_span = event["spans"][1] + chat_span = event["spans"][0] # Usage data should not be present assert SPANDATA.GEN_AI_USAGE_INPUT_TOKENS not in chat_span["data"] @@ -679,7 +640,7 @@ def test_multiple_candidates(sentry_init, capture_events, mock_genai_client): ) (event,) = events - chat_span = event["spans"][1] + chat_span = event["spans"][0] # Should capture all responses # Response text is stored as a JSON string when there are multiple responses @@ -765,7 +726,7 @@ def test_empty_response(sentry_init, capture_events, mock_genai_client): (event,) = events # Should still create spans even with empty candidates - assert len(event["spans"]) == 2 + assert len(event["spans"]) == 1 def test_response_with_different_id_fields( @@ -804,7 +765,7 @@ def test_response_with_different_id_fields( ) (event,) = events - chat_span = event["spans"][1] + chat_span = event["spans"][0] assert chat_span["data"][SPANDATA.GEN_AI_RESPONSE_ID] == "resp-456" assert chat_span["data"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "gemini-1.5-flash-001" @@ -916,7 +877,7 @@ def test_tool_calls_extraction(sentry_init, capture_events, mock_genai_client): ) (event,) = events - chat_span = event["spans"][1] # The chat span + chat_span = event["spans"][0] # The chat span # Check that tool calls are extracted and stored assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS in chat_span["data"]