From 8eacaadc25d152ef45b9c6c3297b3481c4fe1814 Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Thu, 21 May 2026 10:49:32 +0200 Subject: [PATCH 1/6] feat(google-genai): Support span streaming --- .../integrations/google_genai/__init__.py | 341 ++++++++++++------ .../integrations/google_genai/streaming.py | 29 +- sentry_sdk/integrations/google_genai/utils.py | 95 +++-- sentry_sdk/tracing_utils.py | 9 + .../google_genai/test_google_genai.py | 193 ++++++++-- 5 files changed, 492 insertions(+), 175 deletions(-) diff --git a/sentry_sdk/integrations/google_genai/__init__.py b/sentry_sdk/integrations/google_genai/__init__.py index 6dd6e8d0fd..b7d4502f71 100644 --- a/sentry_sdk/integrations/google_genai/__init__.py +++ b/sentry_sdk/integrations/google_genai/__init__.py @@ -11,7 +11,9 @@ from sentry_sdk.ai.utils import get_start_span_function from sentry_sdk.consts import OP, SPANDATA from sentry_sdk.integrations import DidNotEnable, Integration +from sentry_sdk.traces import SpanStatus, StreamedSpan from sentry_sdk.tracing import SPANSTATUS +from sentry_sdk.tracing_utils import has_span_streaming_enabled try: from google.genai.models import AsyncModels, Models @@ -66,23 +68,37 @@ def _wrap_generate_content_stream(f: "Callable[..., Any]") -> "Callable[..., Any def new_generate_content_stream( self: "Any", *args: "Any", **kwargs: "Any" ) -> "Any": - integration = sentry_sdk.get_client().get_integration(GoogleGenAIIntegration) + client = sentry_sdk.get_client() + integration = client.get_integration(GoogleGenAIIntegration) if integration is None: return f(self, *args, **kwargs) _model, contents, model_name = prepare_generate_content_args(args, kwargs) - chat_span = get_start_span_function()( - op=OP.GEN_AI_CHAT, - name=f"chat {model_name}", - origin=ORIGIN, - ) + if has_span_streaming_enabled(client.options): + chat_span = sentry_sdk.traces.start_span( + name=f"chat {model_name}", + attributes={ + "sentry.op": OP.GEN_AI_CHAT, + "sentry.origin": ORIGIN, + }, + ) + + set_on_span = chat_span.set_attribute + else: + chat_span = get_start_span_function()( + op=OP.GEN_AI_CHAT, + name=f"chat {model_name}", + origin=ORIGIN, + ) + + set_on_span = chat_span.set_data chat_span.__enter__() - 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_on_span(SPANDATA.GEN_AI_OPERATION_NAME, "chat") + set_on_span(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM) + set_on_span(SPANDATA.GEN_AI_REQUEST_MODEL, model_name) set_span_data_for_request(chat_span, integration, model_name, contents, kwargs) - chat_span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, True) + set_on_span(SPANDATA.GEN_AI_RESPONSE_STREAMING, True) try: stream = f(self, *args, **kwargs) @@ -96,7 +112,10 @@ def new_iterator() -> "Iterator[Any]": yield chunk except Exception as exc: _capture_exception(exc) - chat_span.set_status(SPANSTATUS.INTERNAL_ERROR) + if isinstance(chat_span, StreamedSpan): + chat_span.status = SpanStatus.ERROR + else: + chat_span.set_status(SPANSTATUS.INTERNAL_ERROR) raise finally: # Accumulate all chunks and set final response data on spans @@ -124,23 +143,37 @@ def _wrap_async_generate_content_stream( async def new_async_generate_content_stream( self: "Any", *args: "Any", **kwargs: "Any" ) -> "Any": - integration = sentry_sdk.get_client().get_integration(GoogleGenAIIntegration) + client = sentry_sdk.get_client() + integration = client.get_integration(GoogleGenAIIntegration) if integration is None: return await f(self, *args, **kwargs) _model, contents, model_name = prepare_generate_content_args(args, kwargs) - chat_span = get_start_span_function()( - op=OP.GEN_AI_CHAT, - name=f"chat {model_name}", - origin=ORIGIN, - ) + if has_span_streaming_enabled(client.options): + chat_span = get_start_span_function()( + name=f"chat {model_name}", + attributes={ + "sentry.op": OP.GEN_AI_CHAT, + "sentry.origin": ORIGIN, + }, + ) + + set_on_span = chat_span.set_attribute + else: + chat_span = get_start_span_function()( + op=OP.GEN_AI_CHAT, + name=f"chat {model_name}", + origin=ORIGIN, + ) + + set_on_span = chat_span.set_data chat_span.__enter__() - 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_on_span(SPANDATA.GEN_AI_OPERATION_NAME, "chat") + set_on_span(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM) + set_on_span(SPANDATA.GEN_AI_REQUEST_MODEL, model_name) set_span_data_for_request(chat_span, integration, model_name, contents, kwargs) - chat_span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, True) + set_on_span(SPANDATA.GEN_AI_RESPONSE_STREAMING, True) try: stream = await f(self, *args, **kwargs) @@ -154,7 +187,10 @@ async def new_async_iterator() -> "AsyncIterator[Any]": yield chunk except Exception as exc: _capture_exception(exc) - chat_span.set_status(SPANSTATUS.INTERNAL_ERROR) + if isinstance(chat_span, StreamedSpan): + chat_span.status = SpanStatus.ERROR + else: + chat_span.set_status(SPANSTATUS.INTERNAL_ERROR) raise finally: # Accumulate all chunks and set final response data on spans @@ -178,34 +214,61 @@ async def new_async_iterator() -> "AsyncIterator[Any]": def _wrap_generate_content(f: "Callable[..., Any]") -> "Callable[..., Any]": @wraps(f) def new_generate_content(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": - integration = sentry_sdk.get_client().get_integration(GoogleGenAIIntegration) + client = sentry_sdk.get_client() + integration = client.get_integration(GoogleGenAIIntegration) if integration is None: return f(self, *args, **kwargs) model, contents, model_name = prepare_generate_content_args(args, kwargs) - with get_start_span_function()( - 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 - ) + if has_span_streaming_enabled(client.options): + with sentry_sdk.traces.start_span( + name=f"chat {model_name}", + attributes={ + "sentry.op": OP.GEN_AI_CHAT, + "sentry.origin": ORIGIN, + }, + ) as chat_span: + chat_span.set_attribute(SPANDATA.GEN_AI_OPERATION_NAME, "chat") + chat_span.set_attribute(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM) + chat_span.set_attribute(SPANDATA.GEN_AI_REQUEST_MODEL, 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.status = SpanStatus.ERROR + raise + + set_span_data_for_response(chat_span, integration, response) + + return response + else: + with get_start_span_function()( + 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 = 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(chat_span, integration, response) - return response + return response return new_generate_content @@ -215,33 +278,59 @@ def _wrap_async_generate_content(f: "Callable[..., Any]") -> "Callable[..., Any] async def new_async_generate_content( self: "Any", *args: "Any", **kwargs: "Any" ) -> "Any": - integration = sentry_sdk.get_client().get_integration(GoogleGenAIIntegration) + client = sentry_sdk.get_client() + integration = client.get_integration(GoogleGenAIIntegration) if integration is None: return await f(self, *args, **kwargs) model, contents, model_name = prepare_generate_content_args(args, kwargs) - with get_start_span_function()( - 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 + if has_span_streaming_enabled(client.options): + with sentry_sdk.traces.start_span( + name=f"chat {model_name}", + attributes={ + "sentry.op": OP.GEN_AI_CHAT, + "sentry.origin": ORIGIN, + }, + ) as chat_span: + chat_span.set_attribute(SPANDATA.GEN_AI_OPERATION_NAME, "chat") + chat_span.set_attribute(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM) + chat_span.set_attribute(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(chat_span, integration, response) + + return response + else: + with get_start_span_function()( + 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 - return response + set_span_data_for_response(chat_span, integration, response) + + return response return new_async_generate_content @@ -249,32 +338,57 @@ async def new_async_generate_content( def _wrap_embed_content(f: "Callable[..., Any]") -> "Callable[..., Any]": @wraps(f) def new_embed_content(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": - integration = sentry_sdk.get_client().get_integration(GoogleGenAIIntegration) + client = sentry_sdk.get_client() + integration = client.get_integration(GoogleGenAIIntegration) if integration is None: return f(self, *args, **kwargs) model_name, contents = prepare_embed_content_args(args, kwargs) - with get_start_span_function()( - op=OP.GEN_AI_EMBEDDINGS, - name=f"embeddings {model_name}", - origin=ORIGIN, - ) as span: - span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "embeddings") - span.set_data(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM) - span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model_name) - set_span_data_for_embed_request(span, integration, contents, kwargs) + if has_span_streaming_enabled(client.options): + with sentry_sdk.traces.start_span( + name=f"embeddings {model_name}", + attributes={ + "sentry.op": OP.GEN_AI_EMBEDDINGS, + "sentry.origin": ORIGIN, + }, + ) as span: + span.set_attribute(SPANDATA.GEN_AI_OPERATION_NAME, "embeddings") + span.set_attribute(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM) + span.set_attribute(SPANDATA.GEN_AI_REQUEST_MODEL, model_name) + set_span_data_for_embed_request(span, integration, contents, kwargs) - try: - response = f(self, *args, **kwargs) - except Exception as exc: - _capture_exception(exc) - span.set_status(SPANSTATUS.INTERNAL_ERROR) - raise + try: + response = f(self, *args, **kwargs) + except Exception as exc: + _capture_exception(exc) + span.status = SpanStatus.ERROR + raise + + set_span_data_for_embed_response(span, integration, response) + + return response + else: + with get_start_span_function()( + op=OP.GEN_AI_EMBEDDINGS, + name=f"embeddings {model_name}", + origin=ORIGIN, + ) as span: + span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "embeddings") + span.set_data(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM) + span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model_name) + set_span_data_for_embed_request(span, integration, contents, kwargs) + + try: + response = f(self, *args, **kwargs) + except Exception as exc: + _capture_exception(exc) + span.set_status(SPANSTATUS.INTERNAL_ERROR) + raise - set_span_data_for_embed_response(span, integration, response) + set_span_data_for_embed_response(span, integration, response) - return response + return response return new_embed_content @@ -284,31 +398,56 @@ def _wrap_async_embed_content(f: "Callable[..., Any]") -> "Callable[..., Any]": async def new_async_embed_content( self: "Any", *args: "Any", **kwargs: "Any" ) -> "Any": - integration = sentry_sdk.get_client().get_integration(GoogleGenAIIntegration) + client = sentry_sdk.get_client() + integration = client.get_integration(GoogleGenAIIntegration) if integration is None: return await f(self, *args, **kwargs) model_name, contents = prepare_embed_content_args(args, kwargs) - with get_start_span_function()( - op=OP.GEN_AI_EMBEDDINGS, - name=f"embeddings {model_name}", - origin=ORIGIN, - ) as span: - span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "embeddings") - span.set_data(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM) - span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model_name) - set_span_data_for_embed_request(span, integration, contents, kwargs) - - try: - response = await f(self, *args, **kwargs) - except Exception as exc: - _capture_exception(exc) - span.set_status(SPANSTATUS.INTERNAL_ERROR) - raise - - set_span_data_for_embed_response(span, integration, response) - - return response + if has_span_streaming_enabled(client.options): + with sentry_sdk.traces.start_span( + name=f"embeddings {model_name}", + attributes={ + "sentry.op": OP.GEN_AI_EMBEDDINGS, + "sentry.origin": ORIGIN, + }, + ) as span: + span.set_attribute(SPANDATA.GEN_AI_OPERATION_NAME, "embeddings") + span.set_attribute(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM) + span.set_attribute(SPANDATA.GEN_AI_REQUEST_MODEL, model_name) + set_span_data_for_embed_request(span, integration, contents, kwargs) + + try: + response = await f(self, *args, **kwargs) + except Exception as exc: + _capture_exception(exc) + span.status = SpanStatus.ERROR + raise + + set_span_data_for_embed_response(span, integration, response) + + return response + else: + with get_start_span_function()( + op=OP.GEN_AI_EMBEDDINGS, + name=f"embeddings {model_name}", + origin=ORIGIN, + ) as span: + span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "embeddings") + span.set_data(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM) + span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model_name) + set_span_data_for_embed_request(span, integration, contents, kwargs) + + try: + response = await f(self, *args, **kwargs) + except Exception as exc: + _capture_exception(exc) + span.set_status(SPANSTATUS.INTERNAL_ERROR) + raise + + set_span_data_for_embed_response(span, integration, response) + + return response return new_async_embed_content diff --git a/sentry_sdk/integrations/google_genai/streaming.py b/sentry_sdk/integrations/google_genai/streaming.py index 7f3f58dc93..b54723e115 100644 --- a/sentry_sdk/integrations/google_genai/streaming.py +++ b/sentry_sdk/integrations/google_genai/streaming.py @@ -1,8 +1,9 @@ -from typing import TYPE_CHECKING, Any, List, Optional, TypedDict +from typing import TYPE_CHECKING, Any, List, Optional, TypedDict, Union 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 ( safe_serialize, ) @@ -97,15 +98,21 @@ def accumulate_streaming_response( def set_span_data_for_streaming_response( - span: "Span", integration: "Any", accumulated_response: "AccumulatedResponse" + span: "Union[Span, StreamedSpan]", + integration: "Any", + accumulated_response: "AccumulatedResponse", ) -> None: """Set span data for accumulated 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") ): - span.set_data( + set_on_span( SPANDATA.GEN_AI_RESPONSE_TEXT, safe_serialize([accumulated_response["text"]]), ) @@ -118,45 +125,45 @@ def set_span_data_for_streaming_response( ) if accumulated_response.get("tool_calls"): - span.set_data( + set_on_span( SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, safe_serialize(accumulated_response["tool_calls"]), ) if accumulated_response.get("id"): - span.set_data(SPANDATA.GEN_AI_RESPONSE_ID, accumulated_response["id"]) + set_on_span(SPANDATA.GEN_AI_RESPONSE_ID, accumulated_response["id"]) if accumulated_response.get("model"): - span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, accumulated_response["model"]) + set_on_span(SPANDATA.GEN_AI_RESPONSE_MODEL, accumulated_response["model"]) if accumulated_response["usage_metadata"] is None: return if accumulated_response["usage_metadata"]["input_tokens"]: - span.set_data( + set_on_span( SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, accumulated_response["usage_metadata"]["input_tokens"], ) if accumulated_response["usage_metadata"]["input_tokens_cached"]: - span.set_data( + set_on_span( SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED, accumulated_response["usage_metadata"]["input_tokens_cached"], ) if accumulated_response["usage_metadata"]["output_tokens"]: - span.set_data( + set_on_span( SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, accumulated_response["usage_metadata"]["output_tokens"], ) if accumulated_response["usage_metadata"]["output_tokens_reasoning"]: - span.set_data( + set_on_span( SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING, accumulated_response["usage_metadata"]["output_tokens_reasoning"], ) if accumulated_response["usage_metadata"]["total_tokens"]: - span.set_data( + set_on_span( SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, accumulated_response["usage_metadata"]["total_tokens"], ) diff --git a/sentry_sdk/integrations/google_genai/utils.py b/sentry_sdk/integrations/google_genai/utils.py index 7a30a56aa7..ee7b8b5e5c 100644 --- a/sentry_sdk/integrations/google_genai/utils.py +++ b/sentry_sdk/integrations/google_genai/utils.py @@ -28,6 +28,8 @@ ) from sentry_sdk.consts import OP, SPANDATA from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.traces import StreamedSpan +from sentry_sdk.tracing_utils import has_span_streaming_enabled from sentry_sdk.utils import ( capture_internal_exceptions, event_from_exception, @@ -664,8 +666,24 @@ def _capture_tool_input( return tool_input -def _create_tool_span(tool_name: str, tool_doc: "Optional[str]") -> "Span": +def _create_tool_span( + tool_name: str, tool_doc: "Optional[str]" +) -> "Union[Span, StreamedSpan]": """Create a span for tool execution.""" + span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options) + if span_streaming: + span = sentry_sdk.traces.start_span( + name=f"execute_tool {tool_name}", + attributes={ + "sentry.op": OP.GEN_AI_EXECUTE_TOOL, + "sentry.origin": ORIGIN, + }, + ) + span.set_attribute(SPANDATA.GEN_AI_TOOL_NAME, tool_name) + if tool_doc: + span.set_attribute(SPANDATA.GEN_AI_TOOL_DESCRIPTION, tool_doc) + return span + span = sentry_sdk.start_span( op=OP.GEN_AI_EXECUTE_TOOL, name=f"execute_tool {tool_name}", @@ -691,21 +709,22 @@ def wrapped_tool(tool: "Tool | Callable[..., Any]") -> "Tool | Callable[..., Any @wraps(tool) async def async_wrapped(*args: "Any", **kwargs: "Any") -> "Any": with _create_tool_span(tool_name, tool_doc) as span: + set_on_span = ( + span.set_attribute + if isinstance(span, StreamedSpan) + else span.set_data + ) # Capture tool input tool_input = _capture_tool_input(args, kwargs, tool) with capture_internal_exceptions(): - span.set_data( - SPANDATA.GEN_AI_TOOL_INPUT, safe_serialize(tool_input) - ) + set_on_span(SPANDATA.GEN_AI_TOOL_INPUT, safe_serialize(tool_input)) try: result = await tool(*args, **kwargs) # Capture tool output with capture_internal_exceptions(): - span.set_data( - SPANDATA.GEN_AI_TOOL_OUTPUT, safe_serialize(result) - ) + set_on_span(SPANDATA.GEN_AI_TOOL_OUTPUT, safe_serialize(result)) return result except Exception as exc: @@ -718,21 +737,22 @@ async def async_wrapped(*args: "Any", **kwargs: "Any") -> "Any": @wraps(tool) def sync_wrapped(*args: "Any", **kwargs: "Any") -> "Any": with _create_tool_span(tool_name, tool_doc) as span: + set_on_span = ( + span.set_attribute + if isinstance(span, StreamedSpan) + else span.set_data + ) # Capture tool input tool_input = _capture_tool_input(args, kwargs, tool) with capture_internal_exceptions(): - span.set_data( - SPANDATA.GEN_AI_TOOL_INPUT, safe_serialize(tool_input) - ) + set_on_span(SPANDATA.GEN_AI_TOOL_INPUT, safe_serialize(tool_input)) try: result = tool(*args, **kwargs) # Capture tool output with capture_internal_exceptions(): - span.set_data( - SPANDATA.GEN_AI_TOOL_OUTPUT, safe_serialize(result) - ) + set_on_span(SPANDATA.GEN_AI_TOOL_OUTPUT, safe_serialize(result)) return result except Exception as exc: @@ -857,18 +877,21 @@ def _transform_system_instructions( def set_span_data_for_request( - span: "Span", + span: "Union[Span, StreamedSpan]", integration: "Any", model: str, contents: "ContentListUnion", kwargs: "dict[str, Any]", ) -> None: """Set span data for the request.""" - span.set_data(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM) - span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model) + set_on_span = ( + span.set_attribute if isinstance(span, StreamedSpan) else span.set_data + ) + set_on_span(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM) + set_on_span(SPANDATA.GEN_AI_REQUEST_MODEL, model) if kwargs.get("stream", False): - span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, True) + set_on_span(SPANDATA.GEN_AI_RESPONSE_STREAMING, True) config: "Optional[GenerateContentConfig]" = kwargs.get("config") @@ -884,7 +907,7 @@ def set_span_data_for_request( system_instructions = config.get("system_instruction") if system_instructions is not None: - span.set_data( + set_on_span( SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, json.dumps(_transform_system_instructions(system_instructions)), ) @@ -923,7 +946,7 @@ def set_span_data_for_request( if hasattr(config, param): value = getattr(config, param) if value is not None: - span.set_data(span_key, value) + set_on_span(span_key, value) # Set tools if available if config is not None and hasattr(config, "tools"): @@ -940,22 +963,27 @@ def set_span_data_for_request( def set_span_data_for_response( - span: "Span", integration: "Any", response: "GenerateContentResponse" + span: "Union[Span, StreamedSpan]", + integration: "Any", + response: "GenerateContentResponse", ) -> None: """Set span data for the response.""" if not response: return + 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 - span.set_data(SPANDATA.GEN_AI_RESPONSE_TEXT, safe_serialize(response_texts)) + 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 - span.set_data(SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, safe_serialize(tool_calls)) + set_on_span(SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, safe_serialize(tool_calls)) finish_reasons = extract_finish_reasons(response) if finish_reasons: @@ -964,33 +992,33 @@ def set_span_data_for_response( ) if getattr(response, "response_id", None): - span.set_data(SPANDATA.GEN_AI_RESPONSE_ID, response.response_id) + set_on_span(SPANDATA.GEN_AI_RESPONSE_ID, response.response_id) if getattr(response, "model_version", None): - span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, response.model_version) + set_on_span(SPANDATA.GEN_AI_RESPONSE_MODEL, response.model_version) usage_data = extract_usage_data(response) if usage_data["input_tokens"]: - span.set_data(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, usage_data["input_tokens"]) + set_on_span(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, usage_data["input_tokens"]) if usage_data["input_tokens_cached"]: - span.set_data( + set_on_span( SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED, usage_data["input_tokens_cached"], ) if usage_data["output_tokens"]: - span.set_data(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, usage_data["output_tokens"]) + set_on_span(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, usage_data["output_tokens"]) if usage_data["output_tokens_reasoning"]: - span.set_data( + set_on_span( SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING, usage_data["output_tokens_reasoning"], ) if usage_data["total_tokens"]: - span.set_data(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, usage_data["total_tokens"]) + set_on_span(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, usage_data["total_tokens"]) def prepare_generate_content_args( @@ -1057,7 +1085,9 @@ def set_span_data_for_embed_request( def set_span_data_for_embed_response( - span: "Span", integration: "Any", response: "EmbedContentResponse" + span: "Union[Span, StreamedSpan]", + integration: "Any", + response: "EmbedContentResponse", ) -> None: """Set span data for embedding response.""" if not response: @@ -1076,4 +1106,7 @@ def set_span_data_for_embed_response( # Set token count if we found any if total_tokens > 0: - span.set_data(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, total_tokens) + if isinstance(span, StreamedSpan): + span.set_attribute(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, total_tokens) + else: + span.set_data(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, total_tokens) diff --git a/sentry_sdk/tracing_utils.py b/sentry_sdk/tracing_utils.py index e6fc8770d6..822114628a 100644 --- a/sentry_sdk/tracing_utils.py +++ b/sentry_sdk/tracing_utils.py @@ -116,6 +116,15 @@ def has_span_streaming_enabled(options: "Optional[dict[str, Any]]") -> bool: return (options.get("_experiments") or {}).get("trace_lifecycle") == "stream" +def should_truncate_gen_ai_input(options: "Optional[dict[str, Any]]") -> bool: + if options is None: + return True + + return not options.get( + "stream_gen_ai_spans", False + ) and not has_span_streaming_enabled(options) + + @contextlib.contextmanager def record_sql_queries( cursor: "Any", diff --git a/tests/integrations/google_genai/test_google_genai.py b/tests/integrations/google_genai/test_google_genai.py index 409e9eabb5..fb50c9cc7d 100644 --- a/tests/integrations/google_genai/test_google_genai.py +++ b/tests/integrations/google_genai/test_google_genai.py @@ -6,6 +6,7 @@ from google.genai import types as genai_types from google.genai.types import Content, Part +import sentry_sdk from sentry_sdk import start_transaction from sentry_sdk._types import BLOB_DATA_SUBSTITUTE from sentry_sdk.consts import OP, SPANDATA @@ -114,6 +115,7 @@ def create_test_config( return genai_types.GenerateContentConfig(**config_dict) +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.parametrize( "send_default_pii, include_prompts", @@ -132,18 +134,20 @@ def test_nonstreaming_generate_content( include_prompts, mock_genai_client, stream_gen_ai_spans, + span_streaming, ): sentry_init( 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, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) # Mock the HTTP response at the _api_client.request() level mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("transaction", "span") with mock.patch.object( @@ -164,8 +168,10 @@ def test_nonstreaming_generate_content( (event,) = (item.payload for item in items if item.type == "transaction") assert event["transaction"] == "google_genai" + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] assert len(spans) == 1 + sentry_sdk.flush() chat_span = next(item.payload for item in items if item.type == "span") # Check chat span @@ -261,6 +267,7 @@ def test_nonstreaming_generate_content( assert chat_span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING] == 3 +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.parametrize("generate_content_config", (False, True)) @pytest.mark.parametrize( @@ -299,17 +306,19 @@ def test_generate_content_with_system_instruction( system_instructions, expected_texts, stream_gen_ai_spans, + span_streaming, ): sentry_init( integrations=[GoogleGenAIIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("span") with mock.patch.object( @@ -329,6 +338,7 @@ def test_generate_content_with_system_instruction( config=config, ) + sentry_sdk.flush() invoke_span = next(item.payload for item in items if item.type == "span") if expected_texts is None: @@ -376,6 +386,7 @@ def test_generate_content_with_system_instruction( ] +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_generate_content_with_tools( sentry_init, @@ -383,11 +394,13 @@ def test_generate_content_with_tools( capture_items, mock_genai_client, stream_gen_ai_spans, + span_streaming, ): sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) # Create a mock tool function @@ -433,7 +446,7 @@ def get_weather(location: str) -> str: mock_http_response = create_mock_http_response(tool_response_json) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("span") with mock.patch.object( @@ -444,6 +457,7 @@ def get_weather(location: str) -> str: model="gemini-1.5-flash", contents="What's the weather?", config=config ) + sentry_sdk.flush() invoke_span = next(item.payload for item in items if item.type == "span") # Check that tools are recorded (data is serialized as a string) @@ -492,18 +506,21 @@ def get_weather(location: str) -> str: assert sorted_tools[1]["description"] == "Get weather information (tool object)" +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_tool_execution( sentry_init, capture_events, capture_items, stream_gen_ai_spans, + span_streaming, ): sentry_init( integrations=[GoogleGenAIIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) # Create a mock tool function @@ -516,7 +533,7 @@ def get_weather(location: str) -> str: wrapped_weather = wrapped_tool(get_weather) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("span") # Execute the wrapped tool @@ -525,8 +542,10 @@ def get_weather(location: str) -> str: assert result == "The weather in San Francisco is sunny" + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] assert len(spans) == 1 + sentry_sdk.flush() tool_span = next(item.payload for item in items if item.type == "span") assert tool_span["attributes"]["sentry.op"] == OP.GEN_AI_EXECUTE_TOOL @@ -558,6 +577,7 @@ def get_weather(location: str) -> str: ) +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_error_handling( sentry_init, @@ -565,13 +585,15 @@ def test_error_handling( capture_items, mock_genai_client, stream_gen_ai_spans, + span_streaming, ): sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("event", "transaction") # Mock an error at the HTTP level @@ -612,6 +634,7 @@ def test_error_handling( assert error_event["exception"]["values"][0]["mechanism"]["type"] == "google_genai" +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_streaming_generate_content( sentry_init, @@ -619,6 +642,7 @@ def test_streaming_generate_content( capture_items, mock_genai_client, stream_gen_ai_spans, + span_streaming, ): """Test streaming with generate_content_stream, verifying chunk accumulation.""" sentry_init( @@ -626,6 +650,7 @@ def test_streaming_generate_content( traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) # Create streaming chunks - simulating a multi-chunk response @@ -690,7 +715,7 @@ def test_streaming_generate_content( stream_chunks = [chunk1_json, chunk2_json, chunk3_json] mock_stream = create_mock_streaming_responses(stream_chunks) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("span") with mock.patch.object( @@ -717,8 +742,10 @@ def test_streaming_generate_content( collected_chunks[2].candidates[0].content.parts[0].text == "help you today?" ) + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] assert len(spans) == 1 + sentry_sdk.flush() chat_span = next(item.payload for item in items if item.type == "span") assert json.loads( @@ -821,6 +848,7 @@ def test_streaming_generate_content( assert chat_span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "gemini-1.5-flash" +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_span_origin( sentry_init, @@ -828,16 +856,18 @@ def test_span_origin( capture_items, mock_genai_client, stream_gen_ai_spans, + span_streaming, ): sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("span", "transaction") with mock.patch.object( @@ -851,6 +881,7 @@ def test_span_origin( (event,) = (item.payload for item in items if item.type == "transaction") assert event["contexts"]["trace"]["origin"] == "manual" + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] for span in spans: assert span["attributes"]["sentry.origin"] == "auto.ai.google_genai" @@ -872,6 +903,7 @@ def test_span_origin( assert span["origin"] == "auto.ai.google_genai" +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_response_without_usage_metadata( sentry_init, @@ -879,12 +911,14 @@ def test_response_without_usage_metadata( capture_items, mock_genai_client, stream_gen_ai_spans, + span_streaming, ): """Test handling of responses without usage metadata""" sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) # Response without usage metadata @@ -902,7 +936,7 @@ def test_response_without_usage_metadata( mock_http_response = create_mock_http_response(response_json) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("span") with mock.patch.object( @@ -913,6 +947,7 @@ def test_response_without_usage_metadata( model="gemini-1.5-flash", contents="Test", config=config ) + sentry_sdk.flush() chat_span = next(item.payload for item in items if item.type == "span") # Usage data should not be present @@ -939,6 +974,7 @@ def test_response_without_usage_metadata( assert SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS not in chat_span["data"] +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_multiple_candidates( sentry_init, @@ -946,6 +982,7 @@ def test_multiple_candidates( capture_items, mock_genai_client, stream_gen_ai_spans, + span_streaming, ): """Test handling of multiple response candidates""" sentry_init( @@ -953,6 +990,7 @@ def test_multiple_candidates( traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) # Response with multiple candidates @@ -982,7 +1020,7 @@ def test_multiple_candidates( mock_http_response = create_mock_http_response(multi_candidate_json) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("span") with mock.patch.object( @@ -993,6 +1031,7 @@ def test_multiple_candidates( model="gemini-1.5-flash", contents="Generate multiple", config=config ) + sentry_sdk.flush() chat_span = next(item.payload for item in items if item.type == "span") # Should capture all responses @@ -1044,6 +1083,7 @@ def test_multiple_candidates( assert finish_reasons == ["STOP", "MAX_TOKENS"] +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_all_configuration_parameters( sentry_init, @@ -1051,17 +1091,19 @@ def test_all_configuration_parameters( capture_items, mock_genai_client, stream_gen_ai_spans, + span_streaming, ): """Test that all configuration parameters are properly recorded""" sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("span") with mock.patch.object( @@ -1080,6 +1122,7 @@ def test_all_configuration_parameters( model="gemini-1.5-flash", contents="Test all params", config=config ) + sentry_sdk.flush() invoke_span = next(item.payload for item in items if item.type == "span") # Check all parameters are recorded @@ -1126,6 +1169,7 @@ def test_all_configuration_parameters( assert invoke_span["data"][SPANDATA.GEN_AI_REQUEST_SEED] == 12345 +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_empty_response( sentry_init, @@ -1133,19 +1177,21 @@ def test_empty_response( capture_items, mock_genai_client, stream_gen_ai_spans, + span_streaming, ): """Test handling of minimal response with no content""" sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) # Minimal response with empty candidates array minimal_response_json = {"candidates": []} mock_http_response = create_mock_http_response(minimal_response_json) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("span") with mock.patch.object( @@ -1160,6 +1206,7 @@ def test_empty_response( assert len(response.candidates) == 0 # Should still create spans even with empty candidates + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] assert len(spans) == 1 else: @@ -1181,6 +1228,7 @@ def test_empty_response( assert len(event["spans"]) == 1 +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_response_with_different_id_fields( sentry_init, @@ -1188,12 +1236,14 @@ def test_response_with_different_id_fields( capture_items, mock_genai_client, stream_gen_ai_spans, + span_streaming, ): """Test handling of different response ID field names""" sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) # Response with response_id and model_version @@ -1213,7 +1263,7 @@ def test_response_with_different_id_fields( mock_http_response = create_mock_http_response(response_json) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("span") with mock.patch.object( @@ -1223,6 +1273,7 @@ def test_response_with_different_id_fields( model="gemini-1.5-flash", contents="Test", config=create_test_config() ) + sentry_sdk.flush() chat_span = next(item.payload for item in items if item.type == "span") assert chat_span["attributes"][SPANDATA.GEN_AI_RESPONSE_ID] == "resp-456" @@ -1270,6 +1321,7 @@ async def async_tool(param: str) -> str: assert hasattr(wrapped_async_tool, "__wrapped__") # Should preserve original +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_contents_as_none( sentry_init, @@ -1277,6 +1329,7 @@ def test_contents_as_none( capture_items, mock_genai_client, stream_gen_ai_spans, + span_streaming, ): """Test handling when contents parameter is None""" sentry_init( @@ -1284,11 +1337,12 @@ def test_contents_as_none( traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("span") with mock.patch.object( @@ -1298,6 +1352,7 @@ def test_contents_as_none( model="gemini-1.5-flash", contents=None, config=create_test_config() ) + sentry_sdk.flush() invoke_span = next(item.payload for item in items if item.type == "span") # Should handle None contents gracefully @@ -1322,6 +1377,7 @@ def test_contents_as_none( assert all(msg["role"] != "user" or msg["content"] is not None for msg in messages) +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_tool_calls_extraction( sentry_init, @@ -1329,12 +1385,14 @@ def test_tool_calls_extraction( capture_items, mock_genai_client, stream_gen_ai_spans, + span_streaming, ): """Test extraction of tool/function calls from response""" sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) # Response with function calls @@ -1374,7 +1432,7 @@ def test_tool_calls_extraction( mock_http_response = create_mock_http_response(function_call_response_json) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("span") with mock.patch.object( @@ -1386,6 +1444,7 @@ def test_tool_calls_extraction( config=create_test_config(), ) + sentry_sdk.flush() chat_span = next( item.payload for item in items if item.type == "span" ) # The chat span @@ -1504,6 +1563,7 @@ def test_google_genai_message_truncation( } +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.parametrize( "send_default_pii, include_prompts", @@ -1522,18 +1582,20 @@ def test_embed_content( include_prompts, mock_genai_client, stream_gen_ai_spans, + span_streaming, ): sentry_init( 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, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) # Mock the HTTP response at the _api_client.request() level mock_http_response = create_mock_http_response(EXAMPLE_EMBED_RESPONSE_JSON) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("transaction", "span") with mock.patch.object( @@ -1554,6 +1616,7 @@ def test_embed_content( assert event["transaction"] == "google_genai_embeddings" # Should have 1 span for embeddings + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] assert len(spans) == 1 (embed_span,) = spans @@ -1635,6 +1698,7 @@ def test_embed_content( assert embed_span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 25 +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_embed_content_string_input( sentry_init, @@ -1642,6 +1706,7 @@ def test_embed_content_string_input( capture_items, mock_genai_client, stream_gen_ai_spans, + span_streaming, ): """Test embed_content with a single string instead of list.""" sentry_init( @@ -1649,6 +1714,7 @@ def test_embed_content_string_input( traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) # Mock response with single embedding @@ -1668,7 +1734,7 @@ def test_embed_content_string_input( } mock_http_response = create_mock_http_response(single_embed_response) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("span") with mock.patch.object( @@ -1679,6 +1745,7 @@ def test_embed_content_string_input( contents="Single text input", ) + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] (embed_span,) = spans @@ -1715,6 +1782,7 @@ def test_embed_content_string_input( assert embed_span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 5 +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_embed_content_error_handling( sentry_init, @@ -1722,14 +1790,16 @@ def test_embed_content_error_handling( capture_items, mock_genai_client, stream_gen_ai_spans, + span_streaming, ): """Test error handling in embed_content.""" sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("transaction", "event") # Mock an error at the HTTP level @@ -1772,6 +1842,7 @@ def test_embed_content_error_handling( assert error_event["exception"]["values"][0]["mechanism"]["type"] == "google_genai" +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_embed_content_without_statistics( sentry_init, @@ -1779,12 +1850,14 @@ def test_embed_content_without_statistics( capture_items, mock_genai_client, stream_gen_ai_spans, + span_streaming, ): """Test embed_content response without statistics (older package versions).""" sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) # Response without statistics (typical for older google-genai versions) @@ -1801,7 +1874,7 @@ def test_embed_content_without_statistics( } mock_http_response = create_mock_http_response(old_version_response) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("span") with mock.patch.object( @@ -1812,6 +1885,7 @@ def test_embed_content_without_statistics( contents=["Test without statistics", "Another test"], ) + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] (embed_span,) = spans @@ -1837,6 +1911,7 @@ def test_embed_content_without_statistics( assert SPANDATA.GEN_AI_USAGE_INPUT_TOKENS not in embed_span["data"] +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_embed_content_span_origin( sentry_init, @@ -1844,16 +1919,18 @@ def test_embed_content_span_origin( capture_items, mock_genai_client, stream_gen_ai_spans, + span_streaming, ): """Test that embed_content spans have correct origin.""" sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) mock_http_response = create_mock_http_response(EXAMPLE_EMBED_RESPONSE_JSON) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("transaction", "span") with mock.patch.object( mock_genai_client._api_client, "request", return_value=mock_http_response @@ -1866,6 +1943,7 @@ def test_embed_content_span_origin( (event,) = (item.payload for item in items if item.type == "transaction") assert event["contexts"]["trace"]["origin"] == "manual" + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] for span in spans: assert span["attributes"]["sentry.origin"] == "auto.ai.google_genai" @@ -1886,6 +1964,7 @@ def test_embed_content_span_origin( assert span["origin"] == "auto.ai.google_genai" +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio @pytest.mark.parametrize( @@ -1905,6 +1984,7 @@ async def test_async_embed_content( include_prompts, mock_genai_client, stream_gen_ai_spans, + span_streaming, ): """Test async embed_content method.""" sentry_init( @@ -1912,12 +1992,13 @@ async def test_async_embed_content( traces_sample_rate=1.0, send_default_pii=send_default_pii, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) # Mock the async HTTP response mock_http_response = create_mock_http_response(EXAMPLE_EMBED_RESPONSE_JSON) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("transaction", "span") with mock.patch.object( @@ -1938,6 +2019,7 @@ async def test_async_embed_content( assert event["transaction"] == "google_genai_embeddings_async" # Should have 1 span for embeddings + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] assert len(spans) == 1 (embed_span,) = spans @@ -2020,6 +2102,7 @@ async def test_async_embed_content( assert embed_span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 25 +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio async def test_async_embed_content_string_input( @@ -2028,6 +2111,7 @@ async def test_async_embed_content_string_input( capture_items, mock_genai_client, stream_gen_ai_spans, + span_streaming, ): """Test async embed_content with a single string instead of list.""" sentry_init( @@ -2035,6 +2119,7 @@ async def test_async_embed_content_string_input( traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) # Mock response with single embedding @@ -2054,7 +2139,7 @@ async def test_async_embed_content_string_input( } mock_http_response = create_mock_http_response(single_embed_response) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("span") with mock.patch.object( @@ -2067,6 +2152,7 @@ async def test_async_embed_content_string_input( contents="Single text input", ) + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] (embed_span,) = spans @@ -2104,6 +2190,7 @@ async def test_async_embed_content_string_input( assert embed_span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 5 +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio async def test_async_embed_content_error_handling( @@ -2112,15 +2199,17 @@ async def test_async_embed_content_error_handling( capture_items, mock_genai_client, stream_gen_ai_spans, + span_streaming, ): """Test error handling in async embed_content.""" sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("transaction", "event") # Mock an error at the HTTP level @@ -2163,6 +2252,7 @@ async def test_async_embed_content_error_handling( assert error_event["exception"]["values"][0]["mechanism"]["type"] == "google_genai" +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio async def test_async_embed_content_without_statistics( @@ -2171,12 +2261,14 @@ async def test_async_embed_content_without_statistics( capture_items, mock_genai_client, stream_gen_ai_spans, + span_streaming, ): """Test async embed_content response without statistics (older package versions).""" sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) # Response without statistics (typical for older google-genai versions) @@ -2193,7 +2285,7 @@ async def test_async_embed_content_without_statistics( } mock_http_response = create_mock_http_response(old_version_response) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("span") with mock.patch.object( @@ -2206,6 +2298,7 @@ async def test_async_embed_content_without_statistics( contents=["Test without statistics", "Another test"], ) + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] (embed_span,) = spans @@ -2233,6 +2326,7 @@ async def test_async_embed_content_without_statistics( assert SPANDATA.GEN_AI_USAGE_INPUT_TOKENS not in embed_span["data"] +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio async def test_async_embed_content_span_origin( @@ -2241,17 +2335,19 @@ async def test_async_embed_content_span_origin( capture_items, mock_genai_client, stream_gen_ai_spans, + span_streaming, ): """Test that async embed_content spans have correct origin.""" sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) mock_http_response = create_mock_http_response(EXAMPLE_EMBED_RESPONSE_JSON) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("transaction", "span") with mock.patch.object( @@ -2267,6 +2363,7 @@ async def test_async_embed_content_span_origin( (event,) = [item.payload for item in items if item.type == "transaction"] assert event["contexts"]["trace"]["origin"] == "manual" + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] for span in spans: assert span["attributes"]["sentry.origin"] == "auto.ai.google_genai" @@ -2291,6 +2388,7 @@ async def test_async_embed_content_span_origin( # Integration tests for generate_content with different input message formats +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_generate_content_with_content_object( sentry_init, @@ -2298,6 +2396,7 @@ def test_generate_content_with_content_object( capture_items, mock_genai_client, stream_gen_ai_spans, + span_streaming, ): """Test generate_content with Content object input.""" sentry_init( @@ -2305,6 +2404,7 @@ def test_generate_content_with_content_object( traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) @@ -2314,7 +2414,7 @@ def test_generate_content_with_content_object( role="user", parts=[genai_types.Part(text="Hello from Content object")] ) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("span") with mock.patch.object( @@ -2324,6 +2424,7 @@ def test_generate_content_with_content_object( model="gemini-1.5-flash", contents=content, config=create_test_config() ) + sentry_sdk.flush() invoke_span = next(item.payload for item in items if item.type == "span") messages = json.loads( @@ -2351,6 +2452,7 @@ def test_generate_content_with_content_object( ] +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_generate_content_with_dict_format( sentry_init, @@ -2358,6 +2460,7 @@ def test_generate_content_with_dict_format( capture_items, mock_genai_client, stream_gen_ai_spans, + span_streaming, ): """Test generate_content with dict format input (ContentDict).""" sentry_init( @@ -2365,6 +2468,7 @@ def test_generate_content_with_dict_format( traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) @@ -2372,7 +2476,7 @@ def test_generate_content_with_dict_format( # Dict format content contents = {"role": "user", "parts": [{"text": "Hello from dict format"}]} - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("span") with mock.patch.object( @@ -2382,6 +2486,7 @@ def test_generate_content_with_dict_format( model="gemini-1.5-flash", contents=contents, config=create_test_config() ) + sentry_sdk.flush() invoke_span = next(item.payload for item in items if item.type == "span") messages = json.loads( @@ -2409,6 +2514,7 @@ def test_generate_content_with_dict_format( ] +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_generate_content_with_file_data( sentry_init, @@ -2416,6 +2522,7 @@ def test_generate_content_with_file_data( capture_items, mock_genai_client, stream_gen_ai_spans, + span_streaming, ): """Test generate_content with file_data (external file reference).""" sentry_init( @@ -2423,6 +2530,7 @@ def test_generate_content_with_file_data( traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) @@ -2439,7 +2547,7 @@ def test_generate_content_with_file_data( ], ) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("span") with mock.patch.object( @@ -2449,6 +2557,7 @@ def test_generate_content_with_file_data( model="gemini-1.5-flash", contents=content, config=create_test_config() ) + sentry_sdk.flush() invoke_span = next(item.payload for item in items if item.type == "span") messages = json.loads( @@ -2482,6 +2591,7 @@ def test_generate_content_with_file_data( assert messages[0]["content"][1]["uri"] == "gs://bucket/image.jpg" +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_generate_content_with_inline_data( sentry_init, @@ -2489,6 +2599,7 @@ def test_generate_content_with_inline_data( capture_items, mock_genai_client, stream_gen_ai_spans, + span_streaming, ): """Test generate_content with inline_data (binary data).""" sentry_init( @@ -2496,6 +2607,7 @@ def test_generate_content_with_inline_data( traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) @@ -2511,7 +2623,7 @@ def test_generate_content_with_inline_data( ], ) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("span") with mock.patch.object( @@ -2521,6 +2633,7 @@ def test_generate_content_with_inline_data( model="gemini-1.5-flash", contents=content, config=create_test_config() ) + sentry_sdk.flush() invoke_span = next(item.payload for item in items if item.type == "span") messages = json.loads( @@ -2650,6 +2763,7 @@ def test_generate_content_with_mixed_string_and_content( assert messages[0]["content"] == [{"text": "Tell me a joke", "type": "text"}] +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_generate_content_with_part_object_directly( sentry_init, @@ -2657,6 +2771,7 @@ def test_generate_content_with_part_object_directly( capture_items, mock_genai_client, stream_gen_ai_spans, + span_streaming, ): """Test generate_content with Part object directly (not wrapped in Content).""" sentry_init( @@ -2664,6 +2779,7 @@ def test_generate_content_with_part_object_directly( traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) @@ -2671,7 +2787,7 @@ def test_generate_content_with_part_object_directly( # Part object directly part = genai_types.Part(text="Direct Part object") - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("span") with mock.patch.object( @@ -2681,6 +2797,7 @@ def test_generate_content_with_part_object_directly( model="gemini-1.5-flash", contents=part, config=create_test_config() ) + sentry_sdk.flush() invoke_span = next(item.payload for item in items if item.type == "span") messages = json.loads( @@ -2749,6 +2866,7 @@ def test_generate_content_with_list_of_dicts( assert messages[0]["content"] == [{"text": "Second user message", "type": "text"}] +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_generate_content_with_dict_inline_data( sentry_init, @@ -2756,6 +2874,7 @@ def test_generate_content_with_dict_inline_data( capture_items, mock_genai_client, stream_gen_ai_spans, + span_streaming, ): """Test generate_content with dict format containing inline_data.""" sentry_init( @@ -2763,6 +2882,7 @@ def test_generate_content_with_dict_inline_data( traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) @@ -2776,7 +2896,7 @@ def test_generate_content_with_dict_inline_data( ], } - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("span") with mock.patch.object( @@ -2786,6 +2906,7 @@ def test_generate_content_with_dict_inline_data( model="gemini-1.5-flash", contents=contents, config=create_test_config() ) + sentry_sdk.flush() invoke_span = next(item.payload for item in items if item.type == "span") messages = json.loads( @@ -2818,6 +2939,7 @@ def test_generate_content_with_dict_inline_data( assert messages[0]["content"][1]["content"] == BLOB_DATA_SUBSTITUTE +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_generate_content_without_parts_property_inline_data( sentry_init, @@ -2825,12 +2947,14 @@ def test_generate_content_without_parts_property_inline_data( capture_items, mock_genai_client, stream_gen_ai_spans, + span_streaming, ): sentry_init( integrations=[GoogleGenAIIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) @@ -2840,7 +2964,7 @@ def test_generate_content_without_parts_property_inline_data( {"inline_data": {"data": b"fake_binary_data", "mime_type": "image/gif"}}, ] - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("span") with mock.patch.object( @@ -2850,6 +2974,7 @@ def test_generate_content_without_parts_property_inline_data( model="gemini-1.5-flash", contents=contents, config=create_test_config() ) + sentry_sdk.flush() invoke_span = next(item.payload for item in items if item.type == "span") messages = json.loads( @@ -2884,6 +3009,7 @@ def test_generate_content_without_parts_property_inline_data( assert messages[0]["content"][1]["inline_data"]["mime_type"] == "image/gif" +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_generate_content_without_parts_property_inline_data_and_binary_data_within_string( sentry_init, @@ -2891,12 +3017,14 @@ def test_generate_content_without_parts_property_inline_data_and_binary_data_wit capture_items, mock_genai_client, stream_gen_ai_spans, + span_streaming, ): sentry_init( integrations=[GoogleGenAIIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) @@ -2911,7 +3039,7 @@ def test_generate_content_without_parts_property_inline_data_and_binary_data_wit }, ] - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("span") with mock.patch.object( @@ -2921,6 +3049,7 @@ def test_generate_content_without_parts_property_inline_data_and_binary_data_wit model="gemini-1.5-flash", contents=contents, config=create_test_config() ) + sentry_sdk.flush() invoke_span = next(item.payload for item in items if item.type == "span") messages = json.loads( From ba2d985d92c1eb3aca97f7d67d662be1d62617bb Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Thu, 21 May 2026 10:55:13 +0200 Subject: [PATCH 2/6] mypy errors --- sentry_sdk/integrations/google_genai/__init__.py | 2 +- sentry_sdk/integrations/google_genai/streaming.py | 11 +++++++---- sentry_sdk/integrations/google_genai/utils.py | 5 ++++- 3 files changed, 12 insertions(+), 6 deletions(-) diff --git a/sentry_sdk/integrations/google_genai/__init__.py b/sentry_sdk/integrations/google_genai/__init__.py index b7d4502f71..83ecdd18c3 100644 --- a/sentry_sdk/integrations/google_genai/__init__.py +++ b/sentry_sdk/integrations/google_genai/__init__.py @@ -303,7 +303,7 @@ async def new_async_generate_content( response = await f(self, *args, **kwargs) except Exception as exc: _capture_exception(exc) - chat_span.set_status(SPANSTATUS.INTERNAL_ERROR) + chat_span.status = SpanStatus.ERROR raise set_span_data_for_response(chat_span, integration, response) diff --git a/sentry_sdk/integrations/google_genai/streaming.py b/sentry_sdk/integrations/google_genai/streaming.py index b54723e115..8414ea4f21 100644 --- a/sentry_sdk/integrations/google_genai/streaming.py +++ b/sentry_sdk/integrations/google_genai/streaming.py @@ -130,10 +130,13 @@ def set_span_data_for_streaming_response( safe_serialize(accumulated_response["tool_calls"]), ) - if accumulated_response.get("id"): - set_on_span(SPANDATA.GEN_AI_RESPONSE_ID, accumulated_response["id"]) - if accumulated_response.get("model"): - set_on_span(SPANDATA.GEN_AI_RESPONSE_MODEL, accumulated_response["model"]) + response_id = accumulated_response.get("id") + if response_id is not None: + set_on_span(SPANDATA.GEN_AI_RESPONSE_ID, response_id) + + response_model = accumulated_response.get("model") + if response_model is not None: + set_on_span(SPANDATA.GEN_AI_RESPONSE_MODEL, response_model) if accumulated_response["usage_metadata"] is None: return diff --git a/sentry_sdk/integrations/google_genai/utils.py b/sentry_sdk/integrations/google_genai/utils.py index ee7b8b5e5c..189758e0e5 100644 --- a/sentry_sdk/integrations/google_genai/utils.py +++ b/sentry_sdk/integrations/google_genai/utils.py @@ -1053,7 +1053,10 @@ def prepare_embed_content_args( def set_span_data_for_embed_request( - span: "Span", integration: "Any", contents: "Any", kwargs: "dict[str, Any]" + span: "Union[Span, StreamedSpan]", + integration: "Any", + contents: "Any", + kwargs: "dict[str, Any]", ) -> None: """Set span data for embedding request.""" # Include input contents if PII is allowed From dc21d8a7558266803c46c4a66bfde5c053495493 Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Thu, 21 May 2026 10:59:27 +0200 Subject: [PATCH 3/6] call enter only in legacy path --- sentry_sdk/integrations/google_genai/__init__.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/sentry_sdk/integrations/google_genai/__init__.py b/sentry_sdk/integrations/google_genai/__init__.py index 83ecdd18c3..cdacf3ffd5 100644 --- a/sentry_sdk/integrations/google_genai/__init__.py +++ b/sentry_sdk/integrations/google_genai/__init__.py @@ -93,7 +93,8 @@ def new_generate_content_stream( ) set_on_span = chat_span.set_data - chat_span.__enter__() + chat_span.__enter__() + set_on_span(SPANDATA.GEN_AI_OPERATION_NAME, "chat") set_on_span(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM) set_on_span(SPANDATA.GEN_AI_REQUEST_MODEL, model_name) @@ -168,7 +169,8 @@ async def new_async_generate_content_stream( ) set_on_span = chat_span.set_data - chat_span.__enter__() + chat_span.__enter__() + set_on_span(SPANDATA.GEN_AI_OPERATION_NAME, "chat") set_on_span(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM) set_on_span(SPANDATA.GEN_AI_REQUEST_MODEL, model_name) From 126f3f54496a1b85b5c5958a91d112f68ae30d59 Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Thu, 21 May 2026 11:02:13 +0200 Subject: [PATCH 4/6] fix(google-genai): Guard against None response ID and response model --- sentry_sdk/integrations/google_genai/streaming.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/sentry_sdk/integrations/google_genai/streaming.py b/sentry_sdk/integrations/google_genai/streaming.py index 7f3f58dc93..05b2318f5b 100644 --- a/sentry_sdk/integrations/google_genai/streaming.py +++ b/sentry_sdk/integrations/google_genai/streaming.py @@ -123,10 +123,13 @@ def set_span_data_for_streaming_response( safe_serialize(accumulated_response["tool_calls"]), ) - if accumulated_response.get("id"): - span.set_data(SPANDATA.GEN_AI_RESPONSE_ID, accumulated_response["id"]) - if accumulated_response.get("model"): - span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, accumulated_response["model"]) + response_id = accumulated_response.get("id") + if response_id is not None: + span.set_data(SPANDATA.GEN_AI_RESPONSE_ID, response_id) + + response_model = accumulated_response.get("model") + if response_model is not None: + span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, response_model) if accumulated_response["usage_metadata"] is None: return From 6dac84813678d2d07fbbd178c278fdce038b5603 Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Thu, 21 May 2026 11:14:30 +0200 Subject: [PATCH 5/6] use should_truncate_gen_ai_input --- sentry_sdk/integrations/google_genai/utils.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/sentry_sdk/integrations/google_genai/utils.py b/sentry_sdk/integrations/google_genai/utils.py index 189758e0e5..a5c3ce9a13 100644 --- a/sentry_sdk/integrations/google_genai/utils.py +++ b/sentry_sdk/integrations/google_genai/utils.py @@ -29,7 +29,10 @@ from sentry_sdk.consts import OP, SPANDATA from sentry_sdk.scope import should_send_default_pii from sentry_sdk.traces import StreamedSpan -from sentry_sdk.tracing_utils import has_span_streaming_enabled +from sentry_sdk.tracing_utils import ( + has_span_streaming_enabled, + should_truncate_gen_ai_input, +) from sentry_sdk.utils import ( capture_internal_exceptions, event_from_exception, @@ -921,9 +924,9 @@ def set_span_data_for_request( client = sentry_sdk.get_client() scope = sentry_sdk.get_current_scope() messages_data = ( - normalized_messages - if client.options.get("stream_gen_ai_spans", False) - else truncate_and_annotate_messages(normalized_messages, span, scope) + truncate_and_annotate_messages(normalized_messages, span, scope) + if should_truncate_gen_ai_input(client.options) + else normalized_messages ) if messages_data is not None: set_data_normalized( From 0d6b7dcc77c412c24cd965747d5cd816c9bf21fb Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Thu, 21 May 2026 11:16:10 +0200 Subject: [PATCH 6/6] use sentry_sdk.traces.start_span directly --- sentry_sdk/integrations/google_genai/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sentry_sdk/integrations/google_genai/__init__.py b/sentry_sdk/integrations/google_genai/__init__.py index cdacf3ffd5..68e0d03271 100644 --- a/sentry_sdk/integrations/google_genai/__init__.py +++ b/sentry_sdk/integrations/google_genai/__init__.py @@ -152,7 +152,7 @@ async def new_async_generate_content_stream( _model, contents, model_name = prepare_generate_content_args(args, kwargs) if has_span_streaming_enabled(client.options): - chat_span = get_start_span_function()( + chat_span = sentry_sdk.traces.start_span( name=f"chat {model_name}", attributes={ "sentry.op": OP.GEN_AI_CHAT,