feat: enable to will_continue-like function response in BaseLlmFlow.run_async method with streaming mode.

Author: Lin-Nikaido

src/google/adk/agents/readonly_context.py

@@ -22,6 +22,7 @@
 if TYPE_CHECKING:
   from google.genai import types
 
+  from ..sessions.session import Session
   from .invocation_context import InvocationContext
 
 
@@ -52,3 +53,8 @@ def agent_name(self) -> str:
   def state(self) -> MappingProxyType[str, Any]:
     """The state of the current session. READONLY field."""
     return MappingProxyType(self._invocation_context.session.state)
+
+  @property
+  def session(self) -> Session:
+    """The current session. READONLY field."""
+    return self._invocation_context.session

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

@@ -24,10 +24,11 @@
 from typing import Optional
 from typing import TYPE_CHECKING
 
-from google.genai import types
 from websockets.exceptions import ConnectionClosed
 from websockets.exceptions import ConnectionClosedOK
 
+from google.genai import types
+
 from . import _output_schema_processor
 from . import functions
 from ...agents.base_agent import BaseAgent
@@ -629,43 +630,51 @@ async def _postprocess_handle_function_calls_async(
       function_call_event: Event,
       llm_request: LlmRequest,
   ) -> AsyncGenerator[Event, None]:
-    if function_response_event := await functions.handle_function_calls_async(
+    # if invocation_context.run_config.streaming_mode == StreamingMode.SSE:
+    #
+    # else:
+    if function_response_event_agen := functions.handle_function_calls_async_gen(
         invocation_context, function_call_event, llm_request.tools_dict
     ):
-      auth_event = functions.generate_auth_event(
-          invocation_context, function_response_event
-      )
-      if auth_event:
-        yield auth_event
-
-      tool_confirmation_event = functions.generate_request_confirmation_event(
-          invocation_context, function_call_event, function_response_event
-      )
-      if tool_confirmation_event:
-        yield tool_confirmation_event
-
-      # Always yield the function response event first
-      yield function_response_event
-
-      # Check if this is a set_model_response function response
-      if json_response := _output_schema_processor.get_structured_model_response(
-          function_response_event
-      ):
-        # Create and yield a final model response event
-        final_event = (
-            _output_schema_processor.create_final_model_response_event(
-                invocation_context, json_response
-            )
+      function_response_event = None
+      async for function_response_event in function_response_event_agen:
+        auth_event = functions.generate_auth_event(
+            invocation_context, function_response_event
         )
-        yield final_event
-      transfer_to_agent = function_response_event.actions.transfer_to_agent
-      if transfer_to_agent:
-        agent_to_run = self._get_agent_to_run(
-            invocation_context, transfer_to_agent
+        if auth_event:
+          yield auth_event
+
+        tool_confirmation_event = functions.generate_request_confirmation_event(
+            invocation_context, function_call_event, function_response_event
         )
-        async with Aclosing(agent_to_run.run_async(invocation_context)) as agen:
-          async for event in agen:
-            yield event
+        if tool_confirmation_event:
+          yield tool_confirmation_event
+
+        # Always yield the function response event first
+        yield function_response_event
+
+        # Check if this is a set_model_response function response
+        if json_response := _output_schema_processor.get_structured_model_response(
+            function_response_event
+        ):
+          # Create and yield a final model response event
+          final_event = (
+              _output_schema_processor.create_final_model_response_event(
+                  invocation_context, json_response
+              )
+          )
+          yield final_event
+      if function_response_event:
+        transfer_to_agent = function_response_event.actions.transfer_to_agent
+        if transfer_to_agent:
+          agent_to_run = self._get_agent_to_run(
+              invocation_context, transfer_to_agent
+          )
+          async with Aclosing(
+              agent_to_run.run_async(invocation_context)
+          ) as agen:
+            async for event in agen:
+              yield event
 
   def _get_agent_to_run(
       self, invocation_context: InvocationContext, agent_name: str

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

@@ -20,18 +20,24 @@
 import copy
 import inspect
 import logging
-import threading
 from typing import Any
 from typing import AsyncGenerator
+from typing import AsyncIterator
 from typing import cast
+from typing import Iterator
+from typing import List
 from typing import Optional
+from typing import Tuple
 from typing import TYPE_CHECKING
 import uuid
 
+from nltk.sem.chat80 import continent
+
 from google.genai import types
 
 from ...agents.active_streaming_tool import ActiveStreamingTool
 from ...agents.invocation_context import InvocationContext
+from ...agents.run_config import StreamingMode
 from ...auth.auth_tool import AuthToolArguments
 from ...events.event import Event
 from ...events.event_actions import EventActions
@@ -184,70 +190,85 @@ async def handle_function_calls_async(
     tools_dict: dict[str, BaseTool],
     filters: Optional[set[str]] = None,
     tool_confirmation_dict: Optional[dict[str, ToolConfirmation]] = None,
+
 ) -> Optional[Event]:
   """Calls the functions and returns the function response event."""
-  function_calls = function_call_event.get_function_calls()
-  return await handle_function_call_list_async(
-      invocation_context,
-      function_calls,
-      tools_dict,
-      filters,
-      tool_confirmation_dict,
-  )
+  async with Aclosing(
+      handle_function_calls_async_gen(
+          invocation_context,
+          function_call_event,
+          tools_dict,
+          filters,
+          tool_confirmation_dict,
+      )
+  ) as agen:
+    last_event = None
+    async for event in agen:
+      last_event = event
+  return last_event
 
 
-async def handle_function_call_list_async(
+async def handle_function_calls_async_gen(
     invocation_context: InvocationContext,
-    function_calls: list[types.FunctionCall],
+    function_call_event: Event,
     tools_dict: dict[str, BaseTool],
     filters: Optional[set[str]] = None,
     tool_confirmation_dict: Optional[dict[str, ToolConfirmation]] = None,
-) -> Optional[Event]:
-  """Calls the functions and returns the function response event."""
+) -> AsyncGenerator[Optional[Event]]:
+  """Calls the functions and returns the function response event as generator."""
   from ...agents.llm_agent import LlmAgent
 
   agent = invocation_context.agent
   if not isinstance(agent, LlmAgent):
-    return None
+    yield None
+    return
+
+  function_calls = function_call_event.get_function_calls()
 
   # Filter function calls
   filtered_calls = [
       fc for fc in function_calls if not filters or fc.id in filters
   ]
 
   if not filtered_calls:
-    return None
+    yield None
+    return
 
-  # Create tasks for parallel execution
-  tasks = [
-      asyncio.create_task(
-          _execute_single_function_call_async(
-              invocation_context,
-              function_call,
-              tools_dict,
-              agent,
-              tool_confirmation_dict[function_call.id]
-              if tool_confirmation_dict
-              else None,
-          )
+  function_call_async_gens = [
+      _execute_single_function_call_async_gen(
+          invocation_context,
+          function_call,
+          tools_dict,
+          agent,
+          tool_confirmation_dict[function_call.id]
+          if tool_confirmation_dict
+          else None,
       )
       for function_call in filtered_calls
   ]
 
-  # Wait for all tasks to complete
-  function_response_events = await asyncio.gather(*tasks)
-
-  # Filter out None results
-  function_response_events = [
-      event for event in function_response_events if event is not None
-  ]
+  merged_event = None
+  result_events: List[Optional[Event]] = [None] * len(function_call_async_gens)
+  function_response_events = []
+  async for idx, event in _concat_function_call_generators(
+      function_call_async_gens
+  ):
+    result_events[idx] = event
+    function_response_events = [
+        event for event in result_events if event is not None
+    ]
+    if function_response_events:
+      merged_event = merge_parallel_function_response_events(
+          function_response_events
+      )
+      if invocation_context.run_config.streaming_mode == StreamingMode.SSE:
+        yield merged_event
+  if invocation_context.run_config.streaming_mode != StreamingMode.SSE:
+    yield merged_event
 
   if not function_response_events:
-    return None
-
-  merged_event = merge_parallel_function_response_events(
-      function_response_events
-  )
+    yield None
+    return
 
   if len(function_response_events) > 1:
     # this is needed for debug traces of parallel calls
@@ -258,16 +279,61 @@ async def handle_function_call_list_async(
           response_event_id=merged_event.id,
           function_response_event=merged_event,
       )
-  return merged_event
 
 
-async def _execute_single_function_call_async(
+async def _concat_function_call_generators(
+    gens: List[AsyncGenerator[Any]],
+) -> AsyncIterator[tuple[int, Any]]:
+  _SENTINEL = object()
+  q: asyncio.Queue[tuple[str, int, Any]] = asyncio.Queue()
+  gens = list(gens)
+  n = len(gens)
+
+  async def __pump(idx: int, agen_: AsyncIterator[Any]):
+    try:
+      async for x in agen_:
+        await q.put(('ITEM', idx, x))
+    except Exception as e:
+      await q.put(('EXC', idx, e))
+    finally:
+      aclose = getattr(agen_, 'aclose', None)
+      if callable(aclose):
+        try:
+          await aclose()
+        except Exception:  # noqa: ignore exception when task canceled.
+          pass
+
+      await q.put(('END', idx, _SENTINEL))
+
+  tasks = [asyncio.create_task(__pump(i, agen)) for i, agen in enumerate(gens)]
+  finished = 0
+  try:
+    while finished < n:
+      kind, i, payload = await q.get()
+      if kind == 'ITEM':
+        yield i, payload
+
+      elif kind == 'EXC':
+        for t in tasks:
+          t.cancel()
+        await asyncio.gather(*tasks, return_exceptions=True)
+        raise payload
+
+      elif kind == 'END':
+        finished += 1
+  finally:
+    for t in tasks:
+      t.cancel()
+    await asyncio.gather(*tasks, return_exceptions=True)
+
+
+async def _execute_single_function_call_async_gen(
     invocation_context: InvocationContext,
     function_call: types.FunctionCall,
     tools_dict: dict[str, BaseTool],
     agent: LlmAgent,
     tool_confirmation: Optional[ToolConfirmation] = None,
-) -> Optional[Event]:
+) -> AsyncGenerator[Optional[Event]]:
   """Execute a single function call with thread safety for state modifications."""
   tool, tool_context = _get_tool_and_context(
       invocation_context,
@@ -310,6 +376,37 @@ async def _execute_single_function_call_async(
         function_response = await __call_tool_async(
             tool, args=function_args, tool_context=tool_context
         )
+        if inspect.isasyncgen(function_response) or isinstance(
+            function_response, AsyncIterator
+        ):
+          res = None
+          async for res in function_response:
+            if inspect.isawaitable(res):
+              res = await res
+            if (
+                invocation_context.run_config.streaming_mode
+                == StreamingMode.SSE
+            ):
+              yield __build_response_event(
+                  tool, res, tool_context, invocation_context
+              )
+          function_response = res
+        elif inspect.isgenerator(function_response) or isinstance(
+            function_response, Iterator
+        ):
+          res = None
+          for res in function_response:
+            if inspect.isawaitable(res):
+              res = await res
+            if (
+                invocation_context.run_config.streaming_mode
+                == StreamingMode.SSE
+            ):
+              yield __build_response_event(
+                  tool, res, tool_context, invocation_context
+              )
+          function_response = res
+
       except Exception as tool_error:
         error_response = (
             await invocation_context.plugin_manager.run_on_tool_error_callback(
@@ -359,7 +456,8 @@ async def _execute_single_function_call_async(
       # Allow long running function to return None to not provide function
       # response.
       if not function_response:
-        return None
+        yield None
+        return
 
     # Note: State deltas are not applied here - they are collected in
     # tool_context.actions.state_delta and applied later when the session
@@ -374,7 +472,7 @@ async def _execute_single_function_call_async(
         args=function_args,
         function_response_event=function_response_event,
     )
-    return function_response_event
+    yield function_response_event
 
 
 async def handle_function_calls_live(