feat: add extension point to add custom attributes to generate_content spans in google genai instrumentor. (#3961)

* feat add ability to add custom attributes to google genai `generate_content {model.name}` spans

* Switch from contextvars to OTel context_api

* Update CHANGELOG.md

---------

Co-authored-by: Aaron Abbott <aaronabbott@google.com>

Author: RKest

instrumentation-genai/opentelemetry-instrumentation-google-genai/CHANGELOG.md

@@ -7,6 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+- Enable the addition of custom attributes to the `generate_content {model.name}` span via the Context API. ([#3961](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3961)).
+
 ## Version 0.5b0 (2025-12-11)
 
 - Ensure log event is written and completion hook is called even when model call results in exception. Put new
@@ -39,4 +41,4 @@ span attribute `gen_ai.response.finish_reasons` is empty ([#3417](https://github
   ([#3298](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3298))
 
 Create an initial version of Open Telemetry instrumentation for github.com/googleapis/python-genai.
-([#3256](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3256)) 
+([#3256](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3256)) 

instrumentation-genai/opentelemetry-instrumentation-google-genai/src/opentelemetry/instrumentation/google_genai/__init__.py

@@ -42,6 +42,7 @@
 """
 
 from .instrumentor import GoogleGenAiSdkInstrumentor
+from .generate_content import GENERATE_CONTENT_EXTRA_ATTRIBUTES_CONTEXT_KEY
 from .version import __version__
 
-__all__ = ["GoogleGenAiSdkInstrumentor", "__version__"]
+__all__ = ["GoogleGenAiSdkInstrumentor", "GENERATE_CONTENT_EXTRA_ATTRIBUTES_CONTEXT_KEY", "__version__"]

instrumentation-genai/opentelemetry-instrumentation-google-genai/src/opentelemetry/instrumentation/google_genai/generate_content.py

@@ -20,7 +20,7 @@
 import logging
 import os
 import time
-from typing import Any, AsyncIterator, Awaitable, Iterator, Optional, Union
+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 +57,9 @@
     MessagePart,
     OutputMessage,
 )
+from opentelemetry.util.types import (
+    AttributeValue,
+)
 from opentelemetry.util.genai.utils import gen_ai_json_dumps
 
 from .allowlist_util import AllowList
@@ -70,6 +73,7 @@
 )
 from .otel_wrapper import OTelWrapper
 from .tool_call_wrapper import wrapped as wrapped_tool
+from opentelemetry import context as context_api
 
 _logger = logging.getLogger(__name__)
 
@@ -80,6 +84,7 @@
 # Constant used for the value of 'gen_ai.operation.name".
 _GENERATE_CONTENT_OP_NAME = "generate_content"
 
+GENERATE_CONTENT_EXTRA_ATTRIBUTES_CONTEXT_KEY = context_api.create_key("generate_content_extra_attributes_context_key")
 
 class _MethodsSnapshot:
     def __init__(self):
@@ -294,6 +299,10 @@ def _create_completion_details_attributes(
     return attributes
 
 
+def _get_extra_generate_content_attributes() -> Optional[Mapping[str, AttributeValue]]:
+    return context_api.get_value(GENERATE_CONTENT_EXTRA_ATTRIBUTES_CONTEXT_KEY)
+
+
 class _GenerateContentInstrumentationHelper:
     def __init__(
         self,
@@ -720,6 +729,8 @@ def instrumented_generate_content(
         with helper.start_span_as_current_span(
             model, "google.genai.Models.generate_content"
         ) as span:
+            if extra_attributes := _get_extra_generate_content_attributes():
+                span.set_attributes(extra_attributes)
             span.set_attributes(request_attributes)
             if helper.sem_conv_opt_in_mode == _StabilityMode.DEFAULT:
                 helper.process_request(contents, config, span)
@@ -793,6 +804,8 @@ def instrumented_generate_content_stream(
         with helper.start_span_as_current_span(
             model, "google.genai.Models.generate_content_stream"
         ) as span:
+            if extra_attributes := _get_extra_generate_content_attributes():
+                span.set_attributes(extra_attributes)
             span.set_attributes(request_attributes)
             if helper.sem_conv_opt_in_mode == _StabilityMode.DEFAULT:
                 helper.process_request(contents, config, span)
@@ -866,6 +879,8 @@ async def instrumented_generate_content(
         with helper.start_span_as_current_span(
             model, "google.genai.AsyncModels.generate_content"
         ) as span:
+            if extra_attributes := _get_extra_generate_content_attributes():
+                span.set_attributes(extra_attributes)
             span.set_attributes(request_attributes)
             if helper.sem_conv_opt_in_mode == _StabilityMode.DEFAULT:
                 helper.process_request(contents, config, span)
@@ -940,6 +955,8 @@ async def instrumented_generate_content_stream(
             "google.genai.AsyncModels.generate_content_stream",
             end_on_exit=False,
         ) as span:
+            if extra_attributes := _get_extra_generate_content_attributes():
+                span.set_attributes(extra_attributes)
             span.set_attributes(request_attributes)
             if (
                 not helper.sem_conv_opt_in_mode

instrumentation-genai/opentelemetry-instrumentation-google-genai/tests/generate_content/nonstreaming_base.py

@@ -31,6 +31,8 @@
 from opentelemetry.util.genai.types import ContentCapturingMode
 
 from .base import TestCase
+from opentelemetry import context as context_api
+from opentelemetry.instrumentation.google_genai import GENERATE_CONTENT_EXTRA_ATTRIBUTES_CONTEXT_KEY
 
 # pylint: disable=too-many-public-methods
 
@@ -99,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 = context_api.attach(context_api.set_value(GENERATE_CONTENT_EXTRA_ATTRIBUTES_CONTEXT_KEY, {"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:
+            context_api.detach(tok)
+
     def test_span_and_event_still_written_when_response_is_exception(self):
         self.configure_exception(ValueError("Uh oh!"))
         patched_environ = patch.dict(

instrumentation-genai/opentelemetry-instrumentation-google-genai/tests/generate_content/streaming_base.py

@@ -15,6 +15,8 @@
 import unittest
 
 from .base import TestCase
+from opentelemetry import context as context_api
+from opentelemetry.instrumentation.google_genai import GENERATE_CONTENT_EXTRA_ATTRIBUTES_CONTEXT_KEY
 
 
 class StreamingTestCase(TestCase):
@@ -42,6 +44,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 = context_api.attach(context_api.set_value(GENERATE_CONTENT_EXTRA_ATTRIBUTES_CONTEXT_KEY, {"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:
+            context_api.detach(tok)
+
     def test_handles_multiple_ressponses(self):
         self.configure_valid_response(text="First response")
         self.configure_valid_response(text="Second response")