130 lines
3.3 KiB
Python
130 lines
3.3 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Annotated
|
|
|
|
from pydantic import BaseModel, Field
|
|
|
|
from wf_authoring import (
|
|
NodeReturn,
|
|
ReducerCatalog,
|
|
WorkflowBuilder,
|
|
node,
|
|
output_to,
|
|
reducer,
|
|
state_field,
|
|
state_path,
|
|
)
|
|
from wf_core import ReducerRef
|
|
|
|
|
|
class ModuloConfig(BaseModel):
|
|
modulus: int = Field(gt=0)
|
|
|
|
|
|
class CounterState(BaseModel):
|
|
total: Annotated[int, state_field(reducer="wf.std.add")] = 0
|
|
|
|
|
|
class EmptyInput(BaseModel):
|
|
pass
|
|
|
|
|
|
class ModuloCounterState(BaseModel):
|
|
total: int = 0
|
|
|
|
|
|
class CounterOutput(BaseModel):
|
|
total: int
|
|
|
|
|
|
@reducer(name="wf.std.add")
|
|
def add(current: int | None, incoming: int) -> int:
|
|
"""Add incoming values into integer state."""
|
|
return (current or 0) + incoming
|
|
|
|
|
|
@reducer(name="wf.std.modulo_add", config_model=ModuloConfig)
|
|
def modulo_add(
|
|
current: int | None,
|
|
incoming: int,
|
|
config: ModuloConfig,
|
|
) -> int:
|
|
"""Add incoming values modulo a configured positive integer."""
|
|
return ((current or 0) + incoming) % config.modulus
|
|
|
|
|
|
def test_reducer_decorator_wraps_plain_callable() -> None:
|
|
assert add.definition.spec.name == "wf.std.add"
|
|
assert add.definition.spec.description == "Add incoming values into integer state."
|
|
|
|
result = add.definition.apply(
|
|
reducer=ReducerRef(name="wf.std.add"),
|
|
current_value=2,
|
|
incoming_value=3,
|
|
destination_path="state.total",
|
|
)
|
|
|
|
assert result == 5
|
|
|
|
|
|
def test_reducer_decorator_wraps_configured_basemodel_callable() -> None:
|
|
schema = modulo_add.definition.spec.config_schema
|
|
|
|
assert schema["properties"]["modulus"]["exclusiveMinimum"] == 0
|
|
|
|
result = modulo_add.definition.apply(
|
|
reducer=ReducerRef(name="wf.std.modulo_add", config={"modulus": 10}),
|
|
current_value=8,
|
|
incoming_value=5,
|
|
destination_path="state.total",
|
|
)
|
|
|
|
assert result == 3
|
|
|
|
|
|
def test_reducer_catalog_exposes_definitions_and_specs() -> None:
|
|
catalog = ReducerCatalog.from_reducers(add, modulo_add)
|
|
|
|
assert set(catalog.definitions) == {"wf.std.add", "wf.std.modulo_add"}
|
|
assert catalog.specs["wf.std.add"].name == "wf.std.add"
|
|
assert catalog.specs["wf.std.modulo_add"].name == "wf.std.modulo_add"
|
|
|
|
|
|
def test_builder_executes_with_custom_reducer_catalog() -> None:
|
|
@node
|
|
def emit(_: EmptyInput) -> NodeReturn[CounterOutput]:
|
|
return NodeReturn(outcome="ok", output=CounterOutput(total=4))
|
|
|
|
builder = WorkflowBuilder(
|
|
name="custom_reducer",
|
|
input_schema=EmptyInput,
|
|
state_schema=CounterState,
|
|
output_schema=CounterOutput,
|
|
reducers=ReducerCatalog.from_reducers(add),
|
|
)
|
|
step = builder.use(
|
|
emit,
|
|
output=[output_to("total", state_path("total"))],
|
|
)
|
|
builder.set_entry_point(step)
|
|
builder.connect(step, "ok", "__end__")
|
|
|
|
run = builder.execute({})
|
|
|
|
assert run.state["total"] == 4
|
|
|
|
|
|
def test_state_field_accepts_configured_reducer_reference() -> None:
|
|
class State(BaseModel):
|
|
total: Annotated[
|
|
int,
|
|
state_field(
|
|
reducer={"name": "wf.std.modulo_add", "config": {"modulus": 10}}
|
|
),
|
|
] = 0
|
|
|
|
field = State.model_fields["total"].metadata[0] # pylint: disable=unsubscriptable-object
|
|
|
|
assert field.reducer.name == "wf.std.modulo_add"
|
|
assert field.reducer.config["modulus"] == 10
|