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

Author: RKest

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)

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(

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")