Add support for async and streaming responses in the Google GenAI instrumentation

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

@@ -5,6 +5,11 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## 0.0.2
+
+Add support for aync and streaming.
+([#3298](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3298))
+
 ## Unreleased
 
 Create an initial version of Open Telemetry instrumentation for github.com/googleapis/python-genai.

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

@@ -4,18 +4,17 @@
 
 Here are some TODO items required to achieve stability for this package:
 
-  1. Add support for streaming interfaces
-  2. Add support for async interfaces
-  3. Add more span-level attributes for request configuration
-  4. Add more span-level attributes for response information
-  5. Verify and correct formatting of events:
+  1. Add more span-level attributes for request configuration
+  2. Add more span-level attributes for response information
+  3. Verify and correct formatting of events:
      - Including the 'role' field for message events
      - Including tool invocation information
-  6. Emit events for safety ratings when they block responses
-  7. Additional cleanup/improvement tasks such as:
+  4. Emit events for safety ratings when they block responses
+  5. Additional cleanup/improvement tasks such as:
      - Adoption of 'wrapt' instead of 'functools.wraps'
      - Bolstering test coverage
-  8. Migrate tests to use VCR.py
+  6. Migrate tests to use VCR.py
+
 ## Future
 
 Beyond the above TODOs, it would also be desirable to extend the

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

@@ -45,15 +45,11 @@
 _logger = logging.getLogger(__name__)
 
 
-# Constant used for the value of 'gen_ai.operation.name".
-_GENERATE_CONTENT_OP_NAME = "generate_content"
-
 # Constant used to make the absence of content more understandable.
 _CONTENT_ELIDED = "<elided>"
 
-# Enable these after these cases are fully vetted and tested
-_INSTRUMENT_STREAMING = False
-_INSTRUMENT_ASYNC = False
+# Constant used for the value of 'gen_ai.operation.name".
+_GENERATE_CONTENT_OP_NAME = "generate_content"
 
 
 class _MethodsSnapshot:
@@ -543,9 +539,6 @@ def _create_instrumented_generate_content_stream(
     snapshot: _MethodsSnapshot, otel_wrapper: OTelWrapper
 ):
     wrapped_func = snapshot.generate_content_stream
-    if not _INSTRUMENT_STREAMING:
-        # TODO: remove once this case has been fully tested
-        return wrapped_func
 
     @functools.wraps(wrapped_func)
     def instrumented_generate_content_stream(
@@ -586,9 +579,6 @@ def _create_instrumented_async_generate_content(
     snapshot: _MethodsSnapshot, otel_wrapper: OTelWrapper
 ):
     wrapped_func = snapshot.async_generate_content
-    if not _INSTRUMENT_ASYNC:
-        # TODO: remove once this case has been fully tested
-        return wrapped_func
 
     @functools.wraps(wrapped_func)
     async def instrumented_generate_content(
@@ -630,9 +620,6 @@ def _create_instrumented_async_generate_content_stream(  # pyright: ignore
     snapshot: _MethodsSnapshot, otel_wrapper: OTelWrapper
 ):
     wrapped_func = snapshot.async_generate_content_stream
-    if not _INSTRUMENT_ASYNC or not _INSTRUMENT_STREAMING:
-        # TODO: remove once this case has been fully tested
-        return wrapped_func
 
     @functools.wraps(wrapped_func)
     async def instrumented_generate_content_stream(
@@ -646,25 +633,33 @@ async def instrumented_generate_content_stream(
         helper = _GenerateContentInstrumentationHelper(
             self, otel_wrapper, model
         )
-        with helper.start_span_as_current_span(
-            model, "google.genai.AsyncModels.generate_content_stream"
-        ):
-            helper.process_request(contents, config)
-            try:
-                async for response in await wrapped_func(
-                    self,
-                    model=model,
-                    contents=contents,
-                    config=config,
-                    **kwargs,
-                ):  # pyright: ignore
-                    helper.process_response(response)
-                    yield response  # pyright: ignore
-            except Exception as error:
-                helper.process_error(error)
-                raise
-            finally:
-                helper.finalize_processing()
+
+        async def _internal_generator():
+            with helper.start_span_as_current_span(
+                model, "google.genai.AsyncModels.generate_content_stream"
+            ):
+                helper.process_request(contents, config)
+                try:
+                    async for response in await wrapped_func(
+                        self,
+                        model=model,
+                        contents=contents,
+                        config=config,
+                        **kwargs,
+                    ):  # pyright: ignore
+                        helper.process_response(response)
+                        yield response  # pyright: ignore
+                except Exception as error:
+                    helper.process_error(error)
+                    raise
+                finally:
+                    helper.finalize_processing()
+
+        class _GeneratorProvider:
+            def __aiter__(self):
+                return _internal_generator()
+
+        return _GeneratorProvider()
 
     return instrumented_generate_content_stream
 

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

@@ -17,4 +17,4 @@
 #   This version should stay below "1.0" until the fundamentals
 #   in "TODOS.md" have been addressed. Please revisit the TODOs
 #   listed there before bumping to a stable version.
-__version__ = "0.0.1.dev"
+__version__ = "0.0.2.dev"

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

@@ -179,6 +179,16 @@ def get_event_named(self, event_name):
                 return event
         return None
 
+    def get_events_named(self, event_name):
+        result = []
+        for event in self.get_finished_logs():
+            event_name_attr = event.attributes.get("event.name")
+            if event_name_attr is None:
+                continue
+            if event_name_attr == event_name:
+                result.append(event)
+        return result
+
     def assert_has_event_named(self, name):
         event = self.get_event_named(name)
         finished_logs = self.get_finished_logs()

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

@@ -37,6 +37,7 @@
 import http.client
 import io
 import json
+from typing import Optional
 
 import requests
 import requests.sessions
@@ -81,7 +82,7 @@ def response(self):
 
 
 def _return_error_status(
-    args: RequestsCallArgs, status_code: int, reason: str = None
+    args: RequestsCallArgs, status_code: int, reason: Optional[str] = None
 ):
     result = requests.Response()
     result.url = args.request.url
@@ -123,6 +124,35 @@ def generate_response_from_dict(args):
     raise ValueError(f"Unsupported response type: {type(response)}")
 
 
+def _to_stream_response_generator(response_generators):
+    if len(response_generators) == 1:
+        return response_generators[0]
+
+    def combined_generator(args):
+        first_response = response_generators[0](args)
+        if first_response.status_code != 200:
+            return first_response
+        result = requests.Response()
+        result.status_code = 200
+        result.headers["content-type"] = "application/json"
+        result.encoding = "utf-8"
+        result.headers["transfer-encoding"] = "chunked"
+        contents = []
+        for generator in response_generators:
+            response = generator(args)
+            if response.status_code != 200:
+                continue
+            response_json = response.json()
+            response_json_str = json.dumps(response_json)
+            contents.append(f"data: {response_json_str}")
+        contents_str = "\r\n".join(contents)
+        full_contents = f"{contents_str}\r\n\r\n"
+        result.raw = io.BytesIO(full_contents.encode())
+        return result
+
+    return combined_generator
+
+
 class RequestsMocker:
     def __init__(self):
         self._original_send = requests.sessions.Session.send
@@ -159,6 +189,38 @@ def _do_send(
         session: requests.sessions.Session,
         request: requests.PreparedRequest,
         **kwargs,
+    ):
+        stream = kwargs.get("stream", False)
+        if not stream:
+            return self._do_send_non_streaming(session, request, **kwargs)
+        return self._do_send_streaming(session, request, **kwargs)
+
+    def _do_send_streaming(
+        self,
+        session: requests.sessions.Session,
+        request: requests.PreparedRequest,
+        **kwargs,
+    ):
+        args = RequestsCallArgs(session, request, **kwargs)
+        response_generators = []
+        for matcher, response_generator in self._handlers:
+            if matcher is None:
+                response_generators.append(response_generator)
+            elif matcher(args):
+                response_generators.append(response_generator)
+        if not response_generators:
+            response_generators.append(_return_404)
+        response_generator = _to_stream_response_generator(response_generators)
+        call = RequestsCall(args, response_generator)
+        result = call.response
+        self._calls.append(call)
+        return result
+
+    def _do_send_non_streaming(
+        self,
+        session: requests.sessions.Session,
+        request: requests.PreparedRequest,
+        **kwargs,
     ):
         args = RequestsCallArgs(session, request, **kwargs)
         response_generator = self._lookup_response_generator(args)

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

@@ -17,31 +17,7 @@
 import unittest
 
 from ..common.base import TestCase
-
-
-def create_valid_response(
-    response_text="The model response", input_tokens=10, output_tokens=20
-):
-    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,
-                        }
-                    ],
-                }
-            }
-        ],
-    }
+from .util import create_valid_response
 
 
 class NonStreamingTestCase(TestCase):
