the second case of success

This commit is contained in:
lda
2026-05-07 04:21:55 +07:00 Verified
parent 0ef6590184
commit b8f56b78e4
8 changed files with 336 additions and 46 deletions
+3
View File
@@ -48,6 +48,7 @@ from .nodes import (
build_registry, build_registry,
node, node,
) )
from .schemas import StateFieldMetadata, state_field
from .subgraph import subgraph_node from .subgraph import subgraph_node
__all__ = [ __all__ = [
@@ -66,6 +67,7 @@ __all__ = [
"AsyncRegistryHandler", "AsyncRegistryHandler",
"SyncRegistryHandler", "SyncRegistryHandler",
"SequenceInput", "SequenceInput",
"StateFieldMetadata",
"TruthyInput", "TruthyInput",
"ValueOutput", "ValueOutput",
"WorkflowBuilder", "WorkflowBuilder",
@@ -94,6 +96,7 @@ __all__ = [
"pick_key", "pick_key",
"node", "node",
"state", "state",
"state_field",
"state_path", "state_path",
"subgraph_node", "subgraph_node",
"truthy", "truthy",
+71 -4
View File
@@ -2,6 +2,7 @@ from __future__ import annotations
from collections.abc import Mapping from collections.abc import Mapping
from dataclasses import dataclass, field from dataclasses import dataclass, field
import re
from typing import Any, Literal, TypeAlias, cast from typing import Any, Literal, TypeAlias, cast
from wf_core import ( from wf_core import (
@@ -49,6 +50,48 @@ def _step_id(ref: StepRef) -> str:
return ref.id return ref.id
def _slug_id(value: str) -> str:
slug = re.sub(r"[^0-9A-Za-z_]+", "_", value).strip("_").lower()
return slug or "step"
def _auto_input_map(
spec: NodeSpec[Any, Any],
*,
input_schema: SchemaRef,
state_schema: StateSchema,
) -> dict[str, str]:
return {
_auto_source_path(field, input_schema=input_schema, state_schema=state_schema): field
for field in spec.input_model.model_json_schema().get("properties", {})
}
def _auto_output_map(
spec: NodeSpec[Any, Any],
*,
state_schema: StateSchema,
) -> dict[str, str]:
return {
field: f"state.{field}"
for field in spec.output_model.model_json_schema().get("properties", {})
if field in state_schema.fields
}
def _auto_source_path(
field: str,
*,
input_schema: SchemaRef,
state_schema: StateSchema,
) -> str:
if field in state_schema.fields:
return f"state.{field}"
if field in input_schema.properties:
return f"input.{field}"
return f"state.{field}"
@dataclass(slots=True) @dataclass(slots=True)
class WorkflowBuilder: class WorkflowBuilder:
name: str name: str
@@ -70,23 +113,47 @@ class WorkflowBuilder:
self, self,
spec: NodeSpec[Any, Any], spec: NodeSpec[Any, Any],
*, *,
id: str, id: str | None = None,
in_map: MapArg | None = None, in_map: MapArg | None = None,
out_map: MapArg | None = None, out_map: MapArg | None = None,
desc: str | None = None, desc: str | None = None,
) -> NodeUse: ) -> NodeUse:
self.node_specs[spec.name] = spec self.node_specs[spec.name] = spec
normalized_input_schema = cast(SchemaRef, self.input_schema)
normalized_state_schema = cast(StateSchema, self.state_schema)
node = NodeUse( node = NodeUse(
id=id, id=id or self._next_step_id(_slug_id(spec.name)),
type="node", type="node",
node=spec.name, node=spec.name,
desc=desc or spec.description, desc=desc or spec.description,
in_map=_normalize_mapping(in_map), in_map=(
out_map=_normalize_mapping(out_map), _auto_input_map(
spec,
input_schema=normalized_input_schema,
state_schema=normalized_state_schema,
)
if in_map is None
else _normalize_mapping(in_map)
),
out_map=(
_auto_output_map(spec, state_schema=normalized_state_schema)
if out_map is None
else _normalize_mapping(out_map)
),
) )
self.nodes.append(node) self.nodes.append(node)
return node return node
def _next_step_id(self, base: str) -> str:
"""Return a stable unused step id based on the requested base name."""
used_ids = {_step_id(node) for node in self.nodes}
if base not in used_ids:
return base
suffix = 2
while f"{base}_{suffix}" in used_ids:
suffix += 1
return f"{base}_{suffix}"
def condition(self, *, id: str, check: CoreCondition | Expr) -> ConditionNode: def condition(self, *, id: str, check: CoreCondition | Expr) -> ConditionNode:
node = ConditionNode( node = ConditionNode(
id=id, id=id,
+54 -2
View File
@@ -1,6 +1,7 @@
from __future__ import annotations from __future__ import annotations
from typing import Any from dataclasses import dataclass
from typing import Any, Literal
from pydantic import BaseModel, TypeAdapter from pydantic import BaseModel, TypeAdapter
@@ -10,6 +11,23 @@ SchemaLike = SchemaRef | type[BaseModel] | type[Any] | dict[str, Any]
StateSchemaLike = StateSchema | type[BaseModel] | type[Any] | dict[str, Any] StateSchemaLike = StateSchema | type[BaseModel] | type[Any] | dict[str, Any]
@dataclass(frozen=True, slots=True)
class StateFieldMetadata:
"""Authoring metadata attached to BaseModel state fields."""
merge_strategy: Literal["replace", "append", "merge_object"] = "replace"
trace: bool = True
def state_field(
*,
merge_strategy: Literal["replace", "append", "merge_object"] = "replace",
trace: bool = True,
) -> StateFieldMetadata:
"""Declare workflow state behavior for an Annotated BaseModel field."""
return StateFieldMetadata(merge_strategy=merge_strategy, trace=trace)
def schema_ref_from(value: SchemaLike) -> SchemaRef: def schema_ref_from(value: SchemaLike) -> SchemaRef:
"""Coerce an authoring schema declaration into a core schema reference.""" """Coerce an authoring schema declaration into a core schema reference."""
if isinstance(value, SchemaRef): if isinstance(value, SchemaRef):
@@ -29,13 +47,34 @@ def state_schema_from(value: StateSchemaLike) -> StateSchema:
return StateSchema.model_validate(value) return StateSchema.model_validate(value)
schema = schema_ref_from(value) schema = schema_ref_from(value)
metadata_by_name = _state_metadata_by_name(value)
fields = { fields = {
name: StateField(type=_state_field_type(property_schema)) name: StateField(
type=_state_field_type(property_schema),
merge_strategy=metadata_by_name.get(
name, StateFieldMetadata()
).merge_strategy,
trace=metadata_by_name.get(name, StateFieldMetadata()).trace,
default=_state_field_default(value, name, property_schema),
)
for name, property_schema in schema.properties.items() for name, property_schema in schema.properties.items()
} }
return StateSchema(fields=fields) return StateSchema(fields=fields)
def _state_metadata_by_name(value: object) -> dict[str, StateFieldMetadata]:
if not isinstance(value, type) or not issubclass(value, BaseModel):
return {}
metadata: dict[str, StateFieldMetadata] = {}
for name, field_info in value.model_fields.items():
for item in field_info.metadata:
if isinstance(item, StateFieldMetadata):
metadata[name] = item
break
return metadata
def _state_field_type(property_schema: object) -> str: def _state_field_type(property_schema: object) -> str:
if not isinstance(property_schema, dict): if not isinstance(property_schema, dict):
return "object" return "object"
@@ -47,3 +86,16 @@ def _state_field_type(property_schema: object) -> str:
if "items" in property_schema: if "items" in property_schema:
return "array" return "array"
return "object" return "object"
def _state_field_default(
value: object,
field_name: str,
property_schema: object,
) -> object:
if isinstance(value, type) and issubclass(value, BaseModel):
field_info = value.model_fields[field_name]
if not field_info.is_required():
return field_info.get_default(call_default_factory=True)
return None
+1
View File
@@ -18,6 +18,7 @@ class StateField(BaseModel):
type: str type: str
merge_strategy: Literal["replace", "append", "merge_object"] = "replace" merge_strategy: Literal["replace", "append", "merge_object"] = "replace"
trace: bool = True trace: bool = True
default: Any = None
class StateSchema(BaseModel): class StateSchema(BaseModel):
+9 -1
View File
@@ -1,15 +1,23 @@
from __future__ import annotations from __future__ import annotations
from copy import deepcopy
from .model import Workflow from .model import Workflow
from .run_state import ExecutionFrame, FrameStatus, RunState, RunStatus from .run_state import ExecutionFrame, FrameStatus, RunState, RunStatus
def create_run_state(workflow: Workflow, workflow_input: dict[str, object]) -> RunState: def create_run_state(workflow: Workflow, workflow_input: dict[str, object]) -> RunState:
state = {
name: deepcopy(field.default)
for name, field in workflow.state_schema.fields.items()
if field.default is not None
}
state.update(dict(workflow_input))
run = RunState( run = RunState(
workflow_name=workflow.name, workflow_name=workflow.name,
status=RunStatus.PENDING, status=RunStatus.PENDING,
workflow_input=dict(workflow_input), workflow_input=dict(workflow_input),
state=dict(workflow_input), state=state,
frames={ frames={
"root": ExecutionFrame( "root": ExecutionFrame(
id="root", id="root",
-1
View File
@@ -1,7 +1,6 @@
from __future__ import annotations from __future__ import annotations
from dataclasses import dataclass from dataclasses import dataclass
from re import L
from fastmcp.server.transforms import Namespace from fastmcp.server.transforms import Namespace
+115 -3
View File
@@ -1,10 +1,11 @@
from __future__ import annotations from __future__ import annotations
from typing import TypedDict from typing import Annotated, TypedDict
from pydantic import BaseModel from pydantic import BaseModel, Field
from wf_authoring import WorkflowBuilder from wf_authoring import WorkflowBuilder, build_registry, node, state_field
from wf_core import RunStatus, execute_workflow
class WorkflowInput(BaseModel): class WorkflowInput(BaseModel):
@@ -25,6 +26,39 @@ class TypedDictInput(TypedDict):
text: str text: str
class AutoBindInput(BaseModel):
text: str
count: int
class AutoBindOutput(BaseModel):
text: str
count: int
class AutoBindState(BaseModel):
text: str
count: int
class AppendState(BaseModel):
items: Annotated[list[str], state_field(merge_strategy="append")] = Field(
default_factory=list
)
class DefaultedState(BaseModel):
items: list[str] = Field(default_factory=list)
metadata: dict[str, str] = Field(default_factory=dict)
explicit: int = 3
@node(name="test.auto_bind")
def auto_bind_node(input: AutoBindInput) -> AutoBindOutput:
"""Return updated fields using automatically mapped state input."""
return AutoBindOutput(text=input.text.upper(), count=input.count + 1)
def test_builder_accepts_basemodel_classes_for_workflow_schemas() -> None: def test_builder_accepts_basemodel_classes_for_workflow_schemas() -> None:
builder = WorkflowBuilder( builder = WorkflowBuilder(
name="model_schema_demo", name="model_schema_demo",
@@ -56,3 +90,81 @@ def test_builder_accepts_typeddict_for_json_schema_refs() -> None:
workflow = builder.compile() workflow = builder.compile()
assert workflow.input_schema.properties["text"]["type"] == "string" assert workflow.input_schema.properties["text"]["type"] == "string"
def test_state_basemodel_can_declare_merge_strategy_with_annotated_metadata() -> None:
builder = WorkflowBuilder(
name="state_metadata_demo",
input_schema=WorkflowInput,
state_schema=AppendState,
output_schema=WorkflowOutput,
start="start",
)
workflow = builder.compile()
assert workflow.state_schema.fields["items"].type == "array"
assert workflow.state_schema.fields["items"].merge_strategy == "append"
def test_state_basemodel_seeds_safe_initial_defaults() -> None:
builder = WorkflowBuilder(
name="state_defaults_demo",
input_schema=WorkflowInput,
state_schema=DefaultedState,
output_schema=WorkflowOutput,
start="start",
)
workflow = builder.compile()
assert workflow.state_schema.fields["items"].default == []
assert workflow.state_schema.fields["metadata"].default == {}
assert workflow.state_schema.fields["explicit"].default == 3
def test_builder_auto_binds_matching_node_inputs_and_outputs_to_state() -> None:
builder = WorkflowBuilder(
name="auto_bind_demo",
input_schema=AutoBindInput,
state_schema=AutoBindState,
output_schema=AutoBindOutput,
start="update",
)
step = builder.use(auto_bind_node, id="update")
builder.connect(step, "ok", "__end__")
workflow = builder.compile()
run = execute_workflow(
workflow,
{"text": "hello", "count": 1},
build_registry(auto_bind_node),
)
assert step.in_map == {
"state.text": "text",
"state.count": "count",
}
assert step.out_map == {
"text": "state.text",
"count": "state.count",
}
assert run.status == RunStatus.COMPLETED
assert run.state["text"] == "HELLO"
assert run.state["count"] == 2
def test_builder_can_auto_id_node_uses_from_spec_name() -> None:
builder = WorkflowBuilder(
name="auto_id_demo",
input_schema=AutoBindInput,
state_schema=AutoBindState,
output_schema=AutoBindOutput,
start="test_auto_bind",
)
first = builder.use(auto_bind_node)
second = builder.use(auto_bind_node)
assert first.id == "test_auto_bind"
assert second.id == "test_auto_bind_2"
+82 -34
View File
@@ -12,17 +12,18 @@ also:
from pprint import pprint from pprint import pprint
import random import random
from typing import Any, Final, Literal, TypedDict from typing import Annotated, Any, Final, Literal, TypedDict
from wf_authoring.builder import WorkflowBuilder from wf_authoring.builder import WorkflowBuilder
from pydantic import BaseModel from pydantic import BaseModel, Field
from wf_authoring import node from wf_authoring import node
from wf_authoring import NodeReturn from wf_authoring import NodeReturn
from wf_authoring.dsl.conditions import compile_condition, expr, state from wf_authoring.dsl.conditions import expr, state
from wf_authoring.nodes.registry import build_registry from wf_authoring.nodes.registry import build_registry
from wf_core.model import Workflow from wf_authoring.schemas import state_field
from wf_core.run_state import RunStatus
from wf_core.runtime.engine import execute_workflow from wf_core.runtime.engine import execute_workflow
from wf_core.tokens import END from wf_core.tokens import END
@@ -47,7 +48,9 @@ class SophisticatedRates(TypedDict):
class Counters(BaseModel): class Counters(BaseModel):
# alot of state["thing"] in the og code, so i think this is the design? # alot of state["thing"] in the og code, so i think this is the design?
counter: SophisticatedCounter # or, that is because of langgraph limitation, counter: SophisticatedCounter = Field(
default_factory=lambda: SophisticatedCounter(c_10=0, c_80=0)
) # or, that is because of langgraph limitation,
# id prefer counter.update with dict.update override (Overwrite, not like that ever worked) # id prefer counter.update with dict.update override (Overwrite, not like that ever worked)
simple_counter: int # add simple_counter: int # add
@@ -59,9 +62,11 @@ class Countdown(BaseModel):
class ContextInput(BaseModel): class ContextInput(BaseModel):
context: Context # final! context: Context # final!
class how_do_i_explain_this(BaseModel): class how_do_i_explain_this(BaseModel):
pity_120_available: bool = True pity_120_available: bool = True
class Input(Counters, ContextInput, Countdown, how_do_i_explain_this): class Input(Counters, ContextInput, Countdown, how_do_i_explain_this):
"input of the graph" "input of the graph"
@@ -101,13 +106,12 @@ class Context(TypedDict): # dataclass support? no. i mean langgraph doesnt.
"still very convoluted logic" "still very convoluted logic"
class Storage(BaseModel): class Storage(BaseModel):
"i NEED to do this?" "i NEED to do this?"
storage: list[Entity] # add! storage: Annotated[list[Entity], state_field(merge_strategy="append")] = Field(
default_factory=list
) # add!
class CurrentRoll(BaseModel): class CurrentRoll(BaseModel):
@@ -117,8 +121,12 @@ class CurrentRoll(BaseModel):
class ThisStorage(Storage, CurrentRoll): ... class ThisStorage(Storage, CurrentRoll): ...
class PartialRates(SophisticatedRates, total=False):
pass
class Rates(BaseModel): class Rates(BaseModel):
rates: SophisticatedRates # or_! rates: Annotated[PartialRates, state_field(merge_strategy="merge_object")] # or_!
class CurrentPools(BaseModel): class CurrentPools(BaseModel):
@@ -217,13 +225,16 @@ def rate_booster(c: Counters) -> NodeReturn[Nothing]:
def s(o: str) -> NodeReturn[Nothing]: def s(o: str) -> NodeReturn[Nothing]:
return NodeReturn(o, Nothing()) return NodeReturn(o, Nothing())
def _popped(storage: list[Entity]) -> bool: def _popped(storage: list[Entity]) -> bool:
return any(c["category"] == "240" for c in reversed(storage)) return any(c["category"] == "240" for c in reversed(storage))
@node @node
def popped(s: Storage) -> how_do_i_explain_this: def popped(s: Storage) -> how_do_i_explain_this:
return how_do_i_explain_this(pity_120_available=_popped(s.storage)) return how_do_i_explain_this(pity_120_available=_popped(s.storage))
@node(outcomes=("240", "80", "10", "1")) @node(outcomes=("240", "80", "10", "1"))
def pre_roll_router(c: CountersContextOutputInputAhhModelType) -> NodeReturn[Nothing]: def pre_roll_router(c: CountersContextOutputInputAhhModelType) -> NodeReturn[Nothing]:
"""another edge. """another edge.
@@ -250,8 +261,17 @@ def pre_roll_router(c: CountersContextOutputInputAhhModelType) -> NodeReturn[Not
class RateChange: class RateChange:
@node(name="force 6* rating") @node(name="force 6* rating")
@staticmethod @staticmethod
def r80(_: Nothing) -> Rates: def r80(r: Rates) -> Rates:
return Rates.model_validate({"rates": {"r_1": 0, "r_10": 0}}) return Rates.model_validate(
{
"rates": {
"r_1": 0,
"r_10": 0,
"r_80": r.rates["r_80"],
"r_240": r.rates["r_240"],
}
}
)
@node(name="force banner rating") @node(name="force banner rating")
@staticmethod @staticmethod
@@ -262,14 +282,23 @@ class RateChange:
@node(name="force 5*+ rating") @node(name="force 5*+ rating")
@staticmethod @staticmethod
def r10(_: Nothing) -> Rates: def r10(r: Rates) -> Rates:
return Rates.model_validate({"rates": {"r_1": 0}}) return Rates.model_validate(
{
"rates": {
"r_1": 0,
"r_10": r.rates["r_10"],
"r_80": r.rates["r_80"],
"r_240": r.rates["r_240"],
}
}
)
@node(name="buff 6* rating") @node(name="buff 6* rating")
@staticmethod @staticmethod
def r65(state: CountersContext) -> Rates: def r65(state: CountersContext) -> Rates:
c = state.counter c = state.counter
assert c["c_80"] > 65, f"routed wrongly, 80 pity currently at {c["c_80"]}" assert c["c_80"] >= 65, f"routed wrongly, 80 pity currently at {c["c_80"]}"
n = c["c_80"] - 64 n = c["c_80"] - 64
rpn = n * 0.05 rpn = n * 0.05
br = state.context["initial_rates"] br = state.context["initial_rates"]
@@ -300,24 +329,24 @@ class RateChange:
class CounterUp: class CounterUp:
@node(name="counter 6* reset") @node(name="counter 6* reset")
@staticmethod @staticmethod
def c80(_: Nothing) -> Counters: def c80(c: Counters) -> Counters:
return Counters.model_validate( return Counters.model_validate(
{ {
"counter": { "counter": {
"c_80": 0, "c_80": 0,
"c_10": 0, "c_10": 0,
}, },
"simple_counter": c.simple_counter,
} }
) )
@node(name="counter 5* reset") @node(name="counter 5* reset")
@staticmethod @staticmethod
def c10(_: Nothing) -> Counters: def c10(c: Counters) -> Counters:
return Counters.model_validate( return Counters.model_validate(
{ {
"counter": { "counter": {"c_10": 0, "c_80": c.counter["c_80"]}, # merge with or_!
"c_10": 0, "simple_counter": c.simple_counter,
}, # merge with or_!
} }
) )
@@ -326,10 +355,8 @@ class CounterUp:
def c1(state: Counters) -> Counters: def c1(state: Counters) -> Counters:
c = state.counter c = state.counter
return Counters( return Counters(
**{ simple_counter=state.simple_counter + 1,
"simple_counter": 1, counter={"c_10": (c["c_10"] + 1) % 10, "c_80": (c["c_80"] + 1) % 80},
"counter": {"c_10": (c["c_10"] + 1) % 10, "c_80": (c["c_80"] + 1) % 80},
}
) )
@@ -340,7 +367,7 @@ class RatesContextInput(Rates, ContextInput): ...
def prep(state: RatesContextInput) -> CurrentPools: def prep(state: RatesContextInput) -> CurrentPools:
r = state.rates r = state.rates
p = state.context["pool"] p = state.context["pool"]
t: Final[tuple[tuple, ...]] = ( # greatest hack t: Final[tuple[tuple[str, str, str], ...]] = ( # greatest hack
("1", "r_1", "n_1"), ("1", "r_1", "n_1"),
("10", "r_10", "n_10"), ("10", "r_10", "n_10"),
("80", "r_80", "n_80"), ("80", "r_80", "n_80"),
@@ -435,7 +462,7 @@ gacha.connect("r_g10", "ok", "prep")
gacha.connect("roll", "ok", "post_roll_router") # missed this! good job gacha.connect("roll", "ok", "post_roll_router") # missed this! good job
gacha.use(post_roll_router, id="post_roll_router") gacha.use(post_roll_router, id="post_roll_router")
gacha.use(popped, id = "reset_avail") gacha.use(popped, id="reset_avail")
gacha.connect("reset_avail", "ok", "c_80") gacha.connect("reset_avail", "ok", "c_80")
for outcome, node_id in { for outcome, node_id in {
@@ -494,7 +521,11 @@ context: Final[Context] = {
def build_input_lite( def build_input_lite(
rolling: int, rolled_previously: int = 0, until_5: int = 10, until_6: int = 80, good_stuff: bool = False rolling: int,
rolled_previously: int = 0,
until_5: int = 10,
until_6: int = 80,
good_stuff: bool = False,
) -> dict[str, Any]: ) -> dict[str, Any]:
if not 0 < until_5 <= 10: if not 0 < until_5 <= 10:
print(f"{until_5 = } invalid, idc") print(f"{until_5 = } invalid, idc")
@@ -507,13 +538,17 @@ def build_input_lite(
"c_10": 10 - until_5, "c_10": 10 - until_5,
"c_80": 80 - until_6, "c_80": 80 - until_6,
}, },
"pity_120_available": not good_stuff "pity_120_available": not good_stuff,
} }
def build_input(context: Context): def build_input(context: Context):
def dec( def dec(
rolling: int, rolled_previously: int = 0, until_5: int = 10, until_6: int = 80, good_stuff: bool = False rolling: int,
rolled_previously: int = 0,
until_5: int = 10,
until_6: int = 80,
good_stuff: bool = False,
) -> Input: ) -> Input:
r = build_input_lite(rolling, rolled_previously, until_5, until_6, good_stuff) r = build_input_lite(rolling, rolled_previously, until_5, until_6, good_stuff)
return Input.model_validate(r | {"context": context}) return Input.model_validate(r | {"context": context})
@@ -525,12 +560,25 @@ def execute(graph: WorkflowBuilder, input: Input):
c = graph.compile() c = graph.compile()
r = build_registry(*(graph.node_specs.values())) r = build_registry(*(graph.node_specs.values()))
i = input.model_dump() i = input.model_dump()
pprint(i) # pprint(i)
pprint(c) # pprint(c)
pprint(r) # pprint(r)
return execute_workflow(c, i, r) return execute_workflow(c, i, r)
execute(gacha, build_input(context)(10, rolled_previously=240-135, until_5=5, until_6=73, good_stuff = True))
# twice in a row! it took 100+ and a miss tho # twice in a row! it took 100+ and a miss tho
# E wf_core.errors.WorkflowExecutionError: node input for init is missing required field 'countdown'
def test():
d = execute(
gacha,
build_input(context)(
20, rolled_previously=240 - 135, until_5=5, until_6=73, good_stuff=False # lets pretend
),
)
assert d.status == RunStatus.COMPLETED, "oops"
state = State.model_validate(d.state)
assert any(i["name"] in context["pool"]["n_240"] for i in state.storage)
pprint(state.storage)
if __name__ == "__main__":
test()