From 1263ed64e30805464fff3391554f65ebbf72746b Mon Sep 17 00:00:00 2001 From: Shangjie Chen Date: Mon, 6 Jul 2026 14:32:19 -0700 Subject: [PATCH] feat: Implement Workflow as Tool core feature Introduce NodeTool, allowing individual Nodes and entire Workflows to be wrapped and executed as standard ADK Tools. This PR implements the core, single-turn execution capability: - Support auto-wrapping of BaseNode (Workflows) directly in Agent.tools. - Wrapping synchronous and asynchronous function nodes as NodeTools. - Providing a complete sample workflow demonstrating how to run a workflow as a tool. Resumption and multi-turn nested HITL support are skipped in this PR and will be fully enabled in the later PR. Co-authored-by: Shangjie Chen PiperOrigin-RevId: 943499058 --- .../samples/workflows/node_as_tool/README.md | 36 + .../samples/workflows/node_as_tool/agent.py | 73 ++ .../workflows/node_as_tool/tests/go.json | 182 +++ src/google/adk/agents/llm_agent.py | 55 +- src/google/adk/tools/_node_tool.py | 156 +++ src/google/adk/workflow/_function_node.py | 25 +- tests/unittests/workflow/test_node_tool.py | 1003 +++++++++++++++++ 7 files changed, 1526 insertions(+), 4 deletions(-) create mode 100644 contributing/samples/workflows/node_as_tool/README.md create mode 100644 contributing/samples/workflows/node_as_tool/agent.py create mode 100644 contributing/samples/workflows/node_as_tool/tests/go.json create mode 100644 src/google/adk/tools/_node_tool.py create mode 100644 tests/unittests/workflow/test_node_tool.py diff --git a/contributing/samples/workflows/node_as_tool/README.md b/contributing/samples/workflows/node_as_tool/README.md new file mode 100644 index 00000000..605050c2 --- /dev/null +++ b/contributing/samples/workflows/node_as_tool/README.md @@ -0,0 +1,36 @@ +# Node as Tool + +## Overview + +Demonstrates wrapping both a regular ADK `Node` (using the `@node` decorator) and a `Workflow` as tools that can be automatically called by a parent `Agent`. + +In this sample: + +1. The parent agent receives an inquiry about a customer's discount. +1. It invokes `customer_lookup_workflow` (a `Workflow` wrapped as a tool) to retrieve customer status. +1. It then invokes `calculate_discount` (a regular `Node` wrapped as a tool) using the retrieved status. + +## Sample Inputs + +- `What discount does customer c123 get?` + + *The parent agent first invokes `customer_lookup_workflow` to verify status, then invokes `calculate_discount` to determine the discount percentage, and summarizes the results.* + +## Agent Topology Graph + +```mermaid +graph TD + customer_service_agent[customer_service_agent] -->|calls| customer_lookup_workflow(customer_lookup_workflow) + customer_service_agent -->|calls| calculate_discount(calculate_discount) +``` + +## How To + +To expose an existing `Node` or `Workflow` as a tool callable by an `Agent`: + +1. Define your `Node` (or `@node`) or `Workflow` and assign both an `input_schema` and a `description`. +1. Pass the node/workflow directly into your parent agent's `tools` list: `Agent(..., tools=[my_node, my_workflow])`. + +## Related Guides + +- [Workflows](../../../../docs/guides/workflows/workflows.md) - Explains building complex multi-step graphs. diff --git a/contributing/samples/workflows/node_as_tool/agent.py b/contributing/samples/workflows/node_as_tool/agent.py new file mode 100644 index 00000000..9b50feca --- /dev/null +++ b/contributing/samples/workflows/node_as_tool/agent.py @@ -0,0 +1,73 @@ +# 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 + +from typing import Generator + +from google.adk import Agent +from google.adk import Event +from google.adk import Workflow +from google.adk.workflow import node +from pydantic import BaseModel +from pydantic import Field + + +# 1. Define schemas +class CustomerLookupArgs(BaseModel): + user_id: str = Field(description="The customer's unique identifier.") + + +# 2. Define a regular Node using the @node decorator. +# This Node is wrapped as a NodeTool automatically by the Agent. +# As a NodeTool, it has the ability to yield intermediate Events during execution. +@node +def calculate_discount(tier: str, ctx) -> Generator[Event | str, None, None]: + """Calculates the discount percentage based on customer tier. + + Args: + tier: The customer's membership tier (e.g., VIP, Standard). + """ + yield Event(message=f"Checking discount rules for tier '{tier}'...") + discount = "20% off" if "VIP" in tier else "5% off" + yield discount + + +# 3. Define a Workflow. +# This Workflow is wrapped as a NodeTool automatically by the Agent. +def lookup_customer_data(node_input: CustomerLookupArgs, ctx) -> dict[str, str]: + return {"user_id": node_input.user_id, "tier": "Verified VIP Member"} + + +customer_lookup_workflow = Workflow( + name="customer_lookup_workflow", + description="Looks up customer status and tier by user_id.", + input_schema=CustomerLookupArgs, + edges=[ + ("START", lookup_customer_data), + ], +) + + +# 4. Define the Agent that uses both Node and Workflow as tools. +root_agent = Agent( + name="customer_service_agent", + instruction=""" + You are a customer service assistant. + 1. First, call `customer_lookup_workflow` using the user_id to get their membership tier. + 2. Then, call `calculate_discount` node with that tier to find out what discount they get. + Summarize these details for the customer. + """, + tools=[customer_lookup_workflow, calculate_discount], +) diff --git a/contributing/samples/workflows/node_as_tool/tests/go.json b/contributing/samples/workflows/node_as_tool/tests/go.json new file mode 100644 index 00000000..10b36485 --- /dev/null +++ b/contributing/samples/workflows/node_as_tool/tests/go.json @@ -0,0 +1,182 @@ +{ + "appName": "node_as_tool", + "events": [ + { + "author": "user", + "content": { + "parts": [ + { + "text": "What discount does customer c123 get?" + } + ], + "role": "user" + }, + "id": "e-1", + "invocationId": "i-1", + "nodeInfo": { + "path": "" + } + }, + { + "author": "customer_service_agent", + "content": { + "parts": [ + { + "functionCall": { + "args": { + "user_id": "c123" + }, + "id": "fc-1", + "name": "customer_lookup_workflow" + } + } + ], + "role": "model" + }, + "id": "e-2", + "invocationId": "i-1", + "longRunningToolIds": [ + "fc-1" + ], + "nodeInfo": { + "path": "customer_service_agent@1" + } + }, + { + "author": "customer_lookup_workflow", + "branch": "customer_lookup_workflow@fc-1", + "id": "e-3", + "invocationId": "i-1", + "nodeInfo": { + "outputFor": [ + "customer_lookup_workflow@1/lookup_customer_data@1", + "customer_lookup_workflow@1" + ], + "path": "customer_lookup_workflow@1/lookup_customer_data@1" + }, + "output": { + "tier": "Verified VIP Member", + "user_id": "c123" + } + }, + { + "author": "customer_service_agent", + "content": { + "parts": [ + { + "functionResponse": { + "id": "fc-1", + "name": "customer_lookup_workflow", + "response": { + "tier": "Verified VIP Member", + "user_id": "c123" + } + } + } + ], + "role": "user" + }, + "id": "e-4", + "invocationId": "i-1", + "nodeInfo": { + "path": "customer_service_agent@1" + } + }, + { + "author": "customer_service_agent", + "content": { + "parts": [ + { + "functionCall": { + "args": { + "tier": "Verified VIP Member" + }, + "id": "fc-2", + "name": "calculate_discount" + } + } + ], + "role": "model" + }, + "id": "e-5", + "invocationId": "i-1", + "longRunningToolIds": [ + "fc-2" + ], + "nodeInfo": { + "path": "customer_service_agent@1" + } + }, + + { + "author": "calculate_discount", + "branch": "calculate_discount@fc-2", + "content": { + "parts": [ + { + "text": "Checking discount rules for tier 'Verified VIP Member'..." + } + ], + "role": "user" + }, + "id": "e-6", + "invocationId": "i-1", + "nodeInfo": { + "path": "calculate_discount@1" + } + }, + { + "author": "calculate_discount", + "branch": "calculate_discount@fc-2", + "id": "e-7", + "invocationId": "i-1", + "nodeInfo": { + "outputFor": [ + "calculate_discount@1" + ], + "path": "calculate_discount@1" + }, + "output": "20% off" + }, + { + "author": "customer_service_agent", + "content": { + "parts": [ + { + "functionResponse": { + "id": "fc-2", + "name": "calculate_discount", + "response": { + "result": "20% off" + } + } + } + ], + "role": "user" + }, + "id": "e-8", + "invocationId": "i-1", + "nodeInfo": { + "path": "customer_service_agent@1" + } + }, + { + "author": "customer_service_agent", + "content": { + "parts": [ + { + "text": "Customer c123 is a Verified VIP Member and gets a 20% discount." + } + ], + "role": "model" + }, + "id": "e-9", + "invocationId": "i-1", + "nodeInfo": { + "path": "customer_service_agent@1" + } + } + ], + "id": "12345678-1234-1234-1234-123456789abc", + "userId": "user" +} diff --git a/src/google/adk/agents/llm_agent.py b/src/google/adk/agents/llm_agent.py index 6f0c2af7..97c91748 100644 --- a/src/google/adk/agents/llm_agent.py +++ b/src/google/adk/agents/llm_agent.py @@ -133,7 +133,6 @@ OnToolErrorCallback: TypeAlias = Union[ InstructionProvider: TypeAlias = Callable[ [ReadonlyContext], Union[str, Awaitable[str]] ] - ToolUnion: TypeAlias = Union[Callable, BaseTool, BaseToolset] @@ -175,6 +174,32 @@ async def _convert_tool_union_to_tools( max_results=vais_tool.max_results, ) ] + from ..workflow._base_node import BaseNode + + if isinstance(tool_union, BaseNode): + from ..tools._node_tool import NodeTool + from .base_agent import BaseAgent + + if isinstance(tool_union, BaseAgent): + raise ValueError( + f"Agent '{tool_union.name}' cannot be wrapped as a NodeTool. Agents" + ' should be invoked as sub-agents.' + ) + + description = tool_union.description + if not description: + raise ValueError( + f"Workflow/Node '{tool_union.name}' must have a description to be" + ' wrapped as a tool.' + ) + + return [ + NodeTool( + node=tool_union, + name=tool_union.name, + description=description, + ) + ] if isinstance(tool_union, BaseTool): return [tool_union] @@ -1011,6 +1036,34 @@ class LlmAgent(BaseAgent, abc.ABC): event.actions.state_delta[self.output_key] = accumulator return accumulator + @model_validator(mode='before') + @classmethod + def _pre_validate_tools(cls, data: Any) -> Any: + if isinstance(data, dict) and 'tools' in data and data['tools']: + from google.adk.agents.base_agent import BaseAgent + from google.adk.tools._node_tool import NodeTool + from google.adk.workflow._base_node import BaseNode + + new_tools = [] + for t in data['tools']: + if isinstance(t, BaseAgent): + raise ValueError( + f"Agent '{t.name}' cannot be wrapped as a NodeTool. Agents should" + ' be invoked as sub-agents.' + ) + elif isinstance(t, BaseNode): + description = t.description + if not description: + raise ValueError( + f"Workflow/Node '{t.name}' must have a description to be" + ' wrapped as a tool.' + ) + new_tools.append(NodeTool(node=t, description=description)) + else: + new_tools.append(t) + data['tools'] = new_tools + return data + @model_validator(mode='after') def __model_validator_after(self) -> LlmAgent: return self diff --git a/src/google/adk/tools/_node_tool.py b/src/google/adk/tools/_node_tool.py new file mode 100644 index 00000000..f69c0ef3 --- /dev/null +++ b/src/google/adk/tools/_node_tool.py @@ -0,0 +1,156 @@ +# 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 + +from typing import Any + +from google.genai import types +from typing_extensions import override + +from ..utils._schema_utils import schema_to_json_schema +from ..workflow._base_node import BaseNode +from ..workflow._errors import NodeInterruptedError +from .base_tool import BaseTool +from .tool_context import ToolContext + + +class NodeTool(BaseTool): + """A tool wrapper that executes a BaseNode (e.g. a Workflow or loop node).""" + + def __init__( + self, + node: BaseNode, + name: str | None = None, + description: str | None = None, + ): + from ..agents.base_agent import BaseAgent + from ..workflow._function_node import FunctionNode + + if isinstance(node, BaseAgent): + raise ValueError( + f"Agent '{node.name}' cannot be wrapped as a NodeTool. Agents should" + ' be invoked as Sub-Agents instead.' + ) + + # Automatically align FunctionNode binding + if ( + isinstance(node, FunctionNode) + and node.parameter_binding != 'node_input' + ): + orig_input_schema = getattr(node, 'input_schema', None) + orig_output_schema = getattr(node, 'output_schema', None) + node = FunctionNode( + func=node._func, + name=node.name, + rerun_on_resume=node.rerun_on_resume, + retry_config=node.retry_config, + timeout=node.timeout, + auth_config=node.auth_config, + parameter_binding='node_input', # Force binding to node_input + state_schema=node.state_schema, + ) + if orig_input_schema is not None: + node.input_schema = orig_input_schema + if orig_output_schema is not None: + node.output_schema = orig_output_schema + + if not getattr(node, 'input_schema', None): + raise ValueError( + f"Node '{node.name}' does not have an input_schema defined." + ' NodeTool requires an explicit Pydantic input_schema on the wrapped' + ' node.' + ) + + self.node = node + super().__init__( + name=name or node.name, + description=description + or node.description + or f'Executes the node: {node.name}', + ) + self.is_long_running = True + + @override + def _get_declaration(self) -> types.FunctionDeclaration: + schema = schema_to_json_schema(self.node.input_schema) + + # The GenAI API strictly requires parameters_json_schema to be an 'object' + # type schema. If the node has a primitive input schema (e.g., str, int), + # we wrap it into an object schema with a 'request' property. + if isinstance(schema, dict) and schema.get('type') != 'object': + schema = { + 'type': 'object', + 'properties': { + 'request': schema, + }, + 'required': ['request'], + } + + decl = types.FunctionDeclaration( + name=self.name, + description=self.description, + parameters_json_schema=schema, + ) + + output_schema = getattr(self.node, 'output_schema', None) + if output_schema: + decl.response_json_schema = schema_to_json_schema(output_schema) + + return decl + + @override + async def run_async( + self, + *, + args: dict[str, Any], + tool_context: ToolContext, + ) -> Any: + import inspect + + from pydantic import BaseModel + + input_schema = self.node.input_schema + node_input: Any + if inspect.isclass(input_schema) and issubclass(input_schema, BaseModel): + try: + # Convert input based on Pydantic schema + node_input = input_schema.model_validate(args) + except Exception as e: + return f'Error validating input for node: {e}' + else: + schema = schema_to_json_schema(input_schema) + if isinstance(schema, dict) and schema.get('type') != 'object': + node_input = args.get('request') + else: + node_input = args + + fc_id = tool_context.function_call_id + base_branch = tool_context.branch + segment = f'{self.name}@{fc_id}' + tool_branch = f'{base_branch}.{segment}' if base_branch else segment + + try: + return await tool_context.run_node( + self.node, + node_input=node_input, + override_branch=tool_branch, + use_sub_branch=False, + raise_on_wait=True, + ) + except NodeInterruptedError as nie: + # Propagates the interrupt up so the runner pauses the invocation + raise nie + except Exception as e: + return f'Error running node {self.name}: {e}' diff --git a/src/google/adk/workflow/_function_node.py b/src/google/adk/workflow/_function_node.py index c6e03ae1..15670634 100644 --- a/src/google/adk/workflow/_function_node.py +++ b/src/google/adk/workflow/_function_node.py @@ -328,9 +328,15 @@ class FunctionNode(BaseNode): is passed through directly and all other non-context parameters are looked up in ``ctx.state``. """ + from pydantic import BaseModel + input_bound = self.parameter_binding == 'node_input' + source: Any if input_bound: - source = node_input if isinstance(node_input, dict) else {} + if isinstance(node_input, (dict, BaseModel)): + source = node_input + else: + source = {} else: source = ctx.state source_name = 'node_input' if input_bound else 'state' @@ -353,8 +359,21 @@ class FunctionNode(BaseNode): kwargs[param_name] = value continue - if param_name in source: - value = source[param_name] + has_param = False + value = None + if isinstance(source, BaseModel): + if hasattr(source, param_name): + has_param = True + value = getattr(source, param_name) + else: + try: + if param_name in source: + has_param = True + value = source[param_name] + except (TypeError, KeyError): + pass + + if has_param: if param_name in self._type_hints: value = self._coerce_param( param_name, diff --git a/tests/unittests/workflow/test_node_tool.py b/tests/unittests/workflow/test_node_tool.py new file mode 100644 index 00000000..82656232 --- /dev/null +++ b/tests/unittests/workflow/test_node_tool.py @@ -0,0 +1,1003 @@ +# 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'}