@@ -56,22 +32,12 @@ def setUp(self):  # pylint: disable=invalid-name
     def generate_content(self, *args, **kwargs):
         raise NotImplementedError("Must implement 'generate_content'.")
 
+    @property
     def expected_function_name(self):
         raise NotImplementedError("Must implement 'expected_function_name'.")
 
-    def configure_valid_response(
-        self,
-        response_text="The model_response",
-        input_tokens=10,
-        output_tokens=20,
-    ):
-        self.requests.add_response(
-            create_valid_response(
-                response_text=response_text,
-                input_tokens=input_tokens,
-                output_tokens=output_tokens,
-            )
-        )
+    def configure_valid_response(self, *args, **kwargs):
+        self.requests.add_response(create_valid_response(*args, **kwargs))
 
     def test_instrumentation_does_not_break_core_functionality(self):
         self.configure_valid_response(response_text="Yep, it works!")

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

@@ -0,0 +1,76 @@
+# Copyright The OpenTelemetry Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+from ..common.base import TestCase
+from .util import create_valid_response
+
+
+class StreamingTestCase(TestCase):
+    # The "setUp" function is defined by "unittest.TestCase" and thus
+    # this name must be used. Uncertain why pylint doesn't seem to
+    # recognize that this is a unit test class for which this is inherited.
+    def setUp(self):  # pylint: disable=invalid-name
+        super().setUp()
+        if self.__class__ == StreamingTestCase:
+            raise unittest.SkipTest("Skipping testcase base.")
+
+    def generate_content(self, *args, **kwargs):
+        raise NotImplementedError("Must implement 'generate_content'.")
+
+    @property
+    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 test_instrumentation_does_not_break_core_functionality(self):
+        self.configure_valid_response(response_text="Yep, it works!")
+        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!")
+
+    def test_handles_multiple_ressponses(self):
+        self.configure_valid_response(response_text="First response")
+        self.configure_valid_response(response_text="Second response")
+        responses = self.generate_content(
+            model="gemini-2.0-flash", contents="Does this work?"
+        )
+        self.assertEqual(len(responses), 2)
+        self.assertEqual(responses[0].text, "First response")
+        self.assertEqual(responses[1].text, "Second response")
+        choice_events = self.otel.get_events_named("gen_ai.choice")
+        self.assertEqual(len(choice_events), 2)
+
+    def test_includes_token_counts_in_span_aggregated_from_responses(self):
+        # Configure multiple responses whose input/output tokens should be
+        # accumulated together when summarizing the end-to-end request.
+        #
+        #   Input: 1 + 3 + 5 => 4 + 5 => 9
+        #   Output: 2 + 4 + 6 => 6 + 6 => 12
+        self.configure_valid_response(input_tokens=1, output_tokens=2)
+        self.configure_valid_response(input_tokens=3, output_tokens=4)
+        self.configure_valid_response(input_tokens=5, output_tokens=6)
+
+        self.generate_content(model="gemini-2.0-flash", contents="Some input")
+
+        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["gen_ai.usage.input_tokens"], 9)
+        self.assertEqual(span.attributes["gen_ai.usage.output_tokens"], 12)