fix: Call all tools in parallel calls during partial authentication

Copybara import of the project:

--
ffd6184d7e402b0787b0fa37fc09cd519adcc7f3 by Calvin Giles <calvin.giles@trademe.co.nz>:

fix: Call all tools in parallel calls during partial authentication

--
c71782a582ba825dbe2246cdb5be3f273ca90dca by seanzhou1023 <seanzhou1023@gmail.com>:

Update auth_preprocessor.py
--
843af6b1bc0bc6291cb9cb23acf11840098ba6dd by seanzhou1023 <seanzhou1023@gmail.com>:

Update test_functions_request_euc.py
--
955e3fa852420ecbf196583caa3cf86b7b80ab56 by seanzhou1023 <seanzhou1023@gmail.com>:

Update test_functions_request_euc.py

COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/853 from calvingiles:fix-parallel-auth-tool-calls f44671e37b9fe44a25c9b1c0c25a26fc634b011c
PiperOrigin-RevId: 765639904
This commit is contained in:
Calvin Giles
2025-05-31 13:13:08 -07:00
committed by Copybara-Service
parent 036f954a2a
commit 0e72efb439
3 changed files with 234 additions and 21 deletions

View File

@@ -100,23 +100,24 @@ class _AuthLlmRequestProcessor(BaseLlmRequestProcessor):
function_calls = event.get_function_calls()
if not function_calls:
continue
for function_call in function_calls:
function_response_event = None
if function_call.id in tools_to_resume:
function_response_event = await functions.handle_function_calls_async(
invocation_context,
event,
{
tool.name: tool
for tool in await agent.canonical_tools(
ReadonlyContext(invocation_context)
)
},
# there could be parallel function calls that require auth
# auth response would be a dict keyed by function call id
tools_to_resume,
)
if function_response_event:
if any([
function_call.id in tools_to_resume
for function_call in function_calls
]):
if function_response_event := await functions.handle_function_calls_async(
invocation_context,
event,
{
tool.name: tool
for tool in await agent.canonical_tools(
ReadonlyContext(invocation_context)
)
},
# there could be parallel function calls that require auth
# auth response would be a dict keyed by function call id
tools_to_resume,
):
yield function_response_event
return
return

View File

@@ -170,10 +170,10 @@ def _rearrange_events_for_latest_function_response(
for idx in range(function_call_event_idx + 1, len(events) - 1):
event = events[idx]
function_responses = event.get_function_responses()
if (
function_responses
and function_responses[0].id in function_responses_ids
):
if function_responses and any([
function_response.id in function_responses_ids
for function_response in function_responses
]):
function_response_events.append(event)
function_response_events.append(events[-1])