From 21eb146aa94b183cb310c6a865d47907b36ac77e Mon Sep 17 00:00:00 2001 From: Jaycee Li Date: Thu, 27 Aug 2026 16:23:35 -0700 Subject: [PATCH] fix: do not run functions once the automatic function calling budget is spent PiperOrigin-RevId: 972215316 --- google/genai/chats.py | 71 ++++-- google/genai/models.py | 50 ++-- .../afc/test_get_max_remote_calls_for_afc.py | 230 +++++++++++++++++- google/genai/tests/chats/test_send_message.py | 24 +- 4 files changed, 316 insertions(+), 59 deletions(-) diff --git a/google/genai/chats.py b/google/genai/chats.py index be10298bc..e638ab0b8 100644 --- a/google/genai/chats.py +++ b/google/genai/chats.py @@ -342,15 +342,20 @@ def send_message( ): break + logger.info(f"AFC remote call {i} is done.") + remaining_remote_calls_afc -= 1 + if remaining_remote_calls_afc == 0: + # No request is left to send a result with, so the functions are not + # called at all. Breaking here leaves the turn to be recorded once, + # below, as the user's message followed by the model's function call. + logger.info("Reached max remote calls for automatic function calling.") + break + func_response_parts = _extra_utils.get_function_response_parts( response, function_map ) if not func_response_parts: break - logger.info(f"AFC remote call {i} is done.") - remaining_remote_calls_afc -= 1 - if remaining_remote_calls_afc == 0: - logger.info("Reached max remote calls for automatic function calling.") func_call_content = response.candidates[0].content func_response_content = types.Content( role="user", parts=func_response_parts @@ -377,7 +382,6 @@ def send_message( ) return response - def send_message_stream( self, message: Union[list[PartUnionDict], PartUnionDict], @@ -489,6 +493,15 @@ def send_message_stream( contents=contents_to_model, # type: ignore[arg-type] config=parsed_config, ) + remaining_remote_calls_afc -= 1 + # No request is left to send a result with, so the functions are not + # called at all. The chunks are still yielded, and the turn is recorded + # once below, ending on the model's unanswered function call. + is_last_remote_call_afc = remaining_remote_calls_afc == 0 + if is_last_remote_call_afc: + logger.info( + "Reached max remote calls for automatic function calling." + ) model_output = [] finish_reason = None @@ -501,7 +514,8 @@ def send_message_stream( is_valid = False if ( - function_map + not is_last_remote_call_afc + and function_map and chunk.candidates and chunk.candidates[0].content and chunk.candidates[0].content.parts @@ -518,15 +532,12 @@ def send_message_stream( finish_reason = chunk.candidates[0].finish_reason yield chunk + if is_last_remote_call_afc: + break if not function_map or not func_response_parts: break logger.info(f"AFC remote call {i} is done.") - remaining_remote_calls_afc -= 1 - if remaining_remote_calls_afc == 0: - logger.info( - "Reached max remote calls for automatic function calling." - ) if chunk and chunk.candidates and chunk.candidates[0].content: func_response_content = types.Content( @@ -802,6 +813,17 @@ async def send_message( ): break + logger.info(f"AFC remote call {i} is done.") + remaining_remote_calls_afc -= 1 + if remaining_remote_calls_afc == 0: + # No request is left to send a result with, so the functions are not + # called at all. Breaking here leaves the turn to be recorded once, + # below, as the user's message followed by the model's function call. + logger.info( + "Reached max remote calls for automatic function calling." + ) + break + func_response_parts = ( await _extra_utils.get_function_response_parts_async( response, function_map @@ -810,13 +832,6 @@ async def send_message( if not func_response_parts: break - logger.info(f"AFC remote call {i} is done.") - remaining_remote_calls_afc -= 1 - if remaining_remote_calls_afc == 0: - logger.info( - "Reached max remote calls for automatic function calling." - ) - func_call_content = response.candidates[0].content func_response_content = types.Content( role="user", parts=func_response_parts @@ -845,7 +860,6 @@ async def send_message( ) return response - async def send_message_stream( self, message: Union[list[PartUnionDict], PartUnionDict], @@ -1034,6 +1048,15 @@ async def async_generator(): # type: ignore[no-untyped-def] contents=contents_to_model, # type: ignore[arg-type] config=final_parsed_config, ) + remaining_remote_calls_afc -= 1 + # No request is left to send a result with, so the functions are not + # called at all. The chunks are still yielded, and the turn is + # recorded once below, ending on the model's unanswered function call. + is_last_remote_call_afc = remaining_remote_calls_afc == 0 + if is_last_remote_call_afc: + logger.info( + "Reached max remote calls for automatic function calling." + ) model_output = [] finish_reason = None @@ -1046,7 +1069,8 @@ async def async_generator(): # type: ignore[no-untyped-def] is_valid = False if ( - function_map + not is_last_remote_call_afc + and function_map and chunk.candidates and chunk.candidates[0].content and chunk.candidates[0].content.parts @@ -1065,15 +1089,12 @@ async def async_generator(): # type: ignore[no-untyped-def] finish_reason = chunk.candidates[0].finish_reason yield chunk + if is_last_remote_call_afc: + break if not function_map or not func_response_parts: break logger.info(f"AFC remote call {i} is done.") - remaining_remote_calls_afc -= 1 - if remaining_remote_calls_afc == 0: - logger.info( - "Reached max remote calls for automatic function calling." - ) func_response_content = types.Content( role="user", parts=func_response_parts diff --git a/google/genai/models.py b/google/genai/models.py index 2577831e6..f9cca3614 100644 --- a/google/genai/models.py +++ b/google/genai/models.py @@ -6281,15 +6281,19 @@ def generate_content( or not response.candidates[0].content.parts ): break + logger.info(f'AFC remote call {i} is done.') + remaining_remote_calls_afc -= 1 + if remaining_remote_calls_afc == 0: + # No request is left to send a result with, so the functions are not + # called at all. The model's function call is returned to the caller to + # run and answer themselves. + logger.info('Reached max remote calls for automatic function calling.') + break func_response_parts = _extra_utils.get_function_response_parts( response, function_map ) if not func_response_parts: break - logger.info(f'AFC remote call {i} is done.') - remaining_remote_calls_afc -= 1 - if remaining_remote_calls_afc == 0: - logger.info('Reached max remote calls for automatic function calling.') func_call_content = response.candidates[0].content func_response_content = types.Content( @@ -6440,6 +6444,13 @@ def generate_content_stream( response = self._generate_content_stream( model=model, contents=contents, config=parsed_config_to_call ) + remaining_remote_calls_afc -= 1 + # No request is left to send a result with, so the functions are not + # called at all. The chunks are still yielded, and the model's function + # call is left for the caller to run and answer themselves. + is_last_remote_call_afc = remaining_remote_calls_afc == 0 + if is_last_remote_call_afc: + logger.info('Reached max remote calls for automatic function calling.') model_output = [] func_response_parts = [] @@ -6455,7 +6466,8 @@ def generate_content_stream( ) if ( - function_map + not is_last_remote_call_afc + and function_map and chunk.candidates and chunk.candidates[0].content and chunk.candidates[0].content.parts @@ -6471,13 +6483,12 @@ def generate_content_stream( yield chunk + if is_last_remote_call_afc: + break if not function_map or not func_response_parts: break logger.info(f'AFC remote call {i} is done.') - remaining_remote_calls_afc -= 1 - if remaining_remote_calls_afc == 0: - logger.info('Reached max remote calls for automatic function calling.') # Append function call and function response parts to contents for the next request. func_response_content = types.Content( @@ -8453,9 +8464,13 @@ async def generate_content( ) remaining_remote_calls_afc -= 1 if remaining_remote_calls_afc == 0: + # No request is left to send a result with, so the functions are not + # called at all. The model's function call is returned to the caller + # to run and answer themselves. logger.info( 'Reached max remote calls for automatic function calling.' ) + break if not function_map: break @@ -8705,6 +8720,15 @@ async def stream_generator(): # type: ignore[no-untyped-def] contents=loop_contents, config=final_parsed_config_to_call, ) + remaining_remote_calls_afc -= 1 + # No request is left to send a result with, so the functions are not + # called at all. The chunks are still yielded, and the model's + # function call is left for the caller to run and answer themselves. + is_last_remote_call_afc = remaining_remote_calls_afc == 0 + if is_last_remote_call_afc: + logger.info( + 'Reached max remote calls for automatic function calling.' + ) model_output = [] func_response_parts = [] @@ -8720,7 +8744,8 @@ async def stream_generator(): # type: ignore[no-untyped-def] ) if ( - function_map + not is_last_remote_call_afc + and function_map and chunk.candidates and chunk.candidates[0].content and chunk.candidates[0].content.parts @@ -8738,15 +8763,12 @@ async def stream_generator(): # type: ignore[no-untyped-def] yield chunk + if is_last_remote_call_afc: + break if not function_map or not func_response_parts: break logger.info(f'AFC remote call {i} is done.') - remaining_remote_calls_afc -= 1 - if remaining_remote_calls_afc == 0: - logger.info( - 'Reached max remote calls for automatic function calling.' - ) # Append function response parts to contents for the next request. func_response_content = types.Content( diff --git a/google/genai/tests/afc/test_get_max_remote_calls_for_afc.py b/google/genai/tests/afc/test_get_max_remote_calls_for_afc.py index 5a72020ed..d9a908575 100644 --- a/google/genai/tests/afc/test_get_max_remote_calls_for_afc.py +++ b/google/genai/tests/afc/test_get_max_remote_calls_for_afc.py @@ -14,11 +14,20 @@ # -"""Tests for get_max_remote_calls_for_afc.""" +"""Tests for the max remote calls budget for AFC. +Covers both reading it off the config, and what the AFC loop does once it is +spent. +""" + +from unittest import mock +import pytest +from ... import _api_client +from ... import _extra_utils +from ... import chats +from ... import models from ... import types from ..._extra_utils import get_max_remote_calls_afc -import pytest def test_config_is_none(): @@ -128,3 +137,220 @@ def test_afc_enabled_max_set_to_float(): ) == 5 ) + + +TEST_FUNCTION_CALL_CONTENT = types.Content( + parts=[ + types.Part( + function_call=types.FunctionCall( + name='get_current_weather', + args={'location': 'San Francisco'}, + ) + ) + ], + role='model', +) + + +TEST_FUNCTION_RESPONSE_PART = types.Part( + function_response=types.FunctionResponse( + name='get_current_weather', + response={'result': 'sunny'}, + ) +) + + +def get_current_weather(location: str) -> str: + """Returns the current weather. + + Args: + location: The city and State, e.g. San Francisco, CA. + """ + return 'sunny' + + +@pytest.fixture +def mock_api_client(): + api_client = mock.MagicMock(spec=_api_client.BaseApiClient) + api_client.api_key = 'TEST_API_KEY' + api_client._host = lambda: 'test_host' + api_client._http_options = {'headers': {}} + api_client.vertexai = False + return api_client + + +def _afc_config(maximum_remote_calls: int) -> types.GenerateContentConfig: + return types.GenerateContentConfig( + tools=[get_current_weather], + automatic_function_calling=types.AutomaticFunctionCallingConfig( + maximum_remote_calls=maximum_remote_calls + ), + ) + + +def test_generate_content_spent_budget_does_not_run_functions(mock_api_client): + """The one allowed request is spent asking, so nothing is run.""" + with mock.patch.object( + models.Models, '_generate_content' + ) as mock_generate_content, mock.patch.object( + _extra_utils, 'get_function_response_parts' + ) as mock_get_function_response_parts: + mock_generate_content.return_value = types.GenerateContentResponse( + candidates=[types.Candidate(content=TEST_FUNCTION_CALL_CONTENT)] + ) + + response = models.Models(api_client_=mock_api_client).generate_content( + model='test_model', + contents='what is the weather in San Francisco?', + config=_afc_config(1), + ) + + assert mock_generate_content.call_count == 1 + # The result could not have been delivered, so the function is never called. + mock_get_function_response_parts.assert_not_called() + assert response.candidates[0].content.parts[0].function_call + + +def test_generate_content_stream_spent_budget_does_not_run_functions( + mock_api_client, +): + """The same over a stream: chunks are yielded, nothing is run.""" + with mock.patch.object( + models.Models, '_generate_content_stream' + ) as mock_generate_content_stream, mock.patch.object( + _extra_utils, 'get_function_response_parts' + ) as mock_get_function_response_parts: + mock_generate_content_stream.return_value = [ + types.GenerateContentResponse( + candidates=[types.Candidate(content=TEST_FUNCTION_CALL_CONTENT)] + ) + ] + + chunks = list( + models.Models(api_client_=mock_api_client).generate_content_stream( + model='test_model', + contents='what is the weather in San Francisco?', + config=_afc_config(1), + ) + ) + + assert mock_generate_content_stream.call_count == 1 + mock_get_function_response_parts.assert_not_called() + assert chunks[0].candidates[0].content.parts[0].function_call + + +def test_send_message_spent_budget_records_the_turn_once(mock_api_client): + """The turn is recorded once, as the message and the unanswered call.""" + with mock.patch.object( + models.Models, '_generate_content' + ) as mock_generate_content, mock.patch.object( + _extra_utils, 'get_function_response_parts' + ) as mock_get_function_response_parts: + mock_generate_content.return_value = types.GenerateContentResponse( + candidates=[types.Candidate(content=TEST_FUNCTION_CALL_CONTENT)] + ) + + chat = chats.Chats(models.Models(api_client_=mock_api_client)).create( + model='test_model', config=_afc_config(1) + ) + chat.send_message('what is the weather in San Francisco?') + history = chat.get_history() + + mock_get_function_response_parts.assert_not_called() + # The model's function call is recorded once, not twice, and the history is + # left ready for the caller to answer that call themselves. + assert [content.role for content in history] == ['user', 'model'] + assert history[1].parts[0].function_call + + +def test_send_message_budget_of_two_answers_the_call(mock_api_client): + """A budget of two is the smallest that lets the model answer.""" + with mock.patch.object( + models.Models, '_generate_content' + ) as mock_generate_content: + mock_generate_content.side_effect = [ + types.GenerateContentResponse( + candidates=[types.Candidate(content=TEST_FUNCTION_CALL_CONTENT)] + ), + types.GenerateContentResponse( + candidates=[ + types.Candidate( + content=types.Content( + parts=[types.Part(text='It is sunny.')], role='model' + ) + ) + ] + ), + ] + + chat = chats.Chats(models.Models(api_client_=mock_api_client)).create( + model='test_model', config=_afc_config(2) + ) + response = chat.send_message('what is the weather in San Francisco?') + history = chat.get_history() + + assert response.text == 'It is sunny.' + assert [content.role for content in history] == [ + 'user', + 'model', + 'user', + 'model', + ] + assert history[2].parts[0].function_response + + +def test_spent_budget_leaves_the_afc_history_empty(mock_api_client): + """Nothing ran, so automatic function calling added no turns to report.""" + with mock.patch.object( + models.Models, '_generate_content' + ) as mock_generate_content, mock.patch.object( + _extra_utils, 'get_function_response_parts' + ): + mock_generate_content.return_value = types.GenerateContentResponse( + candidates=[types.Candidate(content=TEST_FUNCTION_CALL_CONTENT)] + ) + + response = models.Models(api_client_=mock_api_client).generate_content( + model='test_model', + contents='what is the weather in San Francisco?', + config=_afc_config(1), + ) + + assert not response.automatic_function_calling_history + + +def test_afc_history_holds_the_rounds_that_completed(mock_api_client): + """A round that was delivered is reported; the budget stops after it.""" + with mock.patch.object( + models.Models, '_generate_content' + ) as mock_generate_content, mock.patch.object( + _extra_utils, 'get_function_response_parts' + ) as mock_get_function_response_parts: + mock_generate_content.side_effect = [ + types.GenerateContentResponse( + candidates=[types.Candidate(content=TEST_FUNCTION_CALL_CONTENT)] + ), + types.GenerateContentResponse( + candidates=[ + types.Candidate( + content=types.Content( + parts=[types.Part(text='It is sunny.')], role='model' + ) + ) + ] + ), + ] + mock_get_function_response_parts.return_value = [ + TEST_FUNCTION_RESPONSE_PART + ] + + response = models.Models(api_client_=mock_api_client).generate_content( + model='test_model', + contents='what is the weather in San Francisco?', + config=_afc_config(2), + ) + + history = response.automatic_function_calling_history + assert [content.role for content in history] == ['user', 'model', 'user'] + assert history[1].parts[0].function_call + assert history[2].parts[0].function_response diff --git a/google/genai/tests/chats/test_send_message.py b/google/genai/tests/chats/test_send_message.py index 94d451582..15e54ac7c 100644 --- a/google/genai/tests/chats/test_send_message.py +++ b/google/genai/tests/chats/test_send_message.py @@ -335,7 +335,9 @@ def test_with_afc_multiple_remote_calls(client): chat.send_message('Turn this place into a party!') curated_history = chat.get_history() - assert len(curated_history) == 8 + # A budget of 3 buys 3 requests. The third is spent being asked for functions + # that no request is left to answer, so the turn ends on that call. + assert len(curated_history) == 6 assert curated_history[0].role == 'user' assert curated_history[0].parts[0].text == 'Turn this place into a party!' assert curated_history[1].role == 'model' @@ -358,14 +360,6 @@ def test_with_afc_multiple_remote_calls(client): assert len(curated_history[5].parts) == 3 for part in curated_history[5].parts: assert part.function_call - assert curated_history[6].role == 'user' - assert len(curated_history[6].parts) == 3 - for part in curated_history[6].parts: - assert part.function_response - assert curated_history[7].role == 'model' - assert len(curated_history[7].parts) == 3 - for part in curated_history[7].parts: - assert part.function_call @pytest.mark.skipif( @@ -394,7 +388,9 @@ def test_with_afc_multiple_remote_calls_async(client): chat.send_message('Turn this place into a party!') curated_history = chat.get_history() - assert len(curated_history) == 8 + # A budget of 3 buys 3 requests. The third is spent being asked for functions + # that no request is left to answer, so the turn ends on that call. + assert len(curated_history) == 6 assert curated_history[0].role == 'user' assert curated_history[0].parts[0].text == 'Turn this place into a party!' assert curated_history[1].role == 'model' @@ -417,14 +413,6 @@ def test_with_afc_multiple_remote_calls_async(client): assert len(curated_history[5].parts) == 3 for part in curated_history[5].parts: assert part.function_call - assert curated_history[6].role == 'user' - assert len(curated_history[6].parts) == 3 - for part in curated_history[6].parts: - assert part.function_response - assert curated_history[7].role == 'model' - assert len(curated_history[7].parts) == 3 - for part in curated_history[7].parts: - assert part.function_call def test_with_afc_disabled(client): chat = client.chats.create(