Feat(tool): support streaming intermediate results with non live- #4170

src/google/adk/flows/llm_flows/functions.py

@@ -20,13 +20,15 @@
 import base64
 import binascii
 from concurrent.futures import ThreadPoolExecutor
+import contextlib
 import contextvars
 import copy
 import inspect
 import json
 import logging
 import threading
 from typing import Any
+from typing import AsyncGenerator
 from typing import Callable
 from typing import cast
 from typing import Dict
@@ -637,6 +639,12 @@ async def _run_with_trace() -> Event | None:
         function_response = await __call_tool_async(
             tool, args=function_args, tool_context=tool_context
         )
+        function_response = await _drain_tool_result_stream(
+            function_response,
+            tool=tool,
+            tool_context=tool_context,
+            invocation_context=invocation_context,
+        )
       except Exception as tool_error:
         error_response = await _run_on_tool_error_callbacks(
             tool=tool,
@@ -1113,6 +1121,141 @@ async def _emit_streaming_tool_event(
   )
 
 
+async def _iter_tool_yields(
+    generator: object,
+) -> AsyncGenerator[object, None]:
+  """Yields every value a generator tool produces, async or sync.
+
+  Args:
+    generator: The async or sync generator the tool returned.
+
+  Yields:
+    Each value the tool yields, in order.
+  """
+  if inspect.isasyncgen(generator):
+    async with Aclosing(generator) as agen:
+      async for value in agen:
+        yield value
+    return
+
+  with contextlib.closing(cast(Any, generator)) as gen:
+    for value in gen:
+      yield value
+
+
+async def _emit_intermediate_tool_result(
+    function_result: object,
+    *,
+    tool: BaseTool,
+    tool_context: ToolContext,
+    invocation_context: InvocationContext,
+) -> None:
+  """Streams one value a generator tool yielded to the client.
+
+  Marked ``will_continue``, because a further FunctionResponse for the same
+  call always follows: either the tool's next yield, or the result event that
+  answers the call for the model.
+
+  Args:
+    function_result: The value the tool yielded.
+    tool: The tool that yielded it.
+    tool_context: The context of the call, for its function call id.
+    invocation_context: The invocation to enqueue on.
+  """
+  remaining_result, function_response_parts = _extract_multimodal_parts(
+      function_result
+  )
+  content = _build_function_response_content(
+      tool,
+      _normalize_tool_result(remaining_result),
+      tool_context.function_call_id,
+      function_response_parts,
+  )
+  for part in content.parts or []:
+    if part.function_response is not None:
+      part.function_response.will_continue = True
+
+  await invocation_context._enqueue_event(
+      Event(
+          content=content,
+          author=_require_agent_name(invocation_context),
+          invocation_id=invocation_context.invocation_id,
+          branch=invocation_context.branch,
+          partial=True,
+          actions=EventActions(),
+      )
+  )
+
+
+async def _drain_tool_result_stream(
+    function_response: object,
+    *,
+    tool: BaseTool,
+    tool_context: ToolContext,
+    invocation_context: InvocationContext,
+) -> object:
+  """Streams a generator tool's progress and returns the result for the model.
+
+  A tool that yields instead of returning reports progress. Every value it
+  yields is streamed to the client as a partial event the moment it is
+  yielded, and none of them is shown to the model; the last one is also the
+  tool result, and is the only one the model sees. A yielded ``Event`` is a
+  user-facing message rather than a result, so it is streamed and then does
+  not compete to be the last value.
+
+  The last value is delivered twice, once as progress and once as the
+  result; ``will_continue`` tells the two apart, and is true on the progress
+  copy because a further FunctionResponse for that call does follow.
+
+  Args:
+    function_response: The tool's return value. Anything that is not a
+      generator is returned unchanged.
+    tool: The tool being called.
+    tool_context: The context of the call.
+    invocation_context: The invocation to stream on.
+
+  Returns:
+    The last non-Event value the tool yielded, or None if it yielded none.
+  """
+  if not inspect.isasyncgen(function_response) and not inspect.isgenerator(
+      function_response
+  ):
+    return function_response
+
+  can_stream = invocation_context._event_queue is not None
+  if not can_stream:
+    logger.warning(
+        'Tool `%s` yields intermediate results, but this invocation has no'
+        ' event queue to deliver them on, so only its last result is'
+        ' reported. Streaming intermediate results requires running the agent'
+        ' through Runner.',
+        tool.name,
+    )
+
+  result: object = None
+  async with Aclosing(_iter_tool_yields(function_response)) as stream:
+    async for value in stream:
+      if isinstance(value, Event):
+        if can_stream:
+          await _emit_streaming_tool_event(
+              value,
+              tool=tool,
+              tool_context=tool_context,
+              invocation_context=invocation_context,
+          )
+        continue
+      result = value
+      if can_stream:
+        await _emit_intermediate_tool_result(
+            value,
+            tool=tool,
+            tool_context=tool_context,
+            invocation_context=invocation_context,
+        )
+
+  return result
+
+
 async def _process_function_live_helper(
     tool: BaseTool,
     tool_context: ToolContext,