Add tests to verify instrumentation of the tool calls.

Author: michaelsafyan

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

@@ -540,7 +540,7 @@ def _wrapped_tool(otel_wrapper: OTelWrapper, tool: ToolListUnionDict):
     if inspect.iscoroutinefunction(tool):
         return tool
     tool_name = tool.__name__
-    should_record_contents = flags.is_content_recording_enabled()
+    should_record_contents = is_content_recording_enabled()
     @functools.wraps(tool)
     def wrapped_tool(*args, **kwargs):
         with otel_wrapper.start_as_current_span(
@@ -561,7 +561,7 @@ def _wrapped_config_with_tools(
     otel_wrapper: OTelWrapper,
     config: GenerateContentConfig) -> GenerateContentConfig:
     result = copy.copy(config)
-    result.tool = [_wrapped_tool(otel_wrapper, tool) for tool in config.tools]
+    result.tools = [_wrapped_tool(otel_wrapper, tool) for tool in config.tools]
     return result
 
 

instrumentation-genai/opentelemetry-instrumentation-google-genai/tests/common/base.py

@@ -12,10 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Optional
+
 import os
 import unittest
+import unittest.mock
 
 import google.genai
+import google.genai.types as genai_types
+from google.genai.models import Models, AsyncModels
 
 from .instrumentation_context import InstrumentationContext
 from .otel_mocker import OTelMocker
@@ -28,6 +33,7 @@ def refresh(self, request):
 
 
 class TestCase(unittest.TestCase):
+
     def setUp(self):
         self._otel = OTelMocker()
         self._otel.install()
@@ -40,11 +46,31 @@ def setUp(self):
         self._client = None
         self._uses_vertex = False
         self._credentials = _FakeCredentials()
+        self._generate_content_mock = None
+        self._generate_content_stream_mock = None
+        self._original_generate_content = Models.generate_content
+        self._original_generate_content_stream = Models.generate_content_stream
+        self._original_async_generate_content = AsyncModels.generate_content
+        self._original_async_generate_content_stream = (
+            AsyncModels.generate_content_stream
+        )
 
     def _lazy_init(self):
         self._instrumentation_context = InstrumentationContext()
         self._instrumentation_context.install()
 
+    @property
+    def mock_generate_content(self):
+        if self._generate_content_mock is None:
+            self._create_mocks()
+        return self._generate_content_mock
+    
+    @property
+    def mock_generate_content_stream(self):
+        if self._generate_content_stream_mock is None:
+            self._create_mocks()
+        return self._generate_content_stream_mock
+
     @property
     def client(self):
         if self._client is None:
@@ -62,6 +88,81 @@ def otel(self):
     def set_use_vertex(self, use_vertex):
         self._uses_vertex = use_vertex
 
+    def generate_content_response(
+        self,
+        part: Optional[genai_types.Part] = None,
+        parts: Optional[list[genai_types.Part]] = None,
+        content: Optional[genai_types.Content] = None,
+        candidate: Optional[genai_types.Candidate] = None,
+        candidates: Optional[list[genai_types.Candidate]] = None,
+        text: Optional[str] = None):
+        if text is None:
+            text = 'Some response text'
+        if part is None:
+            part = genai_types.Part(text=text)
+        if parts is None:
+            parts = [part]
+        if content is None:
+            content = genai_types.Content(parts=parts, role='model')
+        if candidate is None:
+            candidate = genai_types.Candidate(content=content)
+        if candidates is None:
+            candidates = [candidate]
+        return genai_types.GenerateContentResponse(candidates=candidates)
+
+    def _create_mocks(self):
+        print("Initializing mocks.")
+        if self._client is not None:
+            self._client = None
+        if self._instrumentation_context is not None:
+            self._instrumentation_context.uninstall()
+            self._instrumentation_context = None
+        self._generate_content_mock = unittest.mock.MagicMock()
+        self._generate_content_stream_mock = unittest.mock.MagicMock()
+
+        def convert_response(arg):
+            if isinstance(arg, genai_types.GenerateContentResponse):
+                return arg
+            if isinstance(arg, str):
+                return self.generate_content_response(text=arg)
+            if isinstance(arg, dict):
+                try:
+                    return genai_types.GenerateContentResponse(**arg)
+                except Exception:
+                    return self.generate_content_response(**arg)
+            return arg
+        
+        def default_stream(*args, **kwargs):
+            result = self._generate_content_mock(*args, **kwargs)
+            yield result
+        self._generate_content_stream_mock.side_effect = default_stream
+
+        def sync_variant(*args, **kwargs):
+            return convert_response(self._generate_content_mock(*args, **kwargs))
+        
+        def sync_stream_variant(*args, **kwargs):
+            print("Calling sync stream variant")
+            for result in self._generate_content_stream_mock(*args, **kwargs):
+                yield convert_response(result)
+
+        async def async_variant(*args, **kwargs):
+            print("Calling async non-streaming variant")
+            return sync_variant(*args, **kwargs)
+
+        async def async_stream_variant(*args, **kwargs):
+            print("Calling async stream variant")
+            async def gen():
+                for result in sync_stream_variant(*args, **kwargs):
+                    yield result
+            class GeneratorProvider:
+                def __aiter__(self):
+                    return gen()
+            return GeneratorProvider()
+        Models.generate_content = sync_variant
+        Models.generate_content_stream = sync_stream_variant
+        AsyncModels.generate_content = async_variant
+        AsyncModels.generate_content_stream = async_stream_variant
+
     def _create_client(self):
         self._lazy_init()
         if self._uses_vertex:
@@ -77,5 +178,16 @@ def _create_client(self):
     def tearDown(self):
         if self._instrumentation_context is not None:
             self._instrumentation_context.uninstall()
+        if self._generate_content_mock is None:
+            assert Models.generate_content == self._original_generate_content
+            assert Models.generate_content_stream == self._original_generate_content_stream
+            assert AsyncModels.generate_content == self._original_async_generate_content
+            assert AsyncModels.generate_content_stream == self._original_async_generate_content_stream
         self._requests.uninstall()
         self._otel.uninstall()
+        Models.generate_content = self._original_generate_content
+        Models.generate_content_stream = self._original_generate_content_stream
+        AsyncModels.generate_content = self._original_async_generate_content
+        AsyncModels.generate_content_stream = (
+            self._original_async_generate_content_stream
+        )

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

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from google.genai import types as genai_types
 import json
 import os
 import unittest
@@ -36,8 +37,10 @@ def generate_content(self, *args, **kwargs):
     def expected_function_name(self):
         raise NotImplementedError("Must implement 'expected_function_name'.")
 
-    def configure_valid_response(self, *args, **kwargs):
-        self.requests.add_response(create_valid_response(*args, **kwargs))
+    def configure_valid_response(self, *args, if_matches=None, **kwargs):
+        self.requests.add_response(
+            create_valid_response(*args, **kwargs),
+            if_matches=if_matches)
 
     def test_instrumentation_does_not_break_core_functionality(self):
         self.configure_valid_response(response_text="Yep, it works!")
@@ -197,3 +200,31 @@ def test_records_metrics_data(self):
         self.otel.assert_has_metrics_data_named(
             "gen_ai.client.operation.duration"
         )
+
+    def test_autoinstruments_tools(self):
+        def factorial(n: int):
+            result = 1
+            while n > 1:
+                result *= n
+                n -= 1
+            return result
+
+        def generate_content_impl(*args, **kwargs):
+            config = kwargs['config']
+            tools = config.tools
+            assert len(tools) == 1
+            factorial_tool = tools[0]
+            return str(factorial_tool(5))
+        
+        self.mock_generate_content.side_effect = generate_content_impl
+
+        response=self.generate_content(
+            model="gemini-2.0-flash",
+            contents="Compute 5 factorial",
+            config=genai_types.GenerateContentConfig(tools=[factorial]))
+        self.assertEqual(response.text, '120')
+        self.otel.assert_has_span_named("generate_content gemini-2.0-flash")
+        self.otel.assert_has_span_named("tool_call factorial")
+        generate_content_span = self.otel.get_span_named("generate_content gemini-2.0-flash")
+        tool_call_span = self.otel.get_span_named("tool_call factorial")
+        self.assertEqual(tool_call_span.parent.span_id, generate_content_span.context.span_id)

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

@@ -14,25 +14,33 @@
 
 
 def create_valid_response(
-    response_text="The model response", input_tokens=10, output_tokens=20
+    response_text="The model response",
+    input_tokens=10,
+    output_tokens=20,
+    part=None,
+    parts=None,
+    candidate=None,
+    candidates=None
 ):
+    if part is None:
+        part = {"text": response_text}
+    if parts is None:
+        parts = [part]
+    if candidate is None:
+        candidate = {
+            "content": {
+                "role": "model",
+                "parts": parts,
+            }
+        }
+    if candidates is None:
+        candidates = [candidate]
     return {
         "modelVersion": "gemini-2.0-flash-test123",
         "usageMetadata": {
             "promptTokenCount": input_tokens,
             "candidatesTokenCount": output_tokens,
             "totalTokenCount": input_tokens + output_tokens,
         },
-        "candidates": [
-            {
-                "content": {
-                    "role": "model",
-                    "parts": [
-                        {
-                            "text": response_text,
-                        }
-                    ],
-                }
-            }
-        ],
+        "candidates": candidates,
     }