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

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

@@ -7,5 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+- Add support for async and streaming.
+  ([#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)) 

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:
-     - 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:
-     - Adoption of 'wrapt' instead of 'functools.wraps'
-     - Bolstering test coverage
-  8. Migrate tests to use VCR.py
+ - Add more span-level attributes for request configuration
+ - Add more span-level attributes for response information
+ - Verify and correct formatting of events:
+   - Including the 'role' field for message events
+   - Including tool invocation information
+ - Emit events for safety ratings when they block responses
+ - Additional cleanup/improvement tasks such as:
+   - Adoption of 'wrapt' instead of 'functools.wraps'
+   - Bolstering test coverage
+ - 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:
@@ -220,7 +216,9 @@ def __init__(
         self._response_index = 0
         self._candidate_index = 0
 
-    def start_span_as_current_span(self, model_name, function_name):
+    def start_span_as_current_span(
+        self, model_name, function_name, end_on_exit=True
+    ):
         return self._otel_wrapper.start_as_current_span(
             f"{_GENERATE_CONTENT_OP_NAME} {model_name}",
             start_time=self._start_time,
@@ -230,6 +228,7 @@ def start_span_as_current_span(self, model_name, function_name):
                 gen_ai_attributes.GEN_AI_REQUEST_MODEL: self._genai_request_model,
                 gen_ai_attributes.GEN_AI_OPERATION_NAME: _GENERATE_CONTENT_OP_NAME,
             },
+            end_on_exit=end_on_exit,
         )
 
     def process_request(
@@ -543,9 +542,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 +582,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 +623,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(
@@ -647,24 +637,38 @@ async def instrumented_generate_content_stream(
             self, otel_wrapper, model
         )
         with helper.start_span_as_current_span(
-            model, "google.genai.AsyncModels.generate_content_stream"
-        ):
+            model,
+            "google.genai.AsyncModels.generate_content_stream",
+            end_on_exit=False,
+        ) as span:
             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)
+        try:
+            response_async_generator = await wrapped_func(
+                self,
+                model=model,
+                contents=contents,
+                config=config,
+                **kwargs,
+            )
+        except Exception as error:  # pylint: disable=broad-exception-caught
+            helper.process_error(error)
+            helper.finalize_processing()
+            with trace.use_span(span, end_on_exit=True):
                 raise
-            finally:
-                helper.finalize_processing()
+
+        async def _response_async_generator_wrapper():
+            with trace.use_span(span, end_on_exit=True):
+                try:
+                    async for response in response_async_generator:
+                        helper.process_response(response)
+                        yield response
+                except Exception as error:
+                    helper.process_error(error)
+                    raise
+                finally:
+                    helper.finalize_processing()
+
+        return _response_async_generator_wrapper()
 
     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!")