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

feat: Make resumable llm agents yield checkpoint events

Browse Source 
PiperOrigin-RevId: 813001108
This commit is contained in:
Xinran (Sherry) Tang
2025-09-29 17:08:58 -07:00
committed by Copybara-Service
parent 609a2358eb
commit f005414895
10 changed files with 1133 additions and 206 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/unittests/runners/test_pause_invocation.py
+74 -67
View File
@@ -23,6 +23,7 @@ from google.adk.agents.llm_agent import LlmAgent
from google.adk.agents.loop_agent import LoopAgent
from google.adk.agents.parallel_agent import ParallelAgent
from google.adk.agents.sequential_agent import SequentialAgent
from google.adk.agents.sequential_agent import SequentialAgentState
from google.adk.apps.app import App
from google.adk.apps.app import ResumabilityConfig
from google.adk.events.event import Event
@@ -71,6 +72,7 @@ class _TestingAgent(BaseAgent):
_TRANSFER_RESPONSE_PART = Part.from_function_response(
name="transfer_to_agent", response={"result": None}
)
END_OF_AGENT = testing_utils.END_OF_AGENT
class BasePauseInvocationTest:
@@ -85,15 +87,15 @@ class BasePauseInvocationTest:
def app(self, agent: BaseAgent) -> App:
"""Provides an App for the test."""
return App(
name="InMemoryRunner", # Required for using TestInMemoryRunner.
name="test_app",
root_agent=agent,
resumability_config=ResumabilityConfig(is_resumable=True),
)
@pytest.fixture
def runner(self, app: App) -> testing_utils.TestInMemoryRunner:
def runner(self, app: App) -> testing_utils.InMemoryRunner:
"""Provides an in-memory runner for the agent."""
return testing_utils.TestInMemoryRunner(app=app, app_name=None)
return testing_utils.InMemoryRunner(app=app)
@staticmethod
def mock_model(responses: list[Part]) -> testing_utils.MockModel:
@@ -107,10 +109,6 @@ class TestPauseInvocationWithSingleLlmAgent(BasePauseInvocationTest):
@pytest.fixture
def agent(self) -> BaseAgent:
"""Provides a BaseAgent for the test."""
def test_tool() -> str:
return ""
return LlmAgent(
name="root_agent",
model=self.mock_model(
@@ -120,14 +118,12 @@ class TestPauseInvocationWithSingleLlmAgent(BasePauseInvocationTest):
)
@pytest.mark.asyncio
async def test_pause_on_long_running_function_call(
def test_pause_on_long_running_function_call(
self,
runner: testing_utils.TestInMemoryRunner,
runner: testing_utils.InMemoryRunner,
):
"""Tests that a single LlmAgent pauses on long running function call."""
assert testing_utils.simplify_events(
await runner.run_async_with_new_session("test")
) == [
assert testing_utils.simplify_resumable_app_events(runner.run("test")) == [
("root_agent", Part.from_function_call(name="test_tool", args={})),
]
@@ -158,34 +154,42 @@ class TestPauseInvocationWithSequentialAgent(BasePauseInvocationTest):
)
@pytest.mark.asyncio
async def test_pause_first_agent_on_long_running_function_call(
def test_pause_first_agent_on_long_running_function_call(
self,
runner: testing_utils.TestInMemoryRunner,
runner: testing_utils.InMemoryRunner,
):
"""Tests that a single LlmAgent pauses on long running function call."""
assert testing_utils.simplify_events(
await runner.run_async_with_new_session("test")
) == [
"""Tests that a SequentialAgent pauses on the first sub-agent."""
assert testing_utils.simplify_resumable_app_events(runner.run("test")) == [
(
"root_agent",
SequentialAgentState(current_sub_agent="sub_agent_1").model_dump(
mode="json"
),
),
("sub_agent_1", Part.from_function_call(name="test_tool", args={})),
]
@pytest.mark.asyncio
async def test_pause_second_agent_on_long_running_function_call(
def test_pause_second_agent_on_long_running_function_call(
self,
runner: testing_utils.TestInMemoryRunner,
runner: testing_utils.InMemoryRunner,
):
"""Tests that a single LlmAgent pauses on long running function call."""
# Change the base sequential agent, so that the first agent does not pause.
runner.agent.sub_agents[0].tools = [FunctionTool(func=test_tool)]
runner.agent.sub_agents[0].model = self.mock_model(
runner.root_agent.sub_agents[0].tools = [FunctionTool(func=test_tool)]
runner.root_agent.sub_agents[0].model = self.mock_model(
responses=[
Part.from_function_call(name="test_tool", args={}),
Part.from_text(text="model response after tool call"),
]
)
assert testing_utils.simplify_events(
await runner.run_async_with_new_session("test")
) == [
assert testing_utils.simplify_resumable_app_events(runner.run("test")) == [
(
"root_agent",
SequentialAgentState(current_sub_agent="sub_agent_1").model_dump(
mode="json"
),
),
("sub_agent_1", Part.from_function_call(name="test_tool", args={})),
(
"sub_agent_1",
@@ -194,6 +198,13 @@ class TestPauseInvocationWithSequentialAgent(BasePauseInvocationTest):
),
),
("sub_agent_1", "model response after tool call"),
("sub_agent_1", END_OF_AGENT),
(
"root_agent",
SequentialAgentState(current_sub_agent="sub_agent_2").model_dump(
mode="json"
),
),
("sub_agent_2", Part.from_function_call(name="test_tool", args={})),
]
@@ -221,17 +232,19 @@ class TestPauseInvocationWithParallelAgent(BasePauseInvocationTest):
)
@pytest.mark.asyncio
async def test_pause_on_long_running_function_call(
def test_pause_on_long_running_function_call(
self,
runner: testing_utils.TestInMemoryRunner,
runner: testing_utils.InMemoryRunner,
):
"""Tests that a ParallelAgent pauses on long running function call."""
assert testing_utils.simplify_events(
await runner.run_async_with_new_session("test")
) == [
("sub_agent_1", Part.from_function_call(name="test_tool", args={})),
("sub_agent_2", "Delayed message"),
]
simplified_event_parts = testing_utils.simplify_resumable_app_events(
runner.run("test")
)
assert (
"sub_agent_1",
Part.from_function_call(name="test_tool", args={}),
) in simplified_event_parts
assert ("sub_agent_2", "Delayed message") in simplified_event_parts
class TestPauseInvocationWithNestedParallelAgent(BasePauseInvocationTest):
@@ -265,50 +278,49 @@ class TestPauseInvocationWithNestedParallelAgent(BasePauseInvocationTest):
)
@pytest.mark.asyncio
async def test_pause_on_long_running_function_call_in_nested_agent(
def test_pause_on_long_running_function_call(
self,
runner: testing_utils.TestInMemoryRunner,
runner: testing_utils.InMemoryRunner,
):
"""Tests that a nested ParallelAgent pauses on long running function call."""
assert testing_utils.simplify_events(
await runner.run_async_with_new_session("test")
) == [
(
"nested_sub_agent_1",
Part.from_function_call(name="test_tool", args={}),
),
("sub_agent_1", "Delayed message"),
("nested_sub_agent_2", "Delayed message"),
]
simplified_event_parts = testing_utils.simplify_resumable_app_events(
runner.run("test")
)
assert (
"nested_sub_agent_1",
Part.from_function_call(name="test_tool", args={}),
) in simplified_event_parts
assert ("sub_agent_1", "Delayed message") in simplified_event_parts
assert ("nested_sub_agent_2", "Delayed message") in simplified_event_parts
@pytest.mark.asyncio
async def test_pause_on_multiple_long_running_function_calls(
def test_pause_on_multiple_long_running_function_calls(
self,
runner: testing_utils.TestInMemoryRunner,
runner: testing_utils.InMemoryRunner,
):
"""Tests that a ParallelAgent pauses on long running function calls."""
runner.agent.sub_agents[0] = LlmAgent(
runner.root_agent.sub_agents[0] = LlmAgent(
name="sub_agent_1",
model=self.mock_model(
responses=[
Part.from_function_call(name="test_tool", args={}),
Part.from_function_call(name="test_tool", args={}),
]
),
tools=[LongRunningFunctionTool(func=test_tool)],
)
simplified_events = testing_utils.simplify_events(
await runner.run_async_with_new_session("test")
simplified_events = testing_utils.simplify_resumable_app_events(
runner.run("test")
)
assert len(simplified_events) == 3
assert (
"sub_agent_1",
Part.from_function_call(name="test_tool", args={}),
) in simplified_events
assert ("sub_agent_1", END_OF_AGENT) not in simplified_events
assert (
"nested_sub_agent_1",
Part.from_function_call(name="test_tool", args={}),
) in simplified_events
assert ("nested_sub_agent_1", END_OF_AGENT) not in simplified_events
class TestPauseInvocationWithLoopAgent(BasePauseInvocationTest):
@@ -350,15 +362,14 @@ class TestPauseInvocationWithLoopAgent(BasePauseInvocationTest):
)
@pytest.mark.asyncio
async def test_pause_on_long_running_function_call_in_loop(
def test_pause_on_long_running_function_call(
self,
runner: testing_utils.TestInMemoryRunner,
runner: testing_utils.InMemoryRunner,
):
"""Tests that a LoopAgent pauses on long running function call."""
assert testing_utils.simplify_events(
await runner.run_async_with_new_session("test")
) == [
assert testing_utils.simplify_resumable_app_events(runner.run("test")) == [
("sub_agent_1", "sub agent 1 response"),
("sub_agent_1", END_OF_AGENT),
("sub_agent_2", Part.from_function_call(name="test_tool", args={})),
]
@@ -400,14 +411,12 @@ class TestPauseInvocationWithLlmAgentTree(BasePauseInvocationTest):
)
@pytest.mark.asyncio
async def test_pause_on_transfer_call_part(
def test_pause_on_long_running_function_call(
self,
runner: testing_utils.TestInMemoryRunner,
runner: testing_utils.InMemoryRunner,
):
"""Tests that a tree of resumable LlmAgents yields checkpoint events."""
assert testing_utils.simplify_events(
await runner.run_async_with_new_session("test")
) == [
assert testing_utils.simplify_resumable_app_events(runner.run("test")) == [
("root_agent", _transfer_call_part("sub_llm_agent_1")),
("root_agent", _TRANSFER_RESPONSE_PART),
("sub_llm_agent_1", _transfer_call_part("sub_llm_agent_2")),
@@ -454,14 +463,12 @@ class TestPauseInvocationWithWithTransferLoop(BasePauseInvocationTest):
)
@pytest.mark.asyncio
async def test_agent_tree_yields_checkpoints(
def test_pause_on_long_running_function_call(
self,
runner: testing_utils.TestInMemoryRunner,
runner: testing_utils.InMemoryRunner,
):
"""Tests that a tree of resumable LlmAgents yields checkpoint events."""
assert testing_utils.simplify_events(
await runner.run_async_with_new_session("test")
) == [
assert testing_utils.simplify_resumable_app_events(runner.run("test")) == [
("root_agent", _transfer_call_part("sub_llm_agent_1")),
("root_agent", _TRANSFER_RESPONSE_PART),
("sub_llm_agent_1", _transfer_call_part("sub_llm_agent_2")),
tests/unittests/runners/test_run_tool_confirmation.py
+9 -7
View File
@@ -400,29 +400,31 @@ class TestHITLConfirmationFlowWithResumableApp:
return LlmAgent(name="root_agent", model=mock_model, tools=tools)
@pytest.fixture
def runner(self, agent: LlmAgent) -> testing_utils.TestInMemoryRunner:
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.
app = App(
name="InMemoryRunner", # Required for using TestInMemoryRunner.
name="test_app",
resumability_config=ResumabilityConfig(is_resumable=True),
root_agent=agent,
)
return testing_utils.TestInMemoryRunner(app=app, app_name=None)
return testing_utils.InMemoryRunner(app=app)
@pytest.mark.asyncio
async def test_pause_on_request_confirmation(
def test_pause_on_request_confirmation(
self,
runner: testing_utils.TestInMemoryRunner,
runner: testing_utils.InMemoryRunner,
agent: LlmAgent,
):
"""Tests HITL flow where all tool calls are confirmed."""
events = await runner.run_async_with_new_session("test user query")
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.
assert testing_utils.simplify_events(copy.deepcopy(events)) == [
assert testing_utils.simplify_resumable_app_events(
copy.deepcopy(events)
) == [
(
agent.name,
Part(function_call=FunctionCall(name=agent.tools[0].name, args={})),