wehub-resource-sync
ec2b666284
Continuous Integration / Pre-commit Linter (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.10) (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.11) (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.12) (push) Waiting to run
Continuous Integration / Mypy Check (Python 3.13) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.10) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.11) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.12) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.13) (push) Waiting to run
Continuous Integration / Unit Tests (Python 3.14) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.10) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.11) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.12) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.13) (push) Waiting to run
Continuous Integration / A2A v0.3 Tests (Python 3.14) (push) Waiting to run
Copybara PR Handler / close-imported-pr (push) Waiting to run
842 lines
24 KiB
Python
842 lines
24 KiB
Python
# Copyright 2026 Google LLC
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""Tests for Workflow schema validation (input_schema, output_schema)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from google.adk.apps.app import App
|
|
from google.adk.events.event import Event
|
|
from google.adk.runners import Runner
|
|
from google.adk.sessions.in_memory_session_service import InMemorySessionService
|
|
from google.adk.workflow import BaseNode
|
|
from google.adk.workflow import JoinNode
|
|
from google.adk.workflow import START
|
|
from google.adk.workflow._workflow import Workflow
|
|
from google.genai import types
|
|
from pydantic import BaseModel
|
|
from pydantic import ValidationError
|
|
import pytest
|
|
|
|
from .. import testing_utils
|
|
from .workflow_testing_utils import create_parent_invocation_context
|
|
|
|
|
|
class _OutputModel(BaseModel):
|
|
name: str
|
|
value: int
|
|
|
|
|
|
class _OtherModel(BaseModel):
|
|
name: str
|
|
value: int
|
|
extra: str = 'default'
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_workflow_output_schema_validates_terminal(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Workflow.output_schema validates when downstream reads the output."""
|
|
|
|
def produce() -> dict:
|
|
return {'name': 'result', 'value': 10}
|
|
|
|
def consume(node_input: dict) -> str:
|
|
return f"got {node_input['name']}"
|
|
|
|
inner = Workflow(
|
|
name='wf',
|
|
edges=[(START, produce)],
|
|
output_schema=_OutputModel,
|
|
)
|
|
outer = Workflow(
|
|
name='outer',
|
|
edges=[(START, inner, consume)],
|
|
)
|
|
app = App(name=request.function.__name__, root_agent=outer)
|
|
runner = testing_utils.InMemoryRunner(app=app)
|
|
events = await runner.run_async(testing_utils.get_user_content('start'))
|
|
data_events = [e for e in events if isinstance(e, Event) and e.output]
|
|
assert any(e.output == 'got result' for e in data_events)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_workflow_output_schema_rejects_invalid(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Workflow.output_schema rejects invalid output when downstream reads."""
|
|
|
|
def produce_bad() -> dict:
|
|
return {'wrong_field': 'oops'}
|
|
|
|
def consume(node_input: dict) -> str:
|
|
return 'should not reach'
|
|
|
|
inner = Workflow(
|
|
name='wf',
|
|
edges=[(START, produce_bad)],
|
|
output_schema=_OutputModel,
|
|
)
|
|
outer = Workflow(
|
|
name='outer',
|
|
edges=[(START, inner, consume)],
|
|
)
|
|
app = App(name=request.function.__name__, root_agent=outer)
|
|
runner = testing_utils.InMemoryRunner(app=app)
|
|
with pytest.raises(ValueError):
|
|
await runner.run_async(testing_utils.get_user_content('start'))
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_workflow_output_schema_coerces_defaults(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Workflow.output_schema coerces terminal output (fills defaults)."""
|
|
|
|
def produce() -> dict:
|
|
return {'name': 'x', 'value': 1}
|
|
|
|
def consume(node_input: dict) -> dict:
|
|
return node_input
|
|
|
|
inner = Workflow(
|
|
name='wf',
|
|
edges=[(START, produce)],
|
|
output_schema=_OtherModel,
|
|
)
|
|
outer = Workflow(
|
|
name='outer',
|
|
edges=[(START, inner, consume)],
|
|
)
|
|
app = App(name=request.function.__name__, root_agent=outer)
|
|
runner = testing_utils.InMemoryRunner(app=app)
|
|
events = await runner.run_async(testing_utils.get_user_content('start'))
|
|
data_events = [e for e in events if isinstance(e, Event) and e.output]
|
|
consume_events = [e for e in data_events if e.node_info.name == 'consume']
|
|
assert len(consume_events) == 1
|
|
assert consume_events[0].output == {
|
|
'name': 'x',
|
|
'value': 1,
|
|
'extra': 'default',
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_nested_workflow_output_schema(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Nested workflow's output_schema validates before passing to parent."""
|
|
|
|
def inner_produce() -> dict:
|
|
return {'name': 'inner', 'value': 3}
|
|
|
|
inner_workflow = Workflow(
|
|
name='inner_wf',
|
|
edges=[(START, inner_produce)],
|
|
output_schema=_OutputModel,
|
|
)
|
|
|
|
def outer_consume(node_input: dict) -> str:
|
|
return f"got {node_input['name']}"
|
|
|
|
outer_workflow = Workflow(
|
|
name='outer_wf',
|
|
edges=[
|
|
(START, inner_workflow),
|
|
(inner_workflow, outer_consume),
|
|
],
|
|
)
|
|
app = App(name=request.function.__name__, root_agent=outer_workflow)
|
|
runner = testing_utils.InMemoryRunner(app=app)
|
|
events = await runner.run_async(testing_utils.get_user_content('start'))
|
|
|
|
data_events = [e for e in events if isinstance(e, Event) and e.output]
|
|
assert any(e.output == 'got inner' for e in data_events)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_workflow_output_schema_validates_multiple_terminals(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Each terminal output is validated when downstream reads."""
|
|
|
|
class _JoinedModel(BaseModel):
|
|
branch_a: _OtherModel
|
|
branch_b: _OtherModel
|
|
|
|
def branch_a() -> dict:
|
|
return {'name': 'from_a', 'value': 1}
|
|
|
|
def branch_b() -> dict:
|
|
return {'name': 'from_b', 'value': 2}
|
|
|
|
join = JoinNode(name='join')
|
|
|
|
inner = Workflow(
|
|
name='wf',
|
|
edges=[
|
|
(START, branch_a),
|
|
(START, branch_b),
|
|
(branch_a, join),
|
|
(branch_b, join),
|
|
],
|
|
output_schema=_JoinedModel,
|
|
)
|
|
|
|
def consume(node_input: dict) -> dict:
|
|
return node_input
|
|
|
|
outer = Workflow(
|
|
name='outer',
|
|
edges=[(START, inner, consume)],
|
|
)
|
|
app = App(name=request.function.__name__, root_agent=outer)
|
|
runner = testing_utils.InMemoryRunner(app=app)
|
|
events = await runner.run_async(testing_utils.get_user_content('start'))
|
|
consume_events = [
|
|
e
|
|
for e in events
|
|
if isinstance(e, Event) and e.output and e.node_info.name == 'consume'
|
|
]
|
|
assert len(consume_events) == 1
|
|
|
|
# Both terminal outputs should have 'extra' filled by _OtherModel default.
|
|
output = consume_events[0].output
|
|
assert output['branch_a'] == {
|
|
'name': 'from_a',
|
|
'value': 1,
|
|
'extra': 'default',
|
|
}
|
|
assert output['branch_b'] == {
|
|
'name': 'from_b',
|
|
'value': 2,
|
|
'extra': 'default',
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_workflow_output_schema_rejects_invalid_among_multiple_terminals(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""One invalid terminal among multiple raises validation error."""
|
|
|
|
class _JoinedModel(BaseModel):
|
|
branch_good: _OutputModel
|
|
branch_bad: _OutputModel
|
|
|
|
def branch_good() -> dict:
|
|
return {'name': 'ok', 'value': 1}
|
|
|
|
def branch_bad() -> dict:
|
|
return {'wrong_field': 'oops'}
|
|
|
|
join = JoinNode(name='join')
|
|
|
|
inner = Workflow(
|
|
name='wf',
|
|
edges=[
|
|
(START, branch_good),
|
|
(START, branch_bad),
|
|
(branch_good, join),
|
|
(branch_bad, join),
|
|
],
|
|
output_schema=_JoinedModel,
|
|
)
|
|
|
|
def consume(node_input: dict) -> str:
|
|
return 'should not reach'
|
|
|
|
outer = Workflow(
|
|
name='outer',
|
|
edges=[(START, inner, consume)],
|
|
)
|
|
app = App(name=request.function.__name__, root_agent=outer)
|
|
runner = testing_utils.InMemoryRunner(app=app)
|
|
with pytest.raises(ValueError):
|
|
await runner.run_async(testing_utils.get_user_content('start'))
|
|
|
|
|
|
# ── Primitive and generic type output_schema ─────────────────────────
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_workflow_output_schema_int_coerces(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Workflow output_schema=int coerces string to int at read time."""
|
|
|
|
def produce() -> str:
|
|
return '42'
|
|
|
|
def consume(node_input: int) -> int:
|
|
return node_input
|
|
|
|
inner = Workflow(
|
|
name='wf',
|
|
edges=[(START, produce)],
|
|
output_schema=int,
|
|
)
|
|
outer = Workflow(name='outer', edges=[(START, inner, consume)])
|
|
app = App(name=request.function.__name__, root_agent=outer)
|
|
runner = testing_utils.InMemoryRunner(app=app)
|
|
events = await runner.run_async(testing_utils.get_user_content('start'))
|
|
consume_events = [
|
|
e
|
|
for e in events
|
|
if isinstance(e, Event)
|
|
and e.output is not None
|
|
and e.node_info.name == 'consume'
|
|
]
|
|
assert len(consume_events) == 1
|
|
assert consume_events[0].output == 42
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_workflow_output_schema_int_rejects_invalid(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Workflow output_schema=int rejects non-coercible value at read time."""
|
|
|
|
def produce() -> dict:
|
|
return {'key': 'value'}
|
|
|
|
def consume(node_input: int) -> int:
|
|
return node_input
|
|
|
|
inner = Workflow(
|
|
name='wf',
|
|
edges=[(START, produce)],
|
|
output_schema=int,
|
|
)
|
|
outer = Workflow(name='outer', edges=[(START, inner, consume)])
|
|
app = App(name=request.function.__name__, root_agent=outer)
|
|
runner = testing_utils.InMemoryRunner(app=app)
|
|
with pytest.raises(ValueError):
|
|
await runner.run_async(testing_utils.get_user_content('start'))
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_workflow_output_schema_list_of_str(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Workflow output_schema=list[str] validates list output at read time."""
|
|
|
|
def produce() -> list:
|
|
return ['a', 'b', 'c']
|
|
|
|
def consume(node_input: list) -> list:
|
|
return node_input
|
|
|
|
inner = Workflow(
|
|
name='wf',
|
|
edges=[(START, produce)],
|
|
output_schema=list[str],
|
|
)
|
|
outer = Workflow(name='outer', edges=[(START, inner, consume)])
|
|
app = App(name=request.function.__name__, root_agent=outer)
|
|
runner = testing_utils.InMemoryRunner(app=app)
|
|
events = await runner.run_async(testing_utils.get_user_content('start'))
|
|
consume_events = [
|
|
e
|
|
for e in events
|
|
if isinstance(e, Event)
|
|
and e.output is not None
|
|
and e.node_info.name == 'consume'
|
|
]
|
|
assert len(consume_events) == 1
|
|
assert consume_events[0].output == ['a', 'b', 'c']
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_workflow_output_schema_list_of_basemodel(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Workflow output_schema=list[BaseModel] validates and serializes."""
|
|
|
|
def produce() -> list:
|
|
return [
|
|
{'name': 'x', 'value': 1},
|
|
{'name': 'y', 'value': 2},
|
|
]
|
|
|
|
def consume(node_input: list) -> list:
|
|
return node_input
|
|
|
|
inner = Workflow(
|
|
name='wf',
|
|
edges=[(START, produce)],
|
|
output_schema=list[_OutputModel],
|
|
)
|
|
outer = Workflow(name='outer', edges=[(START, inner, consume)])
|
|
app = App(name=request.function.__name__, root_agent=outer)
|
|
runner = testing_utils.InMemoryRunner(app=app)
|
|
events = await runner.run_async(testing_utils.get_user_content('start'))
|
|
consume_events = [
|
|
e
|
|
for e in events
|
|
if isinstance(e, Event)
|
|
and e.output is not None
|
|
and e.node_info.name == 'consume'
|
|
]
|
|
assert len(consume_events) == 1
|
|
assert consume_events[0].output == [
|
|
{'name': 'x', 'value': 1},
|
|
{'name': 'y', 'value': 2},
|
|
]
|
|
|
|
|
|
# ── End-to-end: input_schema + output_schema combined ──────────────
|
|
|
|
|
|
class _TaskOutput(BaseModel):
|
|
result: str
|
|
score: int
|
|
|
|
|
|
class _ReviewResult(BaseModel):
|
|
result: str
|
|
score: int
|
|
reviewer: str = 'auto'
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_e2e_input_and_output_schema_pipeline(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""output_schema on producer + input_schema on consumer validates both."""
|
|
|
|
def produce() -> _TaskOutput:
|
|
return {'result': 'done', 'score': 88}
|
|
|
|
def consume(node_input: _TaskOutput) -> str:
|
|
return f'result={node_input.result}, score={node_input.score}'
|
|
|
|
agent = Workflow(
|
|
name='wf',
|
|
edges=[(START, produce), (produce, consume)],
|
|
)
|
|
app = App(name=request.function.__name__, root_agent=agent)
|
|
runner = testing_utils.InMemoryRunner(app=app)
|
|
events = await runner.run_async(testing_utils.get_user_content('start'))
|
|
|
|
data_events = [e for e in events if isinstance(e, Event) and e.output]
|
|
assert any(e.output == 'result=done, score=88' for e in data_events)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_e2e_output_schema_fails_before_input_schema(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Producer output_schema failure prevents consumer from running."""
|
|
|
|
def produce_bad() -> _TaskOutput:
|
|
return {'wrong': 'shape'}
|
|
|
|
def consume(node_input: _TaskOutput) -> str:
|
|
return 'should not reach'
|
|
|
|
agent = Workflow(
|
|
name='wf',
|
|
edges=[(START, produce_bad), (produce_bad, consume)],
|
|
)
|
|
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('start'))
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_e2e_fan_out_join_with_schemas(
|
|
request: pytest.FixtureRequest,
|
|
):
|
|
"""Fan-out -> join -> consume with input/output schemas."""
|
|
|
|
class _JoinedResult(BaseModel):
|
|
analyzer: dict
|
|
summarizer: dict
|
|
|
|
def analyzer() -> _TaskOutput:
|
|
return {'result': 'analysis complete', 'score': 92}
|
|
|
|
def summarizer() -> _TaskOutput:
|
|
return {'result': 'summary complete', 'score': 78}
|
|
|
|
join = JoinNode(name='join')
|
|
|
|
def reviewer(node_input: dict) -> _ReviewResult:
|
|
a = node_input['analyzer']
|
|
s = node_input['summarizer']
|
|
return {
|
|
'result': f"{a['result']} + {s['result']}",
|
|
'score': (a['score'] + s['score']) // 2,
|
|
}
|
|
|
|
agent = Workflow(
|
|
name='wf',
|
|
edges=[
|
|
(START, analyzer),
|
|
(START, summarizer),
|
|
(analyzer, join),
|
|
(summarizer, join),
|
|
(join, reviewer),
|
|
],
|
|
)
|
|
app = App(name=request.function.__name__, root_agent=agent)
|
|
runner = testing_utils.InMemoryRunner(app=app)
|
|
events = await runner.run_async(testing_utils.get_user_content('start'))
|
|
|
|
terminal = [
|
|
e
|
|
for e in events
|
|
if isinstance(e, Event)
|
|
and e.output is not None
|
|
and e.node_info.name == 'reviewer'
|
|
]
|
|
assert len(terminal) == 1
|
|
assert terminal[0].output == {
|
|
'result': 'analysis complete + summary complete',
|
|
'score': 85,
|
|
'reviewer': 'auto',
|
|
}
|
|
assert 'wf' in terminal[0].node_info.path
|
|
assert 'reviewer' in terminal[0].node_info.path
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_start_node_with_str_input_schema():
|
|
"""input_schema=str parses user text."""
|
|
|
|
class _AssertingNode(BaseNode):
|
|
|
|
async def _run_impl(
|
|
self, *, ctx: Context, node_input: Any
|
|
) -> AsyncGenerator[Any, None]:
|
|
assert node_input == 'hello'
|
|
yield 'done'
|
|
|
|
node = _AssertingNode(name='node', input_schema=str)
|
|
wf = Workflow(name='wf', edges=[(START, node)])
|
|
|
|
ss = InMemorySessionService()
|
|
runner = Runner(app_name='test', node=wf, session_service=ss)
|
|
session = await ss.create_session(app_name='test', user_id='u')
|
|
|
|
msg = types.Content(parts=[types.Part(text='hello')], role='user')
|
|
events = []
|
|
async for event in runner.run_async(
|
|
user_id='u', session_id=session.id, new_message=msg
|
|
):
|
|
events.append(event)
|
|
|
|
data_events = [e for e in events if isinstance(e, Event) and e.output]
|
|
assert any(e.output == 'done' for e in data_events)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_start_node_with_int_input_schema():
|
|
"""input_schema=int parses user text to int."""
|
|
|
|
class _AssertingNode(BaseNode):
|
|
|
|
async def _run_impl(
|
|
self, *, ctx: Context, node_input: Any
|
|
) -> AsyncGenerator[Any, None]:
|
|
assert node_input == 42
|
|
yield 'done'
|
|
|
|
node = _AssertingNode(name='node', input_schema=int)
|
|
wf = Workflow(name='wf', edges=[(START, node)])
|
|
|
|
ss = InMemorySessionService()
|
|
runner = Runner(app_name='test', node=wf, session_service=ss)
|
|
session = await ss.create_session(app_name='test', user_id='u')
|
|
|
|
msg = types.Content(parts=[types.Part(text='42')], role='user')
|
|
events = []
|
|
|
|
async for event in runner.run_async(
|
|
user_id='u', session_id=session.id, new_message=msg
|
|
):
|
|
events.append(event)
|
|
|
|
data_events = [e for e in events if isinstance(e, Event) and e.output]
|
|
assert any(e.output == 'done' for e in data_events)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_start_node_with_int_list_input_schema():
|
|
"""input_schema=list[int] parses JSON list."""
|
|
|
|
class _AssertingNode(BaseNode):
|
|
|
|
async def _run_impl(
|
|
self, *, ctx: Context, node_input: Any
|
|
) -> AsyncGenerator[Any, None]:
|
|
assert node_input == [1, 2, 3]
|
|
yield 'done'
|
|
|
|
node = _AssertingNode(name='node', input_schema=list[int])
|
|
wf = Workflow(name='wf', edges=[(START, node)])
|
|
|
|
ss = InMemorySessionService()
|
|
runner = Runner(app_name='test', node=wf, session_service=ss)
|
|
session = await ss.create_session(app_name='test', user_id='u')
|
|
|
|
msg = types.Content(parts=[types.Part(text='[1, 2, 3]')], role='user')
|
|
events = []
|
|
|
|
async for event in runner.run_async(
|
|
user_id='u', session_id=session.id, new_message=msg
|
|
):
|
|
events.append(event)
|
|
|
|
data_events = [e for e in events if isinstance(e, Event) and e.output]
|
|
assert any(e.output == 'done' for e in data_events)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_start_node_with_invalid_input_schema():
|
|
"""Invalid input against schema raises error."""
|
|
|
|
class _MyModel(BaseModel):
|
|
age: int
|
|
|
|
class _AssertingNode(BaseNode):
|
|
|
|
async def _run_impl(
|
|
self, *, ctx: Context, node_input: Any
|
|
) -> AsyncGenerator[Any, None]:
|
|
yield 'done'
|
|
|
|
node = _AssertingNode(name='node', input_schema=_MyModel)
|
|
wf = Workflow(name='wf', edges=[(START, node)])
|
|
|
|
ss = InMemorySessionService()
|
|
runner = Runner(app_name='test', node=wf, session_service=ss)
|
|
session = await ss.create_session(app_name='test', user_id='u')
|
|
|
|
# Pass invalid input (Content instead of dict with age)
|
|
msg = types.Content(parts=[types.Part(text='hello')], role='user')
|
|
|
|
# We expect it to raise ValidationError
|
|
with pytest.raises(ValidationError):
|
|
async for event in runner.run_async(
|
|
user_id='u', session_id=session.id, new_message=msg
|
|
):
|
|
pass
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_start_node_receives_parsed_user_content_with_schema():
|
|
"""Parsed input replaces raw Content for first node."""
|
|
|
|
class _AssertingNode(BaseNode):
|
|
|
|
async def _run_impl(
|
|
self, *, ctx: Context, node_input: Any
|
|
) -> AsyncGenerator[Any, None]:
|
|
assert isinstance(node_input, str)
|
|
assert node_input == 'hello'
|
|
yield 'done'
|
|
|
|
node = _AssertingNode(name='node', input_schema=str)
|
|
wf = Workflow(name='wf', edges=[(START, node)])
|
|
|
|
ss = InMemorySessionService()
|
|
runner = Runner(app_name='test', node=wf, session_service=ss)
|
|
session = await ss.create_session(app_name='test', user_id='u')
|
|
|
|
msg = types.Content(parts=[types.Part(text='hello')], role='user')
|
|
events = []
|
|
async for event in runner.run_async(
|
|
user_id='u', session_id=session.id, new_message=msg
|
|
):
|
|
events.append(event)
|
|
|
|
data_events = [e for e in events if isinstance(e, Event) and e.output]
|
|
assert any(e.output == 'done' for e in data_events)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_workflow_with_invalid_output_schema():
|
|
"""Workflow raises ValidationError if terminal output doesn't match output_schema."""
|
|
|
|
from pydantic import ValidationError
|
|
|
|
class _MyModel(BaseModel):
|
|
name: str
|
|
|
|
class _MyNode(BaseNode):
|
|
|
|
async def _run_impl(
|
|
self, *, ctx: Context, node_input: Any
|
|
) -> AsyncGenerator[Any, None]:
|
|
yield {'age': 10}
|
|
|
|
node = _MyNode(name='node')
|
|
wf = Workflow(name='wf', edges=[(START, node)], output_schema=_MyModel)
|
|
|
|
ss = InMemorySessionService()
|
|
runner = Runner(app_name='test', node=wf, session_service=ss)
|
|
session = await ss.create_session(app_name='test', user_id='u')
|
|
|
|
msg = types.Content(parts=[types.Part(text='hello')], role='user')
|
|
|
|
with pytest.raises(ValidationError):
|
|
async for event in runner.run_async(
|
|
user_id='u', session_id=session.id, new_message=msg
|
|
):
|
|
pass
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_node_returns_content_json_parsed():
|
|
"""Node output as types.Content containing JSON is parsed if output_schema is defined."""
|
|
|
|
class _MyModel(BaseModel):
|
|
name: str
|
|
age: int
|
|
|
|
class _MyNode(BaseNode):
|
|
|
|
async def _run_impl(
|
|
self, *, ctx: Context, node_input: Any
|
|
) -> AsyncGenerator[Any, None]:
|
|
yield types.Content(
|
|
parts=[types.Part(text='{"name": "Alice", "age": 30}')]
|
|
)
|
|
|
|
node = _MyNode(name='node', output_schema=_MyModel)
|
|
wf = Workflow(name='wf', edges=[(START, node)])
|
|
|
|
ss = InMemorySessionService()
|
|
runner = Runner(app_name='test', node=wf, session_service=ss)
|
|
session = await ss.create_session(app_name='test', user_id='u')
|
|
|
|
msg = types.Content(parts=[types.Part(text='hello')], role='user')
|
|
events = []
|
|
|
|
async for event in runner.run_async(
|
|
user_id='u', session_id=session.id, new_message=msg
|
|
):
|
|
events.append(event)
|
|
|
|
data_events = [e for e in events if isinstance(e, Event) and e.output]
|
|
|
|
assert len(data_events) == 1
|
|
assert data_events[0].output == {'name': 'Alice', 'age': 30}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_node_returns_raw_string_not_parsed():
|
|
"""Node output as raw JSON string is NOT parsed if output_schema is defined."""
|
|
from pydantic import ValidationError
|
|
|
|
class _MyModel(BaseModel):
|
|
name: str
|
|
age: int
|
|
|
|
class _MyNode(BaseNode):
|
|
|
|
async def _run_impl(
|
|
self, *, ctx: Context, node_input: Any
|
|
) -> AsyncGenerator[Any, None]:
|
|
# This should fail validation because it's a string, not a dict/model
|
|
yield '{"name": "Alice", "age": 30}'
|
|
|
|
node = _MyNode(name='node', output_schema=_MyModel)
|
|
wf = Workflow(name='wf', edges=[(START, node)])
|
|
|
|
ss = InMemorySessionService()
|
|
runner = Runner(app_name='test', node=wf, session_service=ss)
|
|
session = await ss.create_session(app_name='test', user_id='u')
|
|
|
|
msg = types.Content(parts=[types.Part(text='hello')], role='user')
|
|
|
|
with pytest.raises(ValidationError):
|
|
async for _ in runner.run_async(
|
|
user_id='u', session_id=session.id, new_message=msg
|
|
):
|
|
pass
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_output_schema_enforced_by_runtime_without_manual_validation():
|
|
"""Runtime enforces output_schema even when _run_impl doesn't call _validate_output_data."""
|
|
from pydantic import ValidationError
|
|
|
|
class _MyModel(BaseModel):
|
|
name: str
|
|
age: int
|
|
|
|
class _NaiveNode(BaseNode):
|
|
"""Node that yields raw data without any manual validation."""
|
|
|
|
async def _run_impl(
|
|
self, *, ctx: Context, node_input: Any
|
|
) -> AsyncGenerator[Any, None]:
|
|
yield {'color': 'red'} # Does NOT match _MyModel
|
|
|
|
node = _NaiveNode(name='node', output_schema=_MyModel)
|
|
wf = Workflow(name='wf', edges=[(START, node)])
|
|
|
|
ss = InMemorySessionService()
|
|
runner = Runner(app_name='test', node=wf, session_service=ss)
|
|
session = await ss.create_session(app_name='test', user_id='u')
|
|
|
|
msg = types.Content(parts=[types.Part(text='hello')], role='user')
|
|
|
|
with pytest.raises(ValidationError):
|
|
async for _ in runner.run_async(
|
|
user_id='u', session_id=session.id, new_message=msg
|
|
):
|
|
pass
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_output_schema_enforced_for_valid_raw_yield():
|
|
"""Runtime validates and coerces valid raw yields against output_schema."""
|
|
|
|
class _MyModel(BaseModel):
|
|
name: str
|
|
age: int
|
|
|
|
class _NaiveNode(BaseNode):
|
|
"""Node that yields valid raw data without manual validation."""
|
|
|
|
async def _run_impl(
|
|
self, *, ctx: Context, node_input: Any
|
|
) -> AsyncGenerator[Any, None]:
|
|
yield {'name': 'Alice', 'age': 30}
|
|
|
|
node = _NaiveNode(name='node', output_schema=_MyModel)
|
|
wf = Workflow(name='wf', edges=[(START, node)])
|
|
|
|
ss = InMemorySessionService()
|
|
runner = Runner(app_name='test', node=wf, session_service=ss)
|
|
session = await ss.create_session(app_name='test', user_id='u')
|
|
|
|
msg = types.Content(parts=[types.Part(text='hello')], role='user')
|
|
events = []
|
|
|
|
async for event in runner.run_async(
|
|
user_id='u', session_id=session.id, new_message=msg
|
|
):
|
|
events.append(event)
|
|
|
|
data_events = [e for e in events if isinstance(e, Event) and e.output]
|
|
assert len(data_events) == 1
|
|
assert data_events[0].output == {'name': 'Alice', 'age': 30}
|