diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 523da9f1..25e3e2f2 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -1,5 +1,5 @@ # Architecture @@ -161,14 +161,34 @@ compatibility details: | File adapter binding and isolation | [File runtime](src/graphon/file/runtime.py), [engine](src/graphon/engine/engine.py) | [Execution file isolation](tests/engine/test_file_runtime.py), [file scopes](tests/file/test_runtime.py) | | Container ownership and factory preflight | [Graph and NodeFactory](src/graphon/graph/graph.py), [scoping](src/graphon/graph/scoping.py), [validation](src/graphon/graph/validation.py) | [Graph scoping and validation](tests/graph/test_graph_scoping.py) | | Branch joins and readiness | [Scheduler](src/graphon/engine/scheduler.py) | [Scheduling](tests/engine/test_scheduler.py) | -| Container suspension and variable propagation | [Container effects](src/graphon/nodes/container_effects.py), [handlers](src/graphon/engine/container_handler/builtin/) | [Cooperative execution](tests/engine/test_cooperative_container_execution.py), [container variable scoping](tests/workflows/test_full_engine_events.py), [Iteration variable ownership](tests/engine/test_iteration_variable_scopes.py) | +| Container suspension and variable propagation | [Container effects](src/graphon/nodes/container_effects.py), [handlers](src/graphon/engine/container_handler/builtin/) | [Container values](tests/nodes/test_container_effects.py), [cooperative execution](tests/engine/test_cooperative_container_execution.py), [container variable scoping](tests/workflows/test_full_engine_events.py), [Iteration variable ownership](tests/engine/test_iteration_variable_scopes.py) | | Pause, resume, and snapshot quiescence | [Dispatcher](src/graphon/engine/dispatcher.py), [execution guard](src/graphon/runtime/execution.py) | [Serialization](tests/engine/test_runtime_state_serialization.py) | | Snapshot versions and migrations | [Snapshot loading](src/graphon/runtime/runtime_state/snapshot.py) | [Runtime snapshots](tests/runtime/test_runtime_state.py), [version isolation](tests/test_snapshot_version_isolation.py) | +[Concrete segment models](src/graphon/variables/segments.py) own fixed `value_type` tags and value validation; serialized values select their model by that tag. +[Factories](src/graphon/variables/factory.py) infer types for raw values or validate explicit declarations, preserving declared empty arrays and `array[any]` values. +See [factory tests](tests/variables/test_factory.py) for declared types and serialized tag selection. +`array[any]` restores serialized [marked file values](src/graphon/file/constants.py) recursively inside arrays and objects using [File validation](src/graphon/file/models.py), preserving file values in container task and variable pool snapshots. +Unmarked objects retain their data, and invalid marked files fail validation; see [container tests](tests/nodes/test_container_effects.py) and [workflow tests](tests/workflows/test_full_engine_events.py). +[VariableBase.copy_with_value](src/graphon/variables/variables.py) validates replacements for assignment nodes, retaining identity and metadata and promoting integer variables when the new value is a float. +The [v1 assignment tests](tests/nodes/variable_assigner/test_v1_node.py) and [v2 assignment tests](tests/nodes/variable_assigner/test_v2_node.py) cover invalid updates, fractional results, and clearing arrays while preserving their declared type, including `array[file]`. + +[Container values](src/graphon/nodes/container_effects.py) preserve declared segment types through construction, snapshots, and completion, including secret tags. +Conversion removes source variable identity so the parent pool binds completed outputs to the container's selectors. +Pool conversion also assigns a new identity when storing an existing variable under a different selector; see [destination binding tests](tests/runtime/test_variable_pool.py). +Stored outputs remain typed; container node-result inputs and Loop/Iteration presentation events expose plain values. +The [runtime JSON converter](src/graphon/workflow_type_encoder.py) masks typed secrets using their log representation; hosts must apply redaction separately to plain values in logs and presentation. +The [container secret tests](tests/nodes/test_container_effects.py) cover this boundary; [Loop workflow tests](tests/workflows/test_full_engine_events.py) also cover v2 `SET`, `OVER_WRITE`, and `CLEAR` preserving the secret tag and masking. + +[Iteration](src/graphon/nodes/iteration/iteration_node.py) accepts array sequences, including tuples. +When its [handler](src/graphon/engine/container_handler/builtin/iteration.py) flattens outputs, matching declared array types survive in both serial and parallel runs; mixed array types become `array[any]`. +Empty iterations preserve the input array type and bind a separate output variable; see [workflow tests](tests/workflows/test_full_engine_events.py). + For host rendering, bind the engine's adapter with [use_workflow_file_runtime](src/graphon/file/runtime.py); use [EngineEventFilterContext.from_engine](src/graphon/engine/filter/protocol.py) to wire response filters. Supply host adapters again when rebuilding from snapshots. +Runtime JSON conversion renders tuple file arrays as JSON lists within the same adapter scope; see [host file conversion tests](tests/variables/test_workflow_type_encoder.py). Persist after fully consuming the run iterator or from a quiescent `on_graph_end` hook. Hosts must also coordinate direct mutation of exposed runtime objects: the diff --git a/docs/development.md b/docs/development.md index 9ee34371..e67d8901 100644 --- a/docs/development.md +++ b/docs/development.md @@ -1,5 +1,5 @@ # Development navigation @@ -16,15 +16,19 @@ Choose boundaries around state ownership and callers. Extract a responsibility when it has independent ownership or repeated churn; file length alone does not justify splitting a module. +For runtime values, let concrete types own their tags, validation, and value replacement. +Preserve typed values across state transfers and convert to plain values at presentation boundaries, where redaction is a separate responsibility. +Follow the [variable and container contracts](../ARCHITECTURE.md#state-and-execution-invariants) when changing those boundaries. + | Task | Start in source | Relevant checks | | --- | --- | --- | | Change package dependencies or queue ownership | [Import boundaries](../ARCHITECTURE.md#import-boundaries), [runtime queues](../src/graphon/runtime/ready_queue/) | `just imports`, [runtime isolation](../tests/runtime/test_runtime_imports.py), [public contract imports](../tests/test_protocols_exports.py) | | Parse or validate a graph | [graph](../src/graphon/graph/), [graph config](../src/graphon/entities/graph_config.py) | [graph tests](../tests/graph/), especially validation and scoping | | Import Dify DSL or wire node dependencies | [importer](../src/graphon/dsl/importer.py), [node factory](../src/graphon/dsl/node_factory.py) | [DSL tests](../tests/dsl/), especially importer, node factory, and app bootstrap | | Change a built-in node | [nodes](../src/graphon/nodes/), [base node](../src/graphon/nodes/base/node.py) | Matching [node tests](../tests/nodes/); [workflow events](../tests/workflows/test_full_engine_events.py) for graph-visible changes | -| Change dispatch, scheduling, or containers | [engine](../src/graphon/engine/), [container effects](../src/graphon/nodes/container_effects.py) | [engine tests](../tests/engine/), [container dispatch](../tests/nodes/test_container_dispatch.py), [workflow events](../tests/workflows/test_full_engine_events.py) | +| Change dispatch, scheduling, or containers | [engine](../src/graphon/engine/), [container effects](../src/graphon/nodes/container_effects.py) | [engine tests](../tests/engine/), [container values](../tests/nodes/test_container_effects.py), [container dispatch](../tests/nodes/test_container_dispatch.py), [workflow events](../tests/workflows/test_full_engine_events.py) | | Change event payloads or response streaming | [raw frame observations](../ARCHITECTURE.md#raw-frame-observations), [node events](../src/graphon/node_events/), [engine events](../src/graphon/engine_events/), [filters](../src/graphon/engine/filter/) | [frame observation tests](../tests/engine/test_frame_observations.py), [frame failure tests](../tests/engine/test_frame_failures.py), [node event tests](../tests/node_events/), [engine event tests](../tests/engine_events/), [response filter tests](../tests/engine/test_response_stream_filter.py) | -| Change variables or secrets | [variable pool](../src/graphon/runtime/variable_pool.py), [variables](../src/graphon/variables/) | [variable pool tests](../tests/runtime/test_variable_pool.py), [Iteration variable ownership](../tests/engine/test_iteration_variable_scopes.py), [variable tests](../tests/variables/), [HTTP secret masking](../tests/http/test_executor_secret_masking.py) | +| Change variables or secrets | [variable pool](../src/graphon/runtime/variable_pool.py), [segment models](../src/graphon/variables/segments.py), [variable models](../src/graphon/variables/variables.py), [conversion](../src/graphon/variables/factory.py), [variable assignment](../src/graphon/nodes/variable_assigner/) | [variable pool tests](../tests/runtime/test_variable_pool.py), [Iteration variable ownership](../tests/engine/test_iteration_variable_scopes.py), [variable tests](../tests/variables/), [assignment tests](../tests/nodes/variable_assigner/), [HTTP secret masking](../tests/http/test_executor_secret_masking.py) | | Change persistence or resume | [runtime state](../src/graphon/runtime/runtime_state/), [container state](../src/graphon/runtime/container_state.py) | [runtime tests](../tests/runtime/), [engine serialization](../tests/engine/test_runtime_state_serialization.py), [frame failure restoration](../tests/engine/test_frame_failures.py), [version isolation](../tests/test_snapshot_version_isolation.py) | | Change model invocation | [graph-facing LLM protocol](../src/graphon/nodes/llm/runtime_protocols.py), [Slim adapter](../src/graphon/dsl/slim/llm.py), [model runtime](../src/graphon/model_runtime/) | [LLM node tests](../tests/nodes/llm/), [Slim tests](../tests/dsl/test_slim_llm.py), [model dispatch](../tests/model_runtime/test_model_dispatch.py) | | Change file or HTTP integration | [file](../src/graphon/file/), [engine adapter binding](../src/graphon/engine/engine.py), [HTTP](../src/graphon/http/), [node ports](../src/graphon/nodes/protocols.py) | [file tests](../tests/file/), [execution file isolation](../tests/engine/test_file_runtime.py), [HTTP tests](../tests/http/), [HTTP node tests](../tests/nodes/http_request/), [extractor tests](../tests/nodes/document_extractor/) | diff --git a/src/graphon/engine/container_handler/builtin/iteration.py b/src/graphon/engine/container_handler/builtin/iteration.py index d1034038..54f5fe06 100644 --- a/src/graphon/engine/container_handler/builtin/iteration.py +++ b/src/graphon/engine/container_handler/builtin/iteration.py @@ -3,7 +3,7 @@ from __future__ import annotations from datetime import UTC, datetime -from typing import cast, final +from typing import final from graphon.engine_events.base import NodeEvent from graphon.engine_events.node import NodeRunFailedEvent, NodeRunVariableUpdatedEvent @@ -29,7 +29,7 @@ from graphon.runtime.execution import ROOT_FRAME_ID from graphon.runtime.ready_queue import ResumeTask from graphon.runtime.runtime_state import RuntimeState -from graphon.variables.segments import NoneSegment, SerializableSegment +from graphon.variables.segments import ArrayAnySegment, ArraySegment, NoneSegment from graphon.variables.variables import Variable from ...frame import ExecutionFrame, FrameRegistry @@ -206,13 +206,12 @@ def _complete_ready_iteration_frame( result = frame.state.variable_pool.get( run_state.output_selector, ) - output = NoneSegment() if result is None else cast(SerializableSegment, result) run_state = self._complete_iteration_step( frame=frame, frame_state=frame_state, parent_frame=parent_frame, run_state=run_state, - output=output, + output=build_container_value(result), store_output=True, ) self._continue_or_complete_iteration( @@ -266,7 +265,7 @@ def _complete_iteration_step( frame_state: IterationFrameState, parent_frame: ExecutionFrame, run_state: IterationRunState, - output: SerializableSegment, + output: ContainerValue, store_output: bool, ) -> IterationRunState: duration_map = dict(run_state.duration_map) @@ -455,14 +454,7 @@ def _build_iteration_result( status: WorkflowNodeExecutionStatus, error: str = "", ) -> ContainerExecutionResult: - outputs: dict[str, ContainerValue] = { - "output": build_container_value( - self._flatten_outputs_if_needed( - self._ordered_iteration_outputs(run_state), - flatten_output=run_state.flatten_output, - ), - ), - } + outputs = {"output": self._collect_iteration_outputs(run_state)} metadata = self._iteration_metadata(run_state) return ContainerExecutionResult( metadata=self._event_metadata(metadata), @@ -492,36 +484,29 @@ def _iteration_metadata( ), } - def _flatten_outputs_if_needed( - self, - outputs: list[object], - *, - flatten_output: bool, - ) -> list[object]: - if not flatten_output: - return outputs - non_empty_outputs = [output for output in outputs if output is not None] - if not non_empty_outputs: - return outputs - if not all(isinstance(output, list) for output in non_empty_outputs): - return outputs - - flattened: list[object] = [] - for output in outputs: - if isinstance(output, list): - flattened.extend(output) - elif output is not None: - flattened.append(output) - return flattened - - def _ordered_iteration_outputs( + def _collect_iteration_outputs( self, run_state: IterationRunState, - ) -> list[object]: - return [ - run_state.outputs[key].to_object() - for key in sorted(run_state.outputs, key=int) - ] + ) -> ContainerValue: + outputs = [run_state.outputs[key] for key in sorted(run_state.outputs, key=int)] + arrays = [output for output in outputs if isinstance(output, ArraySegment)] + if ( + run_state.flatten_output + and arrays + and all( + isinstance(output, (ArraySegment, NoneSegment)) for output in outputs + ) + ): + flattened = [item for output in arrays for item in output.value] + if len({output.value_type for output in arrays}) == 1: + return build_container_value(arrays[0]).model_copy( + update={"value": flattened} + ) + return ArrayAnySegment(value=flattened) + values = [output.to_object() for output in outputs] + return ( + ArrayAnySegment(value=values) if arrays else build_container_value(values) + ) def _enqueue_container_result( self, diff --git a/src/graphon/engine/event/processor.py b/src/graphon/engine/event/processor.py index b24c876c..ed548403 100644 --- a/src/graphon/engine/event/processor.py +++ b/src/graphon/engine/event/processor.py @@ -46,6 +46,7 @@ ) from graphon.runtime.execution import ROOT_FRAME_ID, GraphExecution from graphon.runtime.ready_queue import ResumeTask, StartTask +from graphon.variables.segments import Segment from ..container_handler import ContainerHandler from ..frame import ExecutionFrame, FrameRegistry @@ -203,11 +204,15 @@ def _( frame: ExecutionFrame, ) -> None: if isinstance(event, NodeRunLoopSucceededEvent): - event.outputs = self._current_loop_outputs( + outputs = self._current_loop_outputs( frame=frame, node_id=event.node_id, outputs=event.outputs, ) + event.outputs = { + key: value.to_object() if isinstance(value, Segment) else value + for key, value in outputs.items() + } self._collect(frame=frame, event=event) @_dispatch.register @@ -403,7 +408,7 @@ def _current_loop_outputs( for key in outputs: value = frame.state.variable_pool.get((node_id, key)) if value is not None: - current_outputs[key] = value.to_object() + current_outputs[key] = value return current_outputs def _store_node_outputs( diff --git a/src/graphon/nodes/container_effects.py b/src/graphon/nodes/container_effects.py index e72c8d44..80f3fb3e 100644 --- a/src/graphon/nodes/container_effects.py +++ b/src/graphon/nodes/container_effects.py @@ -1,9 +1,9 @@ from __future__ import annotations from collections.abc import Mapping -from typing import Annotated, Literal, cast +from typing import Annotated, Literal -from pydantic import BaseModel, BeforeValidator, ConfigDict, Field +from pydantic import BaseModel, BeforeValidator, ConfigDict, Field, TypeAdapter from graphon.enums import ErrorHandleMode from graphon.node_events.base import NodeRunResult @@ -11,9 +11,7 @@ from graphon.variables.segments import Segment, SerializableSegment -def _to_serializable_segment(value: object) -> object: - if isinstance(value, Segment): - return build_segment(value.to_object()) +def _parse_container_value(value: object) -> Segment | Mapping[str, object]: if isinstance(value, Mapping) and "value_type" in value: return value return build_segment(value) @@ -21,13 +19,14 @@ def _to_serializable_segment(value: object) -> object: type ContainerValue = Annotated[ SerializableSegment, - BeforeValidator(_to_serializable_segment), + BeforeValidator(_parse_container_value), ] +_CONTAINER_VALUE_ADAPTER = TypeAdapter(ContainerValue) + def build_container_value(value: object) -> ContainerValue: - raw_value = value.to_object() if isinstance(value, Segment) else value - return cast(ContainerValue, build_segment(raw_value)) + return _CONTAINER_VALUE_ADAPTER.validate_python(build_segment(value)) class ContainerNodeRunResult(NodeRunResult): @@ -38,7 +37,7 @@ def to_node_run_result(self) -> NodeRunResult: return NodeRunResult( status=self.status, inputs={key: value.to_object() for key, value in self.inputs.items()}, - outputs={key: value.to_object() for key, value in self.outputs.items()}, + outputs=self.outputs, metadata=self.metadata, llm_usage=self.llm_usage, error=self.error, diff --git a/src/graphon/nodes/iteration/iteration_node.py b/src/graphon/nodes/iteration/iteration_node.py index 0b508bbc..7a1f2e90 100644 --- a/src/graphon/nodes/iteration/iteration_node.py +++ b/src/graphon/nodes/iteration/iteration_node.py @@ -1,6 +1,6 @@ from collections.abc import Generator, Mapping, Sequence from datetime import UTC, datetime -from typing import Any, cast, override +from typing import Any, override from graphon.enums import ( BuiltinNodeTypes, @@ -9,7 +9,7 @@ WorkflowNodeExecutionStatus, ) from graphon.graph.scoping import resolve_container_id -from graphon.node_events.base import NodeEventPayload, NodeRunResult +from graphon.node_events.base import NodeEventPayload from graphon.node_events.iteration import ( IterationFailedEvent, IterationNextEvent, @@ -20,6 +20,7 @@ from graphon.nodes.base.node import Node from graphon.nodes.container_effects import ( ContainerExecutionResult, + ContainerNodeRunResult, ContainerRunResult, IterationFrameRequest, build_container_value, @@ -118,11 +119,15 @@ def _resume_container_events( if isinstance(result, ContainerExecutionResult): node_run_result = result.node_run_result.to_node_run_result() + outputs = { + key: value.to_object() + for key, value in result.node_run_result.outputs.items() + } if node_run_result.status == WorkflowNodeExecutionStatus.SUCCEEDED: yield IterationSucceededEvent( start_at=self._start_at, inputs=node_run_result.inputs, - outputs=node_run_result.outputs, + outputs=outputs, metadata=result.metadata, steps=result.steps, ) @@ -130,7 +135,7 @@ def _resume_container_events( yield IterationFailedEvent( start_at=self._start_at, inputs=node_run_result.inputs, - outputs=node_run_result.outputs, + outputs=outputs, metadata=result.metadata, steps=result.steps, error=node_run_result.error, @@ -149,10 +154,10 @@ def _resolve_iterator_value(self, variable: object) -> list[object]: msg = f"invalid iterator value: {variable}, please provide a list." raise TypeError(msg) iterator_value = variable.to_object() - if not isinstance(iterator_value, list): + if not isinstance(iterator_value, Sequence): msg = f"Invalid iterator value: {iterator_value}, please provide a list." raise TypeError(msg) - return cast(list[object], iterator_value) + return list(iterator_value) def _run_empty_iteration( self, @@ -160,10 +165,17 @@ def _run_empty_iteration( variable: NoneSegment | ArraySegment, started_at: datetime, ) -> Generator[NodeEventPayload, None, None]: - outputs = {"output": ArrayAnySegment(value=[])} - if isinstance(variable, ArraySegment): - outputs = {"output": variable.model_copy(update={"value": []})} + outputs = { + "output": build_container_value( + variable if isinstance(variable, ArraySegment) else [] + ) + } inputs: dict[str, object] = {"iterator_selector": []} + result = ContainerNodeRunResult( + status=WorkflowNodeExecutionStatus.SUCCEEDED, + outputs=outputs, + inputs={"iterator_selector": ArrayAnySegment(value=[])}, + ) yield IterationStartedEvent( start_at=started_at, inputs=inputs, @@ -172,16 +184,12 @@ def _run_empty_iteration( yield IterationSucceededEvent( start_at=started_at, inputs=inputs, - outputs=outputs, + outputs={key: value.to_object() for key, value in result.outputs.items()}, metadata={}, steps=0, ) yield StreamCompletedEvent( - node_run_result=NodeRunResult( - status=WorkflowNodeExecutionStatus.SUCCEEDED, - outputs=outputs, - inputs=inputs, - ), + node_run_result=result.to_node_run_result(), ) @classmethod diff --git a/src/graphon/nodes/llm/node.py b/src/graphon/nodes/llm/node.py index a36e756e..30384122 100644 --- a/src/graphon/nodes/llm/node.py +++ b/src/graphon/nodes/llm/node.py @@ -81,6 +81,7 @@ ArraySegment, NoneSegment, ObjectSegment, + SecretSegment, StringSegment, ) @@ -1197,7 +1198,7 @@ def _fetch_context( if context_value_variable is None: return - if isinstance(context_value_variable, StringSegment): + if isinstance(context_value_variable, (StringSegment, SecretSegment)): yield RunRetrieverResourceEvent( retriever_resources=[], context=context_value_variable.value, diff --git a/src/graphon/nodes/loop/loop_node.py b/src/graphon/nodes/loop/loop_node.py index c2bcaf21..2afde8a1 100644 --- a/src/graphon/nodes/loop/loop_node.py +++ b/src/graphon/nodes/loop/loop_node.py @@ -104,11 +104,15 @@ def _resume_container_events( if isinstance(result, ContainerExecutionResult): node_run_result = result.node_run_result.to_node_run_result() + outputs = { + key: value.to_object() + for key, value in result.node_run_result.outputs.items() + } if node_run_result.status == WorkflowNodeExecutionStatus.SUCCEEDED: yield LoopSucceededEvent( start_at=self._start_at, inputs=node_run_result.inputs, - outputs=node_run_result.outputs, + outputs=outputs, metadata=result.metadata, steps=result.steps, ) @@ -116,7 +120,7 @@ def _resume_container_events( yield LoopFailedEvent( start_at=self._start_at, inputs=node_run_result.inputs, - outputs=node_run_result.outputs, + outputs=outputs, metadata=result.metadata, steps=result.steps, error=node_run_result.error, @@ -164,7 +168,10 @@ def _initialize_loop_variables( raise ValueError(msg) variable_selector = (self._node_id, loop_variable.label) - initialized_pool.add(variable_selector, processed_segment.value) + initialized_pool.add( + variable_selector, + build_container_value(processed_segment), + ) variable = initialized_pool.variable_dictionary[self._node_id][ loop_variable.label ] diff --git a/src/graphon/nodes/parameter_extractor/parameter_extractor_node.py b/src/graphon/nodes/parameter_extractor/parameter_extractor_node.py index 3e6f45c5..e6f44fe1 100644 --- a/src/graphon/nodes/parameter_extractor/parameter_extractor_node.py +++ b/src/graphon/nodes/parameter_extractor/parameter_extractor_node.py @@ -840,16 +840,16 @@ def _transform_array_parameter_value( if not isinstance(param_value, list): return _TRANSFORM_RESULT_UNSET - segment_value = build_segment_with_type(segment_type=parameter_type, value=[]) + values = [] handler_name = _ARRAY_ITEM_TRANSFORMER_NAMES.get(element_type) if handler_name is None: - return segment_value + return build_segment_with_type(parameter_type, values) for item in param_value: transformed_item = getattr(self, handler_name)(item) if transformed_item is not _TRANSFORM_RESULT_UNSET: - segment_value.value.append(transformed_item) - return segment_value + values.append(transformed_item) + return build_segment_with_type(parameter_type, values) def _transform_number_value(self, value: Any) -> Any: transformed = self._transform_number(value) diff --git a/src/graphon/nodes/variable_assigner/v1/node.py b/src/graphon/nodes/variable_assigner/v1/node.py index cc72250b..49202fdb 100644 --- a/src/graphon/nodes/variable_assigner/v1/node.py +++ b/src/graphon/nodes/variable_assigner/v1/node.py @@ -107,9 +107,7 @@ def _run(self) -> Generator[NodeEventPayload, None, None]: if not income_value: msg = "input value not found" raise VariableOperatorNodeError(msg) - updated_variable = original_variable.model_copy( - update={"value": income_value.value}, - ) + updated_value = income_value.value case WriteMode.APPEND: income_value = self.runtime_state.variable_pool.get( @@ -131,18 +129,13 @@ def _run(self) -> Generator[NodeEventPayload, None, None]: case _: msg = "append mode requires an array variable" raise VariableOperatorNodeError(msg) - updated_variable = original_variable.model_copy( - update={"value": updated_value}, - ) - case WriteMode.CLEAR: income_value = SegmentType.get_zero_value(original_variable.value_type) - updated_variable = original_variable.model_copy( - update={"value": income_value.to_object()}, - ) + updated_value = income_value.to_object() case _: assert_never(write_mode) + updated_variable = original_variable.copy_with_value(updated_value) updated_variables = [ common_helpers.variable_to_processed_data( assigned_variable_selector, diff --git a/src/graphon/nodes/variable_assigner/v2/helpers.py b/src/graphon/nodes/variable_assigner/v2/helpers.py index 8b1ea009..c077ce68 100644 --- a/src/graphon/nodes/variable_assigner/v2/helpers.py +++ b/src/graphon/nodes/variable_assigner/v2/helpers.py @@ -17,6 +17,7 @@ )) _SCALAR_INPUT_TYPES: dict[SegmentType, type[Any]] = { SegmentType.STRING: str, + SegmentType.SECRET: str, SegmentType.BOOLEAN: bool, SegmentType.OBJECT: dict, } @@ -66,6 +67,7 @@ def is_operation_supported( return variable_type in frozenset(( SegmentType.OBJECT, SegmentType.STRING, + SegmentType.SECRET, SegmentType.NUMBER, SegmentType.INTEGER, SegmentType.FLOAT, @@ -107,7 +109,12 @@ def is_constant_input_supported( operation: Operation, ) -> bool: match variable_type: - case SegmentType.STRING | SegmentType.OBJECT | SegmentType.BOOLEAN: + case ( + SegmentType.STRING + | SegmentType.SECRET + | SegmentType.OBJECT + | SegmentType.BOOLEAN + ): return operation in frozenset((Operation.OVER_WRITE, Operation.SET)) case SegmentType.NUMBER | SegmentType.INTEGER | SegmentType.FLOAT: return operation in frozenset(( diff --git a/src/graphon/nodes/variable_assigner/v2/node.py b/src/graphon/nodes/variable_assigner/v2/node.py index 73001989..2326857a 100644 --- a/src/graphon/nodes/variable_assigner/v2/node.py +++ b/src/graphon/nodes/variable_assigner/v2/node.py @@ -232,7 +232,7 @@ def _apply_item( operation=item.operation, value=input_value, ) - updated_variable = variable.model_copy(update={"value": updated_value}) + updated_variable = variable.copy_with_value(updated_value) working_variable_pool.add(updated_variable.selector, updated_variable) return updated_variable.selector diff --git a/src/graphon/variables/__init__.py b/src/graphon/variables/__init__.py index 07b6081d..82a72fbb 100644 --- a/src/graphon/variables/__init__.py +++ b/src/graphon/variables/__init__.py @@ -19,6 +19,7 @@ IntegerSegment, NoneSegment, ObjectSegment, + SecretSegment, Segment, StringSegment, ) @@ -65,6 +66,7 @@ "NoneVariable", "ObjectSegment", "ObjectVariable", + "SecretSegment", "SecretVariable", "Segment", "SegmentGroup", diff --git a/src/graphon/variables/factory.py b/src/graphon/variables/factory.py index b51739b6..cb7caa8e 100644 --- a/src/graphon/variables/factory.py +++ b/src/graphon/variables/factory.py @@ -9,6 +9,8 @@ from typing import Any from uuid import uuid4 +from pydantic import ValidationError + from graphon.file.models import File from .segments import ( @@ -25,7 +27,9 @@ IntegerSegment, NoneSegment, ObjectSegment, + SecretSegment, Segment, + SerializableSegment, StringSegment, ) from .types import SegmentType @@ -88,29 +92,6 @@ def _build_uniform_array_segment( return result -def _build_empty_array_segment_for_type( - *, - segment_type: SegmentType, - value: list[Any], -) -> Segment | None: - match segment_type: - case SegmentType.ARRAY_ANY: - result: Segment | None = ArrayAnySegment(value=value) - case SegmentType.ARRAY_STRING: - result = ArrayStringSegment(value=value) - case SegmentType.ARRAY_BOOLEAN: - result = ArrayBooleanSegment(value=value) - case SegmentType.ARRAY_NUMBER: - result = ArrayNumberSegment(value=value) - case SegmentType.ARRAY_OBJECT: - result = ArrayObjectSegment(value=value) - case SegmentType.ARRAY_FILE: - result = ArrayFileSegment(value=value) - case _: - result = None - return result - - def _existing_variable_from_segment(segment: VariableBase) -> Variable: if isinstance( segment, @@ -261,6 +242,14 @@ def _build_scalar_variable( value=segment.value, selector=selector, ) + case SecretSegment(): + result = SecretVariable( + id=variable_id, + name=name, + description=description, + value=segment.value, + selector=selector, + ) case StringSegment(): result = StringVariable( id=variable_id, @@ -314,46 +303,21 @@ def _build_list_segment(value: list[Any]) -> Segment: return _build_uniform_array_segment(element_type=types.pop(), value=value) -def _build_empty_array_segment( - *, - segment_type: SegmentType, - value: list[Any], -) -> Segment | None: - return _build_empty_array_segment_for_type( - segment_type=segment_type, - value=value, - ) - - -def _resolve_segment_class_for_type_match( - *, - segment_type: SegmentType, - inferred_type: SegmentType, -) -> type[Segment] | None: - if inferred_type == segment_type: - return _SEGMENT_FACTORY[segment_type] - if segment_type == SegmentType.NUMBER and inferred_type in frozenset(( - SegmentType.INTEGER, - SegmentType.FLOAT, - )): - return _SEGMENT_FACTORY[inferred_type] - return None - - def build_segment(value: Any, /) -> Segment: """Build a runtime segment from a Python value.""" segment = _build_non_list_segment(value) if segment is not None: return segment - if isinstance(value, list): - return _build_list_segment(value) + if isinstance(value, (list, tuple)): + return _build_list_segment(list(value)) msg = f"not supported value {value}" raise ValueError(msg) -_SEGMENT_FACTORY: Mapping[SegmentType, type[Segment]] = { +_SEGMENT_FACTORY: Mapping[SegmentType, type[SerializableSegment]] = { SegmentType.NONE: NoneSegment, SegmentType.STRING: StringSegment, + SegmentType.SECRET: SecretSegment, SegmentType.INTEGER: IntegerSegment, SegmentType.FLOAT: FloatSegment, SegmentType.FILE: FileSegment, @@ -368,46 +332,23 @@ def build_segment(value: Any, /) -> Segment: } -def build_segment_with_type(segment_type: SegmentType, value: Any) -> Segment: +def build_segment_with_type( + segment_type: SegmentType, value: Any +) -> SerializableSegment: """Build a segment while enforcing compatibility with the expected runtime type.""" - if value is None: - if segment_type == SegmentType.NONE: - return NoneSegment() - msg = f"Type mismatch: expected {segment_type}, but got None" - raise TypeMismatchError(msg) - - if isinstance(value, list) and len(value) == 0: - empty_segment = _build_empty_array_segment( - segment_type=segment_type, - value=value, + if segment_type == SegmentType.NUMBER: + segment_type = ( + SegmentType.INTEGER if isinstance(value, int) else SegmentType.FLOAT ) - if empty_segment is not None: - return empty_segment - msg = f"Type mismatch: expected {segment_type}, but got empty list" - raise TypeMismatchError(msg) - - inferred_type = SegmentType.infer_segment_type(value) - if inferred_type is None: - msg = ( - f"Type mismatch: expected {segment_type}, but got python object, " - f"type={type(value)}, value={value}" - ) - raise TypeMismatchError(msg) - - segment_class = _resolve_segment_class_for_type_match( - segment_type=segment_type, - inferred_type=inferred_type, - ) - if segment_class is not None: - value_type = ( - inferred_type if segment_type == SegmentType.NUMBER else segment_type - ) - return segment_class(value_type=value_type, value=value) - msg = ( - f"Type mismatch: expected {segment_type}, but got {inferred_type}, " - f"value={value}" - ) - raise TypeMismatchError(msg) + segment_class = _SEGMENT_FACTORY.get(segment_type) + if segment_class is None: + msg = f"Unsupported segment type: {segment_type}" + raise UnsupportedSegmentTypeError(msg) + try: + return segment_class.model_validate({"value": value}, strict=True) + except ValidationError as error: + msg = f"Value does not match {segment_type}" + raise TypeMismatchError(msg) from error def segment_to_variable( @@ -419,7 +360,7 @@ def segment_to_variable( description: str = "", ) -> Variable: """Convert a runtime segment into a runtime variable for storage in the pool.""" - if isinstance(segment, VariableBase): + if isinstance(segment, VariableBase) and tuple(segment.selector) == tuple(selector): return _existing_variable_from_segment(segment) name = name or selector[-1] resolved_variable_id = variable_id or str(uuid4()) diff --git a/src/graphon/variables/segment_group.py b/src/graphon/variables/segment_group.py index e3198870..ad8a7010 100644 --- a/src/graphon/variables/segment_group.py +++ b/src/graphon/variables/segment_group.py @@ -1,9 +1,11 @@ +from typing import Literal + from .segments import Segment from .types import SegmentType class SegmentGroup(Segment): - value_type: SegmentType = SegmentType.GROUP + value_type: Literal[SegmentType.GROUP] = SegmentType.GROUP value: list[Segment] @property diff --git a/src/graphon/variables/segments.py b/src/graphon/variables/segments.py index e053c5d2..defe8c77 100644 --- a/src/graphon/variables/segments.py +++ b/src/graphon/variables/segments.py @@ -1,14 +1,35 @@ import json import sys from collections.abc import Mapping, Sequence -from typing import Annotated, Any +from typing import Annotated, Any, Literal -from pydantic import BaseModel, ConfigDict, Discriminator, Tag, field_validator +from pydantic import AfterValidator, BaseModel, ConfigDict, Field +from graphon.file.constants import maybe_file_object from graphon.file.models import File from .types import SegmentType +SHORT_TOKEN_MAX_LENGTH = 8 + + +def _mask_token(token: str) -> str: + if not token: + return token + if len(token) <= SHORT_TOKEN_MAX_LENGTH: + return "*" * 20 + return token[:6] + "*" * 12 + token[-2:] + + +def _restore_files(value: Any) -> Any: + if maybe_file_object(value): + return File.model_validate(value) + if isinstance(value, (list, tuple)): + return type(value)(_restore_files(item) for item in value) + if isinstance(value, dict): + return {key: _restore_files(item) for key, item in value.items()} + return value + class Segment(BaseModel): """Segment is runtime type used during the execution of workflow. @@ -21,25 +42,6 @@ class Segment(BaseModel): value_type: SegmentType value: Any - @field_validator("value_type") - @classmethod - def validate_value_type(cls, value: SegmentType) -> SegmentType: - """Validate that the provided value is equal to the default value - of the `'value_type'` field. If the value is different, a ValueError is - raised. - - Returns: - The validated `value_type` when it matches the segment class default. - - Raises: - ValueError: If `value` differs from the segment class default `value_type`. - - """ - if value != cls.model_fields["value_type"].default: - msg = "Cannot modify 'value_type'" - raise ValueError(msg) - return value - @property def text(self) -> str: return str(self.value) @@ -62,7 +64,9 @@ def to_object(self) -> object: class NoneSegment(Segment): - value_type: SegmentType = SegmentType.NONE + model_config = ConfigDict(revalidate_instances="always") + + value_type: Literal[SegmentType.NONE] = SegmentType.NONE value: None = None @property @@ -79,12 +83,27 @@ def markdown(self) -> str: class StringSegment(Segment): - value_type: SegmentType = SegmentType.STRING + model_config = ConfigDict(revalidate_instances="always") + + value_type: Literal[SegmentType.STRING] = SegmentType.STRING + value: str + + +class SecretSegment(Segment): + model_config = ConfigDict(revalidate_instances="always") + + value_type: Literal[SegmentType.SECRET] = SegmentType.SECRET value: str + @property + def log(self) -> str: + return _mask_token(self.value) + class FloatSegment(Segment): - value_type: SegmentType = SegmentType.FLOAT + model_config = ConfigDict(revalidate_instances="always") + + value_type: Literal[SegmentType.FLOAT] = SegmentType.FLOAT value: float # NOTE(QuantumGhost): seems that the equality for FloatSegment with `NaN` # value has some problems. @@ -104,12 +123,16 @@ class FloatSegment(Segment): class IntegerSegment(Segment): - value_type: SegmentType = SegmentType.INTEGER + model_config = ConfigDict(revalidate_instances="always") + + value_type: Literal[SegmentType.INTEGER] = SegmentType.INTEGER value: int class ObjectSegment(Segment): - value_type: SegmentType = SegmentType.OBJECT + model_config = ConfigDict(revalidate_instances="always") + + value_type: Literal[SegmentType.OBJECT] = SegmentType.OBJECT value: Mapping[str, Any] @property @@ -126,6 +149,8 @@ def markdown(self) -> str: class ArraySegment(Segment): + model_config = ConfigDict(revalidate_instances="always") + @property def text(self) -> str: # Return empty string for empty arrays instead of "[]" @@ -139,7 +164,9 @@ def markdown(self) -> str: class FileSegment(Segment): - value_type: SegmentType = SegmentType.FILE + model_config = ConfigDict(revalidate_instances="always") + + value_type: Literal[SegmentType.FILE] = SegmentType.FILE value: File @property @@ -156,17 +183,19 @@ def text(self) -> str: class BooleanSegment(Segment): - value_type: SegmentType = SegmentType.BOOLEAN + model_config = ConfigDict(revalidate_instances="always") + + value_type: Literal[SegmentType.BOOLEAN] = SegmentType.BOOLEAN value: bool class ArrayAnySegment(ArraySegment): - value_type: SegmentType = SegmentType.ARRAY_ANY - value: Sequence[Any] + value_type: Literal[SegmentType.ARRAY_ANY] = SegmentType.ARRAY_ANY + value: Annotated[Sequence[Any], AfterValidator(_restore_files)] class ArrayStringSegment(ArraySegment): - value_type: SegmentType = SegmentType.ARRAY_STRING + value_type: Literal[SegmentType.ARRAY_STRING] = SegmentType.ARRAY_STRING value: Sequence[str] @property @@ -178,12 +207,12 @@ def text(self) -> str: class ArrayNumberSegment(ArraySegment): - value_type: SegmentType = SegmentType.ARRAY_NUMBER + value_type: Literal[SegmentType.ARRAY_NUMBER] = SegmentType.ARRAY_NUMBER value: Sequence[float | int] class ArrayObjectSegment(ArraySegment): - value_type: SegmentType = SegmentType.ARRAY_OBJECT + value_type: Literal[SegmentType.ARRAY_OBJECT] = SegmentType.ARRAY_OBJECT value: Sequence[Mapping[str, Any]] @property @@ -195,7 +224,7 @@ def text(self) -> str: class ArrayFileSegment(ArraySegment): - value_type: SegmentType = SegmentType.ARRAY_FILE + value_type: Literal[SegmentType.ARRAY_FILE] = SegmentType.ARRAY_FILE value: Sequence[File] @property @@ -212,41 +241,26 @@ def text(self) -> str: class ArrayBooleanSegment(ArraySegment): - value_type: SegmentType = SegmentType.ARRAY_BOOLEAN + value_type: Literal[SegmentType.ARRAY_BOOLEAN] = SegmentType.ARRAY_BOOLEAN value: Sequence[bool] -def get_segment_discriminator(v: Any) -> SegmentType | None: - if isinstance(v, Segment): - return v.value_type - if isinstance(v, dict): - value_type = v.get("value_type") - if value_type is None: - return None - try: - seg_type = SegmentType(value_type) - except ValueError: - return None - return seg_type - # return None if the discriminator value isn't found - return None - - type SerializableSegment = Annotated[ ( - Annotated[NoneSegment, Tag(SegmentType.NONE)] - | Annotated[StringSegment, Tag(SegmentType.STRING)] - | Annotated[FloatSegment, Tag(SegmentType.FLOAT)] - | Annotated[IntegerSegment, Tag(SegmentType.INTEGER)] - | Annotated[ObjectSegment, Tag(SegmentType.OBJECT)] - | Annotated[FileSegment, Tag(SegmentType.FILE)] - | Annotated[BooleanSegment, Tag(SegmentType.BOOLEAN)] - | Annotated[ArrayAnySegment, Tag(SegmentType.ARRAY_ANY)] - | Annotated[ArrayStringSegment, Tag(SegmentType.ARRAY_STRING)] - | Annotated[ArrayNumberSegment, Tag(SegmentType.ARRAY_NUMBER)] - | Annotated[ArrayObjectSegment, Tag(SegmentType.ARRAY_OBJECT)] - | Annotated[ArrayFileSegment, Tag(SegmentType.ARRAY_FILE)] - | Annotated[ArrayBooleanSegment, Tag(SegmentType.ARRAY_BOOLEAN)] + NoneSegment + | StringSegment + | SecretSegment + | FloatSegment + | IntegerSegment + | ObjectSegment + | FileSegment + | BooleanSegment + | ArrayAnySegment + | ArrayStringSegment + | ArrayNumberSegment + | ArrayObjectSegment + | ArrayFileSegment + | ArrayBooleanSegment ), - Discriminator(get_segment_discriminator), + Field(discriminator="value_type"), ] diff --git a/src/graphon/variables/sensitivity.py b/src/graphon/variables/sensitivity.py index 4379a51e..10535c26 100644 --- a/src/graphon/variables/sensitivity.py +++ b/src/graphon/variables/sensitivity.py @@ -1,6 +1,5 @@ from .segment_group import SegmentGroup -from .segments import Segment -from .variables import SecretVariable +from .segments import SecretSegment, Segment def is_sensitive(segment: Segment) -> bool: @@ -10,10 +9,10 @@ def is_sensitive(segment: Segment) -> bool: ``.log`` representation instead of the plaintext ``.value``. Returns: - True if the segment is a SecretVariable or a SegmentGroup that - contains at least one SecretVariable (at any depth), False otherwise. + True if the segment is a SecretSegment or a SegmentGroup that + contains at least one SecretSegment (at any depth), False otherwise. """ - if isinstance(segment, SecretVariable): + if isinstance(segment, SecretSegment): return True if isinstance(segment, SegmentGroup): return any(is_sensitive(child) for child in segment.value) diff --git a/src/graphon/variables/types.py b/src/graphon/variables/types.py index 5b8fc7eb..3196e480 100644 --- a/src/graphon/variables/types.py +++ b/src/graphon/variables/types.py @@ -105,7 +105,9 @@ def infer_segment_type(cls, value: Any) -> SegmentType | None: internal exhaustive match. """ - if isinstance(value, list): + if isinstance(value, (list, tuple)): + if not value: + return SegmentType.ARRAY_ANY elem_types: set[SegmentType] = set() for item in value: segment_type = cls.infer_segment_type(item) @@ -255,7 +257,6 @@ def get_zero_value(t: SegmentType) -> Segment: # Lazy import to avoid circular dependency between segment types # and factory helpers. from .factory import ( # ruff:ignore[import-outside-top-level] - build_segment, build_segment_with_type, ) @@ -266,7 +267,7 @@ def get_zero_value(t: SegmentType) -> Segment: if zero_value is None: msg = f"unsupported variable type: {t}" raise ValueError(msg) - return build_segment(zero_value) + return build_segment_with_type(t, zero_value) _ARRAY_SEGMENT_TYPE_BY_ELEMENT_TYPE: Mapping[SegmentType, SegmentType] = { @@ -280,6 +281,7 @@ def get_zero_value(t: SegmentType) -> Segment: SegmentType.BOOLEAN: SegmentType.ARRAY_BOOLEAN, } _EMPTY_ARRAY_ZERO_VALUE_TYPES = frozenset(( + SegmentType.ARRAY_FILE, SegmentType.ARRAY_OBJECT, SegmentType.ARRAY_ANY, SegmentType.ARRAY_STRING, @@ -292,6 +294,7 @@ def get_zero_value(t: SegmentType) -> Segment: ] = { SegmentType.OBJECT: {}, SegmentType.STRING: "", + SegmentType.SECRET: "", SegmentType.INTEGER: 0, SegmentType.FLOAT: 0.0, SegmentType.NUMBER: 0, diff --git a/src/graphon/variables/variables.py b/src/graphon/variables/variables.py index 83babc66..ce670b04 100644 --- a/src/graphon/variables/variables.py +++ b/src/graphon/variables/variables.py @@ -1,8 +1,10 @@ +from __future__ import annotations + from collections.abc import Sequence -from typing import Annotated, Any +from typing import Annotated, Any, Self from uuid import uuid4 -from pydantic import BaseModel, Discriminator, Field, Tag +from pydantic import BaseModel, ConfigDict, Field from .segments import ( ArrayAnySegment, @@ -18,21 +20,10 @@ IntegerSegment, NoneSegment, ObjectSegment, + SecretSegment, Segment, StringSegment, - get_segment_discriminator, ) -from .types import SegmentType - -SHORT_TOKEN_MAX_LENGTH = 8 - - -def _obfuscated_token(token: str) -> str: - if not token: - return token - if len(token) <= SHORT_TOKEN_MAX_LENGTH: - return "*" * 20 - return token[:6] + "*" * 12 + token[-2:] class VariableBase(Segment): @@ -43,6 +34,9 @@ class VariableBase(Segment): Note: this class is abstract, you should use subclasses of this class instead. """ + # Ancestor write-back uses variable instance identity to recognize committed writes. + model_config = ConfigDict(revalidate_instances="never") + id: str = Field( default_factory=lambda: str(uuid4()), description="Unique identity for variable.", @@ -51,6 +45,23 @@ class VariableBase(Segment): description: str = Field(default="", description="Description of the variable.") selector: Sequence[str] = Field(default_factory=list) + def copy_with_value(self, value: object) -> Self | FloatVariable: + variable_class = ( + FloatVariable + if isinstance(self, IntegerVariable) and isinstance(value, float) + else type(self) + ) + return variable_class.model_validate( + { + "id": self.id, + "name": self.name, + "description": self.description, + "selector": self.selector, + "value": value, + }, + strict=True, + ) + class StringVariable(StringSegment, VariableBase): pass @@ -88,17 +99,12 @@ class ArrayObjectVariable(ArrayObjectSegment, ArrayVariable): pass -class SecretVariable(StringVariable): - value_type: SegmentType = SegmentType.SECRET - - @property - def log(self) -> str: - return _obfuscated_token(self.value) +class SecretVariable(SecretSegment, VariableBase): + pass class NoneVariable(NoneSegment, VariableBase): - value_type: SegmentType = SegmentType.NONE - value: None = None + pass class FileVariable(FileSegment, VariableBase): @@ -166,20 +172,20 @@ class RAGPipelineVariableInput(BaseModel): # - The union must include all non-abstract subclasses of `VariableBase`. type Variable = Annotated[ ( - Annotated[NoneVariable, Tag(SegmentType.NONE)] - | Annotated[StringVariable, Tag(SegmentType.STRING)] - | Annotated[FloatVariable, Tag(SegmentType.FLOAT)] - | Annotated[IntegerVariable, Tag(SegmentType.INTEGER)] - | Annotated[ObjectVariable, Tag(SegmentType.OBJECT)] - | Annotated[FileVariable, Tag(SegmentType.FILE)] - | Annotated[BooleanVariable, Tag(SegmentType.BOOLEAN)] - | Annotated[ArrayAnyVariable, Tag(SegmentType.ARRAY_ANY)] - | Annotated[ArrayStringVariable, Tag(SegmentType.ARRAY_STRING)] - | Annotated[ArrayNumberVariable, Tag(SegmentType.ARRAY_NUMBER)] - | Annotated[ArrayObjectVariable, Tag(SegmentType.ARRAY_OBJECT)] - | Annotated[ArrayFileVariable, Tag(SegmentType.ARRAY_FILE)] - | Annotated[ArrayBooleanVariable, Tag(SegmentType.ARRAY_BOOLEAN)] - | Annotated[SecretVariable, Tag(SegmentType.SECRET)] + NoneVariable + | StringVariable + | FloatVariable + | IntegerVariable + | ObjectVariable + | FileVariable + | BooleanVariable + | ArrayAnyVariable + | ArrayStringVariable + | ArrayNumberVariable + | ArrayObjectVariable + | ArrayFileVariable + | ArrayBooleanVariable + | SecretVariable ), - Discriminator(get_segment_discriminator), + Field(discriminator="value_type"), ] diff --git a/src/graphon/workflow_type_encoder.py b/src/graphon/workflow_type_encoder.py index 4845181e..5538eb56 100644 --- a/src/graphon/workflow_type_encoder.py +++ b/src/graphon/workflow_type_encoder.py @@ -46,7 +46,7 @@ def value_to_json_encodable_recursive(self, value: Any) -> Any: for key, item in value.items(): encoded_mapping[key] = self.value_to_json_encodable_recursive(item) result = encoded_mapping - case list(): + case list() | tuple(): result = [ self.value_to_json_encodable_recursive(item) for item in value ] diff --git a/tests/http/test_executor_secret_masking.py b/tests/http/test_executor_secret_masking.py index 66734e40..5f3f1ca1 100644 --- a/tests/http/test_executor_secret_masking.py +++ b/tests/http/test_executor_secret_masking.py @@ -25,7 +25,7 @@ from graphon.variables import SecretVariable SECRET = "sk-supersecret-123456" # ruff: ignore[hardcoded-password-string] -# The obfuscated form per _obfuscated_token(): first 6 chars + 12 stars + last 2 chars. +# The obfuscated form per _mask_token(): first 6 chars + 12 stars + last 2 chars. # len(SECRET) = 22 > 8, so: "sk-sup" + "************" + "56" OBFUSCATED = "sk-sup" + "*" * 12 + "56" # When the obfuscated token is percent-encoded (e.g. inside urlencode for params/body) diff --git a/tests/nodes/test_container_dispatch.py b/tests/nodes/test_container_dispatch.py index 533433f5..a0b4e216 100644 --- a/tests/nodes/test_container_dispatch.py +++ b/tests/nodes/test_container_dispatch.py @@ -10,8 +10,38 @@ LoopFrameRequest, build_container_value, ) +from graphon.nodes.iteration.entities import IterationNodeData from graphon.nodes.iteration.iteration_node import IterationNode from graphon.nodes.loop.loop_node import LoopNode +from graphon.runtime import RuntimeState, VariablePool +from graphon.variables.segments import ArrayNumberSegment +from tests.helpers import build_init_params + + +def test_iteration_accepts_a_typed_sequence_input() -> None: + pool = VariablePool() + pool.add(("host", "items"), ArrayNumberSegment(value=(1, 2))) + node = IterationNode( + node_id="iteration", + data=IterationNodeData( + title="Iteration", + iterator_selector=["host", "items"], + output_selector=["iteration", "item"], + start_node_id="iteration-start", + ), + init_params=build_init_params(), + runtime_state=RuntimeState( + workflow_id="workflow", variable_pool=pool, start_at=1 + ), + ) + node.bind_execution_id("iteration-run") + + requests = [ + event for event in node.run() if isinstance(event, IterationFrameRequest) + ] + + assert len(requests) == 1 + assert [item.value for item in requests[0].items] == [1, 2] def test_container_await_requests_have_intrinsic_kind_tags() -> None: diff --git a/tests/nodes/test_container_effects.py b/tests/nodes/test_container_effects.py new file mode 100644 index 00000000..f271c359 --- /dev/null +++ b/tests/nodes/test_container_effects.py @@ -0,0 +1,175 @@ +import pytest +from pydantic import ValidationError + +from graphon.engine.ready_queue.entities import ResumeTask +from graphon.enums import WorkflowNodeExecutionStatus +from graphon.file.enums import FileTransferMethod, FileType +from graphon.file.models import File +from graphon.nodes.container_effects import ( + ContainerExecutionResult, + ContainerNodeRunResult, + build_container_value, +) +from graphon.runtime.runtime_state import RuntimeState +from graphon.runtime.variable_pool import VariablePool +from graphon.variables.segments import ( + ArrayAnySegment, + ArrayBooleanSegment, + ArrayFileSegment, + ArrayNumberSegment, + ArrayObjectSegment, + ArrayStringSegment, + IntegerSegment, + SerializableSegment, +) +from graphon.variables.variables import ArrayFileVariable, SecretVariable +from graphon.workflow_type_encoder import WorkflowRuntimeTypeConverter + + +@pytest.mark.parametrize("use_builder", [False, True], ids=["direct", "builder"]) +@pytest.mark.parametrize( + "segment", + [ + ArrayFileSegment(value=[]), + ArrayStringSegment(value=[]), + ArrayNumberSegment(value=[]), + ArrayObjectSegment(value=[]), + ArrayBooleanSegment(value=[]), + ArrayAnySegment(value=["text"]), + ArrayAnySegment( + value=[ + { + "file": File( + file_type=FileType.DOCUMENT, + transfer_method=FileTransferMethod.LOCAL_FILE, + reference="example-record", + ), + "metadata": {"type": "document", "reference": "plain-value"}, + } + ] + ), + ArrayFileVariable(name="files", selector=["child", "files"], value=[]), + SecretVariable(name="token", selector=["env", "token"], value="example-value"), + ], + ids=[ + "file", + "string", + "number", + "object", + "boolean", + "any", + "nested-file", + "variable", + "secret", + ], +) +def test_container_outputs_keep_their_type_after_snapshot_and_resume( + segment: SerializableSegment, + use_builder: bool, +) -> None: + value = build_container_value(segment) if use_builder else segment + state = RuntimeState( + workflow_id="workflow", variable_pool=VariablePool(), start_at=1 + ) + state.enqueue_ready_task( + ResumeTask( + invocation_id="container-run", + result=ContainerExecutionResult( + metadata={}, + steps=1, + node_run_result=ContainerNodeRunResult( + status=WorkflowNodeExecutionStatus.SUCCEEDED, + inputs={"input": value}, + outputs={"output": value}, + ), + ), + ), + ) + + restored = RuntimeState.from_snapshot(state.dumps()) + task = restored.ready_queue.get(timeout=0.01) + assert isinstance(task, ResumeTask) + assert isinstance(task.result, ContainerExecutionResult) + result = task.result.node_run_result.to_node_run_result() + restored.variable_pool.add(("container", "output"), result.outputs["output"]) + + output = restored.variable_pool.get_variable(("container", "output")) + assert output is not None + assert output.value_type == segment.value_type + assert output.value == segment.value + assert list(output.selector) == ["container", "output"] + assert result.inputs == {"input": segment.to_object()} + + +@pytest.mark.parametrize("use_builder", [False, True], ids=["direct", "builder"]) +def test_live_container_outputs_rebind_variable_selectors(use_builder: bool) -> None: + variable = ArrayFileVariable(name="files", selector=["child", "files"], value=[]) + value = build_container_value(variable) if use_builder else variable + result = ContainerNodeRunResult(outputs={"files": value}).to_node_run_result() + pool = VariablePool() + + pool.add(("container", "files"), result.outputs["files"]) + + files = pool.get_variable(("container", "files")) + assert isinstance(files, ArrayFileSegment) + assert files.value == [] + assert list(files.selector) == ["container", "files"] + + +def test_container_values_preserve_secret_masking() -> None: + secret = SecretVariable(name="token", value="example-value") + + result = ContainerNodeRunResult.model_validate({ + "inputs": {"token": secret}, + "outputs": {"token": build_container_value(secret)}, + }).to_node_run_result() + pool = VariablePool() + pool.add(("container", "token"), result.outputs["token"]) + + token = pool.get(("container", "token")) + assert token is not None + assert token.value_type == "secret" + assert token.value == "example-value" + assert WorkflowRuntimeTypeConverter().to_json_encodable(result.outputs) == { + "token": secret.log, + } + assert token.log == secret.log + assert result.inputs == {"token": "example-value"} + + +def test_container_any_array_rejects_invalid_serialized_file() -> None: + with pytest.raises(ValidationError, match="Missing file reference"): + ContainerNodeRunResult.model_validate({ + "outputs": { + "files": { + "value_type": "array[any]", + "value": [ + { + "dify_model_identity": "__dify__file__", + "type": "document", + "transfer_method": "local_file", + }, + ], + }, + }, + }) + + +@pytest.mark.parametrize("use_builder", [False, True], ids=["direct", "builder"]) +@pytest.mark.parametrize( + "segment", + [ + IntegerSegment(value=1).model_copy(update={"value": "invalid"}), + ArrayNumberSegment(value=[1]).model_copy(update={"value": ["invalid"]}), + ], + ids=["integer", "array"], +) +def test_container_outputs_reject_invalid_typed_values( + segment: SerializableSegment, use_builder: bool +) -> None: + if use_builder: + with pytest.raises(ValidationError): + build_container_value(segment) + else: + with pytest.raises(ValidationError): + ContainerNodeRunResult(outputs={"output": segment}) diff --git a/tests/nodes/variable_assigner/test_v1_node.py b/tests/nodes/variable_assigner/test_v1_node.py index f7ff3de4..1b762935 100644 --- a/tests/nodes/variable_assigner/test_v1_node.py +++ b/tests/nodes/variable_assigner/test_v1_node.py @@ -1,7 +1,10 @@ import time from collections.abc import Sequence +import pytest + from graphon.engine_events.node import ( + NodeRunFailedEvent, NodeRunSucceededEvent, NodeRunVariableUpdatedEvent, ) @@ -12,6 +15,8 @@ from graphon.runtime.variable_pool import VariablePool from graphon.variables.variables import ( ArrayStringVariable, + FloatVariable, + IntegerVariable, StringVariable, ) @@ -158,3 +163,77 @@ def test_clear_array() -> None: assert updated_variables[0].new_value == [] assert updated_event.variable.value == [] assert succeeded_event.node_run_result.inputs == {"value": []} + + +@pytest.mark.parametrize( + ("original", "write_mode", "input_value"), + [ + ( + IntegerVariable(name="target", value=1), + WriteMode.OVER_WRITE, + "not a number", + ), + ( + ArrayStringVariable(name="target", value=["first"]), + WriteMode.APPEND, + 42, + ), + ], + ids=["invalid-overwrite", "invalid-append"], +) +def test_invalid_value_fails_without_updating_variables( + original: IntegerVariable | ArrayStringVariable, + write_mode: WriteMode, + input_value: str | int, +) -> None: + pool = build_variable_pool( + conversation_variables=[original], + variables=[(("input", "value"), input_value)], + ) + original_before = original.model_copy(deep=True) + pool_before = pool.model_copy(deep=True) + node = _build_node( + variable_pool=pool, + assigned_selector=("conversation", original.name), + write_mode=write_mode, + input_selector=("input", "value"), + ) + + events = list(node.run()) + + assert any(isinstance(event, NodeRunFailedEvent) for event in events) + assert not any(isinstance(event, NodeRunVariableUpdatedEvent) for event in events) + assert original == original_before + assert pool == pool_before + + +def test_overwrite_number_keeps_fractional_value_and_variable_identity() -> None: + original = IntegerVariable( + id="number-id", + name="number", + description="Mutable number", + selector=["conversation", "number"], + value=1, + ) + pool = build_variable_pool( + conversation_variables=[original], + variables=[(("input", "number"), FloatVariable(name="input", value=0.5))], + ) + node = _build_node( + variable_pool=pool, + assigned_selector=original.selector, + write_mode=WriteMode.OVER_WRITE, + input_selector=("input", "number"), + ) + + update = next( + event for event in node.run() if isinstance(event, NodeRunVariableUpdatedEvent) + ) + restored = NodeRunVariableUpdatedEvent.model_validate_json(update.model_dump_json()) + + assert restored.variable.value == 0.5 # ruff: ignore[float-equality-comparison] - Exact binary fraction. + assert restored.variable.value_type == "float" + assert restored.variable.id == original.id + assert restored.variable.name == original.name + assert restored.variable.description == original.description + assert restored.variable.selector == original.selector diff --git a/tests/nodes/variable_assigner/test_v2_node.py b/tests/nodes/variable_assigner/test_v2_node.py index 45da22a8..2e5a42e8 100644 --- a/tests/nodes/variable_assigner/test_v2_node.py +++ b/tests/nodes/variable_assigner/test_v2_node.py @@ -1,6 +1,8 @@ import time from collections.abc import Sequence +import pytest + from graphon.engine_events.node import ( NodeRunFailedEvent, NodeRunSucceededEvent, @@ -16,7 +18,9 @@ from graphon.runtime.runtime_state import RuntimeState from graphon.runtime.variable_pool import VariablePool from graphon.variables.variables import ( + ArrayFileVariable, ArrayStringVariable, + IntegerVariable, StringVariable, ) @@ -106,6 +110,32 @@ def test_remove_last_from_array() -> None: assert updated_event.variable.value == ["first", "second"] +def test_clear_preserves_the_file_array_type() -> None: + variable = ArrayFileVariable( + name="files", selector=["iteration", "output"], value=[] + ) + pool = build_variable_pool(variables=[(variable.selector, variable)]) + node = _build_node( + variable_pool=pool, + items=[ + VariableOperationItem( + variable_selector=variable.selector, + input_type=InputType.VARIABLE, + operation=Operation.CLEAR, + ), + ], + ) + + updates = [ + event for event in node.run() if isinstance(event, NodeRunVariableUpdatedEvent) + ] + + assert len(updates) == 1 + assert updates[0].variable.value_type == "array[file]" + assert updates[0].variable.value == [] + assert updates[0].variable.selector == variable.selector + + def test_multiple_operations_emit_single_final_update_per_selector() -> None: conversation_variable = ArrayStringVariable( name="test_conversation_variable", @@ -183,3 +213,48 @@ def test_invalid_constant_input_returns_failed_event() -> None: ) assert failed_event.node_run_result.status == WorkflowNodeExecutionStatus.FAILED assert failed_event.node_run_result.error == "Invalid input value 123" + + +@pytest.mark.parametrize( + ("operation", "value", "expected_value"), + [ + (Operation.DIVIDE, 2, 0.5), + (Operation.ADD, 0.5, 1.5), + ], +) +def test_numeric_updates_keep_valid_types_and_variable_identity( + operation: Operation, + value: float, + expected_value: float, +) -> None: + original = IntegerVariable( + id="number-id", + name="number", + description="Mutable number", + selector=["conversation", "number"], + value=1, + ) + pool = build_variable_pool(conversation_variables=[original]) + node = _build_node( + variable_pool=pool, + items=[ + VariableOperationItem( + variable_selector=original.selector, + input_type=InputType.CONSTANT, + operation=operation, + value=value, + ), + ], + ) + + update = next( + event for event in node.run() if isinstance(event, NodeRunVariableUpdatedEvent) + ) + restored = NodeRunVariableUpdatedEvent.model_validate_json(update.model_dump_json()) + + assert restored.variable.value == expected_value + assert restored.variable.value_type == "float" + assert restored.variable.id == original.id + assert restored.variable.name == original.name + assert restored.variable.description == original.description + assert restored.variable.selector == original.selector diff --git a/tests/runtime/test_variable_pool.py b/tests/runtime/test_variable_pool.py index d9e67ca7..89ab1166 100644 --- a/tests/runtime/test_variable_pool.py +++ b/tests/runtime/test_variable_pool.py @@ -166,6 +166,26 @@ def test_add_keeps_variable_instances_and_supports_segments(self) -> None: assert segment is not None assert segment.value == "Paris" + def test_add_binds_copied_outputs_to_the_destination(self) -> None: + pool = VariablePool.empty() + original = StringVariable( + name="input", selector=["source", "input"], value="before" + ) + pool.add(original.selector, original) + pool.add(("copy", "output"), original) + output = pool.get_variable(("copy", "output")) + assert output is not None + + pool.add(output.selector, output.model_copy(update={"value": "after"})) + + assert pool.get(("source", "input")) == original + updated = pool.get(("copy", "output")) + assert updated is not None + assert updated.value == "after" + assert list(output.selector) == ["copy", "output"] + assert output.name == "output" + assert output.id != original.id + class TestVariablePoolGetNotModifyVariableDictionary: _NODE_ID = "start" diff --git a/tests/variables/test_factory.py b/tests/variables/test_factory.py new file mode 100644 index 00000000..01ced97d --- /dev/null +++ b/tests/variables/test_factory.py @@ -0,0 +1,58 @@ +from typing import assert_type + +import pytest +from pydantic import TypeAdapter + +from graphon.variables.factory import ( + TypeMismatchError, + build_segment, + build_segment_with_type, +) +from graphon.variables.segments import SerializableSegment, StringSegment +from graphon.variables.types import SegmentType + + +@pytest.mark.parametrize( + ("value_type", "value"), + [ + (SegmentType.ARRAY_ANY, ["text"]), + (SegmentType.ARRAY_FILE, []), + (SegmentType.FLOAT, 1), + (SegmentType.SECRET, "example-value"), + ], +) +def test_declared_type_accepts_compatible_values( + value_type: SegmentType, value: object +) -> None: + segment = build_segment_with_type(value_type, value) + + assert segment.value_type == value_type + assert segment.value == value + + +@pytest.mark.parametrize( + ("value_type", "value"), + [(SegmentType.INTEGER, "1"), (SegmentType.ARRAY_STRING, [1])], +) +def test_declared_type_rejects_incompatible_values( + value_type: SegmentType, value: object +) -> None: + with pytest.raises(TypeMismatchError): + build_segment_with_type(value_type, value) + + +def test_empty_array_inference_is_consistent() -> None: + assert SegmentType.infer_segment_type([]) == build_segment([]).value_type + + +def test_serialized_segment_tag_selects_its_value_type() -> None: + segment = TypeAdapter(SerializableSegment).validate_json( + '{"value_type": "string", "value": "text"}' + ) + + if segment.value_type == SegmentType.STRING: + assert_type(segment, StringSegment) + assert_type(segment.value, str) + assert segment.value == "text" + else: + pytest.fail("String tag did not select a string segment") diff --git a/tests/variables/test_workflow_type_encoder.py b/tests/variables/test_workflow_type_encoder.py new file mode 100644 index 00000000..05590405 --- /dev/null +++ b/tests/variables/test_workflow_type_encoder.py @@ -0,0 +1,26 @@ +import json +from unittest.mock import MagicMock + +from graphon.file.enums import FileTransferMethod, FileType +from graphon.file.models import File +from graphon.file.protocols import WorkflowFileRuntimeProtocol +from graphon.file.runtime import use_workflow_file_runtime +from graphon.nodes.container_effects import build_container_value +from graphon.variables.segments import ArrayFileSegment +from graphon.workflow_type_encoder import WorkflowRuntimeTypeConverter + + +def test_tuple_file_array_can_be_encoded_as_json() -> None: + file = File( + file_type=FileType.DOCUMENT, + transfer_method=FileTransferMethod.REMOTE_URL, + remote_url="https://example.test/file.txt", + ) + value = build_container_value(ArrayFileSegment(value=(file,))) + runtime = MagicMock(spec=WorkflowFileRuntimeProtocol) + runtime.resolve_file_url.return_value = file.remote_url + + with use_workflow_file_runtime(runtime): + result = WorkflowRuntimeTypeConverter().to_json_encodable({"files": value}) + + assert json.loads(json.dumps(result)) == {"files": [file.to_dict()]} diff --git a/tests/variables/test_workflow_type_encoder_secret.py b/tests/variables/test_workflow_type_encoder_secret.py index 52e03461..351da6d7 100644 --- a/tests/variables/test_workflow_type_encoder_secret.py +++ b/tests/variables/test_workflow_type_encoder_secret.py @@ -2,7 +2,7 @@ from graphon.workflow_type_encoder import WorkflowRuntimeTypeConverter SECRET = "sk-1234567890abcdef" # ruff: ignore[hardcoded-password-string] -MASKED = "sk-123" + "*" * 12 + "ef" # _obfuscated_token(SECRET) +MASKED = "sk-123" + "*" * 12 + "ef" # _mask_token(SECRET) def test_secret_variable_is_masked() -> None: diff --git a/tests/workflows/test_full_engine_events.py b/tests/workflows/test_full_engine_events.py index 062acfec..49796d6d 100644 --- a/tests/workflows/test_full_engine_events.py +++ b/tests/workflows/test_full_engine_events.py @@ -3,6 +3,7 @@ from collections.abc import Mapping, Sequence from threading import Event, Thread from typing import Any +from unittest.mock import MagicMock import pytest import yaml @@ -35,9 +36,26 @@ NodeRunSucceededEvent, NodeRunVariableUpdatedEvent, ) +from graphon.file.enums import FileTransferMethod, FileType +from graphon.file.models import File +from graphon.file.protocols import WorkflowFileRuntimeProtocol +from graphon.file.runtime import use_workflow_file_runtime +from graphon.runtime import RuntimeState from graphon.variables.factory import build_segment, segment_to_variable -from graphon.variables.segments import Segment -from graphon.variables.variables import ArrayStringVariable, StringVariable +from graphon.variables.segments import ( + ArrayAnySegment, + ArrayFileSegment, + ArrayNumberSegment, + ArrayStringSegment, + NoneSegment, + Segment, +) +from graphon.variables.variables import ( + ArrayStringVariable, + SecretVariable, + StringVariable, +) +from graphon.workflow_type_encoder import WorkflowRuntimeTypeConverter from tests.helpers.workflow_events import ( event_path, fake_slim_llm, @@ -1149,6 +1167,345 @@ def test_full_loop_graph_stops_at_loop_end_node() -> None: assert succeeded.metadata["completed_reason"] == "loop_break" +def test_loop_keeps_empty_typed_array_outputs_in_the_parent_pool() -> None: + dsl = _graph_dsl( + nodes=[ + _start_node(), + { + "id": "loop", + "data": { + "type": "loop", + "loop_count": 1, + "start_node_id": "loop-start", + "break_conditions": [], + "logical_operator": "and", + "loop_variables": [ + { + "label": "items", + "var_type": "array[string]", + "value_type": "constant", + "value": "[]", + } + ], + }, + }, + { + "id": "loop-start", + "data": {"type": "loop-start", "container_id": "loop"}, + }, + { + "id": "stop", + "data": {"type": "loop-end", "container_id": "loop"}, + }, + _end_node([]), + ], + edges=[ + _edge("start", "loop"), + _edge("loop-start", "stop"), + _edge("loop", "end"), + ], + ) + engine = loads(dsl) + + events = list(engine.run()) + + items = engine.runtime_state.variable_pool.get(("loop", "items")) + assert isinstance(items, ArrayStringSegment) + assert items.value == [] + succeeded = next( + event for event in events if isinstance(event, NodeRunLoopSucceededEvent) + ) + assert succeeded.outputs == {"items": [], "loop_round": 1} + + +def test_loop_completes_after_integer_division_produces_a_fraction() -> None: + loop_nodes = _loop_nodes("loop") + loop_nodes[0]["data"]["loop_variables"] = [ + {"label": "number", "var_type": "number", "value_type": "constant", "value": 1}, + ] + dsl = _graph_dsl( + nodes=[ + _start_node(), + *loop_nodes, + { + "id": "divide", + "data": { + "type": "assigner", + "version": "2", + "container_id": "loop", + "items": [ + { + "variable_selector": ["loop", "number"], + "input_type": "constant", + "operation": "/=", + "value": 2, + }, + ], + }, + }, + _end_node([{"variable": "number", "value_selector": ["loop", "number"]}]), + ], + edges=[ + _edge("start", "loop"), + _edge("loop-start", "divide"), + _edge("loop", "end"), + ], + ) + + events = run_workflow(dsl) + + assert final_outputs(events) == {"number": 0.5} + succeeded = next( + event for event in events if isinstance(event, NodeRunLoopSucceededEvent) + ) + assert succeeded.outputs == {"number": 0.5, "loop_round": 1} + + +@pytest.mark.parametrize( + ("operation", "input_type", "value", "expected_value"), + [ + ("set", "constant", "replacement-token", "replacement-token"), + ("over-write", "variable", ["start", "replacement"], "replacement-token"), + ("clear", "variable", None, ""), + ], +) +def test_loop_secret_can_be_assigned_and_cleared( + operation: str, + input_type: str, + value: object, + expected_value: str, +) -> None: + loop_nodes = _loop_nodes("loop") + loop_nodes[0]["data"]["loop_variables"] = [ + { + "label": "token", + "var_type": "string", + "value_type": "variable", + "value": ["env", "token"], + }, + ] + dsl = _graph_dsl( + nodes=[ + _start_node(), + *loop_nodes, + { + "id": "assign", + "data": { + "type": "assigner", + "version": "2", + "container_id": "loop", + "items": [ + { + "variable_selector": ["loop", "token"], + "input_type": input_type, + "operation": operation, + "value": value, + }, + ], + }, + }, + _end_node([{"variable": "token", "value_selector": ["loop", "token"]}]), + ], + edges=[ + _edge("start", "loop"), + _edge("loop-start", "assign"), + _edge("loop", "end"), + ], + ) + engine = loads(dsl, start_inputs={"replacement": "replacement-token"}) + engine.runtime_state.variable_pool.add( + ("env", "token"), SecretVariable(name="token", value="initial-token") + ) + + events = list(engine.run()) + + assert final_outputs(events) == {"token": expected_value} + output = engine.runtime_state.variable_pool.get(("loop", "token")) + assert output is not None + assert output.value_type == "secret" + assert WorkflowRuntimeTypeConverter().to_json_encodable({"token": output}) == { + "token": "replac************en" if expected_value else "", + } + + +@pytest.mark.parametrize("is_parallel", [False, True], ids=["serial", "parallel"]) +@pytest.mark.parametrize( + ("items", "flatten_output", "expected_type", "expected_value"), + [ + ( + [ArrayFileSegment(value=[]), ArrayFileSegment(value=[])], + True, + "array[file]", + [], + ), + ( + [ArrayAnySegment(value=["first"]), ArrayAnySegment(value=["last"])], + True, + "array[any]", + ["first", "last"], + ), + ( + [ + ArrayStringSegment(value=["first"]), + NoneSegment(), + ArrayStringSegment(value=[]), + ], + True, + "array[string]", + ["first"], + ), + ( + [ArrayStringSegment(value=["first"]), ArrayNumberSegment(value=[])], + True, + "array[any]", + ["first"], + ), + ( + [ArrayFileSegment(value=[]), ArrayFileSegment(value=[])], + False, + "array[any]", + [[], []], + ), + ( + [ArrayNumberSegment(value=(1, 2))], + False, + "array[any]", + [(1, 2)], + ), + ( + [ArrayAnySegment(value=[("nested",)])], + False, + "array[any]", + [[("nested",)]], + ), + ], + ids=[ + "empty-files", + "declared-any", + "missing-output", + "mixed-types", + "unflattened", + "tuple-array", + "opaque-elements", + ], +) +def test_iteration_preserves_array_output_types( + is_parallel: bool, + items: list[Segment], + flatten_output: bool, + expected_type: str, + expected_value: list[object], +) -> None: + iteration_nodes = _iteration_nodes( + "iteration", is_parallel=is_parallel, parallel_nums=2 + ) + iteration_nodes[0]["data"]["flatten_output"] = flatten_output + dsl = _graph_dsl( + nodes=[ + _start_node(), + *iteration_nodes, + _end_node([ + {"variable": "items", "value_selector": ["iteration", "output"]} + ]), + ], + edges=[_edge("start", "iteration"), _edge("iteration", "end")], + ) + engine = loads(dsl, start_inputs={"items": items}) + + events = list(engine.run()) + + output = engine.runtime_state.variable_pool.get(("iteration", "output")) + assert output is not None + assert output.value_type == expected_type + assert final_outputs(events) == {"items": expected_value} + succeeded = next( + event for event in events if isinstance(event, NodeRunIterationSucceededEvent) + ) + assert succeeded.outputs == {"output": expected_value} + + +@pytest.mark.parametrize("flatten_output", [False, True]) +def test_iteration_any_array_files_survive_snapshot(flatten_output: bool) -> None: + file = File( + file_type=FileType.DOCUMENT, + transfer_method=FileTransferMethod.LOCAL_FILE, + reference="example-record", + filename="report.txt", + ) + iteration_nodes = _iteration_nodes("iteration") + iteration_nodes[0]["data"]["flatten_output"] = flatten_output + dsl = _graph_dsl( + nodes=[_start_node(), *iteration_nodes, _end_node([])], + edges=[_edge("start", "iteration"), _edge("iteration", "end")], + ) + engine = loads(dsl, start_inputs={"items": [ArrayAnySegment(value=[file])]}) + list(engine.run()) + + restored = RuntimeState.from_snapshot(engine.runtime_state.dumps()) + output = restored.variable_pool.get(("iteration", "output")) + + assert output is not None + assert output.value_type == "array[any]" + assert output.value == ([file] if flatten_output else [[file]]) + runtime = MagicMock(spec=WorkflowFileRuntimeProtocol) + runtime.resolve_file_url.return_value = "https://host.test/report.txt" + with use_workflow_file_runtime(runtime): + expected_files = [file.to_dict()] + assert WorkflowRuntimeTypeConverter().to_json_encodable({"files": output}) == { + "files": expected_files if flatten_output else [expected_files], + } + + +def test_empty_iteration_output_is_independent_of_its_input() -> None: + iteration_nodes = _iteration_nodes("iteration") + iteration_nodes[0]["data"]["iterator_selector"] = ["host", "items"] + dsl = _graph_dsl( + nodes=[ + _start_node(), + *iteration_nodes, + { + "id": "append", + "data": { + "type": "assigner", + "version": "2", + "items": [ + { + "variable_selector": ["iteration", "output"], + "input_type": "variable", + "operation": "append", + "value": ["start", "item"], + }, + ], + }, + }, + _end_node([ + {"variable": "input", "value_selector": ["host", "items"]}, + {"variable": "output", "value_selector": ["iteration", "output"]}, + ]), + ], + edges=[ + _edge("start", "iteration"), + _edge("iteration", "append"), + _edge("append", "end"), + ], + ) + engine = loads(dsl, start_inputs={"item": "appended"}) + engine.runtime_state.variable_pool.add( + ("host", "items"), ArrayStringSegment(value=[]) + ) + + events = list(engine.run()) + + assert final_outputs(events) == {"input": [], "output": ["appended"]} + output = engine.runtime_state.variable_pool.get_variable(("iteration", "output")) + assert isinstance(output, ArrayStringVariable) + assert list(output.selector) == ["iteration", "output"] + succeeded = next( + event for event in events if isinstance(event, NodeRunIterationSucceededEvent) + ) + assert succeeded.outputs == {"output": []} + + def test_full_loop_graph_uses_custom_container_handler_factory() -> None: dsl = _graph_dsl( nodes=[