feat: support supress_instrumentation in opentelemetry-instrumentation-google-genai

Author: RKest

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

@@ -44,6 +44,7 @@
     _OpenTelemetryStabilitySignalType,
     _StabilityMode,
 )
+from opentelemetry.instrumentation.utils import is_instrumentation_enabled
 from opentelemetry.semconv._incubating.attributes import (
     code_attributes,
     gen_ai_attributes,
@@ -705,6 +706,11 @@ def instrumented_generate_content(
         config: Optional[GenerateContentConfigOrDict] = None,
         **kwargs: Any,
     ) -> GenerateContentResponse:
+        if not is_instrumentation_enabled():
+            return wrapped_func(
+                self, model=model, contents=contents, config=config, **kwargs
+            )
+
         candidates = []
         helper = _GenerateContentInstrumentationHelper(
             self,
@@ -778,6 +784,13 @@ def instrumented_generate_content_stream(
         config: Optional[GenerateContentConfigOrDict] = None,
         **kwargs: Any,
     ) -> Iterator[GenerateContentResponse]:
+        if not is_instrumentation_enabled():
+            for resp in wrapped_func(
+                self, model=model, contents=contents, config=config, **kwargs
+            ):
+                yield resp
+            return
+
         candidates: list[Candidate] = []
         helper = _GenerateContentInstrumentationHelper(
             self,
@@ -851,6 +864,10 @@ async def instrumented_generate_content(
         config: Optional[GenerateContentConfigOrDict] = None,
         **kwargs: Any,
     ) -> GenerateContentResponse:
+        if not is_instrumentation_enabled():
+            return await wrapped_func(
+                self, model=model, contents=contents, config=config, **kwargs
+            )
         helper = _GenerateContentInstrumentationHelper(
             self,
             otel_wrapper,
@@ -924,6 +941,11 @@ async def instrumented_generate_content_stream(
         config: Optional[GenerateContentConfigOrDict] = None,
         **kwargs: Any,
     ) -> Awaitable[AsyncIterator[GenerateContentResponse]]:  # type: ignore
+        if not is_instrumentation_enabled():
+            return await wrapped_func(
+                self, model=model, contents=contents, config=config, **kwargs
+            )
+
         helper = _GenerateContentInstrumentationHelper(
             self,
             otel_wrapper,

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

@@ -29,6 +29,7 @@
     _OpenTelemetryStabilitySignalType,
     _StabilityMode,
 )
+from opentelemetry.instrumentation.utils import is_instrumentation_enabled
 from opentelemetry.semconv._incubating.attributes import (
     code_attributes,
 )
@@ -168,6 +169,9 @@ def _wrap_sync_tool_function(
 ):
     @functools.wraps(tool_function)
     def wrapped_function(*args, **kwargs):
+        if not is_instrumentation_enabled():
+            return tool_function(*args, **kwargs)
+
         span_name = _create_function_span_name(tool_function)
         attributes = _create_function_span_attributes(
             tool_function, args, kwargs, extra_span_attributes
@@ -193,6 +197,9 @@ def _wrap_async_tool_function(
 ):
     @functools.wraps(tool_function)
     async def wrapped_function(*args, **kwargs):
+        if not is_instrumentation_enabled():
+            return await tool_function(*args, **kwargs)
+
         span_name = _create_function_span_name(tool_function)
         attributes = _create_function_span_attributes(
             tool_function, args, kwargs, extra_span_attributes

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

@@ -20,6 +20,7 @@
 from google.genai.types import GenerateContentConfig, Part
 from pydantic import BaseModel, Field
 
+from opentelemetry.instrumentation import utils
 from opentelemetry.instrumentation._semconv import (
     _OpenTelemetrySemanticConventionStability,
     _OpenTelemetryStabilitySignalType,
@@ -499,3 +500,14 @@ def test_records_metrics_data(self):
         self.otel.assert_has_metrics_data_named(
             "gen_ai.client.operation.duration"
         )
+
+    def test_suppress_instrumentation(self):
+        self.configure_valid_response(text="Yep, it works!")
+        with utils.suppress_instrumentation():
+            response = self.generate_content(
+                model="gemini-2.0-flash", contents="Does this work?"
+            )
+        self.assertEqual(response.text, "Yep, it works!")
+        self.otel.assert_does_not_have_span_named(
+            "generate_content gemini-2.0-flash"
+        )

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

@@ -14,6 +14,8 @@
 
 import unittest
 
+from opentelemetry.instrumentation import utils
+
 from .base import TestCase
 
 
@@ -70,3 +72,16 @@ def test_includes_token_counts_in_span_aggregated_from_responses(self):
         span = self.otel.get_span_named("generate_content gemini-2.0-flash")
         self.assertEqual(span.attributes["gen_ai.usage.input_tokens"], 9)
         self.assertEqual(span.attributes["gen_ai.usage.output_tokens"], 12)
+
+    def test_suppress_instrumentation(self):
+        self.configure_valid_response(text="Yep, it works!")
+        with utils.suppress_instrumentation():
+            responses = self.generate_content(
+                model="gemini-2.0-flash", contents="Does this work?"
+            )
+        self.assertEqual(len(responses), 1)
+        response = responses[0]
+        self.assertEqual(response.text, "Yep, it works!")
+        self.otel.assert_does_not_have_span_named(
+            "generate_content gemini-2.0-flash"
+        )

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

@@ -16,6 +16,7 @@
 
 import google.genai.types as genai_types
 
+from opentelemetry.instrumentation import utils
 from opentelemetry.instrumentation._semconv import (
     _OpenTelemetrySemanticConventionStability,
     _OpenTelemetryStabilitySignalType,
@@ -440,3 +441,35 @@ def somefunction(x, y=2):
                             generated_span.attributes,
                         )
                 self.tearDown()
+
+    def test_suppress_tool_call_instrumentation(self):
+        calls = []
+
+        def handle(*args, **kwargs):
+            calls.append((args, kwargs))
+            return "some result"
+
+        def somefunction(somearg):
+            print("somearg=%s", somearg)
+
+        self.mock_generate_content.side_effect = handle
+        self.client.models.generate_content(
+            model="some-model-name",
+            contents="Some content",
+            config={
+                "tools": [somefunction],
+            },
+        )
+        self.assertEqual(len(calls), 1)
+        config = calls[0][1]["config"]
+        tools = config.tools
+        wrapped_somefunction = tools[0]
+
+        self.assertIsNone(
+            self.otel.get_span_named("execute_tool somefunction")
+        )
+
+        with utils.suppress_instrumentation():
+            wrapped_somefunction("someparam")
+
+        self.otel.assert_does_not_have_span_named("execute_tool somefunction")