feat: optional concurrent async automatic function calling

google/genai/_extra_utils.py

@@ -360,55 +360,101 @@ def get_function_response_parts(
   return func_response_parts
 
 
+def should_run_afc_concurrently(
+    config: Optional[types.GenerateContentConfigOrDict] = None,
+) -> bool:
+  """Returns whether async AFC should execute function calls concurrently."""
+  if not config:
+    return False
+  config_model = _create_generate_content_config_model(config)
+  if not config_model.automatic_function_calling:
+    return False
+  return bool(config_model.automatic_function_calling.run_concurrently)
+
+
+async def _execute_function_call_async(
+    part: types.Part,
+    function_map: dict[str, Union[Callable[..., Any], McpToGenAiToolAdapter]],
+) -> Optional[types.Part]:
+  """Executes a single function_call part and returns its response part."""
+  if not part.function_call:
+    return None
+  func_name = part.function_call.name
+  if func_name is None:
+    return None
+  func = function_map[func_name]
+  # Treat None as an empty dictionary for execution
+  raw_args = (
+      part.function_call.args if part.function_call.args is not None else {}
+  )
+  args = convert_number_values_for_dict_function_call_args(raw_args)
+  func_response: _common.StringDict
+  try:
+    if isinstance(func, McpToGenAiToolAdapter):
+      mcp_tool_response = await func.call_tool(
+          types.FunctionCall(name=func_name, args=args)
+      )
+      if mcp_tool_response.isError:
+        func_response = {'error': mcp_tool_response}
+      else:
+        func_response = {'result': mcp_tool_response}
+    elif inspect.iscoroutinefunction(func):
+      func_response = {
+          'result': await invoke_function_from_dict_args_async(args, func)
+      }
+    else:
+      func_response = {
+          'result': await asyncio.to_thread(
+              invoke_function_from_dict_args, args, func
+          )
+      }
+  except Exception as e:  # pylint: disable=broad-except
+    func_response = {'error': str(e)}
+  return types.Part.from_function_response(
+      name=func_name, response=func_response
+  )
+
+
 async def get_function_response_parts_async(
     response: types.GenerateContentResponse,
     function_map: dict[str, Union[Callable[..., Any], McpToGenAiToolAdapter]],
+    config: Optional[types.GenerateContentConfigOrDict] = None,
 ) -> list[types.Part]:
-  """Returns the function response parts from the response."""
-  func_response_parts = []
+  """Returns the function response parts from the response.
+
+  When ``AutomaticFunctionCallingConfig.run_concurrently`` is True, multiple
+  function calls from the same model response are executed concurrently via
+  ``asyncio.gather`` (order of response parts matches the call order).
+  """
   if (
-      response.candidates is not None
-      and isinstance(response.candidates[0].content, types.Content)
-      and response.candidates[0].content.parts is not None
+      response.candidates is None
+      or not isinstance(response.candidates[0].content, types.Content)
+      or response.candidates[0].content.parts is None
   ):
-    for part in response.candidates[0].content.parts:
-      if not part.function_call:
-        continue
-      func_name = part.function_call.name
-      if func_name is not None:
-        func = function_map[func_name]
-        # Treat None as an empty dictionary for execution
-        raw_args = (
-            part.function_call.args
-            if part.function_call.args is not None
-            else {}
-        )
-        args = convert_number_values_for_dict_function_call_args(raw_args)
-        try:
-          if isinstance(func, McpToGenAiToolAdapter):
-            mcp_tool_response = await func.call_tool(
-                types.FunctionCall(name=func_name, args=args)
-            )
-            if mcp_tool_response.isError:
-              func_response = {'error': mcp_tool_response}
-            else:
-              func_response = {'result': mcp_tool_response}
-          elif inspect.iscoroutinefunction(func):
-            func_response = {
-                'result': await invoke_function_from_dict_args_async(args, func)
-            }
-          else:
-            func_response = {
-                'result': await asyncio.to_thread(
-                    invoke_function_from_dict_args, args, func
-                )
-            }
-        except Exception as e:  # pylint: disable=broad-except
-          func_response = {'error': str(e)}  # type: ignore[dict-item]
-        func_response_part = types.Part.from_function_response(
-            name=func_name, response=func_response
-        )
-        func_response_parts.append(func_response_part)
+    return []
+
+  function_call_parts = [
+      part
+      for part in response.candidates[0].content.parts
+      if part.function_call and part.function_call.name is not None
+  ]
+  if not function_call_parts:
+    return []
+
+  if should_run_afc_concurrently(config):
+    results = await asyncio.gather(
+        *[
+            _execute_function_call_async(part, function_map)
+            for part in function_call_parts
+        ]
+    )
+    return [part for part in results if part is not None]
+
+  func_response_parts: list[types.Part] = []
+  for part in function_call_parts:
+    func_response_part = await _execute_function_call_async(part, function_map)
+    if func_response_part is not None:
+      func_response_parts.append(func_response_part)
   return func_response_parts
 
 

google/genai/models.py

@@ -8836,7 +8836,7 @@ async def generate_content(
           break
         func_response_parts = (
             await _extra_utils.get_function_response_parts_async(
-                response, function_map
+                response, function_map, final_parsed_config
             )
         )
         if not func_response_parts:
@@ -9104,7 +9104,7 @@ async def stream_generator():  # type: ignore[no-untyped-def]
                   break
                 func_response_parts = (
                     await _extra_utils.get_function_response_parts_async(
-                        chunk, function_map
+                        chunk, function_map, final_parsed_config
                     )
                 )
                 if not func_response_parts:
@@ -9131,7 +9131,7 @@ async def stream_generator():  # type: ignore[no-untyped-def]
               break
             func_response_parts = (
                 await _extra_utils.get_function_response_parts_async(
-                    chunk, function_map
+                    chunk, function_map, final_parsed_config
                 )
             )
 

google/genai/tests/afc/test_get_function_response_parts.py

@@ -16,11 +16,18 @@
 
 """Tests for get_function_response_parts."""
 
+import asyncio
+import time
 import typing
 from typing import Any
 import pytest
-from ..._extra_utils import get_function_response_parts, get_function_response_parts_async
+from ..._extra_utils import (
+    get_function_response_parts,
+    get_function_response_parts_async,
+    should_run_afc_concurrently,
+)
 from ...errors import UnsupportedFunctionError
+from ... import types
 from ...types import Candidate
 from ...types import Content
 from ...types import FunctionCall
@@ -275,3 +282,151 @@ async def call_tool(self, name: str, arguments: dict[str, Any]) -> Any:
     assert actual_part.model_dump_json(
         exclude_none=True
     ) == expected_part.model_dump_json(exclude_none=True)
+
+
+def test_should_run_afc_concurrently():
+  assert should_run_afc_concurrently(None) is False
+  assert should_run_afc_concurrently(types.GenerateContentConfig()) is False
+  assert (
+      should_run_afc_concurrently(
+          types.GenerateContentConfig(
+              automatic_function_calling=types.AutomaticFunctionCallingConfig()
+          )
+      )
+      is False
+  )
+  assert (
+      should_run_afc_concurrently(
+          types.GenerateContentConfig(
+              automatic_function_calling=types.AutomaticFunctionCallingConfig(
+                  run_concurrently=True
+              )
+          )
+      )
+      is True
+  )
+
+
+def _multi_function_response() -> GenerateContentResponse:
+  return GenerateContentResponse(
+      candidates=[
+          Candidate(
+              content=Content(
+                  parts=[
+                      Part(
+                          function_call=FunctionCall(
+                              name='slow_a',
+                              args={},
+                          )
+                      ),
+                      Part(
+                          function_call=FunctionCall(
+                              name='slow_b',
+                              args={},
+                          )
+                      ),
+                  ]
+              )
+          )
+      ]
+  )
+
+
+@pytest.mark.asyncio
+async def test_async_run_concurrently_preserves_order_and_results():
+  async def slow_a() -> str:
+    await asyncio.sleep(0.05)
+    return 'a'
+
+  async def slow_b() -> str:
+    await asyncio.sleep(0.01)
+    return 'b'
+
+  response = _multi_function_response()
+  function_map = {'slow_a': slow_a, 'slow_b': slow_b}
+  config = types.GenerateContentConfig(
+      automatic_function_calling=types.AutomaticFunctionCallingConfig(
+          run_concurrently=True
+      )
+  )
+  actual_parts = await get_function_response_parts_async(
+      response, function_map, config
+  )
+  assert [p.function_response.name for p in actual_parts] == [
+      'slow_a',
+      'slow_b',
+  ]
+  assert actual_parts[0].function_response is not None
+  assert actual_parts[1].function_response is not None
+  assert actual_parts[0].function_response.response == {'result': 'a'}
+  assert actual_parts[1].function_response.response == {'result': 'b'}
+
+
+@pytest.mark.asyncio
+async def test_async_run_concurrently_is_faster_than_sequential():
+  async def slow_a() -> str:
+    await asyncio.sleep(0.1)
+    return 'a'
+
+  async def slow_b() -> str:
+    await asyncio.sleep(0.1)
+    return 'b'
+
+  response = _multi_function_response()
+  function_map = {'slow_a': slow_a, 'slow_b': slow_b}
+  concurrent_config = types.GenerateContentConfig(
+      automatic_function_calling=types.AutomaticFunctionCallingConfig(
+          run_concurrently=True
+      )
+  )
+  sequential_config = types.GenerateContentConfig(
+      automatic_function_calling=types.AutomaticFunctionCallingConfig(
+          run_concurrently=False
+      )
+  )
+
+  start = time.perf_counter()
+  await get_function_response_parts_async(
+      response, function_map, sequential_config
+  )
+  sequential_elapsed = time.perf_counter() - start
+
+  start = time.perf_counter()
+  await get_function_response_parts_async(
+      response, function_map, concurrent_config
+  )
+  concurrent_elapsed = time.perf_counter() - start
+
+  # Concurrent should finish near one sleep; sequential near the sum.
+  assert concurrent_elapsed < sequential_elapsed
+  assert concurrent_elapsed < 0.18
+  assert sequential_elapsed >= 0.18
+
+
+@pytest.mark.asyncio
+async def test_async_run_concurrently_default_remains_sequential():
+  in_flight = 0
+  max_in_flight = 0
+  lock = asyncio.Lock()
+
+  async def tracked(name: str) -> str:
+    nonlocal in_flight, max_in_flight
+    async with lock:
+      in_flight += 1
+      max_in_flight = max(max_in_flight, in_flight)
+    await asyncio.sleep(0.05)
+    async with lock:
+      in_flight -= 1
+    return name
+
+  async def slow_a() -> str:
+    return await tracked('a')
+
+  async def slow_b() -> str:
+    return await tracked('b')
+
+  response = _multi_function_response()
+  await get_function_response_parts_async(
+      response, {'slow_a': slow_a, 'slow_b': slow_b}
+  )
+  assert max_in_flight == 1

google/genai/types.py

@@ -5592,6 +5592,15 @@ class AutomaticFunctionCallingConfig(_common.BaseModel):
       GenerateContentResponse.automatic_function_calling_history.
       """,
   )
+  run_concurrently: Optional[bool] = Field(
+      default=None,
+      description="""If automatic function calling is enabled on the async
+      client, whether to execute multiple function calls from a single model
+      response concurrently (via ``asyncio.gather``).
+      If not set or set to False, function calls are executed sequentially.
+      This field has no effect on the sync client.
+      """,
+  )
 
 
 class AutomaticFunctionCallingConfigDict(TypedDict, total=False):
@@ -5618,6 +5627,13 @@ class AutomaticFunctionCallingConfigDict(TypedDict, total=False):
       GenerateContentResponse.automatic_function_calling_history.
       """
 
+  run_concurrently: Optional[bool]
+  """If automatic function calling is enabled on the async client, whether to
+      execute multiple function calls from a single model response concurrently
+      (via ``asyncio.gather``). If not set or set to False, function calls are
+      executed sequentially. This field has no effect on the sync client.
+      """
+
 
 AutomaticFunctionCallingConfigOrDict = Union[
     AutomaticFunctionCallingConfig, AutomaticFunctionCallingConfigDict