Skip to content

fix: do not run functions once the automatic function calling budget is spent #2921

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: do not run functions once the automatic function calling budget is spent #2921

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
71 changes: 46 additions & 25 deletions google/genai/chats.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -377,7 +382,6 @@ def send_message(
)
return response


def send_message_stream(
self,
message: Union[list[PartUnionDict], PartUnionDict],
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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(
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -845,7 +860,6 @@ async def send_message(
)
return response


async def send_message_stream(
self,
message: Union[list[PartUnionDict], PartUnionDict],
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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
Expand Down
50 changes: 36 additions & 14 deletions google/genai/models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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 = []
Expand All @@ -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
Expand All @@ -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(
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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 = []
Expand All @@ -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
Expand All @@ -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(
Expand Down
Loading
Loading