diff --git a/sentry_sdk/integrations/google_genai/__init__.py b/sentry_sdk/integrations/google_genai/__init__.py index e00af2aa0a..2aeb2a94a6 100644 --- a/sentry_sdk/integrations/google_genai/__init__.py +++ b/sentry_sdk/integrations/google_genai/__init__.py @@ -8,12 +8,9 @@ ) import sentry_sdk -from sentry_sdk.ai.utils import get_start_span_function from sentry_sdk.consts import OP, SPANDATA from sentry_sdk.integrations import DidNotEnable, Integration, _check_minimum_version -from sentry_sdk.traces import SpanStatus, StreamedSpan -from sentry_sdk.tracing import SPANSTATUS -from sentry_sdk.tracing_utils import has_span_streaming_enabled +from sentry_sdk.traces import SpanStatus from sentry_sdk.utils import parse_version try: @@ -80,30 +77,17 @@ def new_generate_content_stream( _model, contents, model_name = prepare_generate_content_args(args, kwargs) - 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, - SPANDATA.GEN_AI_OPERATION_NAME: "chat", - SPANDATA.GEN_AI_SYSTEM: GEN_AI_SYSTEM, - SPANDATA.GEN_AI_REQUEST_MODEL: model_name, - SPANDATA.GEN_AI_RESPONSE_STREAMING: True, - }, - ) - else: - chat_span = get_start_span_function()( - op=OP.GEN_AI_CHAT, - name=f"chat {model_name}", - origin=ORIGIN, - ) - 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) - chat_span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, True) + chat_span = sentry_sdk.traces.start_span( + name=f"chat {model_name}", + attributes={ + "sentry.op": OP.GEN_AI_CHAT, + "sentry.origin": ORIGIN, + SPANDATA.GEN_AI_OPERATION_NAME: "chat", + SPANDATA.GEN_AI_SYSTEM: GEN_AI_SYSTEM, + SPANDATA.GEN_AI_REQUEST_MODEL: model_name, + SPANDATA.GEN_AI_RESPONSE_STREAMING: True, + }, + ) set_span_data_for_request(chat_span, integration, model_name, contents, kwargs) @@ -119,10 +103,7 @@ def new_iterator() -> "Iterator[Any]": yield chunk except Exception as exc: _capture_exception(exc) - if isinstance(chat_span, StreamedSpan): - chat_span.status = SpanStatus.ERROR - else: - chat_span.set_status(SPANSTATUS.INTERNAL_ERROR) + chat_span.status = SpanStatus.ERROR raise finally: # Accumulate all chunks and set final response data on spans @@ -157,30 +138,17 @@ 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 = sentry_sdk.traces.start_span( - name=f"chat {model_name}", - attributes={ - "sentry.op": OP.GEN_AI_CHAT, - "sentry.origin": ORIGIN, - SPANDATA.GEN_AI_OPERATION_NAME: "chat", - SPANDATA.GEN_AI_SYSTEM: GEN_AI_SYSTEM, - SPANDATA.GEN_AI_REQUEST_MODEL: model_name, - SPANDATA.GEN_AI_RESPONSE_STREAMING: True, - }, - ) - else: - chat_span = get_start_span_function()( - op=OP.GEN_AI_CHAT, - name=f"chat {model_name}", - origin=ORIGIN, - ) - 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) - chat_span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, True) + chat_span = sentry_sdk.traces.start_span( + name=f"chat {model_name}", + attributes={ + "sentry.op": OP.GEN_AI_CHAT, + "sentry.origin": ORIGIN, + SPANDATA.GEN_AI_OPERATION_NAME: "chat", + SPANDATA.GEN_AI_SYSTEM: GEN_AI_SYSTEM, + SPANDATA.GEN_AI_REQUEST_MODEL: model_name, + SPANDATA.GEN_AI_RESPONSE_STREAMING: True, + }, + ) set_span_data_for_request(chat_span, integration, model_name, contents, kwargs) @@ -196,10 +164,7 @@ async def new_async_iterator() -> "AsyncIterator[Any]": yield chunk except Exception as exc: _capture_exception(exc) - if isinstance(chat_span, StreamedSpan): - chat_span.status = SpanStatus.ERROR - else: - chat_span.set_status(SPANSTATUS.INTERNAL_ERROR) + chat_span.status = SpanStatus.ERROR raise finally: # Accumulate all chunks and set final response data on spans @@ -230,54 +195,30 @@ def new_generate_content(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": model, contents, model_name = prepare_generate_content_args(args, 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, - SPANDATA.GEN_AI_OPERATION_NAME: "chat", - SPANDATA.GEN_AI_SYSTEM: GEN_AI_SYSTEM, - SPANDATA.GEN_AI_REQUEST_MODEL: model_name, - }, - ) as chat_span: - 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.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 - ) + with sentry_sdk.traces.start_span( + name=f"chat {model_name}", + attributes={ + "sentry.op": OP.GEN_AI_CHAT, + "sentry.origin": ORIGIN, + SPANDATA.GEN_AI_OPERATION_NAME: "chat", + SPANDATA.GEN_AI_SYSTEM: GEN_AI_SYSTEM, + SPANDATA.GEN_AI_REQUEST_MODEL: model_name, + }, + ) as chat_span: + 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) + set_span_data_for_response(chat_span, integration, response) - return response + return response return new_generate_content @@ -294,52 +235,29 @@ async def new_async_generate_content( model, contents, model_name = prepare_generate_content_args(args, 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, - SPANDATA.GEN_AI_OPERATION_NAME: "chat", - SPANDATA.GEN_AI_SYSTEM: GEN_AI_SYSTEM, - SPANDATA.GEN_AI_REQUEST_MODEL: model_name, - }, - ) as chat_span: - 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.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 = await f(self, *args, **kwargs) - except Exception as exc: - _capture_exception(exc) - chat_span.set_status(SPANSTATUS.INTERNAL_ERROR) - raise + with sentry_sdk.traces.start_span( + name=f"chat {model_name}", + attributes={ + "sentry.op": OP.GEN_AI_CHAT, + "sentry.origin": ORIGIN, + SPANDATA.GEN_AI_OPERATION_NAME: "chat", + SPANDATA.GEN_AI_SYSTEM: GEN_AI_SYSTEM, + SPANDATA.GEN_AI_REQUEST_MODEL: model_name, + }, + ) as chat_span: + 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.status = SpanStatus.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_async_generate_content @@ -354,50 +272,28 @@ def new_embed_content(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": model_name, contents = prepare_embed_content_args(args, 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, - SPANDATA.GEN_AI_OPERATION_NAME: "embeddings", - SPANDATA.GEN_AI_SYSTEM: GEN_AI_SYSTEM, - SPANDATA.GEN_AI_REQUEST_MODEL: model_name, - }, - ) as span: - set_span_data_for_embed_request(span, integration, contents, kwargs) - - 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) - - return response + with sentry_sdk.traces.start_span( + name=f"embeddings {model_name}", + attributes={ + "sentry.op": OP.GEN_AI_EMBEDDINGS, + "sentry.origin": ORIGIN, + SPANDATA.GEN_AI_OPERATION_NAME: "embeddings", + SPANDATA.GEN_AI_SYSTEM: GEN_AI_SYSTEM, + SPANDATA.GEN_AI_REQUEST_MODEL: model_name, + }, + ) as span: + set_span_data_for_embed_request(span, integration, contents, kwargs) + + 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 return new_embed_content @@ -414,49 +310,27 @@ async def new_async_embed_content( model_name, contents = prepare_embed_content_args(args, 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, - SPANDATA.GEN_AI_OPERATION_NAME: "embeddings", - SPANDATA.GEN_AI_SYSTEM: GEN_AI_SYSTEM, - SPANDATA.GEN_AI_REQUEST_MODEL: model_name, - }, - ) as span: - 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 + with sentry_sdk.traces.start_span( + name=f"embeddings {model_name}", + attributes={ + "sentry.op": OP.GEN_AI_EMBEDDINGS, + "sentry.origin": ORIGIN, + SPANDATA.GEN_AI_OPERATION_NAME: "embeddings", + SPANDATA.GEN_AI_SYSTEM: GEN_AI_SYSTEM, + SPANDATA.GEN_AI_REQUEST_MODEL: model_name, + }, + ) as span: + 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 return new_async_embed_content diff --git a/sentry_sdk/integrations/google_genai/streaming.py b/sentry_sdk/integrations/google_genai/streaming.py index 86cdcf29ba..3f44585f82 100644 --- a/sentry_sdk/integrations/google_genai/streaming.py +++ b/sentry_sdk/integrations/google_genai/streaming.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Any, List, Optional, TypedDict, Union +from typing import TYPE_CHECKING, Any, List, Optional, TypedDict import sentry_sdk from sentry_sdk.ai.utils import set_data_normalized @@ -21,8 +21,6 @@ if TYPE_CHECKING: from google.genai.types import GenerateContentResponse - from sentry_sdk.tracing import Span - class AccumulatedResponse(TypedDict): id: "Optional[str]" @@ -100,14 +98,11 @@ def accumulate_streaming_response( def set_span_data_for_streaming_response( - span: "Union[Span, StreamedSpan]", + 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 - ) client = sentry_sdk.get_client() if accumulated_response.get("finish_reasons"): @@ -119,41 +114,41 @@ def set_span_data_for_streaming_response( response_id = accumulated_response.get("id") if response_id is not None: - set_on_span(SPANDATA.GEN_AI_RESPONSE_ID, response_id) + span.set_attribute(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) + span.set_attribute(SPANDATA.GEN_AI_RESPONSE_MODEL, response_model) if accumulated_response["usage_metadata"] is None: return if accumulated_response["usage_metadata"]["input_tokens"]: - set_on_span( + span.set_attribute( SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, accumulated_response["usage_metadata"]["input_tokens"], ) if accumulated_response["usage_metadata"]["input_tokens_cached"]: - set_on_span( + span.set_attribute( SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED, accumulated_response["usage_metadata"]["input_tokens_cached"], ) if accumulated_response["usage_metadata"]["output_tokens"]: - set_on_span( + span.set_attribute( SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, accumulated_response["usage_metadata"]["output_tokens"], ) if accumulated_response["usage_metadata"]["output_tokens_reasoning"]: - set_on_span( + span.set_attribute( SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING, accumulated_response["usage_metadata"]["output_tokens_reasoning"], ) if accumulated_response["usage_metadata"]["total_tokens"]: - set_on_span( + span.set_attribute( SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, accumulated_response["usage_metadata"]["total_tokens"], ) @@ -161,13 +156,13 @@ def set_span_data_for_streaming_response( if accumulated_response.get("tool_calls"): if has_data_collection_enabled(client.options): if client.options["data_collection"]["gen_ai"]["outputs"]: - set_on_span( + span.set_attribute( SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, safe_serialize(accumulated_response["tool_calls"]), ) else: # Before data collection was introduced this was unconditionally set - set_on_span( + span.set_attribute( SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, safe_serialize(accumulated_response["tool_calls"]), ) @@ -175,13 +170,13 @@ def set_span_data_for_streaming_response( if accumulated_response.get("text"): if has_data_collection_enabled(client.options): if client.options["data_collection"]["gen_ai"]["outputs"]: - set_on_span( + span.set_attribute( SPANDATA.GEN_AI_RESPONSE_TEXT, safe_serialize([accumulated_response["text"]]), ) elif should_send_default_pii() and integration.include_prompts: - set_on_span( + span.set_attribute( SPANDATA.GEN_AI_RESPONSE_TEXT, safe_serialize([accumulated_response["text"]]), ) diff --git a/sentry_sdk/integrations/google_genai/utils.py b/sentry_sdk/integrations/google_genai/utils.py index 5120d374e1..cd728836c2 100644 --- a/sentry_sdk/integrations/google_genai/utils.py +++ b/sentry_sdk/integrations/google_genai/utils.py @@ -25,14 +25,10 @@ normalize_message_roles, set_data_normalized, transform_google_content_part, - truncate_and_annotate_messages, ) 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, @@ -56,7 +52,6 @@ ) from sentry_sdk._types import TextPart - from sentry_sdk.tracing import Span _is_PIL_available = False try: @@ -684,32 +679,18 @@ def _capture_tool_input( return tool_input -def _create_tool_span( - tool_name: str, tool_doc: "Optional[str]" -) -> "Union[Span, StreamedSpan]": +def _create_tool_span(tool_name: str, tool_doc: "Optional[str]") -> "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, - 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, + span = sentry_sdk.traces.start_span( name=f"execute_tool {tool_name}", - origin=ORIGIN, + attributes={ + "sentry.op": OP.GEN_AI_EXECUTE_TOOL, + "sentry.origin": ORIGIN, + SPANDATA.GEN_AI_TOOL_NAME: tool_name, + }, ) - span.set_data(SPANDATA.GEN_AI_TOOL_NAME, tool_name) if tool_doc: - span.set_data(SPANDATA.GEN_AI_TOOL_DESCRIPTION, tool_doc) + span.set_attribute(SPANDATA.GEN_AI_TOOL_DESCRIPTION, tool_doc) return span @@ -918,7 +899,7 @@ def _transform_system_instructions( def set_span_data_for_request( - span: "Union[Span, StreamedSpan]", + span: "StreamedSpan", integration: "Any", model: str, contents: "ContentListUnion", @@ -926,14 +907,11 @@ def set_span_data_for_request( ) -> None: """Set span data for the request.""" client = sentry_sdk.get_client() - set_on_span = ( - span.set_attribute if isinstance(span, StreamedSpan) else span.set_data - ) - set_on_span(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM) - set_on_span(SPANDATA.GEN_AI_REQUEST_MODEL, model) + span.set_attribute(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM) + span.set_attribute(SPANDATA.GEN_AI_REQUEST_MODEL, model) if kwargs.get("stream", False): - set_on_span(SPANDATA.GEN_AI_RESPONSE_STREAMING, True) + span.set_attribute(SPANDATA.GEN_AI_RESPONSE_STREAMING, True) config: "Optional[GenerateContentConfig]" = kwargs.get("config") @@ -978,7 +956,7 @@ def set_span_data_for_request( system_instructions = config.get("system_instruction") if system_instructions is not None: - set_on_span( + span.set_attribute( SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, json.dumps(_transform_system_instructions(system_instructions)), ) @@ -990,17 +968,11 @@ def set_span_data_for_request( if messages: normalized_messages = normalize_message_roles(messages) client = sentry_sdk.get_client() - scope = sentry_sdk.get_current_scope() - messages_data = ( - truncate_and_annotate_messages(normalized_messages, span, scope) - if not has_span_streaming_enabled(client.options) - else normalized_messages - ) - if messages_data is not None: + if normalized_messages is not None: set_data_normalized( span, SPANDATA.GEN_AI_REQUEST_MESSAGES, - messages_data, + normalized_messages, unpack=False, ) @@ -1017,11 +989,11 @@ def set_span_data_for_request( if hasattr(config, param): value = getattr(config, param) if value is not None: - set_on_span(span_key, value) + span.set_attribute(span_key, value) def set_span_data_for_response( - span: "Union[Span, StreamedSpan]", + span: "StreamedSpan", integration: "Any", response: "GenerateContentResponse", ) -> None: @@ -1029,9 +1001,6 @@ def set_span_data_for_response( return client = sentry_sdk.get_client() - set_on_span = ( - span.set_attribute if isinstance(span, StreamedSpan) else span.set_data - ) finish_reasons = extract_finish_reasons(response) if finish_reasons: @@ -1041,51 +1010,59 @@ def set_span_data_for_response( response_id = getattr(response, "response_id", None) if response_id is not None: - set_on_span(SPANDATA.GEN_AI_RESPONSE_ID, response_id) + span.set_attribute(SPANDATA.GEN_AI_RESPONSE_ID, response_id) model_version = getattr(response, "model_version", None) if model_version is not None: - set_on_span(SPANDATA.GEN_AI_RESPONSE_MODEL, model_version) + span.set_attribute(SPANDATA.GEN_AI_RESPONSE_MODEL, model_version) usage_data = extract_usage_data(response) if usage_data["input_tokens"]: - set_on_span(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, usage_data["input_tokens"]) + span.set_attribute( + SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, usage_data["input_tokens"] + ) if usage_data["input_tokens_cached"]: - set_on_span( + span.set_attribute( SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED, usage_data["input_tokens_cached"], ) if usage_data["output_tokens"]: - set_on_span(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, usage_data["output_tokens"]) + span.set_attribute( + SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, usage_data["output_tokens"] + ) if usage_data["output_tokens_reasoning"]: - set_on_span( + span.set_attribute( SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING, usage_data["output_tokens_reasoning"], ) if usage_data["total_tokens"]: - set_on_span(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, usage_data["total_tokens"]) + span.set_attribute( + SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, usage_data["total_tokens"] + ) tool_calls = extract_tool_calls(response) if tool_calls: if has_data_collection_enabled(client.options): if client.options["data_collection"]["gen_ai"]["outputs"]: - set_on_span( + span.set_attribute( SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, safe_serialize(tool_calls) ) else: # Before data collection was introduced, this was set unconditionally - set_on_span(SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, safe_serialize(tool_calls)) + span.set_attribute( + SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, safe_serialize(tool_calls) + ) if has_data_collection_enabled(client.options): if client.options["data_collection"]["gen_ai"]["outputs"]: response_texts = _extract_response_text(response) if response_texts: - set_on_span( + span.set_attribute( SPANDATA.GEN_AI_RESPONSE_TEXT, safe_serialize(response_texts) ) elif should_send_default_pii() and integration.include_prompts: @@ -1093,7 +1070,9 @@ def set_span_data_for_response( response_texts = _extract_response_text(response) if response_texts: # Format as JSON string array as per documentation - set_on_span(SPANDATA.GEN_AI_RESPONSE_TEXT, safe_serialize(response_texts)) + span.set_attribute( + SPANDATA.GEN_AI_RESPONSE_TEXT, safe_serialize(response_texts) + ) def prepare_generate_content_args( @@ -1128,7 +1107,7 @@ def prepare_embed_content_args( def set_span_data_for_embed_request( - span: "Union[Span, StreamedSpan]", + span: "StreamedSpan", integration: "Any", contents: "Any", kwargs: "dict[str, Any]", @@ -1171,7 +1150,7 @@ def set_span_data_for_embed_request( def set_span_data_for_embed_response( - span: "Union[Span, StreamedSpan]", + span: "StreamedSpan", integration: "Any", response: "EmbedContentResponse", ) -> None: @@ -1192,7 +1171,4 @@ def set_span_data_for_embed_response( # Set token count if we found any if total_tokens > 0: - 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) + span.set_attribute(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, total_tokens) diff --git a/tests/integrations/google_genai/test_google_genai.py b/tests/integrations/google_genai/test_google_genai.py index 55329ddbe1..3f1c191161 100644 --- a/tests/integrations/google_genai/test_google_genai.py +++ b/tests/integrations/google_genai/test_google_genai.py @@ -7,7 +7,6 @@ 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 from sentry_sdk.integrations.google_genai import GoogleGenAIIntegration @@ -115,7 +114,6 @@ def create_test_config( return genai_types.GenerateContentConfig(**config_dict) -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize( "send_default_pii, include_prompts", [ @@ -127,142 +125,86 @@ def create_test_config( ) def test_nonstreaming_generate_content( sentry_init, - capture_events, capture_items, send_default_pii, include_prompts, mock_genai_client, - span_streaming, ): sentry_init( integrations=[GoogleGenAIIntegration(include_prompts=include_prompts)], traces_sample_rate=1.0, send_default_pii=send_default_pii, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) # Mock the HTTP response at the _api_client.request() level mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) + items = capture_items("span") + + with mock.patch.object( + mock_genai_client._api_client, + "request", + return_value=mock_http_response, + ): + config = create_test_config(temperature=0.7, max_output_tokens=100) + mock_genai_client.models.generate_content( + model="gemini-1.5-flash", + contents=[ + "Message demonstrating the absence of truncation.", + "Tell me a joke", + ], + config=config, + ) + + sentry_sdk.flush() + spans = [item.payload for item in items] + assert len(spans) == 1 + chat_span = next(item.payload for item in items if item.type == "span") + + # Check chat span + assert chat_span["attributes"]["sentry.op"] == OP.GEN_AI_CHAT + assert chat_span["name"] == "chat gemini-1.5-flash" + assert chat_span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" + assert chat_span["attributes"][SPANDATA.GEN_AI_SYSTEM] == "gcp.gemini" + assert chat_span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] == "gemini-1.5-flash" - if span_streaming: - items = capture_items("transaction", "span") - - with mock.patch.object( - mock_genai_client._api_client, - "request", - return_value=mock_http_response, - ), sentry_sdk.traces.start_span(name="google_genai"): - config = create_test_config(temperature=0.7, max_output_tokens=100) - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", - contents=[ - "Message demonstrating the absence of truncation.", - "Tell me a joke", + if send_default_pii and include_prompts: + assert json.loads( + chat_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] + ) == [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Message demonstrating the absence of truncation.", + }, + { + "type": "text", + "text": "Tell me a joke", + }, ], - config=config, - ) + } + ] - sentry_sdk.flush() - spans = [item.payload for item in items if item.type == "span"] - assert len(spans) == 2 - assert spans[1]["name"] == "google_genai" - chat_span = next(item.payload for item in items if item.type == "span") - - # Check chat span - assert chat_span["attributes"]["sentry.op"] == OP.GEN_AI_CHAT - assert chat_span["name"] == "chat gemini-1.5-flash" - assert chat_span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" - assert chat_span["attributes"][SPANDATA.GEN_AI_SYSTEM] == "gcp.gemini" - assert ( - chat_span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] == "gemini-1.5-flash" - ) + # Response text is stored as a JSON array + response_text = chat_span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] - if send_default_pii and include_prompts: - assert json.loads( - chat_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - ) == [ - { - "role": "user", - "content": [ - { - "type": "text", - "text": "Message demonstrating the absence of truncation.", - }, - { - "type": "text", - "text": "Tell me a joke", - }, - ], - } - ] - - # Response text is stored as a JSON array - response_text = chat_span["attributes"][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_RESPONSE_TEXT not in chat_span["attributes"] - - # Check token usage - assert chat_span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 - # Output tokens now include reasoning tokens: candidates_token_count (20) + thoughts_token_count (3) = 23 - assert chat_span["attributes"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 23 - assert chat_span["attributes"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 30 - assert chat_span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 5 - assert ( - chat_span["attributes"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING] == 3 - ) + # Parse the JSON array + response_texts = json.loads(response_text) + assert response_texts == ["Hello! How can I help you today?"] else: - events = capture_events() - - with mock.patch.object( - mock_genai_client._api_client, - "request", - return_value=mock_http_response, - ), start_transaction(name="google_genai"): - config = create_test_config(temperature=0.7, max_output_tokens=100) - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents="Tell me a joke", config=config - ) + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in chat_span["attributes"] + + # Check token usage + assert chat_span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 + # Output tokens now include reasoning tokens: candidates_token_count (20) + thoughts_token_count (3) = 23 + assert chat_span["attributes"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 23 + assert chat_span["attributes"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 30 + assert chat_span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 5 + assert chat_span["attributes"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING] == 3 + - assert len(events) == 1 - (event,) = events - - assert event["type"] == "transaction" - assert event["transaction"] == "google_genai" - - assert len(event["spans"]) == 1 - chat_span = event["spans"][0] - - # Check chat span - assert chat_span["op"] == OP.GEN_AI_CHAT - assert chat_span["description"] == "chat gemini-1.5-flash" - assert chat_span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" - assert chat_span["data"][SPANDATA.GEN_AI_SYSTEM] == "gcp.gemini" - assert chat_span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "gemini-1.5-flash" - - if send_default_pii and include_prompts: - # 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_RESPONSE_TEXT not in chat_span["data"] - - # Check token usage - assert chat_span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 - # Output tokens now include reasoning tokens: candidates_token_count (20) + thoughts_token_count (3) = 23 - assert chat_span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 23 - assert chat_span["data"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 30 - assert chat_span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 5 - assert chat_span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING] == 3 - - -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("generate_content_config", (False, True)) @pytest.mark.parametrize( "system_instructions,expected_texts", @@ -293,105 +235,65 @@ def test_nonstreaming_generate_content( ) def test_generate_content_with_system_instruction( sentry_init, - capture_events, capture_items, mock_genai_client, generate_content_config, system_instructions, expected_texts, - span_streaming, ): sentry_init( integrations=[GoogleGenAIIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - config = { - "system_instruction": system_instructions, - "temperature": 0.5, - } - - if generate_content_config: - config = create_test_config(**config) - - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", - contents="What is 2+2?", - config=config, - ) - - sentry_sdk.flush() - invoke_span = next(item.payload for item in items) + with mock.patch.object( + mock_genai_client._api_client, "request", return_value=mock_http_response + ): + config = { + "system_instruction": system_instructions, + "temperature": 0.5, + } - if expected_texts is None: - assert SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS not in invoke_span["attributes"] - return + if generate_content_config: + config = create_test_config(**config) - # (PII is enabled and include_prompts is True in this test) - system_instructions = json.loads( - invoke_span["attributes"][SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS] + mock_genai_client.models.generate_content( + model="gemini-1.5-flash", + contents="What is 2+2?", + config=config, ) - else: - events = capture_events() - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - config = { - "system_instruction": system_instructions, - "temperature": 0.5, - } - - if generate_content_config: - config = create_test_config(**config) - - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", - contents="What is 2+2?", - config=config, - ) - (event,) = events - invoke_span = event["spans"][0] + sentry_sdk.flush() + invoke_span = next(item.payload for item in items) - if expected_texts is None: - assert SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS not in invoke_span["data"] - return + if expected_texts is None: + assert SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS not in invoke_span["attributes"] + return - # (PII is enabled and include_prompts is True in this test) - system_instructions = json.loads( - invoke_span["data"][SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS] - ) + # (PII is enabled and include_prompts is True in this test) + system_instructions = json.loads( + invoke_span["attributes"][SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS] + ) assert system_instructions == [ {"type": "text", "content": text} for text in expected_texts ] -@pytest.mark.parametrize("span_streaming", [True, False]) def test_generate_content_with_tools( sentry_init, - capture_events, capture_items, mock_genai_client, - span_streaming, ): sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) # Create a mock tool function @@ -436,57 +338,29 @@ def get_weather(location: str) -> str: } mock_http_response = create_mock_http_response(tool_response_json) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - config = create_test_config(tools=[get_weather, mock_tool]) - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents="What's the weather?", config=config - ) - - sentry_sdk.flush() - invoke_span = next(item.payload for item in items) - - # Check that tools are recorded (data is serialized as a string) - tools_data_str = invoke_span["attributes"][ - SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS - ] - # Parse the JSON string to verify content - tools_data = json.loads(tools_data_str) - assert len(tools_data) == 2 - - # The order of tools may not be guaranteed, so sort by name and description for comparison - sorted_tools = sorted( - tools_data, key=lambda t: (t.get("name", ""), t.get("name", "")) + with mock.patch.object( + mock_genai_client._api_client, "request", return_value=mock_http_response + ): + config = create_test_config(tools=[get_weather, mock_tool]) + mock_genai_client.models.generate_content( + model="gemini-1.5-flash", contents="What's the weather?", config=config ) - else: - events = capture_events() - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - config = create_test_config(tools=[get_weather, mock_tool]) - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents="What's the weather?", config=config - ) - (event,) = events - invoke_span = event["spans"][0] + sentry_sdk.flush() + invoke_span = next(item.payload for item in items) - # Check that tools are recorded (data is serialized as a string) - tools_data_str = invoke_span["data"][SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS] - # Parse the JSON string to verify content - tools_data = json.loads(tools_data_str) - assert len(tools_data) == 2 + # Check that tools are recorded (data is serialized as a string) + tools_data_str = invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS] + # Parse the JSON string to verify content + tools_data = json.loads(tools_data_str) + assert len(tools_data) == 2 - # The order of tools may not be guaranteed, so sort by name and description for comparison - sorted_tools = sorted( - tools_data, key=lambda t: (t.get("name", ""), t.get("description", "")) - ) + # The order of tools may not be guaranteed, so sort by name and description for comparison + sorted_tools = sorted( + tools_data, key=lambda t: (t.get("name", ""), t.get("name", "")) + ) # The function tool assert sorted_tools[0]["name"] == "get_weather" @@ -497,19 +371,15 @@ def get_weather(location: str) -> str: assert sorted_tools[1]["description"] == "Get weather information (tool object)" -@pytest.mark.parametrize("span_streaming", [True, False]) def test_tool_execution( sentry_init, - capture_events, capture_items, - span_streaming, ): sentry_init( integrations=[GoogleGenAIIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) # Create a mock tool function @@ -521,99 +391,51 @@ def get_weather(location: str) -> str: from sentry_sdk.integrations.google_genai.utils import wrapped_tool wrapped_weather = wrapped_tool(get_weather) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - # Execute the wrapped tool - with start_transaction(name="test_tool"): - result = wrapped_weather("San Francisco") + # Execute the wrapped tool + result = wrapped_weather("San Francisco") - assert result == "The weather in San Francisco is sunny" + assert result == "The weather in San Francisco is sunny" - sentry_sdk.flush() - spans = [item.payload for item in items] - assert len(spans) == 1 - sentry_sdk.flush() - tool_span = next(item.payload for item in items) + sentry_sdk.flush() + spans = [item.payload for item in items] + assert len(spans) == 1 + sentry_sdk.flush() + tool_span = next(item.payload for item in items) - assert tool_span["attributes"]["sentry.op"] == OP.GEN_AI_EXECUTE_TOOL - assert tool_span["name"] == "execute_tool get_weather" - assert tool_span["attributes"][SPANDATA.GEN_AI_TOOL_NAME] == "get_weather" - assert ( - tool_span["attributes"][SPANDATA.GEN_AI_TOOL_DESCRIPTION] - == "Get the weather for a location" - ) - else: - events = capture_events() - - # Execute the wrapped tool - with start_transaction(name="test_tool"): - result = wrapped_weather("San Francisco") - - assert result == "The weather in San Francisco is sunny" - - (event,) = events - assert len(event["spans"]) == 1 - tool_span = event["spans"][0] - - assert tool_span["op"] == OP.GEN_AI_EXECUTE_TOOL - assert tool_span["description"] == "execute_tool get_weather" - assert tool_span["data"][SPANDATA.GEN_AI_TOOL_NAME] == "get_weather" - assert ( - tool_span["data"][SPANDATA.GEN_AI_TOOL_DESCRIPTION] - == "Get the weather for a location" - ) + assert tool_span["attributes"]["sentry.op"] == OP.GEN_AI_EXECUTE_TOOL + assert tool_span["name"] == "execute_tool get_weather" + assert tool_span["attributes"][SPANDATA.GEN_AI_TOOL_NAME] == "get_weather" + assert ( + tool_span["attributes"][SPANDATA.GEN_AI_TOOL_DESCRIPTION] + == "Get the weather for a location" + ) -@pytest.mark.parametrize("span_streaming", [True, False]) def test_error_handling( sentry_init, - capture_events, capture_items, mock_genai_client, - span_streaming, ): sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) - if span_streaming: - items = capture_items("event") - - # Mock an error at the HTTP level - with mock.patch.object( - mock_genai_client._api_client, "request", side_effect=Exception("API Error") - ), start_transaction(name="google_genai"), pytest.raises( - Exception, match="API Error" - ): - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", - contents="This will fail", - config=create_test_config(), - ) + items = capture_items("event") - (error_event,) = (item.payload for item in items) - else: - events = capture_events() - - # Mock an error at the HTTP level - with mock.patch.object( - mock_genai_client._api_client, "request", side_effect=Exception("API Error") - ), start_transaction(name="google_genai"), pytest.raises( - Exception, match="API Error" - ): - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", - contents="This will fail", - config=create_test_config(), - ) + # Mock an error at the HTTP level + with mock.patch.object( + mock_genai_client._api_client, "request", side_effect=Exception("API Error") + ), pytest.raises(Exception, match="API Error"): + mock_genai_client.models.generate_content( + model="gemini-1.5-flash", + contents="This will fail", + config=create_test_config(), + ) - # Should have both transaction and error events - assert len(events) == 2 - error_event, transaction_event = events + (error_event,) = (item.payload for item in items) assert error_event["level"] == "error" assert error_event["exception"]["values"][0]["type"] == "Exception" @@ -621,21 +443,17 @@ def test_error_handling( assert error_event["exception"]["values"][0]["mechanism"]["type"] == "google_genai" -@pytest.mark.parametrize("span_streaming", [True, False]) def test_streaming_generate_content( sentry_init, - capture_events, capture_items, mock_genai_client, - span_streaming, ): """Test streaming with generate_content_stream, verifying chunk accumulation.""" sentry_init( integrations=[GoogleGenAIIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) # Create streaming chunks - simulating a multi-chunk response @@ -699,208 +517,116 @@ def test_streaming_generate_content( # Create streaming mock responses stream_chunks = [chunk1_json, chunk2_json, chunk3_json] mock_stream = create_mock_streaming_responses(stream_chunks) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - mock_genai_client._api_client, "request_streamed", return_value=mock_stream - ), start_transaction(name="google_genai"): - config = create_test_config() - stream = mock_genai_client.models.generate_content_stream( - model="gemini-1.5-flash", - contents=[ - "Message demonstrating the absence of truncation.", - "Stream me a response", - ], - config=config, - ) - - # Consume the stream (this is what users do with the integration wrapper) - collected_chunks = list(stream) - - # Verify we got all chunks - assert len(collected_chunks) == 3 - assert collected_chunks[0].candidates[0].content.parts[0].text == "Hello! " - assert collected_chunks[1].candidates[0].content.parts[0].text == "How can I " - assert ( - collected_chunks[2].candidates[0].content.parts[0].text == "help you today?" - ) - - sentry_sdk.flush() - spans = [item.payload for item in items] - assert len(spans) == 1 - sentry_sdk.flush() - chat_span = next(item.payload for item in items) - - assert json.loads( - chat_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - ) == [ - { - "role": "user", - "content": [ - { - "type": "text", - "text": "Message demonstrating the absence of truncation.", - }, - { - "type": "text", - "text": "Stream me a response", - }, - ], - } - ] - - # Check that streaming flag is set on both spans - assert chat_span["attributes"][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["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] - ) - - assert chat_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["attributes"] - assert ( - chat_span["attributes"][SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS] == "STOP" - ) - assert chat_span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 - assert chat_span["attributes"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 10 - assert chat_span["attributes"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 25 - assert chat_span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 5 - assert ( - chat_span["attributes"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING] == 3 + with mock.patch.object( + mock_genai_client._api_client, "request_streamed", return_value=mock_stream + ): + config = create_test_config() + stream = mock_genai_client.models.generate_content_stream( + model="gemini-1.5-flash", + contents=[ + "Message demonstrating the absence of truncation.", + "Stream me a response", + ], + config=config, ) - # Verify model name - assert ( - chat_span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] == "gemini-1.5-flash" - ) - else: - events = capture_events() - - with mock.patch.object( - mock_genai_client._api_client, "request_streamed", return_value=mock_stream - ), start_transaction(name="google_genai"): - config = create_test_config() - stream = mock_genai_client.models.generate_content_stream( - model="gemini-1.5-flash", contents="Stream me a response", config=config - ) + # Consume the stream (this is what users do with the integration wrapper) + collected_chunks = list(stream) - # Consume the stream (this is what users do with the integration wrapper) - collected_chunks = list(stream) + # Verify we got all chunks + assert len(collected_chunks) == 3 + assert collected_chunks[0].candidates[0].content.parts[0].text == "Hello! " + assert collected_chunks[1].candidates[0].content.parts[0].text == "How can I " + assert collected_chunks[2].candidates[0].content.parts[0].text == "help you today?" - # Verify we got all chunks - assert len(collected_chunks) == 3 - assert collected_chunks[0].candidates[0].content.parts[0].text == "Hello! " - assert collected_chunks[1].candidates[0].content.parts[0].text == "How can I " - assert ( - collected_chunks[2].candidates[0].content.parts[0].text == "help you today?" - ) + sentry_sdk.flush() + spans = [item.payload for item in items] + assert len(spans) == 1 + sentry_sdk.flush() + chat_span = next(item.payload for item in items) - (event,) = events + assert json.loads(chat_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) == [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Message demonstrating the absence of truncation.", + }, + { + "type": "text", + "text": "Stream me a response", + }, + ], + } + ] - assert len(event["spans"]) == 1 - chat_span = event["spans"][0] + # Check that streaming flag is set on both spans + assert chat_span["attributes"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True - # Check that streaming flag is set on both spans - 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["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] + ) - # 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] - ) - assert chat_response_text == [expected_full_text] + assert chat_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 chat_span["data"][SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS] == "STOP" - assert chat_span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 - assert chat_span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 10 - assert chat_span["data"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 25 - assert chat_span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 5 - assert chat_span["data"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING] == 3 + # 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["attributes"] + assert chat_span["attributes"][SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS] == "STOP" + assert chat_span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 + assert chat_span["attributes"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 10 + assert chat_span["attributes"][SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 25 + assert chat_span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 5 + assert chat_span["attributes"][SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING] == 3 - # Verify model name - assert chat_span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "gemini-1.5-flash" + # Verify model name + assert chat_span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] == "gemini-1.5-flash" -@pytest.mark.parametrize("span_streaming", [True, False]) def test_span_origin( sentry_init, - capture_events, capture_items, mock_genai_client, - span_streaming, ): sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) + items = capture_items("span") - if span_streaming: - items = capture_items("span", "transaction") - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), sentry_sdk.traces.start_span(name="google_genai"): - config = create_test_config() - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents="Test origin", config=config - ) - - sentry_sdk.flush() - spans = [item.payload for item in items if item.type == "span"] - for span in spans: - if span["is_segment"] is True: - assert span["attributes"]["sentry.origin"] == "manual" - continue - - assert span["attributes"]["sentry.origin"] == "auto.ai.google_genai" - else: - events = capture_events() - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - config = create_test_config() - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents="Test origin", config=config - ) - - (event,) = events + with mock.patch.object( + mock_genai_client._api_client, "request", return_value=mock_http_response + ): + config = create_test_config() + mock_genai_client.models.generate_content( + model="gemini-1.5-flash", contents="Test origin", config=config + ) - assert event["contexts"]["trace"]["origin"] == "manual" - for span in event["spans"]: - assert span["origin"] == "auto.ai.google_genai" + sentry_sdk.flush() + spans = [item.payload for item in items] + for span in spans: + assert span["attributes"]["sentry.origin"] == "auto.ai.google_genai" -@pytest.mark.parametrize("span_streaming", [True, False]) def test_response_without_usage_metadata( sentry_init, - capture_events, capture_items, mock_genai_client, - span_streaming, ): """Test handling of responses without usage metadata""" sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) # Response without usage metadata @@ -917,60 +643,36 @@ def test_response_without_usage_metadata( } mock_http_response = create_mock_http_response(response_json) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - config = create_test_config() - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents="Test", config=config - ) - - sentry_sdk.flush() - chat_span = next(item.payload for item in items) - - # Usage data should not be present - assert SPANDATA.GEN_AI_USAGE_INPUT_TOKENS not in chat_span["attributes"] - assert SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS not in chat_span["attributes"] - assert SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS not in chat_span["attributes"] - else: - events = capture_events() - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - config = create_test_config() - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents="Test", config=config - ) + with mock.patch.object( + mock_genai_client._api_client, "request", return_value=mock_http_response + ): + config = create_test_config() + mock_genai_client.models.generate_content( + model="gemini-1.5-flash", contents="Test", config=config + ) - (event,) = events - chat_span = event["spans"][0] + sentry_sdk.flush() + chat_span = next(item.payload for item in items) - # Usage data should not be present - assert SPANDATA.GEN_AI_USAGE_INPUT_TOKENS not in chat_span["data"] - assert SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS not in chat_span["data"] - assert SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS not in chat_span["data"] + # Usage data should not be present + assert SPANDATA.GEN_AI_USAGE_INPUT_TOKENS not in chat_span["attributes"] + assert SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS not in chat_span["attributes"] + assert SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS not in chat_span["attributes"] -@pytest.mark.parametrize("span_streaming", [True, False]) def test_multiple_candidates( sentry_init, - capture_events, capture_items, mock_genai_client, - span_streaming, ): """Test handling of multiple response candidates""" sentry_init( integrations=[GoogleGenAIIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) # Response with multiple candidates @@ -999,225 +701,126 @@ def test_multiple_candidates( } mock_http_response = create_mock_http_response(multi_candidate_json) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - config = create_test_config() - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents="Generate multiple", config=config - ) - - sentry_sdk.flush() - chat_span = next(item.payload for item in items) - - # Should capture all responses - # Response text is stored as a JSON string when there are multiple responses - response_text = chat_span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] - - if isinstance(response_text, str) and response_text.startswith("["): - # It's a JSON array - response_list = json.loads(response_text) - assert response_list == ["Response 1", "Response 2"] - else: - # It's concatenated - assert response_text == "Response 1\nResponse 2" - - # Finish reasons are serialized as JSON - finish_reasons = json.loads( - chat_span["attributes"][SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS] + with mock.patch.object( + mock_genai_client._api_client, "request", return_value=mock_http_response + ): + config = create_test_config() + mock_genai_client.models.generate_content( + model="gemini-1.5-flash", contents="Generate multiple", config=config ) - else: - events = capture_events() - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - config = create_test_config() - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents="Generate multiple", config=config - ) - (event,) = events - chat_span = event["spans"][0] + sentry_sdk.flush() + chat_span = next(item.payload for item in items) - # Should capture all responses - # Response text is stored as a JSON string when there are multiple responses - response_text = chat_span["data"][SPANDATA.GEN_AI_RESPONSE_TEXT] + # Should capture all responses + # Response text is stored as a JSON string when there are multiple responses + response_text = chat_span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] - if isinstance(response_text, str) and response_text.startswith("["): - # It's a JSON array - response_list = json.loads(response_text) - assert response_list == ["Response 1", "Response 2"] - else: - # It's concatenated - assert response_text == "Response 1\nResponse 2" + if isinstance(response_text, str) and response_text.startswith("["): + # It's a JSON array + response_list = json.loads(response_text) + assert response_list == ["Response 1", "Response 2"] + else: + # It's concatenated + assert response_text == "Response 1\nResponse 2" - # Finish reasons are serialized as JSON - finish_reasons = json.loads( - chat_span["data"][SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS] - ) + # Finish reasons are serialized as JSON + finish_reasons = json.loads( + chat_span["attributes"][SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS] + ) assert finish_reasons == ["STOP", "MAX_TOKENS"] -@pytest.mark.parametrize("span_streaming", [True, False]) def test_all_configuration_parameters( sentry_init, - capture_events, capture_items, mock_genai_client, - span_streaming, ): """Test that all configuration parameters are properly recorded""" sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - config = create_test_config( - temperature=0.8, - top_p=0.95, - top_k=40, - max_output_tokens=2048, - presence_penalty=0.1, - frequency_penalty=0.2, - seed=12345, - ) - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents="Test all params", config=config - ) - - sentry_sdk.flush() - invoke_span = next(item.payload for item in items) - - # Check all parameters are recorded - assert invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.8 - assert invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.95 - assert invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_TOP_K] == 40 - assert invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 2048 - assert ( - invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY] == 0.1 + with mock.patch.object( + mock_genai_client._api_client, "request", return_value=mock_http_response + ): + config = create_test_config( + temperature=0.8, + top_p=0.95, + top_k=40, + max_output_tokens=2048, + presence_penalty=0.1, + frequency_penalty=0.2, + seed=12345, ) - assert ( - invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY] == 0.2 + mock_genai_client.models.generate_content( + model="gemini-1.5-flash", contents="Test all params", config=config ) - assert invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_SEED] == 12345 - else: - events = capture_events() - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - config = create_test_config( - temperature=0.8, - top_p=0.95, - top_k=40, - max_output_tokens=2048, - presence_penalty=0.1, - frequency_penalty=0.2, - seed=12345, - ) - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents="Test all params", config=config - ) - (event,) = events - invoke_span = event["spans"][0] + sentry_sdk.flush() + invoke_span = next(item.payload for item in items) - # Check all parameters are recorded - assert invoke_span["data"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.8 - assert invoke_span["data"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.95 - assert invoke_span["data"][SPANDATA.GEN_AI_REQUEST_TOP_K] == 40 - assert invoke_span["data"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 2048 - assert invoke_span["data"][SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY] == 0.1 - assert invoke_span["data"][SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY] == 0.2 - assert invoke_span["data"][SPANDATA.GEN_AI_REQUEST_SEED] == 12345 + # Check all parameters are recorded + assert invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.8 + assert invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.95 + assert invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_TOP_K] == 40 + assert invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 2048 + assert invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY] == 0.1 + assert invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY] == 0.2 + assert invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_SEED] == 12345 -@pytest.mark.parametrize("span_streaming", [True, False]) def test_empty_response( sentry_init, - capture_events, capture_items, mock_genai_client, - span_streaming, ): """Test handling of minimal response with no content""" sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) # Minimal response with empty candidates array minimal_response_json = {"candidates": []} mock_http_response = create_mock_http_response(minimal_response_json) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - response = mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents="Test", config=create_test_config() - ) - - # Response will have an empty candidates list - assert response is not None - assert len(response.candidates) == 0 - - # Should still create spans even with empty candidates - sentry_sdk.flush() - spans = [item.payload for item in items] - assert len(spans) == 1 - else: - events = capture_events() - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - response = mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents="Test", config=create_test_config() - ) + with mock.patch.object( + mock_genai_client._api_client, "request", return_value=mock_http_response + ): + response = mock_genai_client.models.generate_content( + model="gemini-1.5-flash", contents="Test", config=create_test_config() + ) - # Response will have an empty candidates list - assert response is not None - assert len(response.candidates) == 0 + # Response will have an empty candidates list + assert response is not None + assert len(response.candidates) == 0 - (event,) = events - # Should still create spans even with empty candidates - assert len(event["spans"]) == 1 + # Should still create spans even with empty candidates + sentry_sdk.flush() + spans = [item.payload for item in items] + assert len(spans) == 1 -@pytest.mark.parametrize("span_streaming", [True, False]) def test_response_with_different_id_fields( sentry_init, - capture_events, capture_items, mock_genai_client, - span_streaming, ): """Test handling of different response ID field names""" sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) # Response with response_id and model_version @@ -1236,42 +839,23 @@ def test_response_with_different_id_fields( } mock_http_response = create_mock_http_response(response_json) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents="Test", config=create_test_config() - ) - - sentry_sdk.flush() - chat_span = next(item.payload for item in items) - - assert chat_span["attributes"][SPANDATA.GEN_AI_RESPONSE_ID] == "resp-456" - assert ( - chat_span["attributes"][SPANDATA.GEN_AI_RESPONSE_MODEL] - == "gemini-1.5-flash-001" + with mock.patch.object( + mock_genai_client._api_client, "request", return_value=mock_http_response + ): + mock_genai_client.models.generate_content( + model="gemini-1.5-flash", contents="Test", config=create_test_config() ) - else: - events = capture_events() - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents="Test", config=create_test_config() - ) - (event,) = events - chat_span = event["spans"][0] + sentry_sdk.flush() + chat_span = next(item.payload for item in items) - 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" - ) + assert chat_span["attributes"][SPANDATA.GEN_AI_RESPONSE_ID] == "resp-456" + assert ( + chat_span["attributes"][SPANDATA.GEN_AI_RESPONSE_MODEL] + == "gemini-1.5-flash-001" + ) def test_tool_with_async_function(sentry_init): @@ -1279,7 +863,6 @@ def test_tool_with_async_function(sentry_init): sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - stream_gen_ai_spans=False, ) # Create an async tool function @@ -1296,74 +879,49 @@ async def async_tool(param: str) -> str: assert hasattr(wrapped_async_tool, "__wrapped__") # Should preserve original -@pytest.mark.parametrize("span_streaming", [True, False]) def test_contents_as_none( sentry_init, - capture_events, capture_items, mock_genai_client, - span_streaming, ): """Test handling when contents parameter is None""" sentry_init( integrations=[GoogleGenAIIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents=None, config=create_test_config() - ) - - sentry_sdk.flush() - invoke_span = next(item.payload for item in items) - - # Should handle None contents gracefully - messages = invoke_span["attributes"].get(SPANDATA.GEN_AI_REQUEST_MESSAGES, []) - else: - events = capture_events() - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents=None, config=create_test_config() - ) + with mock.patch.object( + mock_genai_client._api_client, "request", return_value=mock_http_response + ): + mock_genai_client.models.generate_content( + model="gemini-1.5-flash", contents=None, config=create_test_config() + ) - (event,) = events - invoke_span = event["spans"][0] + sentry_sdk.flush() + invoke_span = next(item.payload for item in items) - # Should handle None contents gracefully - messages = invoke_span["data"].get(SPANDATA.GEN_AI_REQUEST_MESSAGES, []) + # Should handle None contents gracefully + messages = invoke_span["attributes"].get(SPANDATA.GEN_AI_REQUEST_MESSAGES, []) # Should only have system message if any, not user message assert all(msg["role"] != "user" or msg["content"] is not None for msg in messages) -@pytest.mark.parametrize("span_streaming", [True, False]) def test_tool_calls_extraction( sentry_init, - capture_events, capture_items, mock_genai_client, - span_streaming, ): """Test extraction of tool/function calls from response""" sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) # Response with function calls @@ -1402,49 +960,27 @@ def test_tool_calls_extraction( } mock_http_response = create_mock_http_response(function_call_response_json) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", - contents="What's the weather and time?", - config=create_test_config(), - ) - - sentry_sdk.flush() - chat_span = next(item.payload for item in items) # The chat span - - # Check that tool calls are extracted and stored - assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS in chat_span["attributes"] - - # Parse the JSON string to verify content - tool_calls = json.loads( - chat_span["attributes"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS] + with mock.patch.object( + mock_genai_client._api_client, "request", return_value=mock_http_response + ): + mock_genai_client.models.generate_content( + model="gemini-1.5-flash", + contents="What's the weather and time?", + config=create_test_config(), ) - else: - events = capture_events() - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", - contents="What's the weather and time?", - config=create_test_config(), - ) - (event,) = events - chat_span = event["spans"][0] # The chat span + sentry_sdk.flush() + chat_span = next(item.payload for item in items) # The chat span - # Check that tool calls are extracted and stored - assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS in chat_span["data"] + # Check that tool calls are extracted and stored + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS in chat_span["attributes"] - # Parse the JSON string to verify content - tool_calls = json.loads(chat_span["data"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS]) + # Parse the JSON string to verify content + tool_calls = json.loads( + chat_span["attributes"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS] + ) assert len(tool_calls) == 2 @@ -1464,51 +1000,6 @@ def test_tool_calls_extraction( assert json.loads(tool_calls[1]["arguments"]) == {"timezone": "PST"} -def test_google_genai_message_truncation( - sentry_init, capture_events, mock_genai_client -): - """Test that large messages are truncated properly in Google GenAI integration.""" - sentry_init( - integrations=[GoogleGenAIIntegration(include_prompts=True)], - traces_sample_rate=1.0, - send_default_pii=True, - stream_gen_ai_spans=False, - ) - events = capture_events() - - large_content = ( - "This is a very long message that will exceed our size limits. " * 1000 - ) - small_content = "This is a small user message" - - mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ): - with start_transaction(name="google_genai"): - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", - contents=[large_content, small_content], - config=create_test_config(), - ) - - (event,) = events - invoke_span = event["spans"][0] - assert SPANDATA.GEN_AI_REQUEST_MESSAGES in invoke_span["data"] - - messages_data = invoke_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - assert isinstance(messages_data, str) - - parsed_messages = json.loads(messages_data) - assert isinstance(parsed_messages, list) - assert len(parsed_messages) == 1 - assert parsed_messages[0]["role"] == "user" - - # What "small content" becomes because the large message used the entire character limit - assert "..." in parsed_messages[0]["content"][1]["text"] - - # Sample embed content API response JSON EXAMPLE_EMBED_RESPONSE_JSON = { "embeddings": [ @@ -1533,7 +1024,6 @@ def test_google_genai_message_truncation( } -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize( "send_default_pii, include_prompts", [ @@ -1545,139 +1035,79 @@ def test_google_genai_message_truncation( ) def test_embed_content( sentry_init, - capture_events, capture_items, send_default_pii, include_prompts, mock_genai_client, - span_streaming, ): sentry_init( integrations=[GoogleGenAIIntegration(include_prompts=include_prompts)], traces_sample_rate=1.0, send_default_pii=send_default_pii, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) # Mock the HTTP response at the _api_client.request() level mock_http_response = create_mock_http_response(EXAMPLE_EMBED_RESPONSE_JSON) + items = capture_items("span") - if span_streaming: - items = capture_items("transaction", "span") - - with mock.patch.object( - mock_genai_client._api_client, - "request", - return_value=mock_http_response, - ), sentry_sdk.traces.start_span(name="google_genai_embeddings"): - mock_genai_client.models.embed_content( - model="text-embedding-004", - contents=[ - "What is your name?", - "What is your favorite color?", - ], - ) - - # Should have 1 span for embeddings - sentry_sdk.flush() - spans = [item.payload for item in items if item.type == "span"] - assert len(spans) == 2 - assert spans[1]["name"] == "google_genai_embeddings" - (embed_span, _) = spans - - # Check embeddings span - assert embed_span["attributes"]["sentry.op"] == OP.GEN_AI_EMBEDDINGS - assert embed_span["name"] == "embeddings text-embedding-004" - assert embed_span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "embeddings" - assert embed_span["attributes"][SPANDATA.GEN_AI_SYSTEM] == "gcp.gemini" - assert ( - embed_span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] - == "text-embedding-004" - ) - - # Check input texts if PII is allowed - if send_default_pii and include_prompts: - input_texts = json.loads( - embed_span["attributes"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT] - ) - assert input_texts == [ + with mock.patch.object( + mock_genai_client._api_client, + "request", + return_value=mock_http_response, + ): + mock_genai_client.models.embed_content( + model="text-embedding-004", + contents=[ "What is your name?", "What is your favorite color?", - ] - else: - assert SPANDATA.GEN_AI_EMBEDDINGS_INPUT not in embed_span["attributes"] - - # Check usage data (sum of token counts from statistics: 10 + 15 = 25) - # Note: Only available in newer versions with ContentEmbeddingStatistics - if SPANDATA.GEN_AI_USAGE_INPUT_TOKENS in embed_span["attributes"]: - assert embed_span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 25 - else: - events = capture_events() - - with mock.patch.object( - mock_genai_client._api_client, - "request", - return_value=mock_http_response, - ), start_transaction(name="google_genai_embeddings"): - mock_genai_client.models.embed_content( - model="text-embedding-004", - contents=[ - "What is your name?", - "What is your favorite color?", - ], - ) - - assert len(events) == 1 - (event,) = events - - assert event["type"] == "transaction" - assert event["transaction"] == "google_genai_embeddings" + ], + ) - # Should have 1 span for embeddings - assert len(event["spans"]) == 1 - (embed_span,) = event["spans"] + # Should have 1 span for embeddings + sentry_sdk.flush() + spans = [item.payload for item in items] + assert len(spans) == 1 + (embed_span,) = spans - # Check embeddings span - assert embed_span["op"] == OP.GEN_AI_EMBEDDINGS - assert embed_span["description"] == "embeddings text-embedding-004" - assert embed_span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "embeddings" - assert embed_span["data"][SPANDATA.GEN_AI_SYSTEM] == "gcp.gemini" - assert embed_span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "text-embedding-004" + # Check embeddings span + assert embed_span["attributes"]["sentry.op"] == OP.GEN_AI_EMBEDDINGS + assert embed_span["name"] == "embeddings text-embedding-004" + assert embed_span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "embeddings" + assert embed_span["attributes"][SPANDATA.GEN_AI_SYSTEM] == "gcp.gemini" + assert ( + embed_span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] == "text-embedding-004" + ) - # Check input texts if PII is allowed - if send_default_pii and include_prompts: - input_texts = json.loads( - embed_span["data"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT] - ) - assert input_texts == [ - "What is your name?", - "What is your favorite color?", - ] - else: - assert SPANDATA.GEN_AI_EMBEDDINGS_INPUT not in embed_span["data"] + # Check input texts if PII is allowed + if send_default_pii and include_prompts: + input_texts = json.loads( + embed_span["attributes"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT] + ) + assert input_texts == [ + "What is your name?", + "What is your favorite color?", + ] + else: + assert SPANDATA.GEN_AI_EMBEDDINGS_INPUT not in embed_span["attributes"] - # Check usage data (sum of token counts from statistics: 10 + 15 = 25) - # Note: Only available in newer versions with ContentEmbeddingStatistics - if SPANDATA.GEN_AI_USAGE_INPUT_TOKENS in embed_span["data"]: - assert embed_span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 25 + # Check usage data (sum of token counts from statistics: 10 + 15 = 25) + # Note: Only available in newer versions with ContentEmbeddingStatistics + if SPANDATA.GEN_AI_USAGE_INPUT_TOKENS in embed_span["attributes"]: + assert embed_span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 25 -@pytest.mark.parametrize("span_streaming", [True, False]) def test_embed_content_string_input( sentry_init, - capture_events, capture_items, mock_genai_client, - span_streaming, ): """Test embed_content with a single string instead of list.""" sentry_init( integrations=[GoogleGenAIIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) # Mock response with single embedding @@ -1694,108 +1124,57 @@ def test_embed_content_string_input( "metadata": { "billableCharacterCount": 10, }, - } - mock_http_response = create_mock_http_response(single_embed_response) - - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai_embeddings"): - mock_genai_client.models.embed_content( - model="text-embedding-004", - contents="Single text input", - ) - - sentry_sdk.flush() - spans = [item.payload for item in items] - (embed_span,) = spans + } + mock_http_response = create_mock_http_response(single_embed_response) + items = capture_items("span") - # Check that single string is handled correctly - input_texts = json.loads( - embed_span["attributes"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT] + with mock.patch.object( + mock_genai_client._api_client, "request", return_value=mock_http_response + ): + mock_genai_client.models.embed_content( + model="text-embedding-004", + contents="Single text input", ) - assert input_texts == ["Single text input"] - # Should use token_count from statistics (5), not billable_character_count (10) - # Note: Only available in newer versions with ContentEmbeddingStatistics - if SPANDATA.GEN_AI_USAGE_INPUT_TOKENS in embed_span["attributes"]: - assert embed_span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 5 - else: - events = capture_events() - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai_embeddings"): - mock_genai_client.models.embed_content( - model="text-embedding-004", - contents="Single text input", - ) + sentry_sdk.flush() + spans = [item.payload for item in items] + (embed_span,) = spans - (event,) = events - (embed_span,) = event["spans"] + # Check that single string is handled correctly + input_texts = json.loads(embed_span["attributes"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT]) - # Check that single string is handled correctly - input_texts = json.loads(embed_span["data"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT]) - assert input_texts == ["Single text input"] - # Should use token_count from statistics (5), not billable_character_count (10) - # Note: Only available in newer versions with ContentEmbeddingStatistics - if SPANDATA.GEN_AI_USAGE_INPUT_TOKENS in embed_span["data"]: - assert embed_span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 5 + assert input_texts == ["Single text input"] + # Should use token_count from statistics (5), not billable_character_count (10) + # Note: Only available in newer versions with ContentEmbeddingStatistics + if SPANDATA.GEN_AI_USAGE_INPUT_TOKENS in embed_span["attributes"]: + assert embed_span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 5 -@pytest.mark.parametrize("span_streaming", [True, False]) def test_embed_content_error_handling( sentry_init, - capture_events, capture_items, mock_genai_client, - span_streaming, ): """Test error handling in embed_content.""" sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) - if span_streaming: - items = capture_items("event") - - # Mock an error at the HTTP level - with mock.patch.object( - mock_genai_client._api_client, - "request", - side_effect=Exception("Embedding API Error"), - ), start_transaction(name="google_genai_embeddings"), pytest.raises( - Exception, match="Embedding API Error" - ): - mock_genai_client.models.embed_content( - model="text-embedding-004", - contents=["This will fail"], - ) + items = capture_items("event") - (error_event,) = (item.payload for item in items) - else: - events = capture_events() - - # Mock an error at the HTTP level - with mock.patch.object( - mock_genai_client._api_client, - "request", - side_effect=Exception("Embedding API Error"), - ), start_transaction(name="google_genai_embeddings"), pytest.raises( - Exception, match="Embedding API Error" - ): - mock_genai_client.models.embed_content( - model="text-embedding-004", - contents=["This will fail"], - ) + # Mock an error at the HTTP level + with mock.patch.object( + mock_genai_client._api_client, + "request", + side_effect=Exception("Embedding API Error"), + ), pytest.raises(Exception, match="Embedding API Error"): + mock_genai_client.models.embed_content( + model="text-embedding-004", + contents=["This will fail"], + ) - # Should have both transaction and error events - assert len(events) == 2 - error_event, _ = events + (error_event,) = (item.payload for item in items) assert error_event["level"] == "error" assert error_event["exception"]["values"][0]["type"] == "Exception" @@ -1803,20 +1182,16 @@ def test_embed_content_error_handling( assert error_event["exception"]["values"][0]["mechanism"]["type"] == "google_genai" -@pytest.mark.parametrize("span_streaming", [True, False]) def test_embed_content_without_statistics( sentry_init, - capture_events, capture_items, mock_genai_client, - span_streaming, ): """Test embed_content response without statistics (older package versions).""" sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) # Response without statistics (typical for older google-genai versions) @@ -1832,97 +1207,53 @@ def test_embed_content_without_statistics( ], } mock_http_response = create_mock_http_response(old_version_response) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai_embeddings"): - mock_genai_client.models.embed_content( - model="text-embedding-004", - contents=["Test without statistics", "Another test"], - ) - - sentry_sdk.flush() - spans = [item.payload for item in items] - (embed_span,) = spans - - # No usage tokens since there are no statistics in older versions - # This is expected and the integration should handle it gracefully - assert SPANDATA.GEN_AI_USAGE_INPUT_TOKENS not in embed_span["attributes"] - else: - events = capture_events() - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai_embeddings"): - mock_genai_client.models.embed_content( - model="text-embedding-004", - contents=["Test without statistics", "Another test"], - ) + with mock.patch.object( + mock_genai_client._api_client, "request", return_value=mock_http_response + ): + mock_genai_client.models.embed_content( + model="text-embedding-004", + contents=["Test without statistics", "Another test"], + ) - (event,) = events - (embed_span,) = event["spans"] + sentry_sdk.flush() + spans = [item.payload for item in items] + (embed_span,) = spans - # No usage tokens since there are no statistics in older versions - # This is expected and the integration should handle it gracefully - assert SPANDATA.GEN_AI_USAGE_INPUT_TOKENS not in embed_span["data"] + # No usage tokens since there are no statistics in older versions + # This is expected and the integration should handle it gracefully + assert SPANDATA.GEN_AI_USAGE_INPUT_TOKENS not in embed_span["attributes"] -@pytest.mark.parametrize("span_streaming", [True, False]) def test_embed_content_span_origin( sentry_init, - capture_events, capture_items, mock_genai_client, - span_streaming, ): """Test that embed_content spans have correct origin.""" sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) mock_http_response = create_mock_http_response(EXAMPLE_EMBED_RESPONSE_JSON) - if span_streaming: - items = capture_items("transaction", "span") - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), sentry_sdk.traces.start_span(name="google_genai_embeddings"): - mock_genai_client.models.embed_content( - model="text-embedding-004", - contents=["Test origin"], - ) - - sentry_sdk.flush() - spans = [item.payload for item in items if item.type == "span"] - for span in spans: - if span["is_segment"] is True: - assert span["attributes"]["sentry.origin"] == "manual" - continue - - assert span["attributes"]["sentry.origin"] == "auto.ai.google_genai" - else: - events = capture_events() - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai_embeddings"): - mock_genai_client.models.embed_content( - model="text-embedding-004", - contents=["Test origin"], - ) - - (event,) = events + items = capture_items("span") + with mock.patch.object( + mock_genai_client._api_client, "request", return_value=mock_http_response + ): + mock_genai_client.models.embed_content( + model="text-embedding-004", + contents=["Test origin"], + ) - assert event["contexts"]["trace"]["origin"] == "manual" - for span in event["spans"]: - assert span["origin"] == "auto.ai.google_genai" + sentry_sdk.flush() + spans = [item.payload for item in items] + for span in spans: + assert span["attributes"]["sentry.origin"] == "auto.ai.google_genai" -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio @pytest.mark.parametrize( "send_default_pii, include_prompts", @@ -1935,142 +1266,81 @@ def test_embed_content_span_origin( ) async def test_async_embed_content( sentry_init, - capture_events, capture_items, send_default_pii, include_prompts, mock_genai_client, - span_streaming, ): """Test async embed_content method.""" sentry_init( integrations=[GoogleGenAIIntegration(include_prompts=include_prompts)], traces_sample_rate=1.0, send_default_pii=send_default_pii, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) # Mock the async HTTP response mock_http_response = create_mock_http_response(EXAMPLE_EMBED_RESPONSE_JSON) + items = capture_items("span") - if span_streaming: - items = capture_items("transaction", "span") - - with mock.patch.object( - mock_genai_client._api_client, - "async_request", - return_value=mock_http_response, - ), sentry_sdk.traces.start_span(name="google_genai_embeddings_async"): - await mock_genai_client.aio.models.embed_content( - model="text-embedding-004", - contents=[ - "What is your name?", - "What is your favorite color?", - ], - ) - - # Should have 1 span for embeddings - sentry_sdk.flush() - spans = [item.payload for item in items if item.type == "span"] - assert len(spans) == 2 - assert spans[1]["name"] == "google_genai_embeddings_async" - (embed_span, _) = spans - - # Check embeddings span - assert embed_span["attributes"]["sentry.op"] == OP.GEN_AI_EMBEDDINGS - assert embed_span["name"] == "embeddings text-embedding-004" - assert embed_span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "embeddings" - assert embed_span["attributes"][SPANDATA.GEN_AI_SYSTEM] == "gcp.gemini" - assert ( - embed_span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] - == "text-embedding-004" - ) - - # Check input texts if PII is allowed - if send_default_pii and include_prompts: - input_texts = json.loads( - embed_span["attributes"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT] - ) - assert input_texts == [ + with mock.patch.object( + mock_genai_client._api_client, + "async_request", + return_value=mock_http_response, + ): + await mock_genai_client.aio.models.embed_content( + model="text-embedding-004", + contents=[ "What is your name?", "What is your favorite color?", - ] - else: - assert SPANDATA.GEN_AI_EMBEDDINGS_INPUT not in embed_span["attributes"] - - # Check usage data (sum of token counts from statistics: 10 + 15 = 25) - # Note: Only available in newer versions with ContentEmbeddingStatistics - if SPANDATA.GEN_AI_USAGE_INPUT_TOKENS in embed_span["attributes"]: - assert embed_span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 25 - else: - events = capture_events() - - with mock.patch.object( - mock_genai_client._api_client, - "async_request", - return_value=mock_http_response, - ), start_transaction(name="google_genai_embeddings_async"): - await mock_genai_client.aio.models.embed_content( - model="text-embedding-004", - contents=[ - "What is your name?", - "What is your favorite color?", - ], - ) - - assert len(events) == 1 - (event,) = events - - assert event["type"] == "transaction" - - assert event["transaction"] == "google_genai_embeddings_async" + ], + ) - # Should have 1 span for embeddings - assert len(event["spans"]) == 1 - (embed_span,) = event["spans"] + # Should have 1 span for embeddings + sentry_sdk.flush() + spans = [item.payload for item in items] + assert len(spans) == 1 + (embed_span,) = spans - # Check embeddings span - assert embed_span["op"] == OP.GEN_AI_EMBEDDINGS - assert embed_span["description"] == "embeddings text-embedding-004" - assert embed_span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "embeddings" - assert embed_span["data"][SPANDATA.GEN_AI_SYSTEM] == "gcp.gemini" - assert embed_span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "text-embedding-004" + # Check embeddings span + assert embed_span["attributes"]["sentry.op"] == OP.GEN_AI_EMBEDDINGS + assert embed_span["name"] == "embeddings text-embedding-004" + assert embed_span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "embeddings" + assert embed_span["attributes"][SPANDATA.GEN_AI_SYSTEM] == "gcp.gemini" + assert ( + embed_span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] == "text-embedding-004" + ) - # Check input texts if PII is allowed - if send_default_pii and include_prompts: - input_texts = json.loads( - embed_span["data"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT] - ) - assert input_texts == [ - "What is your name?", - "What is your favorite color?", - ] - else: - assert SPANDATA.GEN_AI_EMBEDDINGS_INPUT not in embed_span["data"] + # Check input texts if PII is allowed + if send_default_pii and include_prompts: + input_texts = json.loads( + embed_span["attributes"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT] + ) + assert input_texts == [ + "What is your name?", + "What is your favorite color?", + ] + else: + assert SPANDATA.GEN_AI_EMBEDDINGS_INPUT not in embed_span["attributes"] - # Check usage data (sum of token counts from statistics: 10 + 15 = 25) - # Note: Only available in newer versions with ContentEmbeddingStatistics - if SPANDATA.GEN_AI_USAGE_INPUT_TOKENS in embed_span["data"]: - assert embed_span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 25 + # Check usage data (sum of token counts from statistics: 10 + 15 = 25) + # Note: Only available in newer versions with ContentEmbeddingStatistics + if SPANDATA.GEN_AI_USAGE_INPUT_TOKENS in embed_span["attributes"]: + assert embed_span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 25 -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio async def test_async_embed_content_string_input( sentry_init, - capture_events, capture_items, mock_genai_client, - span_streaming, ): """Test async embed_content with a single string instead of list.""" sentry_init( integrations=[GoogleGenAIIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) # Mock response with single embedding @@ -2089,111 +1359,57 @@ async def test_async_embed_content_string_input( }, } mock_http_response = create_mock_http_response(single_embed_response) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - mock_genai_client._api_client, - "async_request", - return_value=mock_http_response, - ), start_transaction(name="google_genai_embeddings_async"): - await mock_genai_client.aio.models.embed_content( - model="text-embedding-004", - contents="Single text input", - ) - - sentry_sdk.flush() - spans = [item.payload for item in items] - (embed_span,) = spans - - # Check that single string is handled correctly - input_texts = json.loads( - embed_span["attributes"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT] + with mock.patch.object( + mock_genai_client._api_client, + "async_request", + return_value=mock_http_response, + ): + await mock_genai_client.aio.models.embed_content( + model="text-embedding-004", + contents="Single text input", ) - assert input_texts == ["Single text input"] - # Should use token_count from statistics (5), not billable_character_count (10) - # Note: Only available in newer versions with ContentEmbeddingStatistics - if SPANDATA.GEN_AI_USAGE_INPUT_TOKENS in embed_span["attributes"]: - assert embed_span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 5 - else: - events = capture_events() - - with mock.patch.object( - mock_genai_client._api_client, - "async_request", - return_value=mock_http_response, - ), start_transaction(name="google_genai_embeddings_async"): - await mock_genai_client.aio.models.embed_content( - model="text-embedding-004", - contents="Single text input", - ) - (event,) = events - (embed_span,) = event["spans"] + sentry_sdk.flush() + spans = [item.payload for item in items] + (embed_span,) = spans - # Check that single string is handled correctly - input_texts = json.loads(embed_span["data"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT]) - assert input_texts == ["Single text input"] - # Should use token_count from statistics (5), not billable_character_count (10) - # Note: Only available in newer versions with ContentEmbeddingStatistics - if SPANDATA.GEN_AI_USAGE_INPUT_TOKENS in embed_span["data"]: - assert embed_span["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 5 + # Check that single string is handled correctly + input_texts = json.loads(embed_span["attributes"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT]) + assert input_texts == ["Single text input"] + # Should use token_count from statistics (5), not billable_character_count (10) + # Note: Only available in newer versions with ContentEmbeddingStatistics + if SPANDATA.GEN_AI_USAGE_INPUT_TOKENS in embed_span["attributes"]: + assert embed_span["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 5 -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio async def test_async_embed_content_error_handling( sentry_init, - capture_events, capture_items, mock_genai_client, - span_streaming, ): """Test error handling in async embed_content.""" sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) + items = capture_items("event") - if span_streaming: - items = capture_items("event") - - # Mock an error at the HTTP level - with mock.patch.object( - mock_genai_client._api_client, - "async_request", - side_effect=Exception("Async Embedding API Error"), - ), start_transaction(name="google_genai_embeddings_async"), pytest.raises( - Exception, match="Async Embedding API Error" - ): - await mock_genai_client.aio.models.embed_content( - model="text-embedding-004", - contents=["This will fail"], - ) - - (error_event,) = (item.payload for item in items) - else: - events = capture_events() - - # Mock an error at the HTTP level - with mock.patch.object( - mock_genai_client._api_client, - "async_request", - side_effect=Exception("Async Embedding API Error"), - ), start_transaction(name="google_genai_embeddings_async"), pytest.raises( - Exception, match="Async Embedding API Error" - ): - await mock_genai_client.aio.models.embed_content( - model="text-embedding-004", - contents=["This will fail"], - ) + # Mock an error at the HTTP level + with mock.patch.object( + mock_genai_client._api_client, + "async_request", + side_effect=Exception("Async Embedding API Error"), + ), pytest.raises(Exception, match="Async Embedding API Error"): + await mock_genai_client.aio.models.embed_content( + model="text-embedding-004", + contents=["This will fail"], + ) - # Should have both transaction and error events - assert len(events) == 2 - error_event, _ = events + (error_event,) = (item.payload for item in items) assert error_event["level"] == "error" assert error_event["exception"]["values"][0]["type"] == "Exception" @@ -2201,21 +1417,17 @@ 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.asyncio async def test_async_embed_content_without_statistics( sentry_init, - capture_events, capture_items, mock_genai_client, - span_streaming, ): """Test async embed_content response without statistics (older package versions).""" sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) # Response without statistics (typical for older google-genai versions) @@ -2231,124 +1443,71 @@ async def test_async_embed_content_without_statistics( ], } mock_http_response = create_mock_http_response(old_version_response) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - mock_genai_client._api_client, - "async_request", - return_value=mock_http_response, - ), start_transaction(name="google_genai_embeddings_async"): - await mock_genai_client.aio.models.embed_content( - model="text-embedding-004", - contents=["Test without statistics", "Another test"], - ) - - sentry_sdk.flush() - spans = [item.payload for item in items] - (embed_span,) = spans - - # No usage tokens since there are no statistics in older versions - # This is expected and the integration should handle it gracefully - assert SPANDATA.GEN_AI_USAGE_INPUT_TOKENS not in embed_span["attributes"] - else: - events = capture_events() - - with mock.patch.object( - mock_genai_client._api_client, - "async_request", - return_value=mock_http_response, - ), start_transaction(name="google_genai_embeddings_async"): - await mock_genai_client.aio.models.embed_content( - model="text-embedding-004", - contents=["Test without statistics", "Another test"], - ) + with mock.patch.object( + mock_genai_client._api_client, + "async_request", + return_value=mock_http_response, + ): + await mock_genai_client.aio.models.embed_content( + model="text-embedding-004", + contents=["Test without statistics", "Another test"], + ) - (event,) = events - (embed_span,) = event["spans"] + sentry_sdk.flush() + spans = [item.payload for item in items] + (embed_span,) = spans - # No usage tokens since there are no statistics in older versions - # This is expected and the integration should handle it gracefully - assert SPANDATA.GEN_AI_USAGE_INPUT_TOKENS not in embed_span["data"] + # No usage tokens since there are no statistics in older versions + # This is expected and the integration should handle it gracefully + assert SPANDATA.GEN_AI_USAGE_INPUT_TOKENS not in embed_span["attributes"] -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.asyncio async def test_async_embed_content_span_origin( sentry_init, - capture_events, capture_items, mock_genai_client, - span_streaming, ): """Test that async embed_content spans have correct origin.""" sentry_init( integrations=[GoogleGenAIIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) mock_http_response = create_mock_http_response(EXAMPLE_EMBED_RESPONSE_JSON) + items = capture_items("span") - if span_streaming: - items = capture_items("transaction", "span") - - with mock.patch.object( - mock_genai_client._api_client, - "async_request", - return_value=mock_http_response, - ), sentry_sdk.traces.start_span(name="google_genai_embeddings_async"): - await mock_genai_client.aio.models.embed_content( - model="text-embedding-004", - contents=["Test origin"], - ) - - sentry_sdk.flush() - spans = [item.payload for item in items if item.type == "span"] - for span in spans: - if span["is_segment"] is True: - assert span["attributes"]["sentry.origin"] == "manual" - continue - - assert span["attributes"]["sentry.origin"] == "auto.ai.google_genai" - else: - events = capture_events() - - with mock.patch.object( - mock_genai_client._api_client, - "async_request", - return_value=mock_http_response, - ), start_transaction(name="google_genai_embeddings_async"): - await mock_genai_client.aio.models.embed_content( - model="text-embedding-004", - contents=["Test origin"], - ) - - (event,) = events + with mock.patch.object( + mock_genai_client._api_client, + "async_request", + return_value=mock_http_response, + ): + await mock_genai_client.aio.models.embed_content( + model="text-embedding-004", + contents=["Test origin"], + ) - assert event["contexts"]["trace"]["origin"] == "manual" - for span in event["spans"]: - assert span["origin"] == "auto.ai.google_genai" + sentry_sdk.flush() + spans = [item.payload for item in items] + for span in spans: + assert span["attributes"]["sentry.origin"] == "auto.ai.google_genai" # Integration tests for generate_content with different input message formats -@pytest.mark.parametrize("span_streaming", [True, False]) def test_generate_content_with_content_object( sentry_init, - capture_events, capture_items, mock_genai_client, - span_streaming, ): """Test generate_content with Content object input.""" sentry_init( integrations=[GoogleGenAIIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) @@ -2357,37 +1516,19 @@ def test_generate_content_with_content_object( content = genai_types.Content( role="user", parts=[genai_types.Part(text="Hello from Content object")] ) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents=content, config=create_test_config() - ) - - sentry_sdk.flush() - invoke_span = next(item.payload for item in items) - - messages = json.loads( - invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] + with mock.patch.object( + mock_genai_client._api_client, "request", return_value=mock_http_response + ): + mock_genai_client.models.generate_content( + model="gemini-1.5-flash", contents=content, config=create_test_config() ) - else: - events = capture_events() - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents=content, config=create_test_config() - ) - (event,) = events - invoke_span = event["spans"][0] + sentry_sdk.flush() + invoke_span = next(item.payload for item in items) - messages = json.loads(invoke_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) + messages = json.loads(invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) assert len(messages) == 1 assert messages[0]["role"] == "user" @@ -2396,58 +1537,36 @@ def test_generate_content_with_content_object( ] -@pytest.mark.parametrize("span_streaming", [True, False]) def test_generate_content_with_dict_format( sentry_init, - capture_events, capture_items, mock_genai_client, - span_streaming, ): """Test generate_content with dict format input (ContentDict).""" sentry_init( integrations=[GoogleGenAIIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) # Dict format content contents = {"role": "user", "parts": [{"text": "Hello from dict format"}]} + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents=contents, config=create_test_config() - ) - - sentry_sdk.flush() - invoke_span = next(item.payload for item in items) - - messages = json.loads( - invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] + with mock.patch.object( + mock_genai_client._api_client, "request", return_value=mock_http_response + ): + mock_genai_client.models.generate_content( + model="gemini-1.5-flash", contents=contents, config=create_test_config() ) - else: - events = capture_events() - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents=contents, config=create_test_config() - ) - (event,) = events - invoke_span = event["spans"][0] + sentry_sdk.flush() + invoke_span = next(item.payload for item in items) - messages = json.loads(invoke_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) + messages = json.loads(invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) assert len(messages) == 1 assert messages[0]["role"] == "user" @@ -2456,21 +1575,17 @@ def test_generate_content_with_dict_format( ] -@pytest.mark.parametrize("span_streaming", [True, False]) def test_generate_content_with_file_data( sentry_init, - capture_events, capture_items, mock_genai_client, - span_streaming, ): """Test generate_content with file_data (external file reference).""" sentry_init( integrations=[GoogleGenAIIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) @@ -2486,37 +1601,19 @@ def test_generate_content_with_file_data( genai_types.Part(file_data=file_data), ], ) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents=content, config=create_test_config() - ) - - sentry_sdk.flush() - invoke_span = next(item.payload for item in items) - - messages = json.loads( - invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] + with mock.patch.object( + mock_genai_client._api_client, "request", return_value=mock_http_response + ): + mock_genai_client.models.generate_content( + model="gemini-1.5-flash", contents=content, config=create_test_config() ) - else: - events = capture_events() - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents=content, config=create_test_config() - ) - (event,) = events - invoke_span = event["spans"][0] + sentry_sdk.flush() + invoke_span = next(item.payload for item in items) - messages = json.loads(invoke_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) + messages = json.loads(invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) assert len(messages) == 1 assert messages[0]["role"] == "user" @@ -2531,21 +1628,17 @@ def test_generate_content_with_file_data( assert messages[0]["content"][1]["uri"] == "gs://bucket/image.jpg" -@pytest.mark.parametrize("span_streaming", [True, False]) def test_generate_content_with_inline_data( sentry_init, - capture_events, capture_items, mock_genai_client, - span_streaming, ): """Test generate_content with inline_data (binary data).""" sentry_init( integrations=[GoogleGenAIIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) @@ -2560,37 +1653,19 @@ def test_generate_content_with_inline_data( genai_types.Part(inline_data=blob), ], ) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents=content, config=create_test_config() - ) - - sentry_sdk.flush() - invoke_span = next(item.payload for item in items) - - messages = json.loads( - invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] + with mock.patch.object( + mock_genai_client._api_client, "request", return_value=mock_http_response + ): + mock_genai_client.models.generate_content( + model="gemini-1.5-flash", contents=content, config=create_test_config() ) - else: - events = capture_events() - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents=content, config=create_test_config() - ) - (event,) = events - invoke_span = event["spans"][0] + sentry_sdk.flush() + invoke_span = next(item.payload for item in items) - messages = json.loads(invoke_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) + messages = json.loads(invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) assert len(messages) == 1 assert messages[0]["role"] == "user" @@ -2603,16 +1678,15 @@ def test_generate_content_with_inline_data( def test_generate_content_with_function_response( - sentry_init, capture_events, mock_genai_client + sentry_init, capture_items, mock_genai_client ): """Test generate_content with function_response (tool result).""" sentry_init( integrations=[GoogleGenAIIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) - events = capture_events() mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) @@ -2638,37 +1712,40 @@ def test_generate_content_with_function_response( ), ] + items = capture_items("span") + with mock.patch.object( mock_genai_client._api_client, "request", return_value=mock_http_response ): - with start_transaction(name="google_genai"): - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents=contents, config=create_test_config() - ) + mock_genai_client.models.generate_content( + model="gemini-1.5-flash", contents=contents, config=create_test_config() + ) - (event,) = events - invoke_span = event["spans"][0] + sentry_sdk.flush() + spans = [item.payload for item in items] + invoke_span = spans[0] - messages = json.loads(invoke_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) - assert len(messages) == 1 + messages = json.loads(invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) + assert len(messages) == 2 # First message is user message - assert messages[0]["role"] == "tool" - assert messages[0]["content"]["toolCallId"] == "call_123" - assert messages[0]["content"]["toolName"] == "get_weather" - assert messages[0]["content"]["output"] == "Sunny, 72F" + assert messages[1]["role"] == "tool" + assert messages[1]["content"]["toolCallId"] == "call_123" + assert messages[1]["content"]["toolName"] == "get_weather" + assert messages[1]["content"]["output"] == "Sunny, 72F" def test_generate_content_with_mixed_string_and_content( - sentry_init, capture_events, mock_genai_client + sentry_init, + capture_items, + mock_genai_client, ): """Test generate_content with mixed string and Content objects in list.""" sentry_init( integrations=[GoogleGenAIIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) - events = capture_events() mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) @@ -2685,76 +1762,56 @@ def test_generate_content_with_mixed_string_and_content( ), ] + items = capture_items("span") + with mock.patch.object( mock_genai_client._api_client, "request", return_value=mock_http_response ): - with start_transaction(name="google_genai"): - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents=contents, config=create_test_config() - ) + mock_genai_client.models.generate_content( + model="gemini-1.5-flash", contents=contents, config=create_test_config() + ) - (event,) = events - invoke_span = event["spans"][0] + sentry_sdk.flush() + spans = [item.payload for item in items] + invoke_span = spans[0] - messages = json.loads(invoke_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) - assert len(messages) == 1 + messages = json.loads(invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) + assert len(messages) == 3 # User message - assert messages[0]["role"] == "user" - assert messages[0]["content"] == [{"text": "Tell me a joke", "type": "text"}] + assert messages[2]["role"] == "user" + assert messages[2]["content"] == [{"text": "Tell me a joke", "type": "text"}] -@pytest.mark.parametrize("span_streaming", [True, False]) def test_generate_content_with_part_object_directly( sentry_init, - capture_events, capture_items, mock_genai_client, - span_streaming, ): """Test generate_content with Part object directly (not wrapped in Content).""" sentry_init( integrations=[GoogleGenAIIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) # Part object directly part = genai_types.Part(text="Direct Part object") + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents=part, config=create_test_config() - ) - - sentry_sdk.flush() - invoke_span = next(item.payload for item in items) - - messages = json.loads( - invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] + with mock.patch.object( + mock_genai_client._api_client, "request", return_value=mock_http_response + ): + mock_genai_client.models.generate_content( + model="gemini-1.5-flash", contents=part, config=create_test_config() ) - else: - events = capture_events() - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents=part, config=create_test_config() - ) - (event,) = events - invoke_span = event["spans"][0] + sentry_sdk.flush() + invoke_span = next(item.payload for item in items) - messages = json.loads(invoke_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) + messages = json.loads(invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) assert len(messages) == 1 assert messages[0]["role"] == "user" @@ -2762,7 +1819,7 @@ def test_generate_content_with_part_object_directly( def test_generate_content_with_list_of_dicts( - sentry_init, capture_events, mock_genai_client + sentry_init, capture_items, mock_genai_client ): """ Test generate_content with list of dict format inputs. @@ -2775,9 +1832,8 @@ def test_generate_content_with_list_of_dicts( integrations=[GoogleGenAIIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) - events = capture_events() mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) @@ -2788,38 +1844,36 @@ def test_generate_content_with_list_of_dicts( {"role": "user", "parts": [{"text": "Second user message"}]}, ] + items = capture_items("span") + with mock.patch.object( mock_genai_client._api_client, "request", return_value=mock_http_response ): - with start_transaction(name="google_genai"): - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents=contents, config=create_test_config() - ) + mock_genai_client.models.generate_content( + model="gemini-1.5-flash", contents=contents, config=create_test_config() + ) - (event,) = events - invoke_span = event["spans"][0] + sentry_sdk.flush() + spans = [item.payload for item in items] + invoke_span = spans[0] - messages = json.loads(invoke_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) - assert len(messages) == 1 - assert messages[0]["role"] == "user" - assert messages[0]["content"] == [{"text": "Second user message", "type": "text"}] + messages = json.loads(invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) + assert len(messages) == 3 + assert messages[2]["role"] == "user" + assert messages[2]["content"] == [{"text": "Second user message", "type": "text"}] -@pytest.mark.parametrize("span_streaming", [True, False]) def test_generate_content_with_dict_inline_data( sentry_init, - capture_events, capture_items, mock_genai_client, - span_streaming, ): """Test generate_content with dict format containing inline_data.""" sentry_init( integrations=[GoogleGenAIIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) @@ -2832,37 +1886,19 @@ def test_generate_content_with_dict_inline_data( {"inline_data": {"data": b"fake_binary_data", "mime_type": "image/gif"}}, ], } + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents=contents, config=create_test_config() - ) - - sentry_sdk.flush() - invoke_span = next(item.payload for item in items) - - messages = json.loads( - invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] + with mock.patch.object( + mock_genai_client._api_client, "request", return_value=mock_http_response + ): + mock_genai_client.models.generate_content( + model="gemini-1.5-flash", contents=contents, config=create_test_config() ) - else: - events = capture_events() - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents=contents, config=create_test_config() - ) - (event,) = events - invoke_span = event["spans"][0] + sentry_sdk.flush() + invoke_span = next(item.payload for item in items) - messages = json.loads(invoke_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) + messages = json.loads(invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) assert len(messages) == 1 assert messages[0]["role"] == "user" @@ -2876,20 +1912,16 @@ def test_generate_content_with_dict_inline_data( assert messages[0]["content"][1]["content"] == BLOB_DATA_SUBSTITUTE -@pytest.mark.parametrize("span_streaming", [True, False]) def test_generate_content_without_parts_property_inline_data( sentry_init, - capture_events, capture_items, mock_genai_client, - span_streaming, ): sentry_init( integrations=[GoogleGenAIIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) @@ -2898,37 +1930,19 @@ def test_generate_content_without_parts_property_inline_data( {"text": "What's in this image?"}, {"inline_data": {"data": b"fake_binary_data", "mime_type": "image/gif"}}, ] + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents=contents, config=create_test_config() - ) - - sentry_sdk.flush() - invoke_span = next(item.payload for item in items) - - messages = json.loads( - invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] + with mock.patch.object( + mock_genai_client._api_client, "request", return_value=mock_http_response + ): + mock_genai_client.models.generate_content( + model="gemini-1.5-flash", contents=contents, config=create_test_config() ) - else: - events = capture_events() - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents=contents, config=create_test_config() - ) - (event,) = events - invoke_span = event["spans"][0] + sentry_sdk.flush() + invoke_span = next(item.payload for item in items) - messages = json.loads(invoke_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) + messages = json.loads(invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) assert len(messages) == 1 @@ -2944,20 +1958,16 @@ 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]) def test_generate_content_without_parts_property_inline_data_and_binary_data_within_string( sentry_init, - capture_events, capture_items, mock_genai_client, - span_streaming, ): sentry_init( integrations=[GoogleGenAIIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) @@ -2971,37 +1981,19 @@ def test_generate_content_without_parts_property_inline_data_and_binary_data_wit } }, ] + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents=contents, config=create_test_config() - ) - - sentry_sdk.flush() - invoke_span = next(item.payload for item in items) - - messages = json.loads( - invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] + with mock.patch.object( + mock_genai_client._api_client, "request", return_value=mock_http_response + ): + mock_genai_client.models.generate_content( + model="gemini-1.5-flash", contents=contents, config=create_test_config() ) - else: - events = capture_events() - - with mock.patch.object( - mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): - mock_genai_client.models.generate_content( - model="gemini-1.5-flash", contents=contents, config=create_test_config() - ) - (event,) = events - invoke_span = event["spans"][0] + sentry_sdk.flush() + invoke_span = next(item.payload for item in items) - messages = json.loads(invoke_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) + messages = json.loads(invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) assert len(messages) == 1 assert messages[0]["role"] == "user" @@ -3354,7 +2346,6 @@ def __init__(self): } -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize( "data_collection,send_default_pii,include_prompts,expected_present,expected_absent", [ @@ -3456,14 +2447,12 @@ def test_generate_content_data_collection( include_prompts, expected_present, expected_absent, - span_streaming, ): sentry_init_kwargs = dict( integrations=[GoogleGenAIIntegration(include_prompts=include_prompts)], traces_sample_rate=1.0, send_default_pii=send_default_pii, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) if data_collection is not None: sentry_init_kwargs["_experiments"] = {"data_collection": data_collection} @@ -3471,12 +2460,11 @@ def test_generate_content_data_collection( sentry_init(**sentry_init_kwargs) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) - streamed = span_streaming - captured = capture_items("span") if streamed else capture_events() + captured = capture_items("span") with mock.patch.object( mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): + ): mock_genai_client.models.generate_content( model="gemini-1.5-flash", contents="Tell me a joke", @@ -3487,17 +2475,10 @@ def test_generate_content_data_collection( ), ) - if streamed: - sentry_sdk.flush() - spans = [item.payload for item in captured if item.type == "span"] - (span,) = [ - s for s in spans if s["attributes"].get("sentry.op") == OP.GEN_AI_CHAT - ] - span_data = span["attributes"] - else: - (event,) = captured - (span,) = [s for s in event["spans"] if s["op"] == OP.GEN_AI_CHAT] - span_data = span["data"] + sentry_sdk.flush() + spans = [item.payload for item in captured if item.type == "span"] + (span,) = [s for s in spans if s["attributes"].get("sentry.op") == OP.GEN_AI_CHAT] + span_data = span["attributes"] for key in expected_present: assert key in span_data, f"{key} should have been collected" @@ -3517,7 +2498,6 @@ def test_generate_content_data_collection( assert span_data[SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS] == "STOP" -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("send_default_pii", [True, False]) @pytest.mark.parametrize( "data_collection,expected_present,expected_absent", @@ -3581,14 +2561,12 @@ def test_generate_content_data_collection_tools( send_default_pii, expected_present, expected_absent, - span_streaming, ): sentry_init_kwargs = dict( integrations=[GoogleGenAIIntegration(include_prompts=False)], traces_sample_rate=1.0, send_default_pii=send_default_pii, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) if data_collection is not None: sentry_init_kwargs["_experiments"] = {"data_collection": data_collection} @@ -3633,29 +2611,21 @@ def test_generate_content_data_collection_tools( } ) - streamed = span_streaming - captured = capture_items("span") if streamed else capture_events() + captured = capture_items("span") with mock.patch.object( mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai"): + ): mock_genai_client.models.generate_content( model="gemini-1.5-flash", contents="What's the weather?", config=create_test_config(tools=[weather_tool]), ) - if streamed: - sentry_sdk.flush() - spans = [item.payload for item in captured if item.type == "span"] - (span,) = [ - s for s in spans if s["attributes"].get("sentry.op") == OP.GEN_AI_CHAT - ] - span_data = span["attributes"] - else: - (event,) = captured - (span,) = [s for s in event["spans"] if s["op"] == OP.GEN_AI_CHAT] - span_data = span["data"] + sentry_sdk.flush() + spans = [item.payload for item in captured if item.type == "span"] + (span,) = [s for s in spans if s["attributes"].get("sentry.op") == OP.GEN_AI_CHAT] + span_data = span["attributes"] for key in expected_present: assert key in span_data, f"{key} should have been collected" @@ -3673,7 +2643,6 @@ def test_generate_content_data_collection_tools( assert json.loads(tool_calls[0]["arguments"]) == {"location": "San Francisco"} -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize( "data_collection,send_default_pii,include_prompts,expected_present,expected_absent", [ @@ -3775,14 +2744,12 @@ def test_streaming_generate_content_data_collection( include_prompts, expected_present, expected_absent, - span_streaming, ): sentry_init_kwargs = dict( integrations=[GoogleGenAIIntegration(include_prompts=include_prompts)], traces_sample_rate=1.0, send_default_pii=send_default_pii, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) if data_collection is not None: sentry_init_kwargs["_experiments"] = {"data_collection": data_collection} @@ -3826,12 +2793,11 @@ def test_streaming_generate_content_data_collection( ] ) - streamed = span_streaming - captured = capture_items("span") if streamed else capture_events() + captured = capture_items("span") with mock.patch.object( mock_genai_client._api_client, "request_streamed", return_value=mock_stream - ), start_transaction(name="google_genai"): + ): stream = mock_genai_client.models.generate_content_stream( model="gemini-1.5-flash", contents="Tell me a joke", @@ -3841,17 +2807,10 @@ def test_streaming_generate_content_data_collection( ) list(stream) - if streamed: - sentry_sdk.flush() - spans = [item.payload for item in captured if item.type == "span"] - (span,) = [ - s for s in spans if s["attributes"].get("sentry.op") == OP.GEN_AI_CHAT - ] - span_data = span["attributes"] - else: - (event,) = captured - (span,) = [s for s in event["spans"] if s["op"] == OP.GEN_AI_CHAT] - span_data = span["data"] + sentry_sdk.flush() + spans = [item.payload for item in captured if item.type == "span"] + (span,) = [s for s in spans if s["attributes"].get("sentry.op") == OP.GEN_AI_CHAT] + span_data = span["attributes"] for key in expected_present: assert key in span_data, f"{key} should have been collected" @@ -3867,7 +2826,6 @@ def test_streaming_generate_content_data_collection( assert span_data[SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 25 -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("send_default_pii", [True, False]) @pytest.mark.parametrize( "data_collection,expected_present,expected_absent", @@ -3931,14 +2889,12 @@ def test_streaming_generate_content_data_collection_tools( send_default_pii, expected_present, expected_absent, - span_streaming, ): sentry_init_kwargs = dict( integrations=[GoogleGenAIIntegration(include_prompts=False)], traces_sample_rate=1.0, send_default_pii=send_default_pii, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) if data_collection is not None: sentry_init_kwargs["_experiments"] = {"data_collection": data_collection} @@ -3988,12 +2944,11 @@ def test_streaming_generate_content_data_collection_tools( ] ) - streamed = span_streaming - captured = capture_items("span") if streamed else capture_events() + captured = capture_items("span") with mock.patch.object( mock_genai_client._api_client, "request_streamed", return_value=mock_stream - ), start_transaction(name="google_genai"): + ): stream = mock_genai_client.models.generate_content_stream( model="gemini-1.5-flash", contents="What's the weather?", @@ -4001,17 +2956,10 @@ def test_streaming_generate_content_data_collection_tools( ) list(stream) - if streamed: - sentry_sdk.flush() - spans = [item.payload for item in captured if item.type == "span"] - (span,) = [ - s for s in spans if s["attributes"].get("sentry.op") == OP.GEN_AI_CHAT - ] - span_data = span["attributes"] - else: - (event,) = captured - (span,) = [s for s in event["spans"] if s["op"] == OP.GEN_AI_CHAT] - span_data = span["data"] + sentry_sdk.flush() + spans = [item.payload for item in captured if item.type == "span"] + (span,) = [s for s in spans if s["attributes"].get("sentry.op") == OP.GEN_AI_CHAT] + span_data = span["attributes"] for key in expected_present: assert key in span_data, f"{key} should have been collected" @@ -4028,7 +2976,6 @@ def test_streaming_generate_content_data_collection_tools( assert [tool_call["name"] for tool_call in tool_calls] == ["get_weather"] -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize( "data_collection,send_default_pii,include_prompts,expected_present,expected_absent", [ @@ -4114,14 +3061,12 @@ def test_embed_content_data_collection( include_prompts, expected_present, expected_absent, - span_streaming, ): sentry_init_kwargs = dict( integrations=[GoogleGenAIIntegration(include_prompts=include_prompts)], traces_sample_rate=1.0, send_default_pii=send_default_pii, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) if data_collection is not None: sentry_init_kwargs["_experiments"] = {"data_collection": data_collection} @@ -4129,28 +3074,22 @@ def test_embed_content_data_collection( sentry_init(**sentry_init_kwargs) mock_http_response = create_mock_http_response(EXAMPLE_EMBED_RESPONSE_JSON) - streamed = span_streaming - captured = capture_items("span") if streamed else capture_events() + captured = capture_items("span") with mock.patch.object( mock_genai_client._api_client, "request", return_value=mock_http_response - ), start_transaction(name="google_genai_embeddings"): + ): mock_genai_client.models.embed_content( model="text-embedding-004", contents=["What is your name?", "What is your favorite color?"], ) - if streamed: - sentry_sdk.flush() - spans = [item.payload for item in captured if item.type == "span"] - (span,) = [ - s for s in spans if s["attributes"].get("sentry.op") == OP.GEN_AI_EMBEDDINGS - ] - span_data = span["attributes"] - else: - (event,) = captured - (span,) = [s for s in event["spans"] if s["op"] == OP.GEN_AI_EMBEDDINGS] - span_data = span["data"] + sentry_sdk.flush() + spans = [item.payload for item in captured if item.type == "span"] + (span,) = [ + s for s in spans if s["attributes"].get("sentry.op") == OP.GEN_AI_EMBEDDINGS + ] + span_data = span["attributes"] for key in expected_present: assert key in span_data, f"{key} should have been collected" @@ -4164,7 +3103,6 @@ def test_embed_content_data_collection( @pytest.mark.asyncio -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize( "data_collection,send_default_pii,include_prompts,expected_present,expected_absent", [ @@ -4266,14 +3204,12 @@ async def test_async_generate_content_data_collection( include_prompts, expected_present, expected_absent, - span_streaming, ): sentry_init_kwargs = dict( integrations=[GoogleGenAIIntegration(include_prompts=include_prompts)], traces_sample_rate=1.0, send_default_pii=send_default_pii, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) if data_collection is not None: sentry_init_kwargs["_experiments"] = {"data_collection": data_collection} @@ -4281,12 +3217,11 @@ async def test_async_generate_content_data_collection( sentry_init(**sentry_init_kwargs) mock_http_response = create_mock_http_response(EXAMPLE_API_RESPONSE_JSON) - streamed = span_streaming - captured = capture_items("span") if streamed else capture_events() + captured = capture_items("span") with mock.patch.object( mock_genai_client._api_client, "async_request", return_value=mock_http_response - ), start_transaction(name="google_genai"): + ): await mock_genai_client.aio.models.generate_content( model="gemini-1.5-flash", contents="Tell me a joke", @@ -4295,17 +3230,10 @@ async def test_async_generate_content_data_collection( ), ) - if streamed: - sentry_sdk.flush() - spans = [item.payload for item in captured if item.type == "span"] - (span,) = [ - s for s in spans if s["attributes"].get("sentry.op") == OP.GEN_AI_CHAT - ] - span_data = span["attributes"] - else: - (event,) = captured - (span,) = [s for s in event["spans"] if s["op"] == OP.GEN_AI_CHAT] - span_data = span["data"] + sentry_sdk.flush() + spans = [item.payload for item in captured if item.type == "span"] + (span,) = [s for s in spans if s["attributes"].get("sentry.op") == OP.GEN_AI_CHAT] + span_data = span["attributes"] for key in expected_present: assert key in span_data, f"{key} should have been collected" @@ -4319,7 +3247,6 @@ async def test_async_generate_content_data_collection( @pytest.mark.asyncio -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize( "data_collection,send_default_pii,include_prompts,expected_present,expected_absent", [ @@ -4405,14 +3332,12 @@ async def test_async_embed_content_data_collection( include_prompts, expected_present, expected_absent, - span_streaming, ): sentry_init_kwargs = dict( integrations=[GoogleGenAIIntegration(include_prompts=include_prompts)], traces_sample_rate=1.0, send_default_pii=send_default_pii, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) if data_collection is not None: sentry_init_kwargs["_experiments"] = {"data_collection": data_collection} @@ -4420,28 +3345,22 @@ async def test_async_embed_content_data_collection( sentry_init(**sentry_init_kwargs) mock_http_response = create_mock_http_response(EXAMPLE_EMBED_RESPONSE_JSON) - streamed = span_streaming - captured = capture_items("span") if streamed else capture_events() + captured = capture_items("span") with mock.patch.object( mock_genai_client._api_client, "async_request", return_value=mock_http_response - ), start_transaction(name="google_genai_embeddings"): + ): await mock_genai_client.aio.models.embed_content( model="text-embedding-004", contents=["What is your name?", "What is your favorite color?"], ) - if streamed: - sentry_sdk.flush() - spans = [item.payload for item in captured if item.type == "span"] - (span,) = [ - s for s in spans if s["attributes"].get("sentry.op") == OP.GEN_AI_EMBEDDINGS - ] - span_data = span["attributes"] - else: - (event,) = captured - (span,) = [s for s in event["spans"] if s["op"] == OP.GEN_AI_EMBEDDINGS] - span_data = span["data"] + sentry_sdk.flush() + spans = [item.payload for item in captured if item.type == "span"] + (span,) = [ + s for s in spans if s["attributes"].get("sentry.op") == OP.GEN_AI_EMBEDDINGS + ] + span_data = span["attributes"] for key in expected_present: assert key in span_data, f"{key} should have been collected"