Restructure streaming async logic to begin execution earlier.

Author: michaelsafyan

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

@@ -216,7 +216,7 @@ 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,
@@ -226,6 +226,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(
@@ -633,31 +634,41 @@ async def instrumented_generate_content_stream(
         helper = _GenerateContentInstrumentationHelper(
             self, otel_wrapper, model
         )
-
-        async def _internal_generator():
-            with helper.start_span_as_current_span(
-                model, "google.genai.AsyncModels.generate_content_stream"
-            ):
-                helper.process_request(contents, config)
+        with helper.start_span_as_current_span(
+            model,
+            "google.genai.AsyncModels.generate_content_stream",
+            end_on_exit=False) as span:
+         helper.process_request(contents, config)
+        try:
+            response_async_generator = await wrapped_func(
+                    self,
+                    model=model,
+                    contents=contents,
+                    config=config,
+                    **kwargs,
+                )
+        except Exception as error:
+         helper.process_error(error)
+         helper.finalize_processing()
+         with trace.use_span(span, end_on_exit=True):
+            raise
+
+        async def _response_async_generator_wrapper():
+            with trace.use_span(span, end_on_exit=True):
                 try:
-                    async for response in await wrapped_func(
-                        self,
-                        model=model,
-                        contents=contents,
-                        config=config,
-                        **kwargs,
-                    ):  # pyright: ignore
+                    async for response in response_async_generator:
                         helper.process_response(response)
-                        yield response  # pyright: ignore
+                        yield response
                 except Exception as error:
                     helper.process_error(error)
                     raise
                 finally:
                     helper.finalize_processing()
 
         class _GeneratorProvider:
+
             def __aiter__(self):
-                return _internal_generator()
+                return _response_async_generator_wrapper()
 
         return _GeneratorProvider()