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: Modify runner to support resuming an invocation (optionally with a function response)

Browse Source 
PiperOrigin-RevId: 813008406
This commit is contained in:
Xinran (Sherry) Tang
2025-09-29 17:35:18 -07:00
committed by Copybara-Service
parent f005414895
commit fbf75761bb
5 changed files with 473 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/unittests/agents/test_invocation_context.py
+120
View File
@@ -15,9 +15,11 @@
from unittest.mock import Mock
from google.adk.agents.base_agent import BaseAgent
from google.adk.agents.base_agent import BaseAgentState
from google.adk.agents.invocation_context import InvocationContext
from google.adk.apps import ResumabilityConfig
from google.adk.events.event import Event
from google.adk.events.event_actions import EventActions
from google.adk.sessions.base_session_service import BaseSessionService
from google.adk.sessions.session import Session
from google.genai.types import Content
@@ -227,6 +229,124 @@ class TestInvocationContextWithAppResumablity:
invocation_context = self._create_test_invocation_context(None)
assert not invocation_context.is_resumable
def test_populate_invocation_agent_states_not_resumable(self):
"""Tests that populate_invocation_agent_states does nothing if not resumable."""
invocation_context = self._create_test_invocation_context(
ResumabilityConfig(is_resumable=False)
)
event = Event(
invocation_id='inv_1',
author='agent1',
actions=EventActions(end_of_agent=True, agent_state=None),
)
invocation_context.session.events = [event]
invocation_context.populate_invocation_agent_states()
assert not invocation_context.agent_states
assert not invocation_context.end_of_agents
def test_populate_invocation_agent_states_end_of_agent(self):
"""Tests that populate_invocation_agent_states handles end_of_agent."""
invocation_context = self._create_test_invocation_context(
ResumabilityConfig(is_resumable=True)
)
event = Event(
invocation_id='inv_1',
author='agent1',
actions=EventActions(end_of_agent=True, agent_state=None),
)
invocation_context.session.events = [event]
invocation_context.populate_invocation_agent_states()
assert not invocation_context.agent_states
assert invocation_context.end_of_agents == {'agent1': True}
def test_populate_invocation_agent_states_with_agent_state(self):
"""Tests that populate_invocation_agent_states handles agent_state."""
invocation_context = self._create_test_invocation_context(
ResumabilityConfig(is_resumable=True)
)
event = Event(
invocation_id='inv_1',
author='agent1',
actions=EventActions(
end_of_agent=False,
agent_state=BaseAgentState().model_dump(mode='json'),
),
)
invocation_context.session.events = [event]
invocation_context.populate_invocation_agent_states()
assert invocation_context.agent_states == {'agent1': {}}
assert invocation_context.end_of_agents == {'agent1': False}
def test_populate_invocation_agent_states_with_agent_state_and_end_of_agent(
self,
):
"""Tests that populate_invocation_agent_states handles agent_state and end_of_agent."""
invocation_context = self._create_test_invocation_context(
ResumabilityConfig(is_resumable=True)
)
event = Event(
invocation_id='inv_1',
author='agent1',
actions=EventActions(
end_of_agent=True,
agent_state=BaseAgentState().model_dump(mode='json'),
),
)
invocation_context.session.events = [event]
invocation_context.populate_invocation_agent_states()
# When both agent_state and end_of_agent are set, agent_state should be
# cleared, as end_of_agent is of a higher priority.
assert not invocation_context.agent_states
assert invocation_context.end_of_agents == {'agent1': True}
def test_populate_invocation_agent_states_with_content_no_state(self):
"""Tests that populate_invocation_agent_states creates default state."""
invocation_context = self._create_test_invocation_context(
ResumabilityConfig(is_resumable=True)
)
event = Event(
invocation_id='inv_1',
author='agent1',
actions=EventActions(end_of_agent=False, agent_state=None),
content=Content(role='model', parts=[Part(text='hi')]),
)
invocation_context.session.events = [event]
invocation_context.populate_invocation_agent_states()
assert invocation_context.agent_states == {'agent1': BaseAgentState()}
assert invocation_context.end_of_agents == {'agent1': False}
def test_populate_invocation_agent_states_user_message_event(self):
"""Tests that populate_invocation_agent_states ignores user message events for default state."""
invocation_context = self._create_test_invocation_context(
ResumabilityConfig(is_resumable=True)
)
event = Event(
invocation_id='inv_1',
author='user',
actions=EventActions(end_of_agent=False, agent_state=None),
content=Content(role='user', parts=[Part(text='hi')]),
)
invocation_context.session.events = [event]
invocation_context.populate_invocation_agent_states()
assert not invocation_context.agent_states
assert not invocation_context.end_of_agents
def test_populate_invocation_agent_states_no_content(self):
"""Tests that populate_invocation_agent_states ignores events with no content if no state."""
invocation_context = self._create_test_invocation_context(
ResumabilityConfig(is_resumable=True)
)
event = Event(
invocation_id='inv_1',
author='agent1',
actions=EventActions(end_of_agent=None, agent_state=None),
content=None,
)
invocation_context.session.events = [event]
invocation_context.populate_invocation_agent_states()
assert not invocation_context.agent_states
assert not invocation_context.end_of_agents
class TestFindMatchingFunctionCall:
"""Test suite for find_matching_function_call."""
tests/unittests/runners/test_run_tool_confirmation.py
+195 -4
View File
@@ -19,6 +19,8 @@ from unittest import mock
from google.adk.agents.base_agent import BaseAgent
from google.adk.agents.llm_agent import LlmAgent
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.flows.llm_flows.functions import REQUEST_CONFIRMATION_FUNCTION_CALL_NAME
@@ -186,6 +188,7 @@ class TestHITLConfirmationFlowWithSingleAgent(BaseHITLTest):
ask_for_confirmation_function_call_id = (
events[1].content.parts[0].function_call.id
)
invocation_id = events[1].invocation_id
user_confirmation = testing_utils.UserContent(
Part(
function_response=FunctionResponse(
@@ -209,6 +212,8 @@ class TestHITLConfirmationFlowWithSingleAgent(BaseHITLTest):
),
(agent.name, "test llm response after final tool call"),
]
for event in events:
assert event.invocation_id != invocation_id
assert (
testing_utils.simplify_events(copy.deepcopy(events))
== expected_parts_final
@@ -321,6 +326,7 @@ class TestHITLConfirmationFlowWithCustomPayloadSchema(BaseHITLTest):
ask_for_confirmation_function_call_id = (
events[1].content.parts[0].function_call.id
)
invocation_id = events[1].invocation_id
custom_payload = {
"test_custom_payload": {
"int_field": 123,
@@ -358,6 +364,8 @@ class TestHITLConfirmationFlowWithCustomPayloadSchema(BaseHITLTest):
),
(agent.name, "test llm response after final tool call"),
]
for event in events:
assert event.invocation_id != invocation_id
assert (
testing_utils.simplify_events(copy.deepcopy(events))
== expected_parts_final
@@ -380,9 +388,6 @@ class TestHITLConfirmationFlowWithResumableApp:
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
@@ -412,7 +417,7 @@ class TestHITLConfirmationFlowWithResumableApp:
return testing_utils.InMemoryRunner(app=app)
@pytest.mark.asyncio
def test_pause_on_request_confirmation(
async def test_pause_and_resume_on_request_confirmation(
self,
runner: testing_utils.InMemoryRunner,
agent: LlmAgent,
@@ -449,3 +454,189 @@ class TestHITLConfirmationFlowWithResumableApp:
),
),
]
ask_for_confirmation_function_call_id = (
events[1].content.parts[0].function_call.id
)
invocation_id = events[1].invocation_id
user_confirmation = testing_utils.UserContent(
Part(
function_response=FunctionResponse(
id=ask_for_confirmation_function_call_id,
name=REQUEST_CONFIRMATION_FUNCTION_CALL_NAME,
response={"confirmed": True},
)
)
)
events = await runner.run_async(
user_confirmation, invocation_id=invocation_id
)
expected_parts_final = [
(
agent.name,
Part(
function_response=FunctionResponse(
name=agent.tools[0].name,
response={"result": "confirmed=True"},
)
),
),
(agent.name, "test llm response after tool call"),
(agent.name, testing_utils.END_OF_AGENT),
]
for event in events:
assert event.invocation_id == invocation_id
assert (
testing_utils.simplify_resumable_app_events(copy.deepcopy(events))
== expected_parts_final
)
class TestHITLConfirmationFlowWithSequentialAgentAndResumableApp:
"""Tests the HITL confirmation flow with a resumable sequential agent app."""
@pytest.fixture
def tools(self) -> list[FunctionTool]:
"""Provides the tools for the agent."""
return [FunctionTool(func=_test_request_confirmation_function)]
@pytest.fixture
def llm_responses(
self, tools: list[FunctionTool]
) -> list[GenerateContentResponse]:
"""Provides mock LLM responses for the tests."""
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 from second agent"),
]
@pytest.fixture
def mock_model(
self, llm_responses: list[GenerateContentResponse]
) -> testing_utils.MockModel:
"""Provides a mock model with predefined responses."""
return testing_utils.MockModel(responses=llm_responses)
@pytest.fixture
def agent(
self, mock_model: testing_utils.MockModel, tools: list[FunctionTool]
) -> SequentialAgent:
"""Provides a single LlmAgent for the test."""
return SequentialAgent(
name="root_agent",
sub_agents=[
LlmAgent(name="agent1", model=mock_model, tools=tools),
LlmAgent(name="agent2", model=mock_model, tools=[]),
],
)
@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.
app = App(
name="test_app",
resumability_config=ResumabilityConfig(is_resumable=True),
root_agent=agent,
)
return testing_utils.InMemoryRunner(app=app)
@pytest.mark.asyncio
async def test_pause_and_resume_on_request_confirmation(
self,
runner: testing_utils.InMemoryRunner,
agent: SequentialAgent,
):
"""Tests HITL flow where all tool calls are confirmed."""
events = runner.run("test user query")
sub_agent1 = agent.sub_agents[0]
sub_agent2 = agent.sub_agents[1]
# Verify that the invocation is paused after the long running tool call.
# So that no intermediate function response and llm response is generated.
# And the second sub agent is not started.
assert testing_utils.simplify_resumable_app_events(
copy.deepcopy(events)
) == [
(
agent.name,
SequentialAgentState(current_sub_agent=sub_agent1.name).model_dump(
mode="json"
),
),
(
sub_agent1.name,
Part(
function_call=FunctionCall(
name=sub_agent1.tools[0].name, args={}
)
),
),
(
sub_agent1.name,
Part(
function_call=FunctionCall(
name=REQUEST_CONFIRMATION_FUNCTION_CALL_NAME,
args={
"originalFunctionCall": {
"name": sub_agent1.tools[0].name,
"id": mock.ANY,
"args": {},
},
"toolConfirmation": {
"hint": "test hint for request_confirmation",
"confirmed": False,
},
},
)
),
),
]
ask_for_confirmation_function_call_id = (
events[2].content.parts[0].function_call.id
)
invocation_id = events[2].invocation_id
# Resume the invocation and confirm the tool call from sub_agent1, and
# sub_agent2 will continue.
user_confirmation = testing_utils.UserContent(
Part(
function_response=FunctionResponse(
id=ask_for_confirmation_function_call_id,
name=REQUEST_CONFIRMATION_FUNCTION_CALL_NAME,
response={"confirmed": True},
)
)
)
events = await runner.run_async(
user_confirmation, invocation_id=invocation_id
)
expected_parts_final = [
(
sub_agent1.name,
Part(
function_response=FunctionResponse(
name=sub_agent1.tools[0].name,
response={"result": "confirmed=True"},
)
),
),
(sub_agent1.name, "test llm response after tool call"),
(sub_agent1.name, testing_utils.END_OF_AGENT),
(
agent.name,
SequentialAgentState(current_sub_agent=sub_agent2.name).model_dump(
mode="json"
),
),
(sub_agent2.name, "test llm response from second agent"),
(sub_agent2.name, testing_utils.END_OF_AGENT),
(agent.name, testing_utils.END_OF_AGENT),
]
for event in events:
assert event.invocation_id == invocation_id
assert (
testing_utils.simplify_resumable_app_events(copy.deepcopy(events))
== expected_parts_final
)
tests/unittests/testing_utils.py
+4 -1
View File
@@ -273,11 +273,14 @@ class InMemoryRunner:
)
)
async def run_async(self, new_message: types.ContentUnion) -> list[Event]:
async def run_async(
self, new_message: types.ContentUnion, invocation_id: Optional[str] = None
) -> list[Event]:
events = []
async for event in self.runner.run_async(
user_id=self.session.user_id,
session_id=self.session.id,
invocation_id=invocation_id,
new_message=get_user_content(new_message),
):
events.append(event)