From ff01609a2e8bd622f3af4c3d89297311ce16b110 Mon Sep 17 00:00:00 2001 From: lda Date: Sun, 30 Aug 2026 23:56:57 +0700 Subject: [PATCH] feat: make workflow builder edits lossless --- src/wf_authoring/builder/__init__.py | 3 + src/wf_authoring/builder/core.py | 188 +++++++++++++++++++++++++-- src/wf_authoring/builder/mapping.py | 30 ++++- tests/authoring/test_builder.py | 167 +++++++++++++++++++++++- 4 files changed, 373 insertions(+), 15 deletions(-) diff --git a/src/wf_authoring/builder/__init__.py b/src/wf_authoring/builder/__init__.py index 53554983..d00fb2ab 100644 --- a/src/wf_authoring/builder/__init__.py +++ b/src/wf_authoring/builder/__init__.py @@ -1,4 +1,5 @@ from .core import WorkflowBuilder +from .mapping import auto_input_map_from_schema, auto_output_map_from_schema from .refs import BranchRef, BranchResult, DecisionResult, HandleResult, StepRef __all__ = [ @@ -8,4 +9,6 @@ __all__ = [ "HandleResult", "StepRef", "WorkflowBuilder", + "auto_input_map_from_schema", + "auto_output_map_from_schema", ] diff --git a/src/wf_authoring/builder/core.py b/src/wf_authoring/builder/core.py index 30093a45..8b4a0a8f 100644 --- a/src/wf_authoring/builder/core.py +++ b/src/wf_authoring/builder/core.py @@ -15,6 +15,7 @@ from wf_core import ( ForeachItemErrorPolicy, ForeachNode, InterruptNode, + NodeDef, NodeHandler, NodeUse, PreparedSubgraph, @@ -22,6 +23,7 @@ from wf_core import ( SchemaRef, StateSchema, SubgraphNode, + ValidationReport, Workflow, WorkflowRef, execute_workflow, @@ -31,6 +33,7 @@ from wf_core.errors import WorkflowExecutionError from wf_core.models.conditions import BinaryCondition, ExistsCondition, PathOperand from wf_core.models.conditions import Condition as CoreCondition from wf_core.models.steps import ( + InputBinding, InputPathBinding, InputValueBinding, OutputBinding, @@ -54,7 +57,9 @@ from .mapping import ( OutputBindingArg, StepInputBindingArg, auto_input_map, + auto_input_map_from_schema, auto_output_map, + auto_output_map_from_schema, coerce_path, normalize_input_mapping, normalize_input_values, @@ -136,6 +141,11 @@ def _canonical_output_bindings( ] +def _node_defs_compatible(left: NodeDef, right: NodeDef) -> bool: + """Compare contracts by their serialized canonical content.""" + return left.model_dump(mode="json") == right.model_dump(mode="json") + + def _reject_mixed_binding_styles( *, input: object | None, @@ -213,6 +223,8 @@ class WorkflowBuilder: node_specs: dict[str, NodeSpec[Any, Any]] = field(default_factory=dict) nodes: list[Step] = field(default_factory=list) edges: list[Edge] = field(default_factory=list) + workflow_output: list[InputBinding] = field(default_factory=list) + seeded_node_defs: dict[str, NodeDef] = field(default_factory=dict, repr=False) prepared_subgraphs: dict[str, PreparedSubgraph[NodeHandler]] = field( default_factory=dict ) @@ -392,6 +404,134 @@ class WorkflowBuilder: self.nodes.append(node) return node + def use_contract( + self, + node_def: NodeDef, + *, + id: str | None = None, + input: Sequence[StepInputBindingArg] | None = None, + output: Sequence[OutputBindingArg] | None = None, + desc: str | None = None, + ) -> NodeUse: + """Use a schema-backed external node contract without a local handler.""" + existing = self.seeded_node_defs.get(node_def.name) + if existing is not None and not _node_defs_compatible(existing, node_def): + raise ValueError( + f"incompatible duplicate node definition {node_def.name!r}" + ) + self.seeded_node_defs[node_def.name] = node_def.model_copy(deep=True) + normalized_input_schema = cast(SchemaRef, self.input_schema) + normalized_state_schema = cast(StateSchema, self.state_schema) + node_input = ( + normalize_step_input_bindings(input) + if input is not None + else _canonical_input_bindings( + normalize_input_mapping( + auto_input_map_from_schema( + node_def.input_schema, + input_schema=normalized_input_schema, + state_schema=normalized_state_schema, + ) + ), + {}, + ) + ) + node_output = ( + normalize_output_bindings(output) + if output is not None + else _canonical_output_bindings( + normalize_output_mapping( + auto_output_map_from_schema( + node_def.output_schema, + state_schema=normalized_state_schema, + ) + ) + ) + ) + return self.use_ref( + node_def.name, + id=id, + input=node_input, + output=node_output, + desc=desc, + ) + + @classmethod + def from_workflow(cls, workflow: Workflow) -> WorkflowBuilder: + """Create an independently editable builder from a canonical workflow.""" + return cls( + name=workflow.name, + input_schema=workflow.input_schema.model_copy(deep=True), + state_schema=workflow.state_schema.model_copy(deep=True), + output_schema=workflow.output_schema.model_copy(deep=True), + outcomes=tuple(workflow.outcomes), + start=workflow.start, + nodes=[node.model_copy(deep=True) for node in workflow.nodes], + edges=[edge.model_copy(deep=True) for edge in workflow.edges], + workflow_output=[ + binding.model_copy(deep=True) for binding in workflow.output + ], + seeded_node_defs={ + node_def.name: node_def.model_copy(deep=True) + for node_def in workflow.node_defs + }, + ) + + def set_output(self, bindings: Sequence[StepInputBindingArg]) -> None: + """Replace final workflow output projection bindings.""" + self.workflow_output = cast( + list[InputBinding], normalize_step_input_bindings(bindings) + ) + + def set_route(self, source: StepRef, outcome: str, target: StepRef) -> None: + """Replace the unique route for one source/outcome pair.""" + source_id = step_id(source) + target_id = step_id(target) + self.edges = [ + edge + for edge in self.edges + if not (edge.from_ == source_id and edge.outcome == outcome) + ] + self.edges.append( + Edge.model_validate( + {"from": source_id, "outcome": outcome, "to": target_id} + ) + ) + + def remove_route(self, source: StepRef, outcome: str) -> None: + """Remove one route, raising when the requested route is absent.""" + source_id = step_id(source) + matching = [ + edge + for edge in self.edges + if edge.from_ == source_id and edge.outcome == outcome + ] + if not matching: + raise ValueError( + f"route from step {source_id!r} with outcome {outcome!r} not found" + ) + self.edges = [ + edge + for edge in self.edges + if not (edge.from_ == source_id and edge.outcome == outcome) + ] + + def remove_step(self, step: StepRef) -> None: + """Remove an unreferenced step, rejecting dangling graph references.""" + step_id_value = step_id(step) + if not any(node.id == step_id_value for node in self.nodes): + raise ValueError(f"step {step_id_value!r} not found") + if self.start == step_id_value: + raise ValueError( + f"step {step_id_value!r} is still referenced as workflow start" + ) + if any( + edge.from_ == step_id_value or edge.to == step_id_value + for edge in self.edges + ): + raise ValueError(f"step {step_id_value!r} is still referenced by route") + self.nodes = [node for node in self.nodes if node.id != step_id_value] + def subgraph( self, *, @@ -869,21 +1009,45 @@ class WorkflowBuilder: ) return self.match(value, cases, id=id, default=default) + def _build_workflow(self, *, start: str) -> Workflow: + """Build a canonical workflow snapshot from current builder state.""" + node_defs = [ + node_def.model_copy(deep=True) + for node_def in self.seeded_node_defs.values() + ] + by_name = {node_def.name: node_def for node_def in node_defs} + for spec in self.node_specs.values(): + node_def = spec.to_node_def() + existing = by_name.get(node_def.name) + if existing is not None: + if not _node_defs_compatible(existing, node_def): + raise ValueError( + f"incompatible duplicate node definition {node_def.name!r}" + ) + continue + by_name[node_def.name] = node_def + node_defs.append(node_def) + return Workflow( + name=self.name, + input_schema=cast(SchemaRef, self.input_schema).model_copy(deep=True), + state_schema=cast(StateSchema, self.state_schema).model_copy(deep=True), + output_schema=cast(SchemaRef, self.output_schema).model_copy(deep=True), + outcomes=list(self.outcomes), + node_defs=node_defs, + start=start, + output=[binding.model_copy(deep=True) for binding in self.workflow_output], + nodes=[node.model_copy(deep=True) for node in self.nodes], + edges=[edge.model_copy(deep=True) for edge in self.edges], + ) + + def validate_structure(self) -> ValidationReport: + """Return structural issues, including an unset workflow start.""" + return self._build_workflow(start=self.start or "").validate_structure() + def compile(self) -> Workflow: if self.start is None: raise WorkflowExecutionError( "workflow builder requires an explicit start; " "call set_entry_point(...) or pass start=..." ) - node_defs = [spec.to_node_def() for spec in self.node_specs.values()] - return Workflow( - name=self.name, - input_schema=cast(SchemaRef, self.input_schema), - state_schema=cast(StateSchema, self.state_schema), - output_schema=cast(SchemaRef, self.output_schema), - outcomes=list(self.outcomes), - node_defs=node_defs, - start=self.start, - nodes=self.nodes, - edges=self.edges, - ) + return self._build_workflow(start=self.start) diff --git a/src/wf_authoring/builder/mapping.py b/src/wf_authoring/builder/mapping.py index fe45d516..06ebba7c 100644 --- a/src/wf_authoring/builder/mapping.py +++ b/src/wf_authoring/builder/mapping.py @@ -143,11 +143,25 @@ def auto_input_map( state_schema: StateSchema, ) -> dict[str, str]: """Map node input fields from state first, then workflow input.""" + return auto_input_map_from_schema( + spec.to_node_def().input_schema, + input_schema=input_schema, + state_schema=state_schema, + ) + + +def auto_input_map_from_schema( + capability_input_schema: SchemaRef, + *, + input_schema: SchemaRef, + state_schema: StateSchema, +) -> dict[str, str]: + """Map schema-declared capability inputs from state or workflow input.""" 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", {}) + for field in capability_input_schema.properties } @@ -157,10 +171,22 @@ def auto_output_map( state_schema: StateSchema, ) -> dict[str, str]: """Map node output fields back into matching state fields.""" + return auto_output_map_from_schema( + spec.to_node_def().output_schema, + state_schema=state_schema, + ) + + +def auto_output_map_from_schema( + capability_output_schema: SchemaRef, + *, + state_schema: StateSchema, +) -> dict[str, str]: + """Map schema-declared capability outputs into matching state fields.""" state_fields = state_schema.field_map() return { field: f"state.{field}" - for field in spec.output_model.model_json_schema().get("properties", {}) + for field in capability_output_schema.properties if field in state_fields } diff --git a/tests/authoring/test_builder.py b/tests/authoring/test_builder.py index a8a7e0cc..fa17164e 100644 --- a/tests/authoring/test_builder.py +++ b/tests/authoring/test_builder.py @@ -21,7 +21,15 @@ from wf_authoring import ( state_path, ) from wf_authoring.builder.mapping import normalize_input_mapping -from wf_core import END, EndNode, RunStatus, WorkflowExecutionError +from wf_core import ( + END, + EndNode, + NodeDef, + RunStatus, + ValidationIssueCode, + Workflow, + WorkflowExecutionError, +) from wf_core.models.steps import ( InputExpressionBinding, InputPathBinding, @@ -31,6 +39,163 @@ from wf_core.paths import GraphSourcePath, LocalPath, StatePath from wf_platform import CapabilityRef +def _editable_three_step_builder() -> WorkflowBuilder: + builder = WorkflowBuilder( + name="editable_three_step", + input_schema={"type": "object", "properties": {}}, + state_schema={"type": "object", "properties": {}}, + output_schema={"type": "object", "properties": {}}, + start="first", + ) + builder.use_ref("demo.first", id="first") + builder.use_ref("demo.second", id="second") + builder.use_ref("demo.third", id="third") + builder.connect("first", "ok", "second") + builder.connect("second", "ok", "third") + builder.connect("third", "ok", END) + return builder + + +def test_builder_round_trip_preserves_complete_workflow() -> None: + original = Workflow.model_validate( + { + "name": "round_trip", + "input_schema": { + "type": "object", + "properties": {"topic": {"type": "string"}}, + }, + "state_schema": { + "type": "object", + "properties": {"result": {"type": "string"}}, + }, + "output_schema": { + "type": "object", + "properties": {"result": {"type": "string"}}, + }, + "output": [{"path": "state.result", "target": "result"}], + "node_defs": [ + { + "name": "app.default.search", + "input_schema": {"type": "object", "properties": {}}, + "output_schema": { + "type": "object", + "properties": {"result": {"type": "string"}}, + }, + "outcomes": ["ok"], + } + ], + "outcomes": ["ok"], + "start": "search", + "nodes": [ + { + "id": "search", + "type": "node", + "node": "app.default.search", + "input": [], + "output": [{"source": "result", "target": "state.result"}], + }, + {"id": "end_ok", "type": "end", "outcome": "ok"}, + ], + "edges": [{"from": "search", "outcome": "ok", "to": "end_ok"}], + } + ) + + builder = WorkflowBuilder.from_workflow(original) + rebuilt = builder.compile() + + assert rebuilt.model_dump(mode="json", by_alias=True) == original.model_dump( + mode="json", by_alias=True + ) + builder.set_output([{"value": "changed", "target": "result"}]) + original_output = original.output[0] + assert isinstance(original_output, InputPathBinding) + assert str(original_output.path) == "state.result" + + +def test_set_route_replaces_unique_source_outcome_edge() -> None: + builder = _editable_three_step_builder() + + builder.set_route("first", "ok", "third") + + matching = [ + edge for edge in builder.edges if edge.from_ == "first" and edge.outcome == "ok" + ] + assert [(edge.from_, edge.outcome, edge.to) for edge in matching] == [ + ("first", "ok", "third") + ] + + +def test_remove_route_removes_only_requested_source_outcome_pair() -> None: + builder = _editable_three_step_builder() + builder.connect("first", "error", "third") + + builder.remove_route("first", "ok") + + assert [(edge.from_, edge.outcome) for edge in builder.edges] == [ + ("second", "ok"), + ("third", "ok"), + ("first", "error"), + ] + + +def test_remove_step_rejects_referenced_step() -> None: + builder = _editable_three_step_builder() + + with pytest.raises(ValueError, match="still referenced by route"): + builder.remove_step("second") + + +def test_use_contract_registers_schema_only_external_node() -> None: + builder = WorkflowBuilder( + name="contract_builder", + input_schema={"type": "object", "properties": {"topic": {"type": "string"}}}, + state_schema={ + "type": "object", + "properties": {"result": {"type": "string"}}, + }, + output_schema={"type": "object", "properties": {"result": {"type": "string"}}}, + ) + contract = NodeDef.model_validate( + { + "name": "app.remote.search", + "input_schema": { + "type": "object", + "properties": {"topic": {"type": "string"}}, + }, + "output_schema": { + "type": "object", + "properties": {"result": {"type": "string"}}, + }, + "outcomes": ["ok"], + } + ) + + step = builder.use_contract(contract, id="search") + builder.set_entry_point(step) + builder.connect(step, "ok", END) + + assert step.node == contract.name + assert contract.name in {node_def.name for node_def in builder.compile().node_defs} + assert builder.node_specs == {} + assert builder.registry() == {} + assert isinstance(step.input[0], InputPathBinding) + assert step.input[0].path == GraphSourcePath.input("topic") + assert step.output[0].target == StatePath.of("result") + + +def test_validate_structure_reports_unknown_start_when_unset() -> None: + builder = WorkflowBuilder( + name="missing_start", + input_schema={}, + state_schema={"type": "object"}, + output_schema={}, + ) + + report = builder.validate_structure() + + assert report.errors[0].code == ValidationIssueCode.UNKNOWN_START + + def test_builder_auto_binds_matching_node_inputs_and_outputs_to_state() -> None: builder = WorkflowBuilder( name="auto_bind_demo",