From d8cc094a4ab9e4a4d6f7d98fe55a96f51704cc34 Mon Sep 17 00:00:00 2001 From: Max Ind Date: Thu, 20 Nov 2025 13:06:08 +0000 Subject: [PATCH] feat add ability to add custom attributes to google genai `generate_content {model.name}` spans --- .../google_genai/generate_content.py | 11 ++++++++++- .../tests/generate_content/nonstreaming_base.py | 16 ++++++++++++++++ .../tests/generate_content/streaming_base.py | 16 ++++++++++++++++ 3 files changed, 42 insertions(+), 1 deletion(-) diff --git a/instrumentation-genai/opentelemetry-instrumentation-google-genai/src/opentelemetry/instrumentation/google_genai/generate_content.py b/instrumentation-genai/opentelemetry-instrumentation-google-genai/src/opentelemetry/instrumentation/google_genai/generate_content.py index 82dda55d17..fe03c83407 100644 --- a/instrumentation-genai/opentelemetry-instrumentation-google-genai/src/opentelemetry/instrumentation/google_genai/generate_content.py +++ b/instrumentation-genai/opentelemetry-instrumentation-google-genai/src/opentelemetry/instrumentation/google_genai/generate_content.py @@ -20,7 +20,8 @@ import logging import os import time -from typing import Any, AsyncIterator, Awaitable, Iterator, Optional, Union +import contextvars +from typing import Any, AsyncIterator, Awaitable, Iterator, Optional, Union, Mapping from google.genai.models import AsyncModels, Models from google.genai.models import t as transformers @@ -57,6 +58,9 @@ MessagePart, OutputMessage, ) +from opentelemetry.util.types import ( + AttributeValue, +) from opentelemetry.util.genai.utils import gen_ai_json_dumps from .allowlist_util import AllowList @@ -80,6 +84,7 @@ # Constant used for the value of 'gen_ai.operation.name". _GENERATE_CONTENT_OP_NAME = "generate_content" +GENERATE_CONTENT_EXTRA_ATTRIBUTES = contextvars.ContextVar[Mapping[str, AttributeValue]]("GENERATE_CONTENT_EXTRA_ATTRIBUTES", default={}) class _MethodsSnapshot: def __init__(self): @@ -728,6 +733,7 @@ def instrumented_generate_content( with helper.start_span_as_current_span( model, "google.genai.Models.generate_content" ) as span: + span.set_attributes(GENERATE_CONTENT_EXTRA_ATTRIBUTES.get()) span.set_attributes(request_attributes) if helper.sem_conv_opt_in_mode == _StabilityMode.DEFAULT: helper.process_request(contents, config, span) @@ -803,6 +809,7 @@ def instrumented_generate_content_stream( with helper.start_span_as_current_span( model, "google.genai.Models.generate_content_stream" ) as span: + span.set_attributes(GENERATE_CONTENT_EXTRA_ATTRIBUTES.get()) span.set_attributes(request_attributes) if helper.sem_conv_opt_in_mode == _StabilityMode.DEFAULT: helper.process_request(contents, config, span) @@ -878,6 +885,7 @@ async def instrumented_generate_content( with helper.start_span_as_current_span( model, "google.genai.AsyncModels.generate_content" ) as span: + span.set_attributes(GENERATE_CONTENT_EXTRA_ATTRIBUTES.get()) span.set_attributes(request_attributes) if helper.sem_conv_opt_in_mode == _StabilityMode.DEFAULT: helper.process_request(contents, config, span) @@ -954,6 +962,7 @@ async def instrumented_generate_content_stream( "google.genai.AsyncModels.generate_content_stream", end_on_exit=False, ) as span: + span.set_attributes(GENERATE_CONTENT_EXTRA_ATTRIBUTES.get()) span.set_attributes(request_attributes) if not is_experimental_mode: helper.process_request(contents, config, span) diff --git a/instrumentation-genai/opentelemetry-instrumentation-google-genai/tests/generate_content/nonstreaming_base.py b/instrumentation-genai/opentelemetry-instrumentation-google-genai/tests/generate_content/nonstreaming_base.py index 0520f818f9..b7d3460c9e 100644 --- a/instrumentation-genai/opentelemetry-instrumentation-google-genai/tests/generate_content/nonstreaming_base.py +++ b/instrumentation-genai/opentelemetry-instrumentation-google-genai/tests/generate_content/nonstreaming_base.py @@ -32,6 +32,7 @@ from opentelemetry.util.genai.types import ContentCapturingMode from .base import TestCase +from opentelemetry.instrumentation.google_genai.generate_content import GENERATE_CONTENT_EXTRA_ATTRIBUTES # pylint: disable=too-many-public-methods @@ -100,6 +101,21 @@ def test_generated_span_has_minimal_genai_attributes(self): span.attributes["gen_ai.operation.name"], "generate_content" ) + def test_generated_span_has_extra_genai_attributes(self): + self.configure_valid_response(text="Yep, it works!") + tok = GENERATE_CONTENT_EXTRA_ATTRIBUTES.set({"extra_attribute_key": "extra_attribute_value"}) + try: + self.generate_content( + model="gemini-2.0-flash", contents="Does this work?" + ) + self.otel.assert_has_span_named("generate_content gemini-2.0-flash") + span = self.otel.get_span_named("generate_content gemini-2.0-flash") + self.assertEqual(span.attributes["extra_attribute_key"], "extra_attribute_value") + except: + raise + finally: + GENERATE_CONTENT_EXTRA_ATTRIBUTES.reset(tok) + def test_span_and_event_still_written_when_response_is_exception(self): self.configure_exception(ValueError("Uh oh!")) patched_environ = patch.dict( diff --git a/instrumentation-genai/opentelemetry-instrumentation-google-genai/tests/generate_content/streaming_base.py b/instrumentation-genai/opentelemetry-instrumentation-google-genai/tests/generate_content/streaming_base.py index e5bceb7c79..990c281e6c 100644 --- a/instrumentation-genai/opentelemetry-instrumentation-google-genai/tests/generate_content/streaming_base.py +++ b/instrumentation-genai/opentelemetry-instrumentation-google-genai/tests/generate_content/streaming_base.py @@ -15,6 +15,7 @@ import unittest from .base import TestCase +from opentelemetry.instrumentation.google_genai.generate_content import GENERATE_CONTENT_EXTRA_ATTRIBUTES class StreamingTestCase(TestCase): @@ -42,6 +43,21 @@ def test_instrumentation_does_not_break_core_functionality(self): response = responses[0] self.assertEqual(response.text, "Yep, it works!") + def test_generated_span_has_extra_genai_attributes(self): + self.configure_valid_response(text="Yep, it works!") + tok = GENERATE_CONTENT_EXTRA_ATTRIBUTES.set({"extra_attrbiute_key": "extra_attribute_value"}) + try: + self.generate_content( + model="gemini-2.0-flash", contents="Does this work?" + ) + self.otel.assert_has_span_named("generate_content gemini-2.0-flash") + span = self.otel.get_span_named("generate_content gemini-2.0-flash") + self.assertEqual(span.attributes["extra_attrbiute_key"], "extra_attribute_value") + except: + raise + finally: + GENERATE_CONTENT_EXTRA_ATTRIBUTES.reset(tok) + def test_handles_multiple_ressponses(self): self.configure_valid_response(text="First response") self.configure_valid_response(text="Second response")