diff --git a/google/genai/_extra_utils.py b/google/genai/_extra_utils.py index 8e1445fdd..6872f8e6a 100644 --- a/google/genai/_extra_utils.py +++ b/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 diff --git a/google/genai/models.py b/google/genai/models.py index d77346acf..e0c419ea9 100644 --- a/google/genai/models.py +++ b/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 ) ) diff --git a/google/genai/tests/afc/test_get_function_response_parts.py b/google/genai/tests/afc/test_get_function_response_parts.py index e80ec4493..08da262c3 100644 --- a/google/genai/tests/afc/test_get_function_response_parts.py +++ b/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 diff --git a/google/genai/types.py b/google/genai/types.py index 359c00d45..158dfef4d 100644 --- a/google/genai/types.py +++ b/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