fix(google-genai): Remove agent spans for simple requests

Author: alexander-alderman-webb

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
 

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"]