encounter/adk-python
0
0
Fork 0
mirror of https://github.com/encounter/adk-python.git synced 2026-07-09 18:19:28 -07:00
Code Issues Packages Projects Releases Wiki Activity

fix: Yield the long running tool response before pausing execution

Browse Source 
PiperOrigin-RevId: 825056377
This commit is contained in:
Shangjie Chen
2025-10-28 08:41:47 -07:00
committed by Copybara-Service
parent 86f01550bd
commit 9ab17f2afd
5 changed files with 146 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/google/adk/agents/llm_agent.py
+10 -2
View File
@@ -430,16 +430,24 @@ class LlmAgent(BaseAgent):
yield self._create_agent_state_event(ctx)
return
should_pause = False
async with Aclosing(self._llm_flow.run_async(ctx)) as agen:
async for event in agen:
self.__maybe_save_output_to_state(event)
yield event
if ctx.should_pause_invocation(event):
return
# Do not pause immediately, wait until the long running tool call is
# executed.
should_pause = True
if should_pause:
return
if ctx.is_resumable:
events = ctx._get_events(current_invocation=True, current_branch=True)
if events and ctx.should_pause_invocation(events[-1]):
if events and (
ctx.should_pause_invocation(events[-1])
or ctx.should_pause_invocation(events[-2])
):
return
# Only yield an end state if the last event is no longer a long running
# tool call.
src/google/adk/flows/llm_flows/base_llm_flow.py
+19 -5
View File
@@ -383,16 +383,30 @@ class BaseLlmFlow(ABC):
events = invocation_context._get_events(
current_invocation=True, current_branch=True
)
# Long running tool calls should have been handled before this point.
# If there are still long running tool calls, it means the agent is paused
# before, and its branch hasn't been resumed yet.
if (
invocation_context.is_resumable
and events
and len(events) > 1
# TODO: here we are using the last 2 events to decide whether to pause
# the invocation. But this is just being optmisitic, we should find a
# way to pause when the long running tool call is followed by more than
# one text responses.
and (
invocation_context.should_pause_invocation(events[-1])
or invocation_context.should_pause_invocation(events[-2])
)
):
return
if (
invocation_context.is_resumable
and events
and events[-1].get_function_calls()
):
# Long running tool calls should have been handled before this point.
# If there are still long running tool calls, it means the agent is paused
# before, and its branch hasn't been resumed yet.
if invocation_context.should_pause_invocation(events[-1]):
return
model_response_event = events[-1]
async with Aclosing(
self._postprocess_handle_function_calls_async(
src/google/adk/tools/function_tool.py
+1
View File
@@ -205,6 +205,7 @@ You could retry calling this tool, but it is IMPORTANT for you to provide all th
' ToolConfirmation payload.'
),
)
tool_context.actions.skip_summarization = True
return {
'error': (
'This tool call requires confirmation, please approve or'
tests/unittests/runners/test_pause_invocation.py
+38 -2
View File
@@ -44,7 +44,7 @@ def _transfer_call_part(agent_name: str) -> Part:
def test_tool() -> str:
return ""
return "result"
class _TestingAgent(BaseAgent):
@@ -126,6 +126,12 @@ class TestPauseInvocationWithSingleLlmAgent(BasePauseInvocationTest):
"""Tests that a single LlmAgent pauses on long running function call."""
assert testing_utils.simplify_resumable_app_events(runner.run("test")) == [
("root_agent", Part.from_function_call(name="test_tool", args={})),
(
"root_agent",
Part.from_function_response(
name="test_tool", response={"result": "result"}
),
),
]
@@ -168,6 +174,12 @@ class TestPauseInvocationWithSequentialAgent(BasePauseInvocationTest):
),
),
("sub_agent_1", Part.from_function_call(name="test_tool", args={})),
(
"sub_agent_1",
Part.from_function_response(
name="test_tool", response={"result": "result"}
),
),
]
@pytest.mark.asyncio
@@ -195,7 +207,7 @@ class TestPauseInvocationWithSequentialAgent(BasePauseInvocationTest):
(
"sub_agent_1",
Part.from_function_response(
name="test_tool", response={"result": ""}
name="test_tool", response={"result": "result"}
),
),
("sub_agent_1", "model response after tool call"),
@@ -207,6 +219,12 @@ class TestPauseInvocationWithSequentialAgent(BasePauseInvocationTest):
),
),
("sub_agent_2", Part.from_function_call(name="test_tool", args={})),
(
"sub_agent_2",
Part.from_function_response(
name="test_tool", response={"result": "result"}
),
),
]
@@ -384,6 +402,12 @@ class TestPauseInvocationWithLoopAgent(BasePauseInvocationTest):
),
),
("sub_agent_2", Part.from_function_call(name="test_tool", args={})),
(
"sub_agent_2",
Part.from_function_response(
name="test_tool", response={"result": "result"}
),
),
]
@@ -435,6 +459,12 @@ class TestPauseInvocationWithLlmAgentTree(BasePauseInvocationTest):
("sub_llm_agent_1", _transfer_call_part("sub_llm_agent_2")),
("sub_llm_agent_1", _TRANSFER_RESPONSE_PART),
("sub_llm_agent_2", Part.from_function_call(name="test_tool", args={})),
(
"sub_llm_agent_2",
Part.from_function_response(
name="test_tool", response={"result": "result"}
),
),
]
@@ -489,4 +519,10 @@ class TestPauseInvocationWithWithTransferLoop(BasePauseInvocationTest):
("sub_llm_agent_2", _transfer_call_part("root_agent")),
("sub_llm_agent_2", _TRANSFER_RESPONSE_PART),
("root_agent", Part.from_function_call(name="test_tool", args={})),
(
"root_agent",
Part.from_function_response(
name="test_tool", response={"result": "result"}
),
),
]
tests/unittests/runners/test_run_tool_confirmation.py
+78 -41
View File
@@ -36,6 +36,16 @@ import pytest
from .. import testing_utils
HINT_TEXT = (
"Please approve or reject the tool call _test_function() by"
" responding with a FunctionResponse with an"
" expected ToolConfirmation payload."
)
TOOL_CALL_ERROR_RESPONSE = {
"error": "This tool call requires confirmation, please approve or reject."
}
def _create_llm_response_from_tools(
tools: list[FunctionTool],
@@ -57,13 +67,9 @@ def _create_llm_response_from_text(text: str) -> GenerateContentResponse:
)
def _test_request_confirmation_function(
def _test_function(
tool_context: ToolContext,
) -> dict[str, str]:
"""A test tool function that requests confirmation."""
if not tool_context.tool_confirmation:
tool_context.request_confirmation(hint="test hint for request_confirmation")
return {"error": "test error for request_confirmation"}
return {"result": f"confirmed={tool_context.tool_confirmation.confirmed}"}
@@ -82,7 +88,7 @@ def _test_request_confirmation_function_with_custom_schema(
}
},
)
return {"error": "test error for request_confirmation"}
return TOOL_CALL_ERROR_RESPONSE
return {
"result": f"confirmed={tool_context.tool_confirmation.confirmed}",
"custom_payload": tool_context.tool_confirmation.payload,
@@ -104,7 +110,7 @@ class TestHITLConfirmationFlowWithSingleAgent(BaseHITLTest):
@pytest.fixture
def tools(self) -> list[FunctionTool]:
"""Provides the tools for the agent."""
return [FunctionTool(func=_test_request_confirmation_function)]
return [FunctionTool(func=_test_function, require_confirmation=True)]
@pytest.fixture
def llm_responses(
@@ -114,9 +120,6 @@ class TestHITLConfirmationFlowWithSingleAgent(BaseHITLTest):
return [
_create_llm_response_from_tools(tools),
_create_llm_response_from_text("test llm response after tool call"),
_create_llm_response_from_text(
"test llm response after final tool call"
),
]
@pytest.fixture
@@ -163,7 +166,7 @@ class TestHITLConfirmationFlowWithSingleAgent(BaseHITLTest):
"args": {},
},
"toolConfirmation": {
"hint": "test hint for request_confirmation",
"hint": HINT_TEXT,
"confirmed": False,
},
},
@@ -174,12 +177,10 @@ class TestHITLConfirmationFlowWithSingleAgent(BaseHITLTest):
agent.name,
Part(
function_response=FunctionResponse(
name=tools[0].name,
response={"error": "test error for request_confirmation"},
name=tools[0].name, response=TOOL_CALL_ERROR_RESPONSE
)
),
),
(agent.name, "test llm response after tool call"),
]
simplified = testing_utils.simplify_events(copy.deepcopy(events))
@@ -208,11 +209,13 @@ class TestHITLConfirmationFlowWithSingleAgent(BaseHITLTest):
Part(
function_response=FunctionResponse(
name=tools[0].name,
response={"result": f"confirmed={tool_call_confirmed}"},
response={"result": f"confirmed={tool_call_confirmed}"}
if tool_call_confirmed
else {"error": "This tool call is rejected."},
)
),
),
(agent.name, "test llm response after final tool call"),
(agent.name, "test llm response after tool call"),
]
for event in events:
assert event.invocation_id != invocation_id
@@ -312,8 +315,7 @@ class TestHITLConfirmationFlowWithCustomPayloadSchema(BaseHITLTest):
agent.name,
Part(
function_response=FunctionResponse(
name=tools[0].name,
response={"error": "test error for request_confirmation"},
name=tools[0].name, response=TOOL_CALL_ERROR_RESPONSE
)
),
),
@@ -380,7 +382,7 @@ class TestHITLConfirmationFlowWithResumableApp:
@pytest.fixture
def tools(self) -> list[FunctionTool]:
"""Provides the tools for the agent."""
return [FunctionTool(func=_test_request_confirmation_function)]
return [FunctionTool(func=_test_function, require_confirmation=True)]
@pytest.fixture
def llm_responses(
@@ -409,8 +411,8 @@ class TestHITLConfirmationFlowWithResumableApp:
@pytest.fixture
def runner(self, agent: LlmAgent) -> testing_utils.InMemoryRunner:
"""Provides an in-memory runner for the agent."""
# Mark the app as resumable. So that the invocation will be paused after the
# long running tool call.
# Mark the app as resumable. So that the invocation will be paused when
# tool confirmation is requested.
app = App(
name="test_app",
resumability_config=ResumabilityConfig(is_resumable=True),
@@ -427,8 +429,8 @@ class TestHITLConfirmationFlowWithResumableApp:
"""Tests HITL flow where all tool calls are confirmed."""
events = runner.run("test user query")
# Verify that the invocation is paused after the long running tool call.
# So that no intermediate function response and llm response is generated.
# Verify that the invocation is paused when tool confirmation is requested.
# The tool call returns error response, and summarization was skipped.
assert testing_utils.simplify_resumable_app_events(
copy.deepcopy(events)
) == [
@@ -448,13 +450,21 @@ class TestHITLConfirmationFlowWithResumableApp:
"args": {},
},
"toolConfirmation": {
"hint": "test hint for request_confirmation",
"hint": HINT_TEXT,
"confirmed": False,
},
},
)
),
),
(
agent.name,
Part(
function_response=FunctionResponse(
name=agent.tools[0].name, response=TOOL_CALL_ERROR_RESPONSE
)
),
),
]
ask_for_confirmation_function_call_id = (
events[1].content.parts[0].function_call.id
@@ -499,7 +509,7 @@ class TestHITLConfirmationFlowWithSequentialAgentAndResumableApp:
@pytest.fixture
def tools(self) -> list[FunctionTool]:
"""Provides the tools for the agent."""
return [FunctionTool(func=_test_request_confirmation_function)]
return [FunctionTool(func=_test_function, require_confirmation=True)]
@pytest.fixture
def llm_responses(
@@ -535,8 +545,8 @@ class TestHITLConfirmationFlowWithSequentialAgentAndResumableApp:
@pytest.fixture
def runner(self, agent: SequentialAgent) -> testing_utils.InMemoryRunner:
"""Provides an in-memory runner for the agent."""
# Mark the app as resumable. So that the invocation will be paused after the
# long running tool call.
# Mark the app as resumable. So that the invocation will be paused when
# tool confirmation is requested.
app = App(
name="test_app",
resumability_config=ResumabilityConfig(is_resumable=True),
@@ -558,8 +568,8 @@ class TestHITLConfirmationFlowWithSequentialAgentAndResumableApp:
# - sub_agent1 has a tool call that asks for HITL confirmation.
# - sub_agent2 does not have any tool calls.
# - The test will:
# - Run the query and verify that the invocation is paused after the long
# running tool call, at sub_agent1.
# - Run the query and verify that the invocation is paused when tool
# confirmation is requested, at sub_agent1.
# - Resume the invocation and execute the tool call from sub_agent1.
# - Verify that root_agent continues to run sub_agent2.
@@ -568,8 +578,8 @@ class TestHITLConfirmationFlowWithSequentialAgentAndResumableApp:
sub_agent2 = agent.sub_agents[1]
# Step 1:
# Verify that the invocation is paused after the long running tool call.
# So that no intermediate function response and llm response is generated.
# Verify that the invocation is paused when tool confirmation is requested.
# So that no intermediate llm response is generated.
# And the second sub agent is not started.
assert testing_utils.simplify_resumable_app_events(
copy.deepcopy(events)
@@ -600,13 +610,22 @@ class TestHITLConfirmationFlowWithSequentialAgentAndResumableApp:
"args": {},
},
"toolConfirmation": {
"hint": "test hint for request_confirmation",
"hint": HINT_TEXT,
"confirmed": False,
},
},
)
),
),
(
sub_agent1.name,
Part(
function_response=FunctionResponse(
name=sub_agent1.tools[0].name,
response=TOOL_CALL_ERROR_RESPONSE,
)
),
),
]
ask_for_confirmation_function_call_id = (
events[2].content.parts[0].function_call.id
@@ -664,7 +683,7 @@ class TestHITLConfirmationFlowWithParallelAgentAndResumableApp:
@pytest.fixture
def tools(self) -> list[FunctionTool]:
"""Provides the tools for the agent."""
return [FunctionTool(func=_test_request_confirmation_function)]
return [FunctionTool(func=_test_function, require_confirmation=True)]
@pytest.fixture
def llm_responses(
@@ -702,8 +721,8 @@ class TestHITLConfirmationFlowWithParallelAgentAndResumableApp:
@pytest.fixture
def runner(self, agent: ParallelAgent) -> testing_utils.InMemoryRunner:
"""Provides an in-memory runner for the agent."""
# Mark the app as resumable. So that the invocation will be paused after the
# long running tool call.
# Mark the app as resumable. So that the invocation will be paused when
# tool confirmation is requested.
app = App(
name="test_app",
resumability_config=ResumabilityConfig(is_resumable=True),
@@ -725,15 +744,15 @@ class TestHITLConfirmationFlowWithParallelAgentAndResumableApp:
# sub_agent2.
# - Both sub_agents have a tool call that asks for HITL confirmation.
# - The test will:
# - Run the query and verify that each branch is paused after the long
# running tool call.
# - Run the query and verify that each branch is paused when tool
# confirmation is requested.
# - Resume the invocation and execute the tool call of each branch.
sub_agent1 = agent.sub_agents[0]
sub_agent2 = agent.sub_agents[1]
# Verify that each branch is paused after the long running tool call.
# So that no intermediate function response and llm response is generated.
# So that no intermediate llm response is generated.
root_agent_events = [event for event in events if event.branch is None]
sub_agent1_branch_events = [
event
@@ -776,13 +795,22 @@ class TestHITLConfirmationFlowWithParallelAgentAndResumableApp:
"args": {},
},
"toolConfirmation": {
"hint": "test hint for request_confirmation",
"hint": HINT_TEXT,
"confirmed": False,
},
},
)
),
),
(
sub_agent1.name,
Part(
function_response=FunctionResponse(
name=sub_agent1.tools[0].name,
response=TOOL_CALL_ERROR_RESPONSE,
)
),
),
]
assert testing_utils.simplify_resumable_app_events(
copy.deepcopy(sub_agent2_branch_events)
@@ -807,13 +835,22 @@ class TestHITLConfirmationFlowWithParallelAgentAndResumableApp:
"args": {},
},
"toolConfirmation": {
"hint": "test hint for request_confirmation",
"hint": HINT_TEXT,
"confirmed": False,
},
},
)
),
),
(
sub_agent2.name,
Part(
function_response=FunctionResponse(
name=sub_agent2.tools[0].name,
response=TOOL_CALL_ERROR_RESPONSE,
)
),
),
]
ask_for_confirmation_function_call_ids = [
@@ -893,7 +930,7 @@ class TestHITLConfirmationFlowWithParallelAgentAndResumableApp:
sub_agent2.name,
Part(
function_response=FunctionResponse(
name=sub_agent1.tools[0].name,
name=sub_agent2.tools[0].name,
response={"result": "confirmed=True"},
)
),