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
1785 lines
50 KiB
Python
1785 lines
50 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.
|
|
|
|
"""Testings for the FunctionNode."""
|
|
|
|
import copy
|
|
from typing import Any
|
|
from typing import AsyncGenerator
|
|
from typing import Generator
|
|
from typing import Optional
|
|
from typing import Union
|
|
from unittest import mock
|
|
|
|
from google.adk.agents.context import Context
|
|
from google.adk.apps.app import App
|
|
from google.adk.apps.app import ResumabilityConfig
|
|
from google.adk.events._node_path_builder import _NodePathBuilder
|
|
from google.adk.events.event import Event
|
|
from google.adk.events.event import Event as AdkEvent
|
|
from google.adk.events.request_input import RequestInput
|
|
from google.adk.runners import Runner
|
|
from google.adk.sessions import InMemorySessionService
|
|
from google.adk.workflow import FunctionNode
|
|
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.genai import types
|
|
from pydantic import BaseModel
|
|
import pytest
|
|
|
|
from .. import testing_utils
|
|
from .workflow_testing_utils import create_parent_invocation_context
|
|
from .workflow_testing_utils import get_output_events
|
|
from .workflow_testing_utils import get_request_input_events
|
|
from .workflow_testing_utils import simplify_events_with_node
|
|
from .workflow_testing_utils import simplify_events_with_node_and_agent_state
|
|
|
|
ANY = mock.ANY
|
|
|
|
from .workflow_testing_utils import run_workflow
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_various_function_nodes(request: pytest.FixtureRequest):
|
|
"""Tests that Workflow can run with various function nodes."""
|
|
|
|
async def async_gen_func(ctx: Context) -> AsyncGenerator[Any, None]:
|
|
yield Event(
|
|
output='Hello from AsyncGen',
|
|
)
|
|
|
|
def sync_func_out(ctx: Context) -> str:
|
|
return 'Hello from SyncFunc'
|
|
|
|
async def async_func_out(ctx: Context) -> str:
|
|
return 'Hello from AsyncFunc'
|
|
|
|
def sync_func_no_out(ctx: Context) -> None:
|
|
return None
|
|
|
|
async def async_func_no_out(ctx: Context) -> None:
|
|
return None
|
|
|
|
def sync_gen_func(
|
|
ctx: Context,
|
|
) -> Generator[Any, None, None]:
|
|
yield Event(
|
|
output='Hello from SyncGen',
|
|
)
|
|
|
|
async def async_gen_func_raw_output(
|
|
ctx: Context,
|
|
) -> AsyncGenerator[Any, None]:
|
|
yield 'Hello from AsyncGenRawOutput'
|
|
|
|
def sync_gen_func_raw_output(
|
|
ctx: Context,
|
|
) -> Generator[Any, None, None]:
|
|
yield 'Hello from SyncGenRawOutput'
|
|
|
|
agent = Workflow(
|
|
name='test_workflow_agent_various_function_nodes',
|
|
edges=[
|
|
(START, async_gen_func),
|
|
(async_gen_func, sync_func_out),
|
|
(sync_func_out, async_func_out),
|
|
(async_func_out, sync_func_no_out),
|
|
(sync_func_no_out, async_func_no_out),
|
|
(async_func_no_out, sync_gen_func),
|
|
(sync_gen_func, async_gen_func_raw_output),
|
|
(async_gen_func_raw_output, sync_gen_func_raw_output),
|
|
],
|
|
)
|
|
app_instance = App(name=request.function.__name__, root_agent=agent)
|
|
runner = testing_utils.InMemoryRunner(app=app_instance)
|
|
events = await runner.run_async(testing_utils.get_user_content('start'))
|
|
|
|
# Functions with no output (sync_func_no_out, async_func_no_out)
|
|
# will not produce events.
|
|
assert simplify_events_with_node(events) == [
|
|
(
|
|
'test_workflow_agent_various_function_nodes@1/async_gen_func@1',
|
|
{'output': 'Hello from AsyncGen'},
|
|
),
|
|
(
|
|
'test_workflow_agent_various_function_nodes@1/sync_func_out@1',
|
|
{'output': 'Hello from SyncFunc'},
|
|
),
|
|
(
|
|
'test_workflow_agent_various_function_nodes@1/async_func_out@1',
|
|
{'output': 'Hello from AsyncFunc'},
|
|
),
|
|
(
|
|
'test_workflow_agent_various_function_nodes@1/sync_gen_func@1',
|
|
{'output': 'Hello from SyncGen'},
|
|
),
|
|
(
|
|
'test_workflow_agent_various_function_nodes@1/async_gen_func_raw_output@1',
|
|
{
|
|
'output': 'Hello from AsyncGenRawOutput',
|
|
},
|
|
),
|
|
(
|
|
'test_workflow_agent_various_function_nodes@1/sync_gen_func_raw_output@1',
|
|
{
|
|
'output': 'Hello from SyncGenRawOutput',
|
|
},
|
|
),
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_function_node_state_injection(request: pytest.FixtureRequest):
|
|
"""Tests that FunctionNode can inject parameters from workflow state."""
|
|
|
|
async def set_state_node_fn(
|
|
ctx: Context,
|
|
) -> AsyncGenerator[Any, None]:
|
|
yield Event(
|
|
state={'param1': 'value1'},
|
|
)
|
|
|
|
def check_state_node_fn(param1: str, param2: str = 'default2') -> str:
|
|
return f'param1={param1}, param2={param2}'
|
|
|
|
agent = Workflow(
|
|
name='test_workflow_agent_state_injection',
|
|
edges=[
|
|
(START, set_state_node_fn),
|
|
(set_state_node_fn, check_state_node_fn),
|
|
],
|
|
)
|
|
events, _, _ = await run_workflow(agent)
|
|
assert simplify_events_with_node(events) == [
|
|
(
|
|
'test_workflow_agent_state_injection@1/check_state_node_fn@1',
|
|
{
|
|
'output': 'param1=value1, param2=default2',
|
|
},
|
|
),
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_function_node_state_injection_missing_param(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Tests that FunctionNode raises error for missing param."""
|
|
|
|
def check_state_node_fn(param1: str) -> str:
|
|
return f'param1={param1}'
|
|
|
|
agent = Workflow(
|
|
name='test_workflow_agent_state_injection_missing',
|
|
edges=[
|
|
(START, check_state_node_fn),
|
|
],
|
|
)
|
|
with pytest.raises(ValueError, match='Missing value for parameter "param1"'):
|
|
await run_workflow(agent)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_function_node_type_checking(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Tests that FunctionNode performs type checking."""
|
|
|
|
async def set_state_node_fn(
|
|
ctx: Context,
|
|
) -> AsyncGenerator[Any, None]:
|
|
yield Event(
|
|
state={'p1': 'a string'},
|
|
)
|
|
|
|
def check_type_node_fn(p1: int) -> str:
|
|
return f'p1={p1}'
|
|
|
|
agent = Workflow(
|
|
name='test_type_checking',
|
|
edges=[
|
|
(START, set_state_node_fn),
|
|
(set_state_node_fn, check_type_node_fn),
|
|
],
|
|
)
|
|
with pytest.raises(ValueError):
|
|
await run_workflow(agent)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_function_node_input_injection(request: pytest.FixtureRequest):
|
|
"""Tests that FunctionNode can inject parameters from node_input."""
|
|
|
|
def node1_fn() -> dict[str, Any]:
|
|
return {'p1': 'value1_from_node_input', 'p2': 100}
|
|
|
|
def node2_fn(node_input: dict[str, Any]) -> str:
|
|
return f"p1={node_input['p1']}, p2={node_input['p2']}"
|
|
|
|
agent = Workflow(
|
|
name='test_workflow_agent_input_injection_dict',
|
|
edges=[
|
|
(START, node1_fn),
|
|
(node1_fn, node2_fn),
|
|
],
|
|
)
|
|
events, _, _ = await run_workflow(agent)
|
|
assert simplify_events_with_node(events) == [
|
|
(
|
|
'test_workflow_agent_input_injection_dict@1/node1_fn@1',
|
|
{
|
|
'output': {'p1': 'value1_from_node_input', 'p2': 100},
|
|
},
|
|
),
|
|
(
|
|
'test_workflow_agent_input_injection_dict@1/node2_fn@1',
|
|
{
|
|
'output': 'p1=value1_from_node_input, p2=100',
|
|
},
|
|
),
|
|
]
|
|
|
|
|
|
class MyModel(BaseModel):
|
|
p1: str
|
|
p2: int
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_function_node_input_injection_pydantic(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Tests that FunctionNode can inject dict as pydantic model into node_input."""
|
|
|
|
def node1_fn() -> dict[str, Any]:
|
|
return {'p1': 'value1_from_node_input', 'p2': 100}
|
|
|
|
def node2_fn(node_input: MyModel) -> str:
|
|
return f'p1={node_input.p1}, p2={node_input.p2}'
|
|
|
|
agent = Workflow(
|
|
name='test_workflow_agent_input_injection_pydantic',
|
|
edges=[
|
|
(START, node1_fn),
|
|
(node1_fn, node2_fn),
|
|
],
|
|
)
|
|
events, _, _ = await run_workflow(agent)
|
|
assert simplify_events_with_node(events) == [
|
|
(
|
|
'test_workflow_agent_input_injection_pydantic@1/node1_fn@1',
|
|
{
|
|
'output': {'p1': 'value1_from_node_input', 'p2': 100},
|
|
},
|
|
),
|
|
(
|
|
'test_workflow_agent_input_injection_pydantic@1/node2_fn@1',
|
|
{
|
|
'output': 'p1=value1_from_node_input, p2=100',
|
|
},
|
|
),
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_function_node_input_list_wrong_type(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Tests that FunctionNode raises TypeError for list[T] with non-list input."""
|
|
|
|
def node1_fn() -> int:
|
|
return 123
|
|
|
|
def node2_fn(node_input: list[MyModel]) -> str:
|
|
return f'p1={node_input[0].p1}'
|
|
|
|
agent = Workflow(
|
|
name='test_workflow_agent_input_list_wrong_type',
|
|
edges=[
|
|
(START, node1_fn),
|
|
(node1_fn, node2_fn),
|
|
],
|
|
)
|
|
with pytest.raises(ValueError):
|
|
await run_workflow(agent)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_function_node_input_list_no_item_type(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Tests that FunctionNode handles list without item type."""
|
|
|
|
def node1_fn() -> list[int]:
|
|
return [1, 2]
|
|
|
|
def node2_fn(node_input: list) -> str:
|
|
return f'list={node_input}'
|
|
|
|
agent = Workflow(
|
|
name='test_workflow_agent_input_list_no_item_type',
|
|
edges=[
|
|
(START, node1_fn),
|
|
(node1_fn, node2_fn),
|
|
],
|
|
)
|
|
events, _, _ = await run_workflow(agent)
|
|
assert simplify_events_with_node(events) == [
|
|
(
|
|
'test_workflow_agent_input_list_no_item_type@1/node1_fn@1',
|
|
{
|
|
'output': [1, 2],
|
|
},
|
|
),
|
|
(
|
|
'test_workflow_agent_input_list_no_item_type@1/node2_fn@1',
|
|
{
|
|
'output': 'list=[1, 2]',
|
|
},
|
|
),
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_function_node_input_and_state_injection(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Tests parameter injection from node_input and state."""
|
|
|
|
async def nodea_fn(ctx: Context) -> AsyncGenerator[Any, None]:
|
|
yield Event(
|
|
state={'p_state': 'value_from_state'},
|
|
)
|
|
yield 'value_A'
|
|
|
|
def nodeb_fn(
|
|
node_input: str, p_state: str, p_default: str = 'default2'
|
|
) -> str:
|
|
return f'node_input={node_input}, p_state={p_state}, p_default={p_default}'
|
|
|
|
agent = Workflow(
|
|
name='test_node_param_injection_single_and_state',
|
|
edges=[
|
|
(START, nodea_fn),
|
|
(nodea_fn, nodeb_fn),
|
|
],
|
|
)
|
|
events, _, _ = await run_workflow(agent)
|
|
assert simplify_events_with_node(events) == [
|
|
(
|
|
'test_node_param_injection_single_and_state@1/nodea_fn@1',
|
|
{'output': 'value_A'},
|
|
),
|
|
(
|
|
'test_node_param_injection_single_and_state@1/nodeb_fn@1',
|
|
{
|
|
'output': (
|
|
'node_input=value_A, p_state=value_from_state,'
|
|
' p_default=default2'
|
|
),
|
|
},
|
|
),
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_function_node_state_injection_pydantic(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Tests that FunctionNode can inject dict from state as pydantic model."""
|
|
|
|
async def node1_fn(ctx: Context) -> AsyncGenerator[Any, None]:
|
|
yield Event(
|
|
state={'my_model': {'p1': 'value1_from_state', 'p2': 200}},
|
|
)
|
|
|
|
def node2_fn(my_model: MyModel) -> str:
|
|
return f'p1={my_model.p1}, p2={my_model.p2}'
|
|
|
|
agent = Workflow(
|
|
name='test_workflow_agent_state_injection_pydantic',
|
|
edges=[
|
|
(START, node1_fn),
|
|
(node1_fn, node2_fn),
|
|
],
|
|
)
|
|
events, _, _ = await run_workflow(agent)
|
|
assert simplify_events_with_node(events) == [
|
|
(
|
|
'test_workflow_agent_state_injection_pydantic@1/node2_fn@1',
|
|
{
|
|
'output': 'p1=value1_from_state, p2=200',
|
|
},
|
|
),
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_function_node_hitl(request: pytest.FixtureRequest):
|
|
"""Tests that FunctionNode can trigger HITL.
|
|
|
|
Setup: Workflow with request_input_fn -> process_input_fn.
|
|
request_input_fn yields RequestInput.
|
|
Act:
|
|
- Run 1: start workflow, request_input_fn yields RequestInput.
|
|
- Run 2: resume with user response.
|
|
Assert:
|
|
- Run 1: returns RequestInput event with valid ID.
|
|
- Run 2: process_input_fn receives input and completes.
|
|
"""
|
|
|
|
# Given: A workflow where the first node requests input
|
|
def request_input_fn() -> Generator[Any, None, None]:
|
|
yield RequestInput(message='Provide input')
|
|
|
|
def process_input_fn(node_input: dict[str, Any]) -> str:
|
|
return f"received: {node_input['text']}"
|
|
|
|
agent = Workflow(
|
|
name='test_workflow_agent_hitl',
|
|
edges=[
|
|
(START, request_input_fn),
|
|
(request_input_fn, process_input_fn),
|
|
],
|
|
)
|
|
app = App(
|
|
name=request.function.__name__,
|
|
root_agent=agent,
|
|
resumability_config=ResumabilityConfig(is_resumable=True),
|
|
)
|
|
runner = testing_utils.InMemoryRunner(app=app)
|
|
|
|
# When: Starting the workflow
|
|
user_event = testing_utils.get_user_content('start workflow')
|
|
events1 = await runner.run_async(user_event)
|
|
|
|
# Then: It should yield a RequestInput event with a valid ID
|
|
req_events = get_request_input_events(events1)
|
|
assert len(req_events) == 1
|
|
interrupt_id = get_request_input_interrupt_ids(req_events[0])[0]
|
|
invocation_id = events1[0].invocation_id
|
|
|
|
# Verify that the FunctionCall object actually has the id set
|
|
assert req_events[0].content.parts[0].function_call.id == interrupt_id
|
|
|
|
# When: Resuming with user response
|
|
user_input = create_request_input_response(
|
|
interrupt_id, {'text': 'Hello from user'}
|
|
)
|
|
events2 = await runner.run_async(
|
|
new_message=testing_utils.UserContent(user_input),
|
|
invocation_id=invocation_id,
|
|
)
|
|
|
|
# Then: It should complete and pass output to the next node
|
|
data_events = get_output_events(events2, output='received: Hello from user')
|
|
assert len(data_events) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_function_node_adk_events(request: pytest.FixtureRequest):
|
|
"""Tests that FunctionNode can emit ADK events."""
|
|
|
|
def adk_events_fn() -> Generator[Any, None, None]:
|
|
yield AdkEvent(
|
|
author='some_agent', content=types.Content(parts=[{'text': 'event 1'}])
|
|
)
|
|
yield AdkEvent(
|
|
author='some_agent', content=types.Content(parts=[{'text': 'event 2'}])
|
|
)
|
|
|
|
agent = Workflow(
|
|
name='test_workflow_agent_adk_events',
|
|
edges=[
|
|
(START, adk_events_fn),
|
|
],
|
|
)
|
|
events, _, _ = await run_workflow(agent)
|
|
|
|
assert len(events) == 2
|
|
assert isinstance(events[0], AdkEvent)
|
|
assert events[0].content.parts[0].text == 'event 1'
|
|
assert isinstance(events[1], AdkEvent)
|
|
assert events[1].content.parts[0].text == 'event 2'
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_function_node_list_conversion_pydantic(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Tests that FunctionNode correctly converts list[dict] to list[BaseModel]."""
|
|
|
|
class Section(BaseModel):
|
|
section_name: str
|
|
content: str
|
|
|
|
async def upstream_func(ctx: Context) -> list[dict[str, str]]:
|
|
return [
|
|
{'section_name': 's1', 'content': 'c1'},
|
|
{'section_name': 's2', 'content': 'c2'},
|
|
]
|
|
|
|
received_input = None
|
|
|
|
async def aggregate(node_input: list[Section]) -> str:
|
|
nonlocal received_input
|
|
received_input = node_input
|
|
return 'Done'
|
|
|
|
agent = Workflow(
|
|
name='test_function_node_list_conversion_pydantic',
|
|
edges=[
|
|
(START, upstream_func),
|
|
(upstream_func, aggregate),
|
|
],
|
|
)
|
|
unused_events, _, _ = await run_workflow(agent)
|
|
|
|
# Check if received_input contains Section objects
|
|
assert isinstance(received_input, list)
|
|
assert len(received_input) == 2
|
|
assert isinstance(received_input[0], Section)
|
|
assert received_input[0].section_name == 's1'
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_function_node_dict_conversion_pydantic(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Tests that FunctionNode correctly converts dict[str, dict] to dict[str, BaseModel]."""
|
|
|
|
class Section(BaseModel):
|
|
section_name: str
|
|
content: str
|
|
|
|
async def upstream_func(ctx: Context) -> dict[str, dict[str, str]]:
|
|
return {
|
|
'one': {'section_name': 's1', 'content': 'c1'},
|
|
'two': {'section_name': 's2', 'content': 'c2'},
|
|
}
|
|
|
|
received_input = None
|
|
|
|
async def aggregate(node_input: dict[str, Section]) -> str:
|
|
nonlocal received_input
|
|
received_input = node_input
|
|
return 'Done'
|
|
|
|
agent = Workflow(
|
|
name='test_function_node_dict_conversion_pydantic',
|
|
edges=[
|
|
(START, upstream_func),
|
|
(upstream_func, aggregate),
|
|
],
|
|
)
|
|
unused_events, _, _ = await run_workflow(agent)
|
|
|
|
# Check if received_input contains Section objects
|
|
assert isinstance(received_input, dict)
|
|
assert len(received_input) == 2
|
|
assert isinstance(received_input['one'], Section)
|
|
assert received_input['one'].section_name == 's1'
|
|
assert isinstance(received_input['two'], Section)
|
|
assert received_input['two'].content == 'c2'
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_function_node_no_data_returns_none(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Tests that FunctionNode returns Event with output=None if function returns only control event."""
|
|
|
|
def func_no_data() -> Event:
|
|
return Event(route='some_route')
|
|
|
|
agent = Workflow(
|
|
name='test_function_node_no_data_returns_none',
|
|
edges=[
|
|
(START, func_no_data),
|
|
],
|
|
)
|
|
events, _, _ = await run_workflow(agent)
|
|
|
|
assert len(events) == 1
|
|
assert events[0].output is None
|
|
assert events[0].actions.route == 'some_route'
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_function_node_yield_content(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Tests that yielding types.Content sets output=None and content=Content."""
|
|
|
|
def func_yield_content() -> Generator[Any, None, None]:
|
|
yield types.Content(parts=[types.Part(text='some content')])
|
|
|
|
agent = Workflow(
|
|
name='test_function_node_yield_content',
|
|
edges=[
|
|
(START, func_yield_content),
|
|
],
|
|
)
|
|
events, _, _ = await run_workflow(agent)
|
|
|
|
assert len(events) == 1
|
|
assert events[0].output is None
|
|
assert events[0].content is not None
|
|
assert events[0].content.parts[0].text == 'some content'
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_function_node_yield_event_with_content(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Tests that yielding Event(content=...) retains output=None and content=... ."""
|
|
|
|
def func_yield_event_with_content() -> Generator[Any, None, None]:
|
|
yield Event(content=types.Content(parts=[types.Part(text='some content')]))
|
|
|
|
agent = Workflow(
|
|
name='test_function_node_yield_event_with_content',
|
|
edges=[
|
|
(START, func_yield_event_with_content),
|
|
],
|
|
)
|
|
events, _, _ = await run_workflow(agent)
|
|
|
|
assert len(events) == 1
|
|
assert events[0].output is None
|
|
assert events[0].content is not None
|
|
assert events[0].content.parts[0].text == 'some content'
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_content_to_str_auto_conversion(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Tests that Content is auto-converted to str when function expects str."""
|
|
received_inputs = []
|
|
|
|
def record_input(node_input: str) -> str:
|
|
received_inputs.append(node_input)
|
|
return f'Hello, {node_input}!'
|
|
|
|
agent = Workflow(
|
|
name='test_content_to_str',
|
|
edges=[
|
|
(START, record_input),
|
|
],
|
|
)
|
|
app = App(
|
|
name=request.function.__name__,
|
|
root_agent=agent,
|
|
)
|
|
runner = testing_utils.InMemoryRunner(app=app)
|
|
user_event = testing_utils.get_user_content('start workflow')
|
|
await runner.run_async(user_event)
|
|
|
|
assert len(received_inputs) == 1
|
|
assert isinstance(received_inputs[0], str)
|
|
assert received_inputs[0] == 'start workflow'
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_content_to_str_multi_part(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Tests Content with multiple text parts is concatenated."""
|
|
received_inputs = []
|
|
|
|
def record_input(node_input: str) -> str:
|
|
received_inputs.append(node_input)
|
|
return node_input
|
|
|
|
agent = Workflow(
|
|
name='test_content_to_str_multi_part',
|
|
edges=[
|
|
(START, record_input),
|
|
],
|
|
)
|
|
app = App(
|
|
name=request.function.__name__,
|
|
root_agent=agent,
|
|
)
|
|
runner = testing_utils.InMemoryRunner(app=app)
|
|
user_event = testing_utils.get_user_content(
|
|
types.Content(
|
|
parts=[
|
|
types.Part(text='Hello '),
|
|
types.Part(text='World'),
|
|
],
|
|
role='user',
|
|
)
|
|
)
|
|
await runner.run_async(user_event)
|
|
|
|
assert len(received_inputs) == 1
|
|
assert received_inputs[0] == 'Hello World'
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_content_to_str_warns_on_non_text(
|
|
request: pytest.FixtureRequest,
|
|
caplog,
|
|
):
|
|
"""Tests that non-text parts produce a warning during conversion."""
|
|
import logging
|
|
|
|
received_inputs = []
|
|
|
|
def record_input(node_input: str) -> str:
|
|
received_inputs.append(node_input)
|
|
return node_input
|
|
|
|
agent = Workflow(
|
|
name='test_content_to_str_warns',
|
|
edges=[
|
|
(START, record_input),
|
|
],
|
|
)
|
|
app = App(
|
|
name=request.function.__name__,
|
|
root_agent=agent,
|
|
)
|
|
runner = testing_utils.InMemoryRunner(app=app)
|
|
user_event = testing_utils.get_user_content(
|
|
types.Content(
|
|
parts=[
|
|
types.Part(text='Hello'),
|
|
types.Part(
|
|
inline_data=types.Blob(data=b'img', mime_type='image/png')
|
|
),
|
|
],
|
|
role='user',
|
|
)
|
|
)
|
|
with caplog.at_level(logging.WARNING):
|
|
await runner.run_async(user_event)
|
|
|
|
assert len(received_inputs) == 1
|
|
assert received_inputs[0] == 'Hello'
|
|
assert 'non-text parts' in caplog.text
|
|
|
|
|
|
def _produce_list():
|
|
return [1, 2, 3]
|
|
|
|
|
|
def _produce_dict():
|
|
return {'key': 'value'}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
'produce_value, expected',
|
|
[
|
|
(_produce_list, [1, 2, 3]),
|
|
(_produce_dict, {'key': 'value'}),
|
|
],
|
|
ids=['list', 'dict'],
|
|
)
|
|
async def test_union_type_accepts_matching_member(
|
|
request: pytest.FixtureRequest,
|
|
produce_value,
|
|
expected,
|
|
):
|
|
"""Tests that Union type annotation accepts values matching any member."""
|
|
received_inputs = []
|
|
|
|
def record_input(node_input: Union[list, dict]) -> str:
|
|
received_inputs.append(node_input)
|
|
return 'ok'
|
|
|
|
agent = Workflow(
|
|
name='test_union_accept',
|
|
edges=[
|
|
(START, produce_value),
|
|
(produce_value, record_input),
|
|
],
|
|
)
|
|
app = App(
|
|
name=request.function.__name__,
|
|
root_agent=agent,
|
|
)
|
|
runner = testing_utils.InMemoryRunner(app=app)
|
|
await runner.run_async(testing_utils.get_user_content('go'))
|
|
|
|
assert len(received_inputs) == 1
|
|
assert received_inputs[0] == expected
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_optional_str_with_content_auto_conversion(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Tests that Optional[str] auto-converts Content from START node."""
|
|
received_inputs = []
|
|
|
|
def record_input(node_input: Optional[str]) -> str:
|
|
received_inputs.append(node_input)
|
|
return 'ok'
|
|
|
|
agent = Workflow(
|
|
name='test_optional_content',
|
|
edges=[
|
|
(START, record_input),
|
|
],
|
|
)
|
|
app = App(
|
|
name=request.function.__name__,
|
|
root_agent=agent,
|
|
)
|
|
runner = testing_utils.InMemoryRunner(app=app)
|
|
await runner.run_async(testing_utils.get_user_content('hello'))
|
|
|
|
assert len(received_inputs) == 1
|
|
assert received_inputs[0] == 'hello'
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_union_type_rejects_non_matching(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Tests that Union type annotation rejects values not matching any member."""
|
|
|
|
def bad_input(node_input: Union[str, int]) -> str:
|
|
return 'should not reach'
|
|
|
|
def produce_list() -> list:
|
|
return [1, 2, 3]
|
|
|
|
agent = Workflow(
|
|
name='test_union_reject',
|
|
edges=[
|
|
(START, produce_list),
|
|
(produce_list, bad_input),
|
|
],
|
|
)
|
|
app = App(
|
|
name=request.function.__name__,
|
|
root_agent=agent,
|
|
)
|
|
runner = testing_utils.InMemoryRunner(app=app)
|
|
with pytest.raises(ValueError):
|
|
await runner.run_async(testing_utils.get_user_content('go'))
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_function_node_ctx_state_delta_sync(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Tests that state set via ctx.state in a sync function is persisted."""
|
|
|
|
def set_state_via_ctx(ctx: Context) -> str:
|
|
ctx.state['user_request'] = 'build a tracker app'
|
|
return 'done'
|
|
|
|
def read_state(user_request: str) -> str:
|
|
return f'request={user_request}'
|
|
|
|
agent = Workflow(
|
|
name='test_ctx_state_delta_sync',
|
|
edges=[
|
|
(START, set_state_via_ctx),
|
|
(set_state_via_ctx, read_state),
|
|
],
|
|
)
|
|
events, _, _ = await run_workflow(agent)
|
|
simplified = simplify_events_with_node(events, include_state_delta=True)
|
|
assert simplified == [
|
|
(
|
|
'test_ctx_state_delta_sync@1/set_state_via_ctx@1',
|
|
{
|
|
'output': 'done',
|
|
'state_delta': {'user_request': 'build a tracker app'},
|
|
},
|
|
),
|
|
(
|
|
'test_ctx_state_delta_sync@1/read_state@1',
|
|
{
|
|
'output': 'request=build a tracker app',
|
|
},
|
|
),
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_function_node_ctx_state_delta_async(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Tests that state set via ctx.state in an async function is persisted."""
|
|
|
|
async def set_state_via_ctx(ctx: Context) -> str:
|
|
ctx.state['counter'] = 42
|
|
ctx.state['name'] = 'test'
|
|
return 'set'
|
|
|
|
def read_state(counter: int, name: str) -> str:
|
|
return f'counter={counter}, name={name}'
|
|
|
|
agent = Workflow(
|
|
name='test_ctx_state_delta_async',
|
|
edges=[
|
|
(START, set_state_via_ctx),
|
|
(set_state_via_ctx, read_state),
|
|
],
|
|
)
|
|
events, _, _ = await run_workflow(agent)
|
|
simplified = simplify_events_with_node(events, include_state_delta=True)
|
|
assert simplified == [
|
|
(
|
|
'test_ctx_state_delta_async@1/set_state_via_ctx@1',
|
|
{
|
|
'output': 'set',
|
|
'state_delta': {'counter': 42, 'name': 'test'},
|
|
},
|
|
),
|
|
(
|
|
'test_ctx_state_delta_async@1/read_state@1',
|
|
{
|
|
'output': 'counter=42, name=test',
|
|
},
|
|
),
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_function_node_ctx_state_delta_none_return(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Tests that state is persisted even when function returns None."""
|
|
|
|
def set_state_return_none(ctx: Context) -> None:
|
|
ctx.state['my_key'] = 'my_value'
|
|
|
|
def read_state(my_key: str) -> str:
|
|
return f'my_key={my_key}'
|
|
|
|
agent = Workflow(
|
|
name='test_ctx_state_delta_none_return',
|
|
edges=[
|
|
(START, set_state_return_none),
|
|
(set_state_return_none, read_state),
|
|
],
|
|
)
|
|
events, _, _ = await run_workflow(agent)
|
|
simplified = simplify_events_with_node(events, include_state_delta=True)
|
|
assert simplified == [
|
|
(
|
|
'test_ctx_state_delta_none_return@1/set_state_return_none@1',
|
|
{
|
|
'state_delta': {'my_key': 'my_value'},
|
|
'output': None,
|
|
},
|
|
),
|
|
(
|
|
'test_ctx_state_delta_none_return@1/read_state@1',
|
|
{
|
|
'output': 'my_key=my_value',
|
|
},
|
|
),
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_function_node_ctx_state_delta_with_event_return(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Tests that ctx.state changes merge into a returned Event's state_delta."""
|
|
|
|
def set_state_return_event(ctx: Context) -> Event:
|
|
ctx.state['from_ctx'] = 'ctx_value'
|
|
return Event(
|
|
output='result',
|
|
state={'from_event': 'event_value'},
|
|
)
|
|
|
|
def read_state(from_ctx: str, from_event: str) -> str:
|
|
return f'from_ctx={from_ctx}, from_event={from_event}'
|
|
|
|
agent = Workflow(
|
|
name='test_ctx_state_delta_event_return',
|
|
edges=[
|
|
(START, set_state_return_event),
|
|
(set_state_return_event, read_state),
|
|
],
|
|
)
|
|
events, _, _ = await run_workflow(agent)
|
|
simplified = simplify_events_with_node(events, include_state_delta=True)
|
|
assert simplified == [
|
|
(
|
|
'test_ctx_state_delta_event_return@1/set_state_return_event@1',
|
|
{
|
|
'output': 'result',
|
|
'state_delta': {
|
|
'from_event': 'event_value',
|
|
'from_ctx': 'ctx_value',
|
|
},
|
|
},
|
|
),
|
|
(
|
|
'test_ctx_state_delta_event_return@1/read_state@1',
|
|
{
|
|
'output': 'from_ctx=ctx_value, from_event=event_value',
|
|
},
|
|
),
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_function_node_ctx_state_delta_generator(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Tests that ctx.state changes are captured in generator yields."""
|
|
|
|
def gen_with_state(ctx: Context) -> Generator[Any, None, None]:
|
|
ctx.state['key1'] = 'value1'
|
|
yield Event(state={'key1': 'value1'})
|
|
ctx.state['key2'] = 'value2'
|
|
yield 'done'
|
|
|
|
def read_state(key1: str, key2: str) -> str:
|
|
return f'key1={key1}, key2={key2}'
|
|
|
|
agent = Workflow(
|
|
name='test_ctx_state_delta_generator',
|
|
edges=[
|
|
(START, gen_with_state),
|
|
(gen_with_state, read_state),
|
|
],
|
|
)
|
|
events, _, _ = await run_workflow(agent)
|
|
simplified = simplify_events_with_node(events, include_state_delta=True)
|
|
|
|
# First yield is a state-only Event, second is the data output with
|
|
# accumulated state from both ctx.state assignments.
|
|
assert simplified == [
|
|
(
|
|
'test_ctx_state_delta_generator@1/gen_with_state@1',
|
|
{
|
|
'output': None,
|
|
'state_delta': {'key1': 'value1'},
|
|
},
|
|
),
|
|
(
|
|
'test_ctx_state_delta_generator@1/gen_with_state@1',
|
|
{
|
|
'output': 'done',
|
|
'state_delta': {'key2': 'value2'},
|
|
},
|
|
),
|
|
(
|
|
'test_ctx_state_delta_generator@1/read_state@1',
|
|
{
|
|
'output': 'key1=value1, key2=value2',
|
|
},
|
|
),
|
|
]
|
|
|
|
|
|
# ── FunctionNode output_schema ──────────────────────────────────────
|
|
|
|
|
|
class _OutputModel(BaseModel):
|
|
name: str
|
|
value: int
|
|
|
|
|
|
class _OtherModel(BaseModel):
|
|
name: str
|
|
value: int
|
|
extra: str = 'default'
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_output_schema_inferred_validates_dict(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Inferred output_schema validates dict return from BaseModel function."""
|
|
|
|
def produce() -> _OutputModel:
|
|
return {'name': 'test', 'value': 42}
|
|
|
|
node = FunctionNode(func=produce)
|
|
assert node.output_schema is _OutputModel
|
|
|
|
agent = Workflow(name='wf', edges=[(START, node)])
|
|
events, _, _ = await run_workflow(agent)
|
|
|
|
data_events = [
|
|
e
|
|
for e in events
|
|
if isinstance(e, Event)
|
|
and e.output is not None
|
|
and _NodePathBuilder.from_string(e.node_info.path).is_direct_child_of(
|
|
_NodePathBuilder.from_string('wf@1')
|
|
)
|
|
]
|
|
assert len(data_events) == 1
|
|
assert data_events[0].output == {'name': 'test', 'value': 42}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_output_schema_inferred_rejects_invalid(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Inferred output_schema rejects invalid dict."""
|
|
|
|
def produce() -> _OutputModel:
|
|
return {'name': 'test'} # missing 'value'
|
|
|
|
node = FunctionNode(func=produce)
|
|
agent = Workflow(name='wf', edges=[(START, node)])
|
|
with pytest.raises(ValueError):
|
|
await run_workflow(agent)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_output_schema_inferred_rejects_wrong_type(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Inferred output_schema rejects non-dict, non-BaseModel return."""
|
|
|
|
def produce() -> _OutputModel:
|
|
return 'not a dict'
|
|
|
|
node = FunctionNode(func=produce)
|
|
agent = Workflow(name='wf', edges=[(START, node)])
|
|
with pytest.raises(ValueError):
|
|
await run_workflow(agent)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_output_schema_generator_rejects_invalid_item(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""A generator that yields an invalid item mid-stream raises."""
|
|
|
|
def produce_items() -> Generator[_OutputModel, None, None]:
|
|
yield {'name': 'a', 'value': 1}
|
|
yield {'name': 'bad'} # missing 'value'
|
|
|
|
node = FunctionNode(func=produce_items)
|
|
agent = Workflow(name='wf', edges=[(START, node)])
|
|
with pytest.raises(ValueError):
|
|
await run_workflow(agent)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_output_schema_inferred_coerces_defaults(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Inferred output_schema fills in default fields."""
|
|
|
|
def produce() -> _OtherModel:
|
|
return {'name': 'test', 'value': 5}
|
|
|
|
node = FunctionNode(func=produce)
|
|
assert node.output_schema is _OtherModel
|
|
|
|
agent = Workflow(name='wf', edges=[(START, node)])
|
|
events, _, _ = await run_workflow(agent)
|
|
|
|
data_events = [
|
|
e
|
|
for e in events
|
|
if isinstance(e, Event)
|
|
and e.output is not None
|
|
and _NodePathBuilder.from_string(e.node_info.path).is_direct_child_of(
|
|
_NodePathBuilder.from_string('wf@1')
|
|
)
|
|
]
|
|
assert len(data_events) == 1
|
|
assert data_events[0].output == {
|
|
'name': 'test',
|
|
'value': 5,
|
|
'extra': 'default',
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_output_schema_inferred_from_return_hint(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""output_schema is auto-inferred from -> BaseModel return hint."""
|
|
|
|
def produce() -> _OutputModel:
|
|
return _OutputModel(name='inferred', value=1)
|
|
|
|
node = FunctionNode(func=produce)
|
|
assert node.output_schema is _OutputModel
|
|
|
|
agent = Workflow(name='wf', edges=[(START, node)])
|
|
events, _, _ = await run_workflow(agent)
|
|
|
|
data_events = [
|
|
e
|
|
for e in events
|
|
if isinstance(e, Event)
|
|
and e.output is not None
|
|
and _NodePathBuilder.from_string(e.node_info.path).is_direct_child_of(
|
|
_NodePathBuilder.from_string('wf@1')
|
|
)
|
|
]
|
|
assert len(data_events) == 1
|
|
assert data_events[0].output == {'name': 'inferred', 'value': 1}
|
|
|
|
|
|
def test_output_schema_no_inference_for_non_basemodel():
|
|
"""Non-BaseModel return hints (str, dict, etc.) don't trigger inference."""
|
|
|
|
def produce() -> dict:
|
|
return {'any': 'thing'}
|
|
|
|
node = FunctionNode(func=produce)
|
|
assert node.output_schema is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_output_schema_inferred_type_coercion(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Pydantic coerces compatible types (str '123' -> int 123)."""
|
|
|
|
def produce() -> _OutputModel:
|
|
return {'name': 'coerce', 'value': '42'}
|
|
|
|
node = FunctionNode(func=produce)
|
|
agent = Workflow(name='wf', edges=[(START, node)])
|
|
events, _, _ = await run_workflow(agent)
|
|
|
|
data_events = [
|
|
e
|
|
for e in events
|
|
if isinstance(e, Event)
|
|
and e.output is not None
|
|
and _NodePathBuilder.from_string(e.node_info.path).is_direct_child_of(
|
|
_NodePathBuilder.from_string('wf@1')
|
|
)
|
|
]
|
|
assert len(data_events) == 1
|
|
assert data_events[0].output == {'name': 'coerce', 'value': 42}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_output_schema_none_return(request: pytest.FixtureRequest):
|
|
"""Returning None with inferred output_schema skips validation."""
|
|
|
|
def produce_none() -> _OutputModel:
|
|
return None
|
|
|
|
node = FunctionNode(func=produce_none)
|
|
assert node.output_schema is _OutputModel
|
|
|
|
def downstream(node_input: Any) -> str:
|
|
return f'got: {node_input}'
|
|
|
|
agent = Workflow(name='wf', edges=[(START, node), (node, downstream)])
|
|
events, _, _ = await run_workflow(agent)
|
|
|
|
data_events = [
|
|
e
|
|
for e in events
|
|
if isinstance(e, Event)
|
|
and e.output is not None
|
|
and _NodePathBuilder.from_string(e.node_info.path).is_direct_child_of(
|
|
_NodePathBuilder.from_string('wf@1')
|
|
)
|
|
]
|
|
assert len(data_events) == 1
|
|
assert data_events[0].output == 'got: None'
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_output_schema_validates_returned_event_data(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""When a function returns an Event with data, output_schema validates it."""
|
|
|
|
def produce() -> _OutputModel:
|
|
return Event(output={'name': 'evt', 'value': 7})
|
|
|
|
node = FunctionNode(func=produce)
|
|
assert node.output_schema is _OutputModel
|
|
|
|
agent = Workflow(name='wf', edges=[(START, node)])
|
|
events, _, _ = await run_workflow(agent)
|
|
|
|
data_events = [
|
|
e
|
|
for e in events
|
|
if isinstance(e, Event)
|
|
and e.output is not None
|
|
and _NodePathBuilder.from_string(e.node_info.path).is_direct_child_of(
|
|
_NodePathBuilder.from_string('wf@1')
|
|
)
|
|
]
|
|
assert len(data_events) == 1
|
|
assert data_events[0].output == {'name': 'evt', 'value': 7}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_output_schema_rejects_invalid_returned_event_data(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""When a function returns an Event with invalid data, validation raises."""
|
|
|
|
def produce() -> _OutputModel:
|
|
return Event(output={'wrong_field': 'oops'})
|
|
|
|
node = FunctionNode(func=produce)
|
|
agent = Workflow(name='wf', edges=[(START, node)])
|
|
with pytest.raises(ValueError):
|
|
await run_workflow(agent)
|
|
|
|
|
|
# ── FunctionNode input_schema ──────────────────────────────────────
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_input_schema_validates_dict(request: pytest.FixtureRequest):
|
|
"""Dict input is validated and coerced through inferred input_schema."""
|
|
received = []
|
|
|
|
def process(node_input: _OutputModel) -> str:
|
|
received.append(node_input)
|
|
return 'ok'
|
|
|
|
def produce() -> dict:
|
|
return {'name': 'test', 'value': 42}
|
|
|
|
node = FunctionNode(func=process)
|
|
assert node.input_schema is _OutputModel
|
|
|
|
agent = Workflow(name='wf', edges=[(START, produce), (produce, node)])
|
|
await run_workflow(agent)
|
|
|
|
# input_schema validates before FunctionNode converts dict -> BaseModel
|
|
assert received == [_OutputModel(name='test', value=42)]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_input_schema_rejects_invalid_dict(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Dict missing required fields raises validation error."""
|
|
|
|
def process(node_input: _OutputModel) -> str:
|
|
return 'should not reach'
|
|
|
|
def produce() -> dict:
|
|
return {'name': 'test'} # missing 'value'
|
|
|
|
node = FunctionNode(func=process)
|
|
agent = Workflow(name='wf', edges=[(START, produce), (produce, node)])
|
|
with pytest.raises(ValueError):
|
|
await run_workflow(agent)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_input_schema_coerces_types(request: pytest.FixtureRequest):
|
|
"""Pydantic coerces compatible types in input (str '5' -> int 5)."""
|
|
received = []
|
|
|
|
def process(node_input: _OutputModel) -> str:
|
|
received.append(node_input)
|
|
return 'ok'
|
|
|
|
def produce() -> dict:
|
|
return {'name': 'test', 'value': '5'}
|
|
|
|
node = FunctionNode(func=process)
|
|
agent = Workflow(name='wf', edges=[(START, produce), (produce, node)])
|
|
await run_workflow(agent)
|
|
|
|
assert received == [_OutputModel(name='test', value=5)]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_input_schema_fills_defaults(request: pytest.FixtureRequest):
|
|
"""Inferred input_schema fills default fields."""
|
|
received = []
|
|
|
|
def process(node_input: _OtherModel) -> str:
|
|
received.append(node_input)
|
|
return 'ok'
|
|
|
|
def produce() -> dict:
|
|
return {'name': 'test', 'value': 1}
|
|
|
|
node = FunctionNode(func=process)
|
|
assert node.input_schema is _OtherModel
|
|
|
|
agent = Workflow(name='wf', edges=[(START, produce), (produce, node)])
|
|
await run_workflow(agent)
|
|
|
|
assert received == [_OtherModel(name='test', value=1, extra='default')]
|
|
|
|
|
|
def test_input_schema_no_inference_for_non_basemodel():
|
|
"""Non-BaseModel node_input hints don't trigger inference."""
|
|
|
|
def process(node_input: dict) -> str:
|
|
return 'ok'
|
|
|
|
node = FunctionNode(func=process)
|
|
assert node.input_schema is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_input_schema_none_passthrough(request: pytest.FixtureRequest):
|
|
"""None input with input_schema skips validation."""
|
|
|
|
def produce_none() -> None:
|
|
return None
|
|
|
|
def process(node_input: _OutputModel | None = None) -> str:
|
|
return f'got: {node_input}'
|
|
|
|
node = FunctionNode(func=process)
|
|
agent = Workflow(
|
|
name='wf', edges=[(START, produce_none), (produce_none, node)]
|
|
)
|
|
events, _, _ = await run_workflow(agent)
|
|
|
|
data_events = [
|
|
e
|
|
for e in events
|
|
if isinstance(e, Event)
|
|
and e.output is not None
|
|
and _NodePathBuilder.from_string(e.node_info.path).is_direct_child_of(
|
|
_NodePathBuilder.from_string('wf@1')
|
|
)
|
|
]
|
|
assert any(e.output == 'got: None' for e in data_events)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# auth_config tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestAuthConfig:
|
|
"""Tests for FunctionNode auth_config behavior."""
|
|
|
|
def test_raises_without_rerun_on_resume(self):
|
|
"""auth_config raises ValueError when rerun_on_resume is not True."""
|
|
from fastapi.openapi.models import APIKey
|
|
from fastapi.openapi.models import APIKeyIn
|
|
from google.adk.auth.auth_credential import AuthCredential
|
|
from google.adk.auth.auth_credential import AuthCredentialTypes
|
|
from google.adk.auth.auth_tool import AuthConfig
|
|
|
|
auth_config = AuthConfig(
|
|
auth_scheme=APIKey(**{'in': APIKeyIn.header, 'name': 'X-Api-Key'}),
|
|
raw_auth_credential=AuthCredential(
|
|
auth_type=AuthCredentialTypes.API_KEY,
|
|
api_key='placeholder',
|
|
),
|
|
credential_key='test_key',
|
|
)
|
|
with pytest.raises(ValueError, match='rerun_on_resume=True'):
|
|
FunctionNode(func=lambda: None, name='n', auth_config=auth_config)
|
|
|
|
def test_no_auth_config_default(self):
|
|
"""auth_config defaults to None."""
|
|
node = FunctionNode(func=lambda: None, name='n')
|
|
assert node.auth_config is None
|
|
|
|
def test_rerun_on_resume_explicit_true_with_auth(self):
|
|
"""Explicit rerun_on_resume=True with auth_config is fine."""
|
|
from fastapi.openapi.models import APIKey
|
|
from fastapi.openapi.models import APIKeyIn
|
|
from google.adk.auth.auth_credential import AuthCredential
|
|
from google.adk.auth.auth_credential import AuthCredentialTypes
|
|
from google.adk.auth.auth_tool import AuthConfig
|
|
|
|
auth_config = AuthConfig(
|
|
auth_scheme=APIKey(**{'in': APIKeyIn.header, 'name': 'X-Api-Key'}),
|
|
raw_auth_credential=AuthCredential(
|
|
auth_type=AuthCredentialTypes.API_KEY,
|
|
api_key='placeholder',
|
|
),
|
|
credential_key='test_key',
|
|
)
|
|
node = FunctionNode(
|
|
func=lambda: None,
|
|
name='n',
|
|
auth_config=auth_config,
|
|
rerun_on_resume=True,
|
|
)
|
|
assert node.rerun_on_resume is True
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# parameter_binding='node_input' tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestParameterBindingNodeInput:
|
|
"""Tests for FunctionNode with parameter_binding='node_input'."""
|
|
|
|
def test_schemas_inferred_from_signature(self):
|
|
"""input_schema and output_schema are inferred from func signature."""
|
|
|
|
def add(x: int, y: int) -> int:
|
|
"""Add two numbers."""
|
|
return x + y
|
|
|
|
node = FunctionNode(func=add, name='add', parameter_binding='node_input')
|
|
|
|
assert node.parameter_binding == 'node_input'
|
|
assert node.input_schema is not None
|
|
assert 'properties' in node.input_schema
|
|
assert 'x' in node.input_schema['properties']
|
|
assert 'y' in node.input_schema['properties']
|
|
assert node.output_schema == {'type': 'integer'}
|
|
|
|
def test_ctx_param_excluded_from_schema(self):
|
|
"""Context parameter is excluded from input_schema."""
|
|
|
|
def greet(name: str, ctx: Context) -> str:
|
|
return f'Hello, {name}!'
|
|
|
|
node = FunctionNode(
|
|
func=greet, name='greet', parameter_binding='node_input'
|
|
)
|
|
|
|
assert node.input_schema is not None
|
|
assert 'name' in node.input_schema['properties']
|
|
assert 'ctx' not in node.input_schema.get('properties', {})
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
'producer_output, add_func, expected_output',
|
|
[
|
|
pytest.param(
|
|
{'x': 3, 'y': 4},
|
|
staticmethod(lambda x, y: x + y),
|
|
7,
|
|
id='all_params_provided',
|
|
),
|
|
pytest.param(
|
|
{'x': 5},
|
|
None, # uses default func defined below
|
|
15,
|
|
id='missing_param_uses_default',
|
|
),
|
|
],
|
|
)
|
|
async def test_bind_from_node_input(
|
|
self,
|
|
request: pytest.FixtureRequest,
|
|
producer_output: dict,
|
|
add_func,
|
|
expected_output: int,
|
|
):
|
|
"""Parameters are bound from node_input dict."""
|
|
|
|
if add_func is None:
|
|
|
|
def add_func(x: int, y: int = 10):
|
|
return x + y
|
|
|
|
def produce():
|
|
return producer_output
|
|
|
|
node = FunctionNode(
|
|
func=add_func, name='add', parameter_binding='node_input'
|
|
)
|
|
|
|
agent = Workflow(
|
|
name='test_bind_from_node_input',
|
|
edges=[
|
|
(START, produce),
|
|
(produce, node),
|
|
],
|
|
)
|
|
events, _, _ = await run_workflow(agent)
|
|
assert simplify_events_with_node(events) == [
|
|
(
|
|
'test_bind_from_node_input@1/produce@1',
|
|
{'output': producer_output},
|
|
),
|
|
(
|
|
'test_bind_from_node_input@1/add@1',
|
|
{'output': expected_output},
|
|
),
|
|
]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_bind_from_node_input_missing_required(
|
|
self, request: pytest.FixtureRequest
|
|
):
|
|
"""Missing required param in node_input mode raises ValueError."""
|
|
|
|
def produce():
|
|
return {'x': 5}
|
|
|
|
def add(x: int, y: int):
|
|
return x + y
|
|
|
|
node = FunctionNode(func=add, name='add', parameter_binding='node_input')
|
|
|
|
agent = Workflow(
|
|
name='test_bind_node_input_missing',
|
|
edges=[
|
|
(START, produce),
|
|
(produce, node),
|
|
],
|
|
)
|
|
with pytest.raises(ValueError, match='Missing value for parameter "y"'):
|
|
await run_workflow(agent)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_bind_from_node_input_with_ctx(
|
|
self, request: pytest.FixtureRequest
|
|
):
|
|
"""Context parameter is injected alongside node_input params."""
|
|
received_ctx = []
|
|
|
|
def produce():
|
|
return {'name': 'Alice'}
|
|
|
|
def greet(name: str, ctx: Context):
|
|
received_ctx.append(ctx)
|
|
return f'Hello, {name}!'
|
|
|
|
node = FunctionNode(
|
|
func=greet, name='greet', parameter_binding='node_input'
|
|
)
|
|
|
|
agent = Workflow(
|
|
name='test_bind_node_input_ctx',
|
|
edges=[
|
|
(START, produce),
|
|
(produce, node),
|
|
],
|
|
)
|
|
events, _, _ = await run_workflow(agent)
|
|
|
|
assert len(received_ctx) == 1
|
|
assert isinstance(received_ctx[0], Context)
|
|
assert simplify_events_with_node(events) == [
|
|
(
|
|
'test_bind_node_input_ctx@1/produce@1',
|
|
{'output': {'name': 'Alice'}},
|
|
),
|
|
(
|
|
'test_bind_node_input_ctx@1/greet@1',
|
|
{'output': 'Hello, Alice!'},
|
|
),
|
|
]
|
|
|
|
def test_model_copy_preserves_parameter_binding(self):
|
|
"""model_copy preserves parameter_binding and input_schema."""
|
|
|
|
def add(x: int, y: int) -> int:
|
|
return x + y
|
|
|
|
node = FunctionNode(func=add, name='add', parameter_binding='node_input')
|
|
copied = node.model_copy(update={'name': 'add_copy'})
|
|
|
|
assert copied.parameter_binding == 'node_input'
|
|
assert copied.input_schema is not None
|
|
assert 'x' in copied.input_schema['properties']
|
|
|
|
def test_type_hints_cache(self):
|
|
"""Verifies that type hints are cached and robustly unwrapped."""
|
|
from google.adk.workflow._function_node import _get_type_hints_cached
|
|
from google.adk.workflow._function_node import _get_type_hints_for_unwrapped
|
|
|
|
def my_func(x: int, y: str) -> bool:
|
|
return True
|
|
|
|
# Clear cache first to have predictable results
|
|
_get_type_hints_for_unwrapped.cache_clear()
|
|
|
|
hints1 = _get_type_hints_cached(my_func)
|
|
assert hints1 == {'x': int, 'y': str, 'return': bool}
|
|
|
|
# Call again, should hit cache
|
|
hints2 = _get_type_hints_cached(my_func)
|
|
assert hints2 == {'x': int, 'y': str, 'return': bool}
|
|
assert _get_type_hints_for_unwrapped.cache_info().hits == 1
|
|
|
|
# Test partial
|
|
import functools
|
|
|
|
partial_func = functools.partial(my_func, x=1)
|
|
hints3 = _get_type_hints_cached(partial_func)
|
|
# Partial should unwrap to my_func and hit cache!
|
|
assert hints3 == {'x': int, 'y': str, 'return': bool}
|
|
assert _get_type_hints_for_unwrapped.cache_info().hits == 2
|
|
|
|
# Test callable object
|
|
class MyCallable:
|
|
|
|
def __call__(self, z: float) -> None:
|
|
pass
|
|
|
|
obj = MyCallable()
|
|
hints4 = _get_type_hints_cached(obj)
|
|
assert hints4 == {'z': float, 'return': type(None)}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_function_node_wrapped_partial(request: pytest.FixtureRequest):
|
|
"""Tests that FunctionNode correctly unwraps functools.partial for async/sync generators and coroutines."""
|
|
import functools
|
|
|
|
async def async_gen_fn(
|
|
prefix: str, ctx: Context
|
|
) -> AsyncGenerator[Any, None]:
|
|
yield Event(output=f'{prefix} from AsyncGen')
|
|
|
|
def sync_gen_fn(prefix: str, ctx: Context) -> Generator[Any, None, None]:
|
|
yield Event(output=f'{prefix} from SyncGen')
|
|
|
|
async def async_fn(prefix: str, ctx: Context) -> str:
|
|
return f'{prefix} from AsyncCoro'
|
|
|
|
p_async_gen = functools.partial(async_gen_fn, 'Hello')
|
|
p_sync_gen = functools.partial(sync_gen_fn, 'Hello')
|
|
p_async = functools.partial(async_fn, 'Hello')
|
|
|
|
agent = Workflow(
|
|
name='test_workflow_partial_unwrapping',
|
|
edges=[
|
|
(START, p_async_gen),
|
|
(p_async_gen, p_sync_gen),
|
|
(p_sync_gen, p_async),
|
|
],
|
|
)
|
|
events, _, _ = await run_workflow(agent)
|
|
|
|
assert simplify_events_with_node(events) == [
|
|
(
|
|
'test_workflow_partial_unwrapping@1/async_gen_fn@1',
|
|
{'output': 'Hello from AsyncGen'},
|
|
),
|
|
(
|
|
'test_workflow_partial_unwrapping@1/sync_gen_fn@1',
|
|
{'output': 'Hello from SyncGen'},
|
|
),
|
|
(
|
|
'test_workflow_partial_unwrapping@1/async_fn@1',
|
|
{'output': 'Hello from AsyncCoro'},
|
|
),
|
|
]
|