refactor: Extract a utility for aggregating partial streaming responses and emitting LlmResponses for them as needed

PiperOrigin-RevId: 800521404

src/google/adk/models/gemini_llm_connection.py

@@ -144,6 +144,8 @@ class GeminiLlmConnection(BaseLlmConnection):
 
     text = ''
     async with Aclosing(self._gemini_session.receive()) as agen:
+      # TODO(b/440101573): Reuse StreamingResponseAggregator to accumulate
+      # partial content and emit responses as needed.
       async for message in agen:
         logger.debug('Got LLM Live message: %s', message)
         if message.server_content:

src/google/adk/models/google_llm.py

@@ -33,6 +33,7 @@ from typing_extensions import override
 
 from .. import version
 from ..utils.context_utils import Aclosing
+from ..utils.streaming_utils import StreamingResponseAggregator
 from ..utils.variant_utils import GoogleLLMVariant
 from .base_llm import BaseLlm
 from .base_llm_connection import BaseLlmConnection
@@ -133,68 +134,23 @@ class Gemini(BaseLlm):
           contents=llm_request.contents,
           config=llm_request.config,
       )
-      response = None
-      thought_text = ''
-      text = ''
-      usage_metadata = None
+
       # for sse, similar as bidi (see receive method in gemini_llm_connecton.py),
       # we need to mark those text content as partial and after all partial
       # contents are sent, we send an accumulated event which contains all the
       # previous partial content. The only difference is bidi rely on
       # complete_turn flag to detect end while sse depends on finish_reason.
+      aggregator = StreamingResponseAggregator()
       async with Aclosing(responses) as agen:
         async for response in agen:
           logger.debug(_build_response_log(response))
-          llm_response = LlmResponse.create(response)
-          usage_metadata = llm_response.usage_metadata
-          if (
-              llm_response.content
-              and llm_response.content.parts
-              and llm_response.content.parts[0].text
-          ):
-            part0 = llm_response.content.parts[0]
-            if part0.thought:
-              thought_text += part0.text
-            else:
-              text += part0.text
-            llm_response.partial = True
-          elif (thought_text or text) and (
-              not llm_response.content
-              or not llm_response.content.parts
-              # don't yield the merged text event when receiving audio data
-              or not llm_response.content.parts[0].inline_data
-          ):
-            parts = []
-            if thought_text:
-              parts.append(types.Part(text=thought_text, thought=True))
-            if text:
-              parts.append(types.Part.from_text(text=text))
-            yield LlmResponse(
-                content=types.ModelContent(parts=parts),
-                usage_metadata=llm_response.usage_metadata,
-            )
-            thought_text = ''
-            text = ''
-          yield llm_response
-
-      # generate an aggregated content at the end regardless the
-      # response.candidates[0].finish_reason
-      if (text or thought_text) and response and response.candidates:
-        parts = []
-        if thought_text:
-          parts.append(types.Part(text=thought_text, thought=True))
-        if text:
-          parts.append(types.Part.from_text(text=text))
-        yield LlmResponse(
-            content=types.ModelContent(parts=parts),
-            error_code=None
-            if response.candidates[0].finish_reason == FinishReason.STOP
-            else response.candidates[0].finish_reason,
-            error_message=None
-            if response.candidates[0].finish_reason == FinishReason.STOP
-            else response.candidates[0].finish_message,
-            usage_metadata=usage_metadata,
-        )
+          async with Aclosing(
+              aggregator.process_response(response)
+          ) as aggregator_gen:
+            async for llm_response in aggregator_gen:
+              yield llm_response
+      if (close_result := aggregator.close()) is not None:
+        yield close_result
 
     else:
       response = await self.api_client.aio.models.generate_content(