Files
google--adk-python/tests/unittests/workflow/test_node_tool.py
T
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
chore: import upstream snapshot with attribution
2026-07-13 13:25:13 +08:00

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'}