Checkpoint current state.

Author: michaelsafyan

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

@@ -40,11 +40,6 @@
 _logger = logging.getLogger(__name__)
 
 
-# Enable these after these cases are fully vetted and tested
-_INSTRUMENT_STREAMING = False
-_INSTRUMENT_ASYNC = False
-
-
 class _MethodsSnapshot:
     def __init__(self):
         self._original_generate_content = Models.generate_content
@@ -436,10 +431,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(
         self: Models,
@@ -479,10 +470,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(
         self: AsyncModels,
@@ -523,10 +510,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(
         self: AsyncModels,
@@ -539,25 +522,30 @@ 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/tests/common/requests_mocker.py

@@ -123,6 +123,27 @@ 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
+            # TODO: ...
+    return combined_generator
+
+
 class RequestsMocker:
     def __init__(self):
         self._original_send = requests.sessions.Session.send
@@ -159,6 +180,37 @@ def _do_send(
         session: requests.sessions.Session,
         request: requests.PreparedRequest,
         **kwargs,
+    ):
+        stream=kwargs.get('stream', False)
+        if not stream:
+            return _do_send_non_streaming(session, request, **kwargs)
+        return _do_send_streaming(session, request, **kwargs)
+
+    def _do_send_streaming(
+        self,
+        session: requests.sessions.Session,
+        request: requests.PreparedRequest,
+        **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)
+        args = RequestsCallArgs(session, request, **kwargs)
+        response_generator = _to_stream_response_generator(response_generators)
+        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,17 @@ 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,
+        *args,
+        **kwargs
     ):
         self.requests.add_response(
-            create_valid_response(
-                response_text=response_text,
-                input_tokens=input_tokens,
-                output_tokens=output_tokens,
-            )
-        )
+            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,54 @@
+# 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 json
+import os
+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!")

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

@@ -12,76 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
-# TODO: Once the async non-streaming case has been fully implemented,
-# reimplement this in terms of "nonstreaming_base.py".
-
 import asyncio
-import logging
-import unittest
-
-from ..common.base import TestCase
 
+from .nonstreaming_base import NonStreamingTestCase
 
-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,
-                        }
-                    ],
-                }
-            }
-        ],
-    }
-
-
-# Temporary test fixture just to ensure that the in-progress work to
-# implement this case doesn't break the original code.
-class TestGenerateContentAsyncNonstreaming(TestCase):
-    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,
-            )
-        )
 
+class TestGenerateContentAsyncNonstreaming(NonStreamingTestCase):
+    
     def generate_content(self, *args, **kwargs):
         return asyncio.run(
             self.client.aio.models.generate_content(*args, **kwargs)  # pylint: disable=missing-kwoa
         )
 
-    def test_async_generate_content_not_broken_by_instrumentation(self):
-        self.configure_valid_response(response_text="Yep, it works!")
-        response = self.generate_content(
-            model="gemini-2.0-flash", contents="Does this work?"
-        )
-        self.assertEqual(response.text, "Yep, it works!")
-
-
-def main():
-    logging.basicConfig(level=logging.DEBUG)
-    unittest.main()
-
-
-if __name__ == "__main__":
-    main()
+    @property
+    def expected_function_name(self):
+        return 'google.genai.AsyncModels.generate_content'