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
1004 lines
30 KiB
Python
1004 lines
30 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.
|
|
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
from typing import Any
|
|
|
|
from google.adk.agents.llm_agent import LlmAgent
|
|
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._node_tool import NodeTool
|
|
from google.adk.tools.long_running_tool import LongRunningFunctionTool
|
|
from google.adk.workflow import JoinNode
|
|
from google.adk.workflow import node
|
|
from google.adk.workflow import START
|
|
from google.adk.workflow._node_status import NodeStatus
|
|
from google.adk.workflow._workflow import Workflow
|
|
from google.adk.workflow.utils._workflow_hitl_utils import create_request_input_response
|
|
from google.adk.workflow.utils._workflow_hitl_utils import get_request_input_interrupt_ids
|
|
from google.adk.workflow.utils._workflow_hitl_utils import REQUEST_INPUT_FUNCTION_CALL_NAME
|
|
from google.genai import types
|
|
from pydantic import BaseModel
|
|
from pydantic import Field
|
|
import pytest
|
|
|
|
from . import workflow_testing_utils
|
|
from .. import testing_utils
|
|
from .workflow_testing_utils import RequestInputNode
|
|
|
|
|
|
class UserInfo(BaseModel):
|
|
name: str
|
|
age: int
|
|
|
|
|
|
class DummyRequest(BaseModel):
|
|
request: str = ''
|
|
|
|
|
|
def test_node_tool_requires_input_schema():
|
|
"""NodeTool raises ValueError if wrapped node has no input_schema."""
|
|
wf = Workflow(name='no_schema_wf', edges=[])
|
|
with pytest.raises(ValueError, match='does not have an input_schema defined'):
|
|
NodeTool(node=wf)
|
|
|
|
|
|
@pytest.mark.skip(reason='Requires CL 2 subagent branch refactor')
|
|
@pytest.mark.asyncio
|
|
async def test_workflow_as_tool_hitl_resume(request: pytest.FixtureRequest):
|
|
"""Workflow-as-a-tool suspends on RequestInput and resumes successfully.
|
|
|
|
Setup:
|
|
- LlmAgent 'parent_agent' uses WorkflowTool 'collect_user_info_tool'.
|
|
- The tool wraps 'sub_workflow' which has a RequestInputNode and a
|
|
format_response node.
|
|
Act:
|
|
- Turn 1: Run with 'Start task'. The model calls the tool, which suspends.
|
|
- Turn 2: Resume with the user input response to the interrupt.
|
|
Assert:
|
|
- Turn 1: Event history contains the RequestInput function call.
|
|
- Turn 2: The workflow tool resumes and finishes, and parent agent produces
|
|
final text response.
|
|
"""
|
|
# 1. Define the sub-workflow that has an input interrupt
|
|
input_node = RequestInputNode(
|
|
name='input_node',
|
|
message='What is your name and age?',
|
|
response_schema=UserInfo.model_json_schema(),
|
|
)
|
|
|
|
def format_response(node_input: dict[str, Any]):
|
|
yield Event(
|
|
output=f"User {node_input['name']} is {node_input['age']} years old."
|
|
)
|
|
|
|
sub_workflow = Workflow(
|
|
name='sub_workflow',
|
|
edges=[
|
|
(START, input_node),
|
|
(input_node, format_response),
|
|
],
|
|
)
|
|
sub_workflow.input_schema = DummyRequest
|
|
|
|
# 2. Wrap the sub-workflow as a WorkflowTool
|
|
wf_tool = NodeTool(
|
|
node=sub_workflow,
|
|
name='collect_user_info_tool',
|
|
description='Call this tool to collect customer name and age.',
|
|
)
|
|
|
|
# 3. Define the parent agent that calls this tool
|
|
# In the first turn, the model decides to call the tool.
|
|
# In the second turn, after the tool resumes and returns output, the model replies to the user.
|
|
parent_agent = LlmAgent(
|
|
name='parent_agent',
|
|
model=testing_utils.MockModel.create(
|
|
responses=[
|
|
types.Part.from_function_call(
|
|
name='collect_user_info_tool',
|
|
args={},
|
|
),
|
|
types.Part.from_text(
|
|
text='Thank you! I received the user details.'
|
|
),
|
|
]
|
|
),
|
|
tools=[wf_tool],
|
|
)
|
|
|
|
# 4. Wrap the parent agent in an App with resumability enabled
|
|
app = App(
|
|
name=request.function.__name__,
|
|
root_agent=parent_agent,
|
|
resumability_config=ResumabilityConfig(is_resumable=True),
|
|
)
|
|
runner = testing_utils.InMemoryRunner(app=app)
|
|
|
|
# Turn 1: Run the agent, triggering the tool call.
|
|
# The sub-workflow starts, hits the RequestInputNode, and suspends.
|
|
user_event = testing_utils.get_user_content('Start task')
|
|
events1 = await runner.run_async(user_event)
|
|
|
|
simplified_events1 = (
|
|
workflow_testing_utils.simplify_events_with_node_and_agent_state(
|
|
copy.deepcopy(events1),
|
|
)
|
|
)
|
|
|
|
# Verify that we got a RequestInput event
|
|
request_input_event = workflow_testing_utils.find_function_call_event(
|
|
events1, REQUEST_INPUT_FUNCTION_CALL_NAME
|
|
)
|
|
assert request_input_event is not None
|
|
args = request_input_event.content.parts[0].function_call.args
|
|
assert args['message'] == 'What is your name and age?'
|
|
|
|
interrupt_id = get_request_input_interrupt_ids(request_input_event)[0]
|
|
invocation_id = request_input_event.invocation_id
|
|
|
|
# Turn 2: Resume with the user input resolving the interrupt.
|
|
user_input = create_request_input_response(
|
|
interrupt_id, {'name': 'Alice', 'age': 25}
|
|
)
|
|
events2 = await runner.run_async(
|
|
new_message=testing_utils.UserContent(user_input),
|
|
invocation_id=invocation_id,
|
|
)
|
|
|
|
simplified_events2 = (
|
|
workflow_testing_utils.simplify_events_with_node_and_agent_state(
|
|
copy.deepcopy(events2),
|
|
)
|
|
)
|
|
|
|
# Verify the tool workflow finished executing, returned the output,
|
|
# and the parent agent LLM produced its final response.
|
|
text_responses = [
|
|
event.content.parts[0].text
|
|
for event in events2
|
|
if event.content and event.content.parts and event.content.parts[0].text
|
|
]
|
|
assert 'Thank you! I received the user details.' in text_responses
|
|
|
|
|
|
@pytest.mark.skip(reason='Requires CL 2 subagent branch refactor')
|
|
@pytest.mark.asyncio
|
|
async def test_workflow_as_tool_hitl_resume_non_resumable_app(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Workflow-as-a-tool suspends and resumes successfully even when the App has resumability disabled."""
|
|
# 1. Define the sub-workflow that has an input interrupt
|
|
input_node = RequestInputNode(
|
|
name='input_node',
|
|
message='What is your name and age?',
|
|
response_schema=UserInfo.model_json_schema(),
|
|
)
|
|
|
|
def format_response(node_input: dict[str, Any]):
|
|
yield Event(
|
|
output=f"User {node_input['name']} is {node_input['age']} years old."
|
|
)
|
|
|
|
sub_workflow = Workflow(
|
|
name='sub_workflow',
|
|
edges=[
|
|
(START, input_node),
|
|
(input_node, format_response),
|
|
],
|
|
)
|
|
sub_workflow.input_schema = DummyRequest
|
|
|
|
# 2. Wrap the sub-workflow as a WorkflowTool
|
|
wf_tool = NodeTool(
|
|
node=sub_workflow,
|
|
name='collect_user_info_tool',
|
|
description='Call this tool to collect customer name and age.',
|
|
)
|
|
|
|
# 3. Define the parent agent that calls this tool
|
|
parent_agent = LlmAgent(
|
|
name='parent_agent',
|
|
model=testing_utils.MockModel.create(
|
|
responses=[
|
|
types.Part.from_function_call(
|
|
name='collect_user_info_tool',
|
|
args={},
|
|
),
|
|
types.Part.from_text(
|
|
text='Thank you! I received the user details.'
|
|
),
|
|
]
|
|
),
|
|
tools=[wf_tool],
|
|
)
|
|
|
|
# 4. Wrap the parent agent in an App with resumability disabled
|
|
app = App(
|
|
name=request.function.__name__,
|
|
root_agent=parent_agent,
|
|
resumability_config=None,
|
|
)
|
|
runner = testing_utils.InMemoryRunner(app=app)
|
|
|
|
# Turn 1: Run the agent, triggering the tool call.
|
|
user_event = testing_utils.get_user_content('Start task')
|
|
events1 = await runner.run_async(user_event)
|
|
|
|
# Verify that we got a RequestInput event
|
|
request_input_event = workflow_testing_utils.find_function_call_event(
|
|
events1, REQUEST_INPUT_FUNCTION_CALL_NAME
|
|
)
|
|
assert request_input_event is not None
|
|
args = request_input_event.content.parts[0].function_call.args
|
|
assert args['message'] == 'What is your name and age?'
|
|
|
|
interrupt_id = get_request_input_interrupt_ids(request_input_event)[0]
|
|
invocation_id = request_input_event.invocation_id
|
|
|
|
# Turn 2: Resume with the user input resolving the interrupt.
|
|
user_input = create_request_input_response(
|
|
interrupt_id, {'name': 'Alice', 'age': 25}
|
|
)
|
|
events2 = await runner.run_async(
|
|
new_message=testing_utils.UserContent(user_input),
|
|
invocation_id=invocation_id,
|
|
)
|
|
|
|
# Verify the tool workflow finished executing, returned the output,
|
|
# and the parent agent LLM produced its final response.
|
|
text_responses = [
|
|
event.content.parts[0].text
|
|
for event in events2
|
|
if event.content and event.content.parts and event.content.parts[0].text
|
|
]
|
|
assert 'Thank you! I received the user details.' in text_responses
|
|
|
|
|
|
def test_node_tool_rejects_agent():
|
|
"""NodeTool raises ValueError if initialized with any BaseAgent."""
|
|
agent = LlmAgent(
|
|
name='my_agent',
|
|
instruction='Answer questions',
|
|
)
|
|
with pytest.raises(ValueError, match='cannot be wrapped as a NodeTool'):
|
|
NodeTool(node=agent)
|
|
|
|
|
|
class GreetRequest(BaseModel):
|
|
request: str
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_function_node_wrapped_as_tool_returns_output(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""NodeTool wraps a function node and returns expected output."""
|
|
|
|
@node
|
|
def greet_node(request: str) -> str:
|
|
return f'Hello, {request}!'
|
|
|
|
greet_node.input_schema = GreetRequest
|
|
greet_tool = NodeTool(node=greet_node, name='greet_tool')
|
|
|
|
parent_agent = LlmAgent(
|
|
name='parent_agent',
|
|
model=testing_utils.MockModel.create(
|
|
responses=[
|
|
types.Part.from_function_call(
|
|
name='greet_tool',
|
|
args={'request': 'world'},
|
|
),
|
|
types.Part.from_text(text='Processed greet.'),
|
|
]
|
|
),
|
|
tools=[greet_tool],
|
|
)
|
|
|
|
app = App(
|
|
name=request.function.__name__,
|
|
root_agent=parent_agent,
|
|
)
|
|
runner = testing_utils.InMemoryRunner(app=app)
|
|
events = await runner.run_async(testing_utils.get_user_content('Greet world'))
|
|
|
|
func_response_events = [
|
|
e
|
|
for e in events
|
|
if e.content and e.content.parts and e.content.parts[0].function_response
|
|
]
|
|
assert len(func_response_events) == 1
|
|
assert func_response_events[0].content.parts[
|
|
0
|
|
].function_response.response == {'result': 'Hello, world!'}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_workflow_tool_with_join_node(request: pytest.FixtureRequest):
|
|
"""WorkflowTool containing a JoinNode works correctly when wrapped as a tool."""
|
|
node_a = workflow_testing_utils.TestingNode(name='NodeA', output={'a': 1})
|
|
node_b = workflow_testing_utils.TestingNode(name='NodeB', output={'b': 2})
|
|
node_join = JoinNode(name='NodeJoin')
|
|
|
|
def format_response(node_input: dict[str, Any]):
|
|
yield Event(
|
|
output=(
|
|
f"A is {node_input['NodeA']['a']} and B is"
|
|
f" {node_input['NodeB']['b']}."
|
|
)
|
|
)
|
|
|
|
sub_workflow = Workflow(
|
|
name='sub_workflow',
|
|
edges=[
|
|
(START, node_a),
|
|
(START, node_b),
|
|
(node_a, node_join),
|
|
(node_b, node_join),
|
|
(node_join, format_response),
|
|
],
|
|
)
|
|
sub_workflow.input_schema = DummyRequest
|
|
|
|
wf_tool = NodeTool(
|
|
node=sub_workflow,
|
|
name='my_join_tool',
|
|
description='Collect parallel items.',
|
|
)
|
|
|
|
parent_agent = LlmAgent(
|
|
name='parent_agent',
|
|
model=testing_utils.MockModel.create(
|
|
responses=[
|
|
types.Part.from_function_call(
|
|
name='my_join_tool',
|
|
args={},
|
|
),
|
|
types.Part.from_text(text='Done.'),
|
|
]
|
|
),
|
|
tools=[wf_tool],
|
|
)
|
|
|
|
app = App(
|
|
name=request.function.__name__,
|
|
root_agent=parent_agent,
|
|
)
|
|
runner = testing_utils.InMemoryRunner(app=app)
|
|
events = await runner.run_async(testing_utils.get_user_content('Run join'))
|
|
|
|
func_response_events = [
|
|
e
|
|
for e in events
|
|
if e.content and e.content.parts and e.content.parts[0].function_response
|
|
]
|
|
assert len(func_response_events) == 1
|
|
assert func_response_events[0].content.parts[
|
|
0
|
|
].function_response.response == {'result': 'A is 1 and B is 2.'}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_workflow_tool_with_dynamic_node(request: pytest.FixtureRequest):
|
|
"""WorkflowTool containing a dynamic node schedules and executes it correctly."""
|
|
|
|
@node
|
|
async def child(*, ctx, node_input):
|
|
yield f'child got: {node_input}'
|
|
|
|
@node(rerun_on_resume=True)
|
|
async def parent_node(*, ctx, node_input):
|
|
result = await ctx.run_node(child, node_input='hello')
|
|
yield f'parent got: {result}'
|
|
|
|
sub_workflow = Workflow(
|
|
name='sub_workflow',
|
|
edges=[
|
|
(START, parent_node),
|
|
],
|
|
)
|
|
sub_workflow.input_schema = DummyRequest
|
|
|
|
wf_tool = NodeTool(
|
|
node=sub_workflow,
|
|
name='my_dynamic_tool',
|
|
description='Call dynamic node.',
|
|
)
|
|
|
|
parent_agent = LlmAgent(
|
|
name='parent_agent',
|
|
model=testing_utils.MockModel.create(
|
|
responses=[
|
|
types.Part.from_function_call(
|
|
name='my_dynamic_tool',
|
|
args={},
|
|
),
|
|
types.Part.from_text(text='Done.'),
|
|
]
|
|
),
|
|
tools=[wf_tool],
|
|
)
|
|
|
|
app = App(
|
|
name=request.function.__name__,
|
|
root_agent=parent_agent,
|
|
)
|
|
runner = testing_utils.InMemoryRunner(app=app)
|
|
events = await runner.run_async(testing_utils.get_user_content('Run dynamic'))
|
|
|
|
func_response_events = [
|
|
e
|
|
for e in events
|
|
if e.content and e.content.parts and e.content.parts[0].function_response
|
|
]
|
|
assert len(func_response_events) == 1
|
|
assert func_response_events[0].content.parts[
|
|
0
|
|
].function_response.response == {'result': 'parent got: child got: hello'}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_workflow_tool_with_nested_workflows(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""WorkflowTool wrapping a nested workflow executes successfully."""
|
|
inner_node = workflow_testing_utils.TestingNode(
|
|
name='inner_node', output='inner_output'
|
|
)
|
|
inner_wf = Workflow(
|
|
name='inner_wf',
|
|
edges=[
|
|
(START, inner_node),
|
|
],
|
|
)
|
|
inner_wf.input_schema = None
|
|
|
|
outer_node = workflow_testing_utils.TestingNode(
|
|
name='outer_node', output='outer_output'
|
|
)
|
|
outer_wf = Workflow(
|
|
name='outer_wf',
|
|
edges=[
|
|
(START, outer_node, inner_wf),
|
|
],
|
|
)
|
|
outer_wf.input_schema = DummyRequest
|
|
|
|
wf_tool = NodeTool(
|
|
node=outer_wf,
|
|
name='nested_wf_tool',
|
|
description='Call nested workflow.',
|
|
)
|
|
|
|
parent_agent = LlmAgent(
|
|
name='parent_agent',
|
|
model=testing_utils.MockModel.create(
|
|
responses=[
|
|
types.Part.from_function_call(
|
|
name='nested_wf_tool',
|
|
args={},
|
|
),
|
|
types.Part.from_text(text='Done.'),
|
|
]
|
|
),
|
|
tools=[wf_tool],
|
|
)
|
|
|
|
app = App(
|
|
name=request.function.__name__,
|
|
root_agent=parent_agent,
|
|
)
|
|
runner = testing_utils.InMemoryRunner(app=app)
|
|
events = await runner.run_async(testing_utils.get_user_content('Run nested'))
|
|
|
|
func_response_events = [
|
|
e
|
|
for e in events
|
|
if e.content and e.content.parts and e.content.parts[0].function_response
|
|
]
|
|
assert len(func_response_events) == 1
|
|
assert func_response_events[0].content.parts[
|
|
0
|
|
].function_response.response == {'result': 'inner_output'}
|
|
|
|
|
|
@pytest.mark.skip(reason='Requires CL 2 subagent branch refactor')
|
|
@pytest.mark.asyncio
|
|
async def test_workflow_tool_with_dynamic_node_hitl_resume(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""WorkflowTool with a dynamic node containing HITL suspends and resumes successfully."""
|
|
# 1. Define dynamic node calling a child RequestInputNode
|
|
input_node = RequestInputNode(
|
|
name='input_node',
|
|
message='Enter value:',
|
|
response_schema=UserInfo.model_json_schema(),
|
|
)
|
|
|
|
@node(rerun_on_resume=True)
|
|
async def parent_node(*, ctx, node_input):
|
|
result = await ctx.run_node(input_node)
|
|
yield f'parent got: {result["name"]}'
|
|
|
|
sub_workflow = Workflow(
|
|
name='sub_workflow',
|
|
edges=[
|
|
(START, parent_node),
|
|
],
|
|
)
|
|
sub_workflow.input_schema = DummyRequest
|
|
|
|
# 2. Wrap as WorkflowTool
|
|
wf_tool = NodeTool(
|
|
node=sub_workflow,
|
|
name='my_dynamic_hitl_tool',
|
|
description='Call dynamic HITL node.',
|
|
)
|
|
|
|
# 3. Define parent agent
|
|
parent_agent = LlmAgent(
|
|
name='parent_agent',
|
|
model=testing_utils.MockModel.create(
|
|
responses=[
|
|
types.Part.from_function_call(
|
|
name='my_dynamic_hitl_tool',
|
|
args={},
|
|
),
|
|
types.Part.from_text(text='Task completed.'),
|
|
]
|
|
),
|
|
tools=[wf_tool],
|
|
)
|
|
|
|
# 4. App with resumability enabled
|
|
app = App(
|
|
name=request.function.__name__,
|
|
root_agent=parent_agent,
|
|
resumability_config=ResumabilityConfig(is_resumable=True),
|
|
)
|
|
runner = testing_utils.InMemoryRunner(app=app)
|
|
|
|
# Turn 1: Run the agent, triggering the tool call and dynamic node suspend.
|
|
user_event = testing_utils.get_user_content('Start')
|
|
events1 = await runner.run_async(user_event)
|
|
|
|
request_input_event = workflow_testing_utils.find_function_call_event(
|
|
events1, REQUEST_INPUT_FUNCTION_CALL_NAME
|
|
)
|
|
assert request_input_event is not None
|
|
interrupt_id = get_request_input_interrupt_ids(request_input_event)[0]
|
|
invocation_id = request_input_event.invocation_id
|
|
|
|
# Turn 2: Resume with the user input response.
|
|
user_input = create_request_input_response(
|
|
interrupt_id, {'name': 'Bob', 'age': 30}
|
|
)
|
|
events2 = await runner.run_async(
|
|
new_message=testing_utils.UserContent(user_input),
|
|
invocation_id=invocation_id,
|
|
)
|
|
|
|
# Verify the tool workflow finished executing, returned output,
|
|
# and parent agent replied.
|
|
text_responses = [
|
|
event.content.parts[0].text
|
|
for event in events2
|
|
if event.content and event.content.parts and event.content.parts[0].text
|
|
]
|
|
assert 'Task completed.' in text_responses
|
|
|
|
|
|
@pytest.mark.skip(
|
|
reason='Known framework issue with MockModel nested HITL in sub-workflow'
|
|
)
|
|
@pytest.mark.asyncio
|
|
async def test_workflow_as_tool_nested_hitl(request: pytest.FixtureRequest):
|
|
"""Parent LLM agent -> workflow -> LLM agent -> NodeTool(HITL) propagation."""
|
|
# 1. Define the deepest node that raises RequestInput
|
|
input_node = RequestInputNode(
|
|
name='deep_input_node',
|
|
message='Give me some input:',
|
|
response_schema={
|
|
'type': 'object',
|
|
'properties': {'val': {'type': 'string'}},
|
|
},
|
|
)
|
|
|
|
# 2. Wrap it as a NodeTool
|
|
input_node.input_schema = DummyRequest
|
|
node_tool = NodeTool(node=input_node, name='my_node_tool')
|
|
|
|
# 3. Define the child agent that uses this NodeTool
|
|
child_agent = LlmAgent(
|
|
name='child_agent',
|
|
model=testing_utils.MockModel.create(
|
|
responses=[
|
|
types.Part.from_function_call(
|
|
name='my_node_tool',
|
|
args={},
|
|
),
|
|
types.Part.from_text(text='Child agent processed.'),
|
|
]
|
|
),
|
|
tools=[node_tool],
|
|
)
|
|
|
|
# 4. Define the sub-workflow containing the child agent
|
|
sub_workflow = Workflow(
|
|
name='sub_workflow',
|
|
edges=[
|
|
(START, child_agent),
|
|
],
|
|
)
|
|
sub_workflow.input_schema = DummyRequest
|
|
|
|
# 5. Wrap the sub-workflow as a WorkflowTool
|
|
wf_tool = NodeTool(
|
|
node=sub_workflow,
|
|
name='my_wf_tool',
|
|
description='Call sub workflow.',
|
|
)
|
|
|
|
# 6. Define the parent agent that calls the WorkflowTool
|
|
parent_agent = LlmAgent(
|
|
name='parent_agent',
|
|
model=testing_utils.MockModel.create(
|
|
responses=[
|
|
types.Part.from_function_call(
|
|
name='my_wf_tool',
|
|
args={},
|
|
),
|
|
types.Part.from_text(text='Parent agent finished successfully.'),
|
|
]
|
|
),
|
|
tools=[wf_tool],
|
|
)
|
|
|
|
# 7. Wrap in App and Runner
|
|
app = App(
|
|
name=request.function.__name__,
|
|
root_agent=parent_agent,
|
|
resumability_config=ResumabilityConfig(is_resumable=True),
|
|
)
|
|
runner = testing_utils.InMemoryRunner(app=app)
|
|
|
|
# Turn 1: Run
|
|
events1 = await runner.run_async(testing_utils.get_user_content('Start task'))
|
|
|
|
# Assert Turn 1: Expect RequestInput event
|
|
request_input_event = workflow_testing_utils.find_function_call_event(
|
|
events1, REQUEST_INPUT_FUNCTION_CALL_NAME
|
|
)
|
|
assert request_input_event is not None
|
|
args = request_input_event.content.parts[0].function_call.args
|
|
assert args['message'] == 'Give me some input:'
|
|
|
|
interrupt_id = get_request_input_interrupt_ids(request_input_event)[0]
|
|
invocation_id = request_input_event.invocation_id
|
|
|
|
# Turn 2: Resume
|
|
user_input = create_request_input_response(interrupt_id, {'val': 'hello'})
|
|
events2 = await runner.run_async(
|
|
new_message=testing_utils.UserContent(user_input),
|
|
invocation_id=invocation_id,
|
|
)
|
|
|
|
# Assert Turn 2: Expect completion
|
|
text_responses = [
|
|
event.content.parts[0].text
|
|
for event in events2
|
|
if event.content and event.content.parts and event.content.parts[0].text
|
|
]
|
|
assert 'Parent agent finished successfully.' in text_responses
|
|
|
|
|
|
@pytest.mark.skip(
|
|
reason='Known framework issue with MockModel multi-HITL in nested workflow'
|
|
)
|
|
@pytest.mark.asyncio
|
|
async def test_workflow_as_tool_nested_multi_hitl(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Parent LLM agent -> workflow -> LLM agent -> NodeTool(HITL) twice."""
|
|
# 1. Define the deepest node that raises RequestInput
|
|
input_node = RequestInputNode(
|
|
name='deep_input_node',
|
|
message='Give me some input:',
|
|
response_schema={
|
|
'type': 'object',
|
|
'properties': {'val': {'type': 'string'}},
|
|
},
|
|
)
|
|
|
|
# 2. Wrap it as a NodeTool
|
|
input_node.input_schema = DummyRequest
|
|
node_tool = NodeTool(node=input_node, name='my_node_tool')
|
|
|
|
# 3. Define the child agent that uses this NodeTool twice
|
|
child_agent = LlmAgent(
|
|
name='child_agent',
|
|
model=testing_utils.MockModel.create(
|
|
responses=[
|
|
types.Part.from_function_call(
|
|
name='my_node_tool',
|
|
args={},
|
|
),
|
|
types.Part.from_function_call(
|
|
name='my_node_tool',
|
|
args={},
|
|
),
|
|
types.Part.from_text(text='Child agent finished.'),
|
|
]
|
|
),
|
|
tools=[node_tool],
|
|
)
|
|
|
|
# 4. Define the sub-workflow containing the child agent
|
|
sub_workflow = Workflow(
|
|
name='sub_workflow',
|
|
edges=[
|
|
(START, child_agent),
|
|
],
|
|
)
|
|
sub_workflow.input_schema = DummyRequest
|
|
|
|
# 5. Wrap the sub-workflow as a WorkflowTool
|
|
wf_tool = NodeTool(
|
|
node=sub_workflow,
|
|
name='my_wf_tool',
|
|
description='Call sub workflow.',
|
|
)
|
|
|
|
# 6. Define the parent agent that calls the WorkflowTool
|
|
parent_agent = LlmAgent(
|
|
name='parent_agent',
|
|
model=testing_utils.MockModel.create(
|
|
responses=[
|
|
types.Part.from_function_call(
|
|
name='my_wf_tool',
|
|
args={},
|
|
),
|
|
types.Part.from_text(text='Parent agent finished successfully.'),
|
|
]
|
|
),
|
|
tools=[wf_tool],
|
|
)
|
|
|
|
# 7. Wrap in App and Runner
|
|
app = App(
|
|
name=request.function.__name__,
|
|
root_agent=parent_agent,
|
|
resumability_config=ResumabilityConfig(is_resumable=True),
|
|
)
|
|
runner = testing_utils.InMemoryRunner(app=app)
|
|
|
|
# Turn 1: Run -> triggers first HITL
|
|
events1 = await runner.run_async(testing_utils.get_user_content('Start task'))
|
|
request_input_event1 = workflow_testing_utils.find_function_call_event(
|
|
events1, REQUEST_INPUT_FUNCTION_CALL_NAME
|
|
)
|
|
assert request_input_event1 is not None
|
|
interrupt_id1 = get_request_input_interrupt_ids(request_input_event1)[0]
|
|
invocation_id = request_input_event1.invocation_id
|
|
|
|
# Turn 2: Resume first HITL -> triggers second HITL
|
|
user_input1 = create_request_input_response(interrupt_id1, {'val': 'hello'})
|
|
events2 = await runner.run_async(
|
|
new_message=testing_utils.UserContent(user_input1),
|
|
invocation_id=invocation_id,
|
|
)
|
|
request_input_event2 = workflow_testing_utils.find_function_call_event(
|
|
events2, REQUEST_INPUT_FUNCTION_CALL_NAME
|
|
)
|
|
assert request_input_event2 is not None
|
|
interrupt_id2 = get_request_input_interrupt_ids(request_input_event2)[0]
|
|
assert interrupt_id1 != interrupt_id2
|
|
|
|
# Turn 3: Resume second HITL -> finishes
|
|
user_input2 = create_request_input_response(interrupt_id2, {'val': 'world'})
|
|
events3 = await runner.run_async(
|
|
new_message=testing_utils.UserContent(user_input2),
|
|
invocation_id=invocation_id,
|
|
)
|
|
|
|
# Assert Turn 3: Expect completion
|
|
text_responses = [
|
|
event.content.parts[0].text
|
|
for event in events3
|
|
if event.content and event.content.parts and event.content.parts[0].text
|
|
]
|
|
assert 'Parent agent finished successfully.' in text_responses
|
|
|
|
|
|
@pytest.mark.skip(reason='Requires CL 2 subagent branch refactor')
|
|
@pytest.mark.asyncio
|
|
async def test_workflow_as_tool_nested_lro(request: pytest.FixtureRequest):
|
|
"""Parent LLM agent -> workflow -> LLM agent -> LRO tool."""
|
|
|
|
# 1. Define LRO tool function
|
|
def my_lro_func():
|
|
return None
|
|
|
|
# 2. Define child agent with LRO tool
|
|
child_agent = LlmAgent(
|
|
name='child_agent',
|
|
model=testing_utils.MockModel.create(
|
|
responses=[
|
|
types.Part.from_function_call(
|
|
name='my_lro_func',
|
|
args={},
|
|
),
|
|
types.Part.from_text(text='Child agent finished after LRO.'),
|
|
]
|
|
),
|
|
tools=[LongRunningFunctionTool(func=my_lro_func)],
|
|
)
|
|
|
|
# 3. Define sub-workflow
|
|
sub_workflow = Workflow(
|
|
name='sub_workflow',
|
|
edges=[
|
|
(START, child_agent),
|
|
],
|
|
)
|
|
sub_workflow.input_schema = DummyRequest
|
|
|
|
# 4. Wrap as WorkflowTool
|
|
wf_tool = NodeTool(
|
|
node=sub_workflow,
|
|
name='my_wf_tool',
|
|
description='Call sub workflow.',
|
|
)
|
|
|
|
# 5. Define parent agent
|
|
parent_agent = LlmAgent(
|
|
name='parent_agent',
|
|
model=testing_utils.MockModel.create(
|
|
responses=[
|
|
types.Part.from_function_call(
|
|
name='my_wf_tool',
|
|
args={},
|
|
),
|
|
types.Part.from_text(text='Parent agent finished successfully.'),
|
|
]
|
|
),
|
|
tools=[wf_tool],
|
|
)
|
|
|
|
# 6. Wrap in App and Runner
|
|
app = App(
|
|
name=request.function.__name__,
|
|
root_agent=parent_agent,
|
|
resumability_config=ResumabilityConfig(is_resumable=True),
|
|
)
|
|
runner = testing_utils.InMemoryRunner(app=app)
|
|
|
|
# Turn 1: Run -> should pause on LRO
|
|
events1 = await runner.run_async(testing_utils.get_user_content('Start task'))
|
|
assert any(e.long_running_tool_ids for e in events1)
|
|
|
|
invocation_id = events1[0].invocation_id
|
|
fc_event = workflow_testing_utils.find_function_call_event(
|
|
events1, 'my_lro_func'
|
|
)
|
|
assert fc_event is not None
|
|
function_call_id = fc_event.content.parts[0].function_call.id
|
|
|
|
# Turn 2: Resume with LRO response
|
|
tool_response = testing_utils.UserContent(
|
|
types.Part(
|
|
function_response=types.FunctionResponse(
|
|
id=function_call_id,
|
|
name='my_lro_func',
|
|
response={'result': 'LRO finished'},
|
|
)
|
|
)
|
|
)
|
|
events2 = await runner.run_async(
|
|
new_message=tool_response,
|
|
invocation_id=invocation_id,
|
|
)
|
|
|
|
# Assert Turn 2: Expect completion
|
|
text_responses = [
|
|
event.content.parts[0].text
|
|
for event in events2
|
|
if event.content and event.content.parts and event.content.parts[0].text
|
|
]
|
|
assert 'Parent agent finished successfully.' in text_responses
|
|
|
|
|
|
def test_node_tool_auto_converts_function_node_binding():
|
|
"""NodeTool automatically converts FunctionNode parameter_binding to 'node_input'."""
|
|
|
|
@node
|
|
def my_func_node(request: str) -> str:
|
|
"""A dummy node."""
|
|
return f'Result: {request}'
|
|
|
|
# Originally it is 'state' mode by default
|
|
assert my_func_node.parameter_binding == 'state'
|
|
# input_schema is originally None
|
|
assert getattr(my_func_node, 'input_schema', None) is None
|
|
|
|
# Wrap it
|
|
tool = NodeTool(node=my_func_node)
|
|
|
|
# Check that the wrapped node copy is converted to 'node_input' mode
|
|
assert tool.node.parameter_binding == 'node_input'
|
|
# And input_schema is automatically inferred
|
|
schema = tool.node.input_schema
|
|
assert 'request' in schema['properties']
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_node_tool_primitive_input_schema(request: pytest.FixtureRequest):
|
|
"""NodeTool automatically wraps primitive input_schema to object in declaration and unwraps in run."""
|
|
|
|
def echo_func(node_input: str):
|
|
yield Event(output=f'Echo: {node_input}')
|
|
|
|
sub_workflow = Workflow(
|
|
name='sub_workflow',
|
|
edges=[
|
|
(START, echo_func),
|
|
],
|
|
)
|
|
sub_workflow.input_schema = str
|
|
tool = NodeTool(node=sub_workflow, name='primitive_tool')
|
|
|
|
# 1. Check declaration is wrapped to object schema
|
|
decl = tool._get_declaration()
|
|
assert decl.parameters_json_schema is not None
|
|
assert decl.parameters_json_schema['type'] == 'object'
|
|
assert 'request' in decl.parameters_json_schema['properties']
|
|
assert (
|
|
decl.parameters_json_schema['properties']['request']['type'] == 'string'
|
|
)
|
|
|
|
# 2. Run the tool (passing wrapped argument) and check execution
|
|
parent_agent = LlmAgent(
|
|
name='parent_agent',
|
|
model=testing_utils.MockModel.create(
|
|
responses=[
|
|
types.Part.from_function_call(
|
|
name='primitive_tool',
|
|
args={'request': 'hello_world'},
|
|
),
|
|
types.Part.from_text(text='Finished.'),
|
|
]
|
|
),
|
|
tools=[tool],
|
|
)
|
|
app = App(
|
|
name=request.function.__name__,
|
|
root_agent=parent_agent,
|
|
)
|
|
runner = testing_utils.InMemoryRunner(app=app)
|
|
events = await runner.run_async(testing_utils.get_user_content('Run'))
|
|
|
|
func_response_events = [
|
|
e
|
|
for e in events
|
|
if e.content and e.content.parts and e.content.parts[0].function_response
|
|
]
|
|
assert len(func_response_events) == 1
|
|
assert func_response_events[0].content.parts[
|
|
0
|
|
].function_response.response == {'result': 'Echo: hello_world'}
|