Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
alexander-alderman-webb wants to merge 1 commit into
base: master
Choose a base branch
from
Draft

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
126 changes: 37 additions & 89 deletions sentry_sdk/integrations/google_genai/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}",
Expand Down Expand Up @@ -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
Expand All @@ -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}",
Expand Down Expand Up @@ -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
Expand All @@ -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

Expand All @@ -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

Expand Down
57 changes: 9 additions & 48 deletions tests/integrations/google_genai/test_google_genai.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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(
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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"]
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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"
Expand Down Expand Up @@ -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"]
Expand Down
Loading