wehub-resource-sync
ec2b666284
Continuous Integration / Pre-commit Linter (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.10) (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.11) (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.12) (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.13) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.10) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.11) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.12) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.13) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.14) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.10) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.11) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.12) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.13) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.14) (push) Waiting to run
Copybara PR Handler / close-imported-pr (push) Waiting to run
539 lines
17 KiB
Python
539 lines
17 KiB
Python
# Copyright 2026 Google LLC
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""Tests for the resumption flow with different agent structures."""
|
|
|
|
import asyncio
|
|
from typing import AsyncGenerator
|
|
|
|
from google.adk.agents.base_agent import BaseAgent
|
|
from google.adk.agents.invocation_context import InvocationContext
|
|
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.apps.app import App
|
|
from google.adk.apps.app import ResumabilityConfig
|
|
from google.adk.events.event import Event
|
|
from google.adk.tools.exit_loop_tool import exit_loop
|
|
from google.adk.tools.function_tool import FunctionTool
|
|
from google.adk.tools.long_running_tool import LongRunningFunctionTool
|
|
from google.genai.types import Part
|
|
import pytest
|
|
|
|
from .. import testing_utils
|
|
|
|
|
|
def _transfer_call_part(agent_name: str) -> Part:
|
|
return Part.from_function_call(
|
|
name="transfer_to_agent", args={"agent_name": agent_name}
|
|
)
|
|
|
|
|
|
def test_tool():
|
|
"""A test tool; returns None to simulate a pending long-running operation."""
|
|
return None
|
|
|
|
|
|
class _TestingAgent(BaseAgent):
|
|
"""A testing agent that generates an event after a delay."""
|
|
|
|
delay: float = 0
|
|
"""The delay before the agent generates an event."""
|
|
|
|
def event(self, ctx: InvocationContext):
|
|
return Event(
|
|
author=self.name,
|
|
branch=ctx.branch,
|
|
invocation_id=ctx.invocation_id,
|
|
content=testing_utils.ModelContent(
|
|
parts=[Part.from_text(text="Delayed message")]
|
|
),
|
|
)
|
|
|
|
async def _run_async_impl(
|
|
self, ctx: InvocationContext
|
|
) -> AsyncGenerator[Event, None]:
|
|
await asyncio.sleep(self.delay)
|
|
yield self.event(ctx)
|
|
|
|
|
|
_TRANSFER_RESPONSE_PART = Part.from_function_response(
|
|
name="transfer_to_agent", response={"result": None}
|
|
)
|
|
END_OF_AGENT = testing_utils.END_OF_AGENT
|
|
|
|
|
|
class BasePauseInvocationTest:
|
|
"""Base class for pausing invocation tests with common fixtures."""
|
|
|
|
@pytest.fixture
|
|
def agent(self) -> BaseAgent:
|
|
"""Provides a BaseAgent for the test."""
|
|
return BaseAgent(name="test_agent")
|
|
|
|
@pytest.fixture
|
|
def app(self, agent: BaseAgent) -> App:
|
|
"""Provides an App for the test."""
|
|
return App(
|
|
name="test_app",
|
|
root_agent=agent,
|
|
resumability_config=ResumabilityConfig(is_resumable=True),
|
|
)
|
|
|
|
@pytest.fixture
|
|
def runner(self, app: App) -> testing_utils.InMemoryRunner:
|
|
"""Provides an in-memory runner for the agent."""
|
|
return testing_utils.InMemoryRunner(app=app)
|
|
|
|
@staticmethod
|
|
def mock_model(responses: list[Part]) -> testing_utils.MockModel:
|
|
"""Provides a mock model with predefined responses."""
|
|
return testing_utils.MockModel.create(responses=responses)
|
|
|
|
|
|
class TestPauseInvocationWithSingleLlmAgent(BasePauseInvocationTest):
|
|
"""Tests the resumption flow with a single LlmAgent."""
|
|
|
|
@pytest.fixture
|
|
def agent(self) -> BaseAgent:
|
|
"""Provides a BaseAgent for the test."""
|
|
return LlmAgent(
|
|
name="root_agent",
|
|
model=self.mock_model(
|
|
responses=[Part.from_function_call(name="test_tool", args={})]
|
|
),
|
|
tools=[LongRunningFunctionTool(func=test_tool)],
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
def test_pause_on_long_running_function_call(
|
|
self,
|
|
runner: testing_utils.InMemoryRunner,
|
|
):
|
|
"""Tests that a single LlmAgent pauses on long running function call."""
|
|
actual = testing_utils.simplify_resumable_app_events(runner.run("test"))
|
|
behavioral = [e for e in actual if not isinstance(e[1], dict)]
|
|
assert behavioral == [
|
|
# execute_tools yields the interrupt event with long_running_tool_ids.
|
|
("root_agent", Part.from_function_call(name="test_tool", args={})),
|
|
]
|
|
|
|
|
|
class TestPauseInvocationWithSequentialAgent(BasePauseInvocationTest):
|
|
"""Tests pausing invocation with a SequentialAgent."""
|
|
|
|
@pytest.fixture
|
|
def agent(self) -> BaseAgent:
|
|
"""Provides a BaseAgent for the test."""
|
|
sub_agent1 = LlmAgent(
|
|
name="sub_agent_1",
|
|
model=self.mock_model(
|
|
responses=[Part.from_function_call(name="test_tool", args={})]
|
|
),
|
|
tools=[LongRunningFunctionTool(func=test_tool)],
|
|
)
|
|
sub_agent2 = LlmAgent(
|
|
name="sub_agent_2",
|
|
model=self.mock_model(
|
|
responses=[Part.from_function_call(name="test_tool", args={})]
|
|
),
|
|
tools=[LongRunningFunctionTool(func=test_tool)],
|
|
)
|
|
return SequentialAgent(
|
|
name="root_agent",
|
|
sub_agents=[sub_agent1, sub_agent2],
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
def test_pause_first_agent_on_long_running_function_call(
|
|
self,
|
|
runner: testing_utils.InMemoryRunner,
|
|
):
|
|
"""Tests that a SequentialAgent pauses on the first sub-agent."""
|
|
actual = testing_utils.simplify_resumable_app_events(runner.run("test"))
|
|
behavioral = [e for e in actual if not isinstance(e[1], dict)]
|
|
assert behavioral == [
|
|
("sub_agent_1", Part.from_function_call(name="test_tool", args={})),
|
|
]
|
|
|
|
@pytest.mark.xfail(
|
|
reason=(
|
|
"Tests implementation details that are different in V2 and will be"
|
|
" deprecated."
|
|
)
|
|
)
|
|
@pytest.mark.asyncio
|
|
def test_pause_second_agent_on_long_running_function_call(
|
|
self,
|
|
):
|
|
"""Tests that a single LlmAgent pauses on long running function call."""
|
|
# Construct sub_agent_1 with regular FunctionTool (not long-running).
|
|
sub_agent_1 = LlmAgent(
|
|
name="sub_agent_1",
|
|
model=self.mock_model(
|
|
responses=[
|
|
Part.from_function_call(name="test_tool", args={}),
|
|
Part.from_text(text="model response after tool call"),
|
|
]
|
|
),
|
|
tools=[FunctionTool(func=test_tool)],
|
|
)
|
|
sub_agent_2 = LlmAgent(
|
|
name="sub_agent_2",
|
|
model=self.mock_model(
|
|
responses=[Part.from_function_call(name="test_tool", args={})]
|
|
),
|
|
tools=[LongRunningFunctionTool(func=test_tool)],
|
|
)
|
|
agent = SequentialAgent(
|
|
name="root_agent",
|
|
sub_agents=[sub_agent_1, sub_agent_2],
|
|
)
|
|
app = App(
|
|
name="test_app",
|
|
root_agent=agent,
|
|
resumability_config=ResumabilityConfig(is_resumable=True),
|
|
)
|
|
runner = testing_utils.InMemoryRunner(app=app)
|
|
actual = testing_utils.simplify_resumable_app_events(runner.run("test"))
|
|
behavioral = [e for e in actual if not isinstance(e[1], dict)]
|
|
assert behavioral == [
|
|
("sub_agent_1", Part.from_function_call(name="test_tool", args={})),
|
|
(
|
|
"sub_agent_1",
|
|
Part.from_function_response(
|
|
name="test_tool", response={"result": None}
|
|
),
|
|
),
|
|
("sub_agent_1", "model response after tool call"),
|
|
# Wrapper emits output before END_OF_AGENT.
|
|
("root_agent", "model response after tool call"),
|
|
("sub_agent_1", END_OF_AGENT),
|
|
("sub_agent_2", Part.from_function_call(name="test_tool", args={})),
|
|
]
|
|
|
|
|
|
class TestPauseInvocationWithParallelAgent(BasePauseInvocationTest):
|
|
"""Tests pausing invocation with a ParallelAgent."""
|
|
|
|
@pytest.fixture
|
|
def agent(self) -> BaseAgent:
|
|
"""Provides a BaseAgent for the test."""
|
|
sub_agent1 = LlmAgent(
|
|
name="sub_agent_1",
|
|
model=self.mock_model(
|
|
responses=[Part.from_function_call(name="test_tool", args={})]
|
|
),
|
|
tools=[LongRunningFunctionTool(func=test_tool)],
|
|
)
|
|
sub_agent2 = _TestingAgent(
|
|
name="sub_agent_2",
|
|
delay=0.5,
|
|
)
|
|
return ParallelAgent(
|
|
name="root_agent",
|
|
sub_agents=[sub_agent1, sub_agent2],
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
def test_pause_on_long_running_function_call(
|
|
self,
|
|
runner: testing_utils.InMemoryRunner,
|
|
):
|
|
"""Tests that a ParallelAgent pauses on long running function call."""
|
|
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):
|
|
"""Tests pausing invocation with a nested ParallelAgent."""
|
|
|
|
@pytest.fixture
|
|
def agent(self) -> BaseAgent:
|
|
"""Provides a BaseAgent for the test."""
|
|
nested_sub_agent_1 = LlmAgent(
|
|
name="nested_sub_agent_1",
|
|
model=self.mock_model(
|
|
responses=[Part.from_function_call(name="test_tool", args={})]
|
|
),
|
|
tools=[LongRunningFunctionTool(func=test_tool)],
|
|
)
|
|
nested_sub_agent_2 = _TestingAgent(
|
|
name="nested_sub_agent_2",
|
|
delay=0.5,
|
|
)
|
|
nested_parallel_agent = ParallelAgent(
|
|
name="nested_parallel_agent",
|
|
sub_agents=[nested_sub_agent_1, nested_sub_agent_2],
|
|
)
|
|
sub_agent_1 = _TestingAgent(
|
|
name="sub_agent_1",
|
|
delay=0.5,
|
|
)
|
|
return ParallelAgent(
|
|
name="root_agent",
|
|
sub_agents=[sub_agent_1, nested_parallel_agent],
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
def test_pause_on_long_running_function_call(
|
|
self,
|
|
runner: testing_utils.InMemoryRunner,
|
|
):
|
|
"""Tests that a nested ParallelAgent pauses on long running function call."""
|
|
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
|
|
def test_pause_on_multiple_long_running_function_calls(
|
|
self,
|
|
):
|
|
"""Tests that a ParallelAgent pauses on long running function calls."""
|
|
nested_sub_agent_1 = LlmAgent(
|
|
name="nested_sub_agent_1",
|
|
model=self.mock_model(
|
|
responses=[Part.from_function_call(name="test_tool", args={})]
|
|
),
|
|
tools=[LongRunningFunctionTool(func=test_tool)],
|
|
)
|
|
nested_sub_agent_2 = _TestingAgent(
|
|
name="nested_sub_agent_2",
|
|
delay=0.5,
|
|
)
|
|
nested_parallel_agent = ParallelAgent(
|
|
name="nested_parallel_agent",
|
|
sub_agents=[nested_sub_agent_1, nested_sub_agent_2],
|
|
)
|
|
sub_agent_1 = LlmAgent(
|
|
name="sub_agent_1",
|
|
model=self.mock_model(
|
|
responses=[
|
|
Part.from_function_call(name="test_tool", args={}),
|
|
]
|
|
),
|
|
tools=[LongRunningFunctionTool(func=test_tool)],
|
|
)
|
|
agent = ParallelAgent(
|
|
name="root_agent",
|
|
sub_agents=[sub_agent_1, nested_parallel_agent],
|
|
)
|
|
app = App(
|
|
name="test_app",
|
|
root_agent=agent,
|
|
resumability_config=ResumabilityConfig(is_resumable=True),
|
|
)
|
|
runner = testing_utils.InMemoryRunner(app=app)
|
|
simplified_events = testing_utils.simplify_resumable_app_events(
|
|
runner.run("test")
|
|
)
|
|
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):
|
|
"""Tests pausing invocation with a LoopAgent."""
|
|
|
|
@pytest.fixture
|
|
def agent(self) -> BaseAgent:
|
|
"""Provides a BaseAgent for the test."""
|
|
sub_agent_1 = LlmAgent(
|
|
name="sub_agent_1",
|
|
model=self.mock_model(
|
|
responses=[
|
|
Part.from_text(text="sub agent 1 response"),
|
|
]
|
|
),
|
|
)
|
|
sub_agent_2 = LlmAgent(
|
|
name="sub_agent_2",
|
|
model=self.mock_model(
|
|
responses=[
|
|
Part.from_function_call(name="test_tool", args={}),
|
|
]
|
|
),
|
|
tools=[LongRunningFunctionTool(func=test_tool)],
|
|
)
|
|
sub_agent_3 = LlmAgent(
|
|
name="sub_agent_3",
|
|
model=self.mock_model(
|
|
responses=[
|
|
Part.from_function_call(name="exit_loop", args={}),
|
|
]
|
|
),
|
|
tools=[exit_loop],
|
|
)
|
|
return LoopAgent(
|
|
name="root_agent",
|
|
sub_agents=[sub_agent_1, sub_agent_2, sub_agent_3],
|
|
max_iterations=2,
|
|
)
|
|
|
|
@pytest.mark.xfail(
|
|
reason=(
|
|
"Tests implementation details that are different in V2 and will be"
|
|
" deprecated."
|
|
)
|
|
)
|
|
@pytest.mark.asyncio
|
|
def test_pause_on_long_running_function_call(
|
|
self,
|
|
runner: testing_utils.InMemoryRunner,
|
|
):
|
|
"""Tests that a LoopAgent pauses on long running function call."""
|
|
actual = testing_utils.simplify_resumable_app_events(runner.run("test"))
|
|
behavioral = [e for e in actual if not isinstance(e[1], dict)]
|
|
assert behavioral == [
|
|
("sub_agent_1", "sub agent 1 response"),
|
|
# Wrapper emits output before END_OF_AGENT.
|
|
("root_agent", "sub agent 1 response"),
|
|
("sub_agent_1", END_OF_AGENT),
|
|
("sub_agent_2", Part.from_function_call(name="test_tool", args={})),
|
|
]
|
|
|
|
|
|
class TestPauseInvocationWithLlmAgentTree(BasePauseInvocationTest):
|
|
"""Tests the pausing invocation with a tree of LlmAgents."""
|
|
|
|
@pytest.fixture
|
|
def agent(self) -> LlmAgent:
|
|
"""Provides an LlmAgent with sub-agents for the test."""
|
|
sub_llm_agent_1 = LlmAgent(
|
|
name="sub_llm_agent_1",
|
|
model=self.mock_model(
|
|
responses=[
|
|
_transfer_call_part("sub_llm_agent_2"),
|
|
"llm response not used",
|
|
]
|
|
),
|
|
)
|
|
sub_llm_agent_2 = LlmAgent(
|
|
name="sub_llm_agent_2",
|
|
model=self.mock_model(
|
|
responses=[
|
|
Part.from_function_call(name="test_tool", args={}),
|
|
"llm response not used",
|
|
]
|
|
),
|
|
tools=[LongRunningFunctionTool(func=test_tool)],
|
|
)
|
|
return LlmAgent(
|
|
name="root_agent",
|
|
model=self.mock_model(
|
|
responses=[
|
|
_transfer_call_part("sub_llm_agent_1"),
|
|
"llm response not used",
|
|
]
|
|
),
|
|
sub_agents=[sub_llm_agent_1, sub_llm_agent_2],
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
def test_pause_on_long_running_function_call(
|
|
self,
|
|
runner: testing_utils.InMemoryRunner,
|
|
):
|
|
"""Tests that a tree of resumable LlmAgents yields checkpoint events."""
|
|
actual = testing_utils.simplify_resumable_app_events(runner.run("test"))
|
|
behavioral = [e for e in actual if not isinstance(e[1], dict)]
|
|
assert behavioral == [
|
|
("root_agent", _transfer_call_part("sub_llm_agent_1")),
|
|
("root_agent", _TRANSFER_RESPONSE_PART),
|
|
("root_agent", END_OF_AGENT),
|
|
("sub_llm_agent_1", _transfer_call_part("sub_llm_agent_2")),
|
|
("sub_llm_agent_1", _TRANSFER_RESPONSE_PART),
|
|
("sub_llm_agent_1", END_OF_AGENT),
|
|
("sub_llm_agent_2", Part.from_function_call(name="test_tool", args={})),
|
|
]
|
|
|
|
|
|
class TestPauseInvocationWithWithTransferLoop(BasePauseInvocationTest):
|
|
"""Tests pausing the invocation when the agent transfer forms a loop."""
|
|
|
|
@pytest.fixture
|
|
def agent(self) -> LlmAgent:
|
|
"""Provides an LlmAgent with sub-agents for the test."""
|
|
sub_llm_agent_1 = LlmAgent(
|
|
name="sub_llm_agent_1",
|
|
model=self.mock_model(
|
|
responses=[
|
|
_transfer_call_part("sub_llm_agent_2"),
|
|
"llm response not used",
|
|
]
|
|
),
|
|
)
|
|
sub_llm_agent_2 = LlmAgent(
|
|
name="sub_llm_agent_2",
|
|
model=self.mock_model(
|
|
responses=[
|
|
_transfer_call_part("root_agent"),
|
|
"llm response not used",
|
|
]
|
|
),
|
|
)
|
|
return LlmAgent(
|
|
name="root_agent",
|
|
model=self.mock_model(
|
|
responses=[
|
|
_transfer_call_part("sub_llm_agent_1"),
|
|
Part.from_function_call(name="test_tool", args={}),
|
|
"llm response not used",
|
|
]
|
|
),
|
|
sub_agents=[sub_llm_agent_1, sub_llm_agent_2],
|
|
tools=[LongRunningFunctionTool(func=test_tool)],
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
def test_pause_on_long_running_function_call(
|
|
self,
|
|
runner: testing_utils.InMemoryRunner,
|
|
):
|
|
"""Tests that a tree of resumable LlmAgents yields checkpoint events."""
|
|
actual = testing_utils.simplify_resumable_app_events(runner.run("test"))
|
|
behavioral = [e for e in actual if not isinstance(e[1], dict)]
|
|
assert behavioral == [
|
|
("root_agent", _transfer_call_part("sub_llm_agent_1")),
|
|
("root_agent", _TRANSFER_RESPONSE_PART),
|
|
("root_agent", END_OF_AGENT),
|
|
("sub_llm_agent_1", _transfer_call_part("sub_llm_agent_2")),
|
|
("sub_llm_agent_1", _TRANSFER_RESPONSE_PART),
|
|
("sub_llm_agent_1", END_OF_AGENT),
|
|
("sub_llm_agent_2", _transfer_call_part("root_agent")),
|
|
("sub_llm_agent_2", _TRANSFER_RESPONSE_PART),
|
|
("sub_llm_agent_2", END_OF_AGENT),
|
|
("root_agent", Part.from_function_call(name="test_tool", args={})),
|
|
]
|