fix(google-genai): make streaming responses work (again) (#3421)

Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com>

Author: nirga

packages/opentelemetry-instrumentation-google-generativeai/opentelemetry/instrumentation/google_generativeai/__init__.py

@@ -48,6 +48,18 @@
         "method": "generate_content",
         "span_name": "gemini.generate_content",
     },
+    {
+        "package": "google.genai.models",
+        "object": "Models",
+        "method": "generate_content_stream",
+        "span_name": "gemini.generate_content",
+    },
+    {
+        "package": "google.genai.models",
+        "object": "AsyncModels",
+        "method": "generate_content_stream",
+        "span_name": "gemini.generate_content",
+    },
 ]
 
 
@@ -66,8 +78,10 @@ def _build_from_streaming_response(
     event_logger,
 ):
     complete_response = ""
+    last_chunk = None
     for item in response:
         item_to_yield = item
+        last_chunk = item
         complete_response += str(item.text)
 
         yield item_to_yield
@@ -76,16 +90,18 @@ def _build_from_streaming_response(
         emit_choice_events(response, event_logger)
     else:
         set_response_attributes(span, complete_response, llm_model)
-    set_model_response_attributes(span, response, llm_model)
+    set_model_response_attributes(span, last_chunk or response, llm_model)
     span.end()
 
 
 async def _abuild_from_streaming_response(
     span, response: GenerateContentResponse, llm_model, event_logger
 ):
     complete_response = ""
+    last_chunk = None
     async for item in response:
         item_to_yield = item
+        last_chunk = item
         complete_response += str(item.text)
 
         yield item_to_yield
@@ -94,7 +110,7 @@ async def _abuild_from_streaming_response(
         emit_choice_events(response, event_logger)
     else:
         set_response_attributes(span, complete_response, llm_model)
-    set_model_response_attributes(span, response, llm_model)
+    set_model_response_attributes(span, last_chunk if last_chunk else response, llm_model)
     span.end()
 
 

packages/opentelemetry-instrumentation-google-generativeai/tests/test_new_library_instrumentation.py

@@ -9,23 +9,31 @@ def test_library_instrumentation():
     """Test that the google-genai library gets properly instrumented."""
     # Import the library
     from google import genai
-    from google.genai.models import Models
+    from google.genai.models import AsyncModels, Models
 
     # Set up instrumentor
     instrumentor = GoogleGenerativeAiInstrumentor()
 
     # Verify methods are not wrapped initially
     assert not hasattr(Models.generate_content, '__wrapped__')
+    assert not hasattr(Models.generate_content_stream, '__wrapped__')
+    assert not hasattr(AsyncModels.generate_content, '__wrapped__')
+    assert not hasattr(AsyncModels.generate_content_stream, '__wrapped__')
 
     try:
         instrumentor.instrument()
 
-        # Verify methods are now wrapped
+        # Verify all methods are now wrapped
         assert hasattr(Models.generate_content, '__wrapped__')
+        assert hasattr(Models.generate_content_stream, '__wrapped__')
+        assert hasattr(AsyncModels.generate_content, '__wrapped__')
+        assert hasattr(AsyncModels.generate_content_stream, '__wrapped__')
 
-        # Verify it's our wrapper
-        wrapped_method = Models.generate_content
-        assert callable(wrapped_method)
+        # Verify they're callable
+        assert callable(Models.generate_content)
+        assert callable(Models.generate_content_stream)
+        assert callable(AsyncModels.generate_content)
+        assert callable(AsyncModels.generate_content_stream)
 
         # Test that we can create a client
         client = genai.Client(api_key="test_key")
@@ -38,6 +46,9 @@ def test_library_instrumentation():
 
         # Verify methods are unwrapped
         assert not hasattr(Models.generate_content, '__wrapped__')
+        assert not hasattr(Models.generate_content_stream, '__wrapped__')
+        assert not hasattr(AsyncModels.generate_content, '__wrapped__')
+        assert not hasattr(AsyncModels.generate_content_stream, '__wrapped__')
 
 
 def test_instrumentation_dependencies():
@@ -54,7 +65,7 @@ def test_wrapped_methods():
     instrumentor = GoogleGenerativeAiInstrumentor()
     methods = instrumentor._wrapped_methods()
 
-    assert len(methods) == 2
+    assert len(methods) == 4
 
     # Should be using new library methods
     packages = [method.get("package", "") for method in methods]
@@ -64,3 +75,8 @@ def test_wrapped_methods():
     objects = [method.get("object", "") for method in methods]
     assert "Models" in objects
     assert "AsyncModels" in objects
+
+    # Should wrap both generate_content and generate_content_stream
+    wrapped_methods = [method.get("method", "") for method in methods]
+    assert wrapped_methods.count("generate_content") == 2
+    assert wrapped_methods.count("generate_content_stream") == 2