wehub-resource-sync
cddb07a176
build container image / cpu (push) Waiting to run
build container image / cuda (push) Waiting to run
build container image / rocm (push) Waiting to run
frontend tests / frontend-tests (push) Waiting to run
openapi checks / openapi-checks (push) Waiting to run
python tests / py3.12: macos-default (push) Waiting to run
python tests / py3.11: windows-cpu (push) Waiting to run
python tests / py3.12: windows-cpu (push) Waiting to run
python tests / py3.11: linux-cpu (push) Waiting to run
python tests / py3.12: linux-cpu (push) Waiting to run
typegen checks / typegen-checks (push) Waiting to run
uv lock checks / uv-lock-checks (push) Waiting to run
frontend checks / frontend-checks (push) Waiting to run
lfs checks / lfs-check (push) Waiting to run
python checks / python-checks (push) Waiting to run
python tests / py3.11: macos-default (push) Waiting to run
docs / deploy (push) Has been cancelled
docs / changes (push) Has been cancelled
docs / check-and-build (push) Has been cancelled
1570 lines
66 KiB
Python
1570 lines
66 KiB
Python
from typing import Optional
|
|
from unittest.mock import Mock
|
|
|
|
import pytest
|
|
from pydantic import TypeAdapter
|
|
|
|
from invokeai.app.invocations.baseinvocation import BaseInvocation, BaseInvocationOutput, InvocationContext
|
|
from invokeai.app.invocations.collections import RangeInvocation
|
|
from invokeai.app.invocations.fields import InputField, OutputField
|
|
from invokeai.app.invocations.logic import IfInvocation, IfInvocationOutput
|
|
from invokeai.app.invocations.math import AddInvocation, MultiplyInvocation
|
|
from invokeai.app.invocations.primitives import (
|
|
BooleanCollectionInvocation,
|
|
BooleanCollectionOutput,
|
|
BooleanInvocation,
|
|
BooleanOutput,
|
|
)
|
|
from invokeai.app.services.shared.graph import (
|
|
CollectInvocation,
|
|
Graph,
|
|
GraphExecutionState,
|
|
IterateInvocation,
|
|
WorkflowCallFrame,
|
|
)
|
|
|
|
# This import must happen before other invoke imports or test in other files(!!) break
|
|
from tests.test_nodes import (
|
|
AnyTypeTestInvocation,
|
|
AnyTypeTestInvocationOutput,
|
|
PromptCollectionTestInvocation,
|
|
PromptTestInvocation,
|
|
TextToImageTestInvocation,
|
|
create_edge,
|
|
)
|
|
|
|
|
|
class IntegerCollectionTestInvocationOutput(BaseInvocationOutput):
|
|
collection: list[int] = OutputField(default=[])
|
|
|
|
|
|
class IntegerCollectionFromItemTestInvocation(BaseInvocation):
|
|
value: int = InputField(default=0)
|
|
|
|
def invoke(self, context: InvocationContext) -> IntegerCollectionTestInvocationOutput:
|
|
base = self.value * 10
|
|
return IntegerCollectionTestInvocationOutput(collection=[base, base + 1])
|
|
|
|
|
|
class IntegerCollectionPassthroughTestInvocation(BaseInvocation):
|
|
collection: list[int] = InputField(default=[])
|
|
|
|
def invoke(self, context: InvocationContext) -> IntegerCollectionTestInvocationOutput:
|
|
return IntegerCollectionTestInvocationOutput(collection=self.collection.copy())
|
|
|
|
|
|
class TwoIntegerCollectionsTestInvocation(BaseInvocation):
|
|
first: list[int] = InputField(default=[])
|
|
second: list[int] = InputField(default=[])
|
|
|
|
def invoke(self, context: InvocationContext) -> IntegerCollectionTestInvocationOutput:
|
|
return IntegerCollectionTestInvocationOutput(collection=self.first + self.second)
|
|
|
|
|
|
@pytest.fixture
|
|
def simple_graph() -> Graph:
|
|
g = Graph()
|
|
g.add_node(PromptTestInvocation(id="1", prompt="Banana sushi"))
|
|
g.add_node(TextToImageTestInvocation(id="2"))
|
|
g.add_edge(create_edge("1", "prompt", "2", "prompt"))
|
|
return g
|
|
|
|
|
|
def invoke_next(g: GraphExecutionState) -> tuple[Optional[BaseInvocation], Optional[BaseInvocationOutput]]:
|
|
n = g.next()
|
|
if n is None:
|
|
return (None, None)
|
|
|
|
print(f"invoking {n.id}: {type(n)}")
|
|
o = n.invoke(Mock(InvocationContext))
|
|
g.complete(n.id, o)
|
|
|
|
return (n, o)
|
|
|
|
|
|
def execute_all_nodes(g: GraphExecutionState) -> list[str]:
|
|
"""Execute the graph to completion and return source node ids in execution order."""
|
|
|
|
executed_source_ids: list[str] = []
|
|
while True:
|
|
invocation, _output = invoke_next(g)
|
|
if invocation is None:
|
|
break
|
|
executed_source_ids.append(g.prepared_source_mapping[invocation.id])
|
|
|
|
return executed_source_ids
|
|
|
|
|
|
def test_graph_state_executes_in_order(simple_graph: Graph):
|
|
g = GraphExecutionState(graph=simple_graph)
|
|
|
|
n1 = invoke_next(g)
|
|
n2 = invoke_next(g)
|
|
n3 = g.next()
|
|
|
|
assert g.prepared_source_mapping[n1[0].id] == "1"
|
|
assert g.prepared_source_mapping[n2[0].id] == "2"
|
|
assert n3 is None
|
|
assert g.results[n1[0].id].prompt == n1[0].prompt
|
|
assert n2[0].prompt == n1[0].prompt
|
|
|
|
|
|
def test_graph_is_complete(simple_graph: Graph):
|
|
g = GraphExecutionState(graph=simple_graph)
|
|
_ = invoke_next(g)
|
|
_ = invoke_next(g)
|
|
_ = g.next()
|
|
|
|
assert g.is_complete()
|
|
|
|
|
|
def test_graph_is_not_complete(simple_graph: Graph):
|
|
g = GraphExecutionState(graph=simple_graph)
|
|
_ = invoke_next(g)
|
|
_ = g.next()
|
|
|
|
assert not g.is_complete()
|
|
|
|
|
|
def test_graph_waiting_on_workflow_call_blocks_other_ready_nodes():
|
|
graph = Graph()
|
|
graph.add_node(PromptTestInvocation(id="prompt_a", prompt="a"))
|
|
graph.add_node(PromptTestInvocation(id="prompt_b", prompt="b"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
|
|
first = g.next()
|
|
assert first is not None
|
|
|
|
waiting_frame = g.build_workflow_call_frame(exec_node_id=first.id, workflow_id="workflow-a")
|
|
g.begin_waiting_on_workflow_call(waiting_frame)
|
|
|
|
assert g.next() is None
|
|
assert not g.is_complete()
|
|
assert g.is_waiting_on_workflow_call()
|
|
|
|
|
|
def test_graph_build_workflow_call_frame_uses_prepared_and_source_ids():
|
|
g = GraphExecutionState(graph=Graph())
|
|
g.execution_graph.add_node(PromptTestInvocation(id="prepared-call", prompt="a"))
|
|
g.prepared_source_mapping["prepared-call"] = "source-call"
|
|
|
|
frame = g.build_workflow_call_frame(exec_node_id="prepared-call", workflow_id="workflow-a")
|
|
|
|
assert frame.prepared_call_node_id == "prepared-call"
|
|
assert frame.source_call_node_id == "source-call"
|
|
assert frame.workflow_id == "workflow-a"
|
|
assert frame.depth == 1
|
|
|
|
|
|
def test_graph_build_workflow_call_frame_rejects_depth_over_limit():
|
|
graph = Graph()
|
|
graph.add_node(PromptTestInvocation(id="source-call", prompt="a"))
|
|
g = GraphExecutionState(
|
|
graph=graph,
|
|
workflow_call_stack=[
|
|
WorkflowCallFrame(
|
|
prepared_call_node_id=f"prepared-{i}",
|
|
source_call_node_id=f"source-{i}",
|
|
workflow_id=f"workflow-{i}",
|
|
depth=i + 1,
|
|
)
|
|
for i in range(4)
|
|
],
|
|
)
|
|
g.execution_graph.add_node(PromptTestInvocation(id="prepared-call", prompt="a"))
|
|
g.prepared_source_mapping["prepared-call"] = "source-call"
|
|
|
|
with pytest.raises(ValueError, match="Maximum workflow call depth"):
|
|
g.build_workflow_call_frame(exec_node_id="prepared-call", workflow_id="workflow-a")
|
|
|
|
|
|
def test_graph_execution_state_serializes_workflow_call_state():
|
|
graph = Graph()
|
|
graph.add_node(PromptTestInvocation(id="source-call", prompt="a"))
|
|
g = GraphExecutionState(graph=graph)
|
|
g.execution_graph.add_node(PromptTestInvocation(id="prepared-call", prompt="a"))
|
|
g.prepared_source_mapping["prepared-call"] = "source-call"
|
|
|
|
frame = g.build_workflow_call_frame(exec_node_id="prepared-call", workflow_id="workflow-a")
|
|
g.workflow_call_stack.append(frame)
|
|
g.begin_waiting_on_workflow_call(frame)
|
|
|
|
restored = GraphExecutionState.model_validate(g.model_dump(warnings=False))
|
|
|
|
assert restored.workflow_call_stack == [frame]
|
|
assert restored.waiting_workflow_call == frame
|
|
assert restored.max_workflow_call_depth == 4
|
|
|
|
|
|
def test_graph_waiting_on_workflow_call_blocks_until_suspended_node_is_completed():
|
|
graph = Graph()
|
|
graph.add_node(PromptTestInvocation(id="prompt_a", prompt="a"))
|
|
graph.add_node(PromptTestInvocation(id="prompt_b", prompt="b"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
|
|
first = g.next()
|
|
assert first is not None
|
|
|
|
waiting_frame = g.build_workflow_call_frame(exec_node_id=first.id, workflow_id="workflow-a")
|
|
g.begin_waiting_on_workflow_call(waiting_frame)
|
|
assert g.next() is None
|
|
|
|
g.end_waiting_on_workflow_call()
|
|
g.complete(first.id, first.invoke(Mock(InvocationContext)))
|
|
|
|
resumed = g.next()
|
|
assert resumed is not None
|
|
assert resumed.id != first.id
|
|
assert g.prepared_source_mapping[resumed.id] == "prompt_b"
|
|
|
|
|
|
def test_graph_begin_waiting_on_workflow_call_rejects_double_entry():
|
|
g = GraphExecutionState(graph=Graph())
|
|
g.execution_graph.add_node(PromptTestInvocation(id="prepared-call", prompt="a"))
|
|
g.prepared_source_mapping["prepared-call"] = "source-call"
|
|
|
|
first_frame = g.build_workflow_call_frame(exec_node_id="prepared-call", workflow_id="workflow-a")
|
|
g.begin_waiting_on_workflow_call(first_frame)
|
|
|
|
with pytest.raises(ValueError, match="already waiting"):
|
|
g.begin_waiting_on_workflow_call(first_frame)
|
|
|
|
|
|
def test_graph_build_workflow_call_frame_rejects_missing_execution_node():
|
|
g = GraphExecutionState(graph=Graph())
|
|
|
|
with pytest.raises(Exception, match="not found in execution graph"):
|
|
g.build_workflow_call_frame(exec_node_id="missing-node", workflow_id="workflow-a")
|
|
|
|
|
|
def test_graph_build_workflow_call_frame_rejects_unprepared_execution_node():
|
|
g = GraphExecutionState(graph=Graph())
|
|
g.execution_graph.add_node(PromptTestInvocation(id="prepared-call", prompt="a"))
|
|
|
|
with pytest.raises(ValueError, match="not a prepared execution node"):
|
|
g.build_workflow_call_frame(exec_node_id="prepared-call", workflow_id="workflow-a")
|
|
|
|
|
|
def test_graph_child_workflow_execution_state_inherits_stack_and_isolates_runtime_state():
|
|
parent_graph = Graph()
|
|
child_graph = Graph()
|
|
|
|
parent = GraphExecutionState(graph=parent_graph)
|
|
parent.execution_graph.add_node(PromptTestInvocation(id="prepared-parent", prompt="a"))
|
|
parent.prepared_source_mapping["prepared-parent"] = "source-parent"
|
|
parent.results["prepared-parent"] = PromptTestInvocation(id="result-node", prompt="existing").invoke(
|
|
Mock(InvocationContext)
|
|
)
|
|
parent.executed.add("prepared-parent")
|
|
|
|
root_frame = parent.build_workflow_call_frame(exec_node_id="prepared-parent", workflow_id="workflow-a")
|
|
parent.workflow_call_stack.append(root_frame)
|
|
|
|
parent.execution_graph.add_node(PromptTestInvocation(id="prepared-child", prompt="b"))
|
|
parent.prepared_source_mapping["prepared-child"] = "source-child"
|
|
child_frame = parent.build_workflow_call_frame(exec_node_id="prepared-child", workflow_id="workflow-b")
|
|
|
|
child_state = parent.create_child_workflow_execution_state(graph=child_graph, frame=child_frame)
|
|
|
|
assert child_state.graph == child_graph
|
|
assert child_state.workflow_call_stack == [root_frame, child_frame]
|
|
assert child_state.max_workflow_call_depth == parent.max_workflow_call_depth
|
|
assert child_state.waiting_workflow_call is None
|
|
assert child_state.results == {}
|
|
assert child_state.executed == set()
|
|
|
|
|
|
def test_graph_waiting_workflow_call_tracks_parent_child_metadata():
|
|
parent = GraphExecutionState(graph=Graph())
|
|
parent.execution_graph.add_node(PromptTestInvocation(id="prepared-parent", prompt="a"))
|
|
parent.prepared_source_mapping["prepared-parent"] = "source-parent"
|
|
frame = parent.build_workflow_call_frame(exec_node_id="prepared-parent", workflow_id="workflow-a")
|
|
|
|
child = parent.create_child_workflow_execution_state(graph=Graph(), frame=frame)
|
|
parent.begin_waiting_on_workflow_call(frame)
|
|
parent.attach_waiting_workflow_call_child_session(child)
|
|
|
|
assert parent.waiting_workflow_call_execution is not None
|
|
assert parent.waiting_workflow_call_execution.parent_session_id == parent.id
|
|
assert parent.waiting_workflow_call_execution.child_session_id == child.id
|
|
assert parent.waiting_workflow_call_execution.status == "running_child"
|
|
assert child.workflow_call_parent is not None
|
|
assert child.workflow_call_parent.workflow_call_id == parent.waiting_workflow_call_execution.id
|
|
assert child.workflow_call_parent.parent_session_id == parent.id
|
|
|
|
|
|
def test_graph_attach_waiting_workflow_call_child_sessions_tracks_fan_out_metadata():
|
|
parent = GraphExecutionState(graph=Graph())
|
|
parent.execution_graph.add_node(AddInvocation(id="prepared-parent", a=1, b=2))
|
|
parent.prepared_source_mapping["prepared-parent"] = "source-parent"
|
|
|
|
frame = parent.build_workflow_call_frame(exec_node_id="prepared-parent", workflow_id="workflow-a")
|
|
child_a = parent.create_child_workflow_execution_state(Graph(), frame)
|
|
child_b = parent.create_child_workflow_execution_state(Graph(), frame)
|
|
|
|
parent.begin_waiting_on_workflow_call(frame)
|
|
parent.attach_waiting_workflow_call_child_sessions([child_a, child_b])
|
|
|
|
assert parent.waiting_workflow_call_execution is not None
|
|
assert parent.waiting_workflow_call_execution.child_session_ids == [child_a.id, child_b.id]
|
|
assert parent.waiting_workflow_call_execution.expected_child_count == 2
|
|
assert parent.waiting_workflow_call_child_session is None
|
|
assert child_a.workflow_call_parent is not None
|
|
assert child_b.workflow_call_parent is not None
|
|
|
|
|
|
def test_graph_record_waiting_workflow_call_child_completion_aggregates_named_values():
|
|
parent = GraphExecutionState(graph=Graph())
|
|
parent.execution_graph.add_node(AddInvocation(id="prepared-parent", a=1, b=2))
|
|
parent.prepared_source_mapping["prepared-parent"] = "source-parent"
|
|
|
|
frame = parent.build_workflow_call_frame(exec_node_id="prepared-parent", workflow_id="workflow-a")
|
|
child_a = parent.create_child_workflow_execution_state(Graph(), frame)
|
|
child_b = parent.create_child_workflow_execution_state(Graph(), frame)
|
|
|
|
parent.begin_waiting_on_workflow_call(frame)
|
|
parent.attach_waiting_workflow_call_child_sessions([child_a, child_b])
|
|
|
|
is_complete, aggregated_values = parent.record_waiting_workflow_call_child_completion(
|
|
101, {"sum": 3, "images": "image-a"}
|
|
)
|
|
assert is_complete is False
|
|
assert aggregated_values == {"sum": [3], "images": ["image-a"]}
|
|
|
|
is_complete, aggregated_values = parent.record_waiting_workflow_call_child_completion(
|
|
102, {"sum": 7, "images": "image-b"}
|
|
)
|
|
assert is_complete is True
|
|
assert aggregated_values == {"sum": [3, 7], "images": ["image-a", "image-b"]}
|
|
assert parent.waiting_workflow_call_execution is not None
|
|
assert parent.waiting_workflow_call_execution.completed_child_item_ids == [101, 102]
|
|
|
|
|
|
def test_graph_record_waiting_workflow_call_child_completion_preserves_enqueue_order():
|
|
parent = GraphExecutionState(graph=Graph())
|
|
parent.execution_graph.add_node(AddInvocation(id="prepared-parent", a=1, b=2))
|
|
parent.prepared_source_mapping["prepared-parent"] = "source-parent"
|
|
|
|
frame = parent.build_workflow_call_frame(exec_node_id="prepared-parent", workflow_id="workflow-a")
|
|
child_a = parent.create_child_workflow_execution_state(Graph(), frame)
|
|
child_b = parent.create_child_workflow_execution_state(Graph(), frame)
|
|
parent.begin_waiting_on_workflow_call(frame)
|
|
parent.attach_waiting_workflow_call_child_sessions([child_a, child_b])
|
|
parent.set_waiting_workflow_call_child_item_ids([101, 102])
|
|
|
|
parent.record_waiting_workflow_call_child_completion(102, {"sum": 7})
|
|
is_complete, aggregated_values = parent.record_waiting_workflow_call_child_completion(101, {"sum": 3})
|
|
|
|
assert is_complete is True
|
|
assert aggregated_values == {"sum": [3, 7]}
|
|
|
|
|
|
def test_graph_end_waiting_on_workflow_call_records_lifecycle_history():
|
|
parent = GraphExecutionState(graph=Graph())
|
|
parent.execution_graph.add_node(PromptTestInvocation(id="prepared-parent", prompt="a"))
|
|
parent.prepared_source_mapping["prepared-parent"] = "source-parent"
|
|
frame = parent.build_workflow_call_frame(exec_node_id="prepared-parent", workflow_id="workflow-a")
|
|
|
|
child = parent.create_child_workflow_execution_state(graph=Graph(), frame=frame)
|
|
parent.begin_waiting_on_workflow_call(frame)
|
|
parent.attach_waiting_workflow_call_child_session(child)
|
|
parent.end_waiting_on_workflow_call(status="failed", error_message="child failed")
|
|
|
|
assert parent.waiting_workflow_call is None
|
|
assert parent.waiting_workflow_call_execution is None
|
|
assert parent.waiting_workflow_call_child_session is None
|
|
assert len(parent.workflow_call_history) == 1
|
|
assert parent.workflow_call_history[0].status == "failed"
|
|
assert parent.workflow_call_history[0].error_message == "child failed"
|
|
assert parent.workflow_call_history[0].parent_session_id == parent.id
|
|
assert parent.workflow_call_history[0].child_session_id == child.id
|
|
|
|
|
|
def test_graph_execution_state_serializes_recursive_workflow_call_stack():
|
|
g = GraphExecutionState(
|
|
graph=Graph(),
|
|
workflow_call_stack=[
|
|
WorkflowCallFrame(
|
|
prepared_call_node_id="prepared-a",
|
|
source_call_node_id="source-a",
|
|
workflow_id="workflow-a",
|
|
depth=1,
|
|
),
|
|
WorkflowCallFrame(
|
|
prepared_call_node_id="prepared-b",
|
|
source_call_node_id="source-b",
|
|
workflow_id="workflow-b",
|
|
depth=2,
|
|
),
|
|
WorkflowCallFrame(
|
|
prepared_call_node_id="prepared-a-2",
|
|
source_call_node_id="source-a-2",
|
|
workflow_id="workflow-a",
|
|
depth=3,
|
|
),
|
|
],
|
|
)
|
|
|
|
restored = GraphExecutionState.model_validate(g.model_dump(warnings=False))
|
|
|
|
assert restored.workflow_call_stack == g.workflow_call_stack
|
|
|
|
|
|
# TODO: test completion with iterators/subgraphs
|
|
|
|
|
|
def test_graph_state_expands_iterator():
|
|
graph = Graph()
|
|
graph.add_node(RangeInvocation(id="0", start=0, stop=3, step=1))
|
|
graph.add_node(IterateInvocation(id="1"))
|
|
graph.add_node(MultiplyInvocation(id="2", b=10))
|
|
graph.add_node(AddInvocation(id="3", b=1))
|
|
graph.add_edge(create_edge("0", "collection", "1", "collection"))
|
|
graph.add_edge(create_edge("1", "item", "2", "a"))
|
|
graph.add_edge(create_edge("2", "value", "3", "a"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
while not g.is_complete():
|
|
invoke_next(g)
|
|
|
|
prepared_add_nodes = g.source_prepared_mapping["3"]
|
|
results = {g.results[n].value for n in prepared_add_nodes}
|
|
expected = {1, 11, 21}
|
|
assert results == expected
|
|
|
|
|
|
def test_graph_state_collects():
|
|
graph = Graph()
|
|
test_prompts = ["Banana sushi", "Cat sushi"]
|
|
graph.add_node(PromptCollectionTestInvocation(id="1", collection=list(test_prompts)))
|
|
graph.add_node(IterateInvocation(id="2"))
|
|
graph.add_node(PromptTestInvocation(id="3"))
|
|
graph.add_node(CollectInvocation(id="4"))
|
|
graph.add_edge(create_edge("1", "collection", "2", "collection"))
|
|
graph.add_edge(create_edge("2", "item", "3", "prompt"))
|
|
graph.add_edge(create_edge("3", "prompt", "4", "item"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
_ = invoke_next(g)
|
|
_ = invoke_next(g)
|
|
_ = invoke_next(g)
|
|
_ = invoke_next(g)
|
|
_ = invoke_next(g)
|
|
n6 = invoke_next(g)
|
|
|
|
assert isinstance(n6[0], CollectInvocation)
|
|
|
|
assert sorted(g.results[n6[0].id].collection) == sorted(test_prompts)
|
|
|
|
|
|
def test_graph_state_resumes_partially_executed_session_after_json_round_trip():
|
|
graph = Graph()
|
|
graph.add_node(RangeInvocation(id="c", start=1, stop=5, step=1))
|
|
graph.add_node(IterateInvocation(id="iter"))
|
|
graph.add_node(AddInvocation(id="add", b=1))
|
|
graph.add_node(CollectInvocation(id="collect"))
|
|
|
|
graph.add_edge(create_edge("c", "collection", "iter", "collection"))
|
|
graph.add_edge(create_edge("iter", "item", "add", "a"))
|
|
graph.add_edge(create_edge("add", "value", "collect", "item"))
|
|
|
|
state = GraphExecutionState(graph=graph)
|
|
|
|
for _ in range(4):
|
|
invocation, output = invoke_next(state)
|
|
assert invocation is not None
|
|
assert output is not None
|
|
|
|
raw = state.model_dump_json(warnings=False, exclude_none=True)
|
|
resumed = TypeAdapter(GraphExecutionState).validate_json(raw, strict=False)
|
|
registry = resumed._prepared_registry()
|
|
|
|
assert all(
|
|
registry.get_iteration_path(exec_node_id) is not None for exec_node_id in resumed.prepared_source_mapping
|
|
)
|
|
|
|
executed_source_ids = execute_all_nodes(resumed)
|
|
|
|
assert executed_source_ids
|
|
assert "add" in executed_source_ids
|
|
assert "collect" in resumed.source_prepared_mapping
|
|
|
|
prepared_collect_id = next(iter(resumed.source_prepared_mapping["collect"]))
|
|
assert resumed.results[prepared_collect_id].collection == [2, 3, 4, 5]
|
|
|
|
|
|
def test_if_graph_state_resumes_resolved_branch_after_json_round_trip():
|
|
graph = Graph()
|
|
graph.add_node(BooleanInvocation(id="condition", value=True))
|
|
graph.add_node(PromptTestInvocation(id="true_value", prompt="true branch"))
|
|
graph.add_node(PromptTestInvocation(id="false_value", prompt="false branch"))
|
|
graph.add_node(IfInvocation(id="if"))
|
|
graph.add_node(PromptTestInvocation(id="selected_output"))
|
|
|
|
graph.add_edge(create_edge("condition", "value", "if", "condition"))
|
|
graph.add_edge(create_edge("true_value", "prompt", "if", "true_input"))
|
|
graph.add_edge(create_edge("false_value", "prompt", "if", "false_input"))
|
|
graph.add_edge(create_edge("if", "value", "selected_output", "prompt"))
|
|
|
|
state = GraphExecutionState(graph=graph)
|
|
|
|
for _ in range(2):
|
|
invocation, output = invoke_next(state)
|
|
assert invocation is not None
|
|
assert output is not None
|
|
|
|
raw = state.model_dump_json(warnings=False, exclude_none=True)
|
|
resumed = TypeAdapter(GraphExecutionState).validate_json(raw, strict=False)
|
|
|
|
executed_source_ids = execute_all_nodes(resumed)
|
|
|
|
prepared_selected_output_id = next(iter(resumed.source_prepared_mapping["selected_output"]))
|
|
assert resumed.results[prepared_selected_output_id].prompt == "true branch"
|
|
assert set(executed_source_ids) == {"if", "selected_output"}
|
|
assert "false_value" not in executed_source_ids
|
|
|
|
|
|
def test_graph_state_prepares_eagerly():
|
|
"""Tests that all prepareable nodes are prepared"""
|
|
graph = Graph()
|
|
|
|
test_prompts = ["Banana sushi", "Cat sushi"]
|
|
graph.add_node(PromptCollectionTestInvocation(id="prompt_collection", collection=list(test_prompts)))
|
|
graph.add_node(IterateInvocation(id="iterate"))
|
|
graph.add_node(PromptTestInvocation(id="prompt_iterated"))
|
|
graph.add_edge(create_edge("prompt_collection", "collection", "iterate", "collection"))
|
|
graph.add_edge(create_edge("iterate", "item", "prompt_iterated", "prompt"))
|
|
|
|
# separated, fully-preparable chain of nodes
|
|
graph.add_node(PromptTestInvocation(id="prompt_chain_1", prompt="Dinosaur sushi"))
|
|
graph.add_node(PromptTestInvocation(id="prompt_chain_2"))
|
|
graph.add_node(PromptTestInvocation(id="prompt_chain_3"))
|
|
graph.add_edge(create_edge("prompt_chain_1", "prompt", "prompt_chain_2", "prompt"))
|
|
graph.add_edge(create_edge("prompt_chain_2", "prompt", "prompt_chain_3", "prompt"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
g.next()
|
|
|
|
assert "prompt_collection" in g.source_prepared_mapping
|
|
assert "prompt_chain_1" in g.source_prepared_mapping
|
|
assert "prompt_chain_2" in g.source_prepared_mapping
|
|
assert "prompt_chain_3" in g.source_prepared_mapping
|
|
assert "iterate" not in g.source_prepared_mapping
|
|
assert "prompt_iterated" not in g.source_prepared_mapping
|
|
|
|
|
|
def test_graph_executes_depth_first():
|
|
"""Tests that the graph executes depth-first, executing a branch as far as possible before moving to the next branch"""
|
|
|
|
def assert_topo_order_and_all_executed(state: GraphExecutionState, order: list[str]):
|
|
"""
|
|
Validates:
|
|
1) Every materialized exec node executed exactly once.
|
|
2) Execution order respects all exec-graph dependencies (u→v ⇒ u before v).
|
|
"""
|
|
# order must be EXEC node ids in run order
|
|
exec_nodes = set(state.execution_graph.nodes.keys())
|
|
|
|
# 1) coverage: all exec nodes ran, and no duplicates
|
|
pos = {nid: i for i, nid in enumerate(order)}
|
|
assert set(pos.keys()) == exec_nodes, (
|
|
f"Executed {len(pos)} of {len(exec_nodes)} nodes. Missing: {sorted(exec_nodes - set(pos))[:10]}"
|
|
)
|
|
assert len(pos) == len(order), "Duplicate execution detected"
|
|
|
|
# 2) topo order: parents before children
|
|
for e in state.execution_graph.edges:
|
|
u = e.source.node_id
|
|
v = e.destination.node_id
|
|
assert pos[u] < pos[v], f"child {v} ran before parent {u}"
|
|
|
|
graph = Graph()
|
|
|
|
test_prompts = ["Banana sushi", "Cat sushi"]
|
|
graph.add_node(PromptCollectionTestInvocation(id="prompt_collection", collection=list(test_prompts)))
|
|
graph.add_node(IterateInvocation(id="iterate"))
|
|
graph.add_node(PromptTestInvocation(id="prompt_iterated"))
|
|
graph.add_node(PromptTestInvocation(id="prompt_successor"))
|
|
graph.add_edge(create_edge("prompt_collection", "collection", "iterate", "collection"))
|
|
graph.add_edge(create_edge("iterate", "item", "prompt_iterated", "prompt"))
|
|
graph.add_edge(create_edge("prompt_iterated", "prompt", "prompt_successor", "prompt"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
order: list[str] = []
|
|
|
|
while True:
|
|
n = g.next()
|
|
if n is None:
|
|
break
|
|
o = n.invoke(Mock(InvocationContext))
|
|
g.complete(n.id, o)
|
|
order.append(n.id)
|
|
|
|
assert_topo_order_and_all_executed(g, order)
|
|
|
|
|
|
def test_graph_scheduler_drains_active_class_before_switching():
|
|
graph = Graph()
|
|
graph.add_node(PromptTestInvocation(id="prompt_a", prompt="a"))
|
|
graph.add_node(PromptTestInvocation(id="prompt_b", prompt="b"))
|
|
graph.add_node(TextToImageTestInvocation(id="image"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
g.set_ready_order([PromptTestInvocation, TextToImageTestInvocation])
|
|
|
|
first = invoke_next(g)[0]
|
|
second = invoke_next(g)[0]
|
|
third = invoke_next(g)[0]
|
|
|
|
assert first is not None
|
|
assert g.prepared_source_mapping[first.id] == "prompt_a"
|
|
assert g.prepared_source_mapping[second.id] == "prompt_b"
|
|
assert g.prepared_source_mapping[third.id] == "image"
|
|
|
|
|
|
def test_graph_scheduler_skips_stale_ready_entries():
|
|
graph = Graph()
|
|
graph.add_node(PromptTestInvocation(id="prompt_a", prompt="a"))
|
|
graph.add_node(PromptTestInvocation(id="prompt_b", prompt="b"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
g.set_ready_order([PromptTestInvocation])
|
|
|
|
first = invoke_next(g)[0]
|
|
assert first is not None
|
|
|
|
prompt_queue = g._queue_for(PromptTestInvocation.__name__)
|
|
prompt_queue.appendleft(first.id)
|
|
|
|
second = g.next()
|
|
|
|
assert second is not None
|
|
assert second.id != first.id
|
|
assert g.prepared_source_mapping[second.id] == "prompt_b"
|
|
|
|
|
|
def test_graph_scheduler_falls_back_to_non_priority_ready_classes():
|
|
graph = Graph()
|
|
graph.add_node(TextToImageTestInvocation(id="image"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
g.set_ready_order([PromptTestInvocation])
|
|
|
|
next_node = g.next()
|
|
|
|
assert next_node is not None
|
|
assert g.prepared_source_mapping[next_node.id] == "image"
|
|
|
|
|
|
# Because this tests deterministic ordering, we run it multiple times
|
|
@pytest.mark.parametrize("execution_number", range(5))
|
|
def test_graph_iterate_execution_order(execution_number: int):
|
|
"""Tests that iterate nodes execution is ordered by the order of the collection"""
|
|
|
|
graph = Graph()
|
|
|
|
test_prompts = ["Banana sushi", "Cat sushi", "Strawberry Sushi", "Dinosaur Sushi"]
|
|
graph.add_node(PromptCollectionTestInvocation(id="prompt_collection", collection=list(test_prompts)))
|
|
graph.add_node(IterateInvocation(id="iterate"))
|
|
graph.add_node(PromptTestInvocation(id="prompt_iterated"))
|
|
graph.add_edge(create_edge("prompt_collection", "collection", "iterate", "collection"))
|
|
graph.add_edge(create_edge("iterate", "item", "prompt_iterated", "prompt"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
_ = invoke_next(g)
|
|
_ = invoke_next(g)
|
|
assert _[1].item == "Banana sushi"
|
|
_ = invoke_next(g)
|
|
assert _[1].item == "Cat sushi"
|
|
_ = invoke_next(g)
|
|
assert _[1].item == "Strawberry Sushi"
|
|
_ = invoke_next(g)
|
|
assert _[1].item == "Dinosaur Sushi"
|
|
_ = invoke_next(g)
|
|
|
|
|
|
# Because this tests deterministic ordering, we run it multiple times
|
|
@pytest.mark.parametrize("execution_number", range(5))
|
|
def test_graph_nested_iterate_execution_order(execution_number: int):
|
|
"""
|
|
Validates best-effort in-order execution for nodes expanded under nested iterators.
|
|
Expected lexicographic order by (outer_index, inner_index), subject to readiness.
|
|
"""
|
|
graph = Graph()
|
|
|
|
# Outer iterator: [0, 1]
|
|
graph.add_node(RangeInvocation(id="outer_range", start=0, stop=2, step=1))
|
|
graph.add_node(IterateInvocation(id="outer_iter"))
|
|
|
|
# Inner iterator is derived from the outer item:
|
|
# start = outer_item * 10
|
|
# stop = start + 2 => yields 2 items per outer item
|
|
graph.add_node(MultiplyInvocation(id="mul10", b=10))
|
|
graph.add_node(AddInvocation(id="stop_plus2", b=2))
|
|
graph.add_node(RangeInvocation(id="inner_range", start=0, stop=1, step=1))
|
|
graph.add_node(IterateInvocation(id="inner_iter"))
|
|
|
|
# Observe inner items (they encode outer via start=outer*10)
|
|
graph.add_node(AddInvocation(id="sum", b=0))
|
|
|
|
graph.add_edge(create_edge("outer_range", "collection", "outer_iter", "collection"))
|
|
graph.add_edge(create_edge("outer_iter", "item", "mul10", "a"))
|
|
graph.add_edge(create_edge("mul10", "value", "stop_plus2", "a"))
|
|
graph.add_edge(create_edge("mul10", "value", "inner_range", "start"))
|
|
graph.add_edge(create_edge("stop_plus2", "value", "inner_range", "stop"))
|
|
graph.add_edge(create_edge("inner_range", "collection", "inner_iter", "collection"))
|
|
graph.add_edge(create_edge("inner_iter", "item", "sum", "a"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
sum_values: list[int] = []
|
|
|
|
while True:
|
|
n, o = invoke_next(g)
|
|
if n is None:
|
|
break
|
|
if g.prepared_source_mapping[n.id] == "sum":
|
|
sum_values.append(o.value)
|
|
|
|
assert sum_values == [0, 1, 10, 11]
|
|
|
|
|
|
def test_graph_collector_nested_under_outer_iterator_collects_only_current_outer_iteration_items():
|
|
graph = Graph()
|
|
|
|
graph.add_node(RangeInvocation(id="outer_range", start=0, stop=2, step=1))
|
|
graph.add_node(IterateInvocation(id="outer_iter"))
|
|
graph.add_node(IntegerCollectionFromItemTestInvocation(id="inner_collection"))
|
|
graph.add_node(IterateInvocation(id="inner_iter"))
|
|
graph.add_node(AddInvocation(id="inner_item", b=0))
|
|
graph.add_node(CollectInvocation(id="collect"))
|
|
graph.add_node(IntegerCollectionPassthroughTestInvocation(id="per_outer_consumer"))
|
|
|
|
graph.add_edge(create_edge("outer_range", "collection", "outer_iter", "collection"))
|
|
graph.add_edge(create_edge("outer_iter", "item", "inner_collection", "value"))
|
|
graph.add_edge(create_edge("inner_collection", "collection", "inner_iter", "collection"))
|
|
graph.add_edge(create_edge("inner_iter", "item", "inner_item", "a"))
|
|
graph.add_edge(create_edge("inner_item", "value", "collect", "item"))
|
|
graph.add_edge(create_edge("collect", "collection", "per_outer_consumer", "collection"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
execute_all_nodes(g)
|
|
|
|
prepared_consumer_ids = g.source_prepared_mapping["per_outer_consumer"]
|
|
consumer_collections = sorted(g.results[node_id].collection for node_id in prepared_consumer_ids)
|
|
|
|
assert consumer_collections == [[0, 1], [10, 11]]
|
|
|
|
|
|
def test_graph_collector_reuses_outer_collection_input_for_each_nested_iterator_group():
|
|
graph = Graph()
|
|
|
|
graph.add_node(RangeInvocation(id="base_collection", start=100, stop=101, step=1))
|
|
graph.add_node(RangeInvocation(id="outer_range", start=0, stop=2, step=1))
|
|
graph.add_node(IterateInvocation(id="outer_iter"))
|
|
graph.add_node(IntegerCollectionFromItemTestInvocation(id="inner_collection"))
|
|
graph.add_node(IterateInvocation(id="inner_iter"))
|
|
graph.add_node(AddInvocation(id="inner_item", b=0))
|
|
graph.add_node(CollectInvocation(id="collect"))
|
|
|
|
graph.add_edge(create_edge("base_collection", "collection", "collect", "collection"))
|
|
graph.add_edge(create_edge("outer_range", "collection", "outer_iter", "collection"))
|
|
graph.add_edge(create_edge("outer_iter", "item", "inner_collection", "value"))
|
|
graph.add_edge(create_edge("inner_collection", "collection", "inner_iter", "collection"))
|
|
graph.add_edge(create_edge("inner_iter", "item", "inner_item", "a"))
|
|
graph.add_edge(create_edge("inner_item", "value", "collect", "item"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
execute_all_nodes(g)
|
|
|
|
prepared_collect_ids = g.source_prepared_mapping["collect"]
|
|
collect_results = sorted(g.results[node_id].collection for node_id in prepared_collect_ids)
|
|
|
|
assert collect_results == [[100, 0, 1], [100, 10, 11]]
|
|
|
|
|
|
def test_graph_collector_nested_under_three_iterators_preserves_outer_iteration_paths():
|
|
graph = Graph()
|
|
|
|
graph.add_node(RangeInvocation(id="outer_range", start=0, stop=2, step=1))
|
|
graph.add_node(IterateInvocation(id="outer_iter"))
|
|
graph.add_node(IntegerCollectionFromItemTestInvocation(id="middle_collection"))
|
|
graph.add_node(IterateInvocation(id="middle_iter"))
|
|
graph.add_node(IntegerCollectionFromItemTestInvocation(id="inner_collection"))
|
|
graph.add_node(IterateInvocation(id="inner_iter"))
|
|
graph.add_node(AddInvocation(id="inner_item", b=0))
|
|
graph.add_node(CollectInvocation(id="collect"))
|
|
graph.add_node(IntegerCollectionPassthroughTestInvocation(id="per_middle_consumer"))
|
|
|
|
graph.add_edge(create_edge("outer_range", "collection", "outer_iter", "collection"))
|
|
graph.add_edge(create_edge("outer_iter", "item", "middle_collection", "value"))
|
|
graph.add_edge(create_edge("middle_collection", "collection", "middle_iter", "collection"))
|
|
graph.add_edge(create_edge("middle_iter", "item", "inner_collection", "value"))
|
|
graph.add_edge(create_edge("inner_collection", "collection", "inner_iter", "collection"))
|
|
graph.add_edge(create_edge("inner_iter", "item", "inner_item", "a"))
|
|
graph.add_edge(create_edge("inner_item", "value", "collect", "item"))
|
|
graph.add_edge(create_edge("collect", "collection", "per_middle_consumer", "collection"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
execute_all_nodes(g)
|
|
|
|
prepared_consumer_ids = g.source_prepared_mapping["per_middle_consumer"]
|
|
consumer_collections = sorted(g.results[node_id].collection for node_id in prepared_consumer_ids)
|
|
|
|
assert consumer_collections == [[0, 1], [10, 11], [100, 101], [110, 111]]
|
|
|
|
|
|
def test_graph_collector_with_mixed_depth_item_inputs_keeps_outer_iterations_separate():
|
|
graph = Graph()
|
|
|
|
graph.add_node(RangeInvocation(id="outer_range", start=0, stop=2, step=1))
|
|
graph.add_node(IterateInvocation(id="outer_iter"))
|
|
graph.add_node(AddInvocation(id="outer_item", b=100))
|
|
graph.add_node(IntegerCollectionFromItemTestInvocation(id="inner_collection"))
|
|
graph.add_node(IterateInvocation(id="inner_iter"))
|
|
graph.add_node(AddInvocation(id="inner_item", b=0))
|
|
graph.add_node(CollectInvocation(id="collect"))
|
|
|
|
graph.add_edge(create_edge("outer_range", "collection", "outer_iter", "collection"))
|
|
graph.add_edge(create_edge("outer_iter", "item", "outer_item", "a"))
|
|
graph.add_edge(create_edge("outer_item", "value", "collect", "item"))
|
|
graph.add_edge(create_edge("outer_iter", "item", "inner_collection", "value"))
|
|
graph.add_edge(create_edge("inner_collection", "collection", "inner_iter", "collection"))
|
|
graph.add_edge(create_edge("inner_iter", "item", "inner_item", "a"))
|
|
graph.add_edge(create_edge("inner_item", "value", "collect", "item"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
execute_all_nodes(g)
|
|
|
|
collect_results = sorted(g.results[node_id].collection for node_id in g.source_prepared_mapping["collect"])
|
|
|
|
assert collect_results == [[100, 0, 1], [101, 10, 11]]
|
|
|
|
|
|
def test_graph_consumer_matches_collector_parents_at_different_iteration_depths():
|
|
graph = Graph()
|
|
|
|
graph.add_node(RangeInvocation(id="outer_range", start=0, stop=2, step=1))
|
|
graph.add_node(IterateInvocation(id="outer_iter"))
|
|
graph.add_node(IntegerCollectionFromItemTestInvocation(id="middle_collection"))
|
|
graph.add_node(IterateInvocation(id="middle_iter"))
|
|
graph.add_node(AddInvocation(id="shallow_item", b=0))
|
|
graph.add_node(CollectInvocation(id="shallow_collect"))
|
|
graph.add_node(IntegerCollectionFromItemTestInvocation(id="inner_collection"))
|
|
graph.add_node(IterateInvocation(id="inner_iter"))
|
|
graph.add_node(AddInvocation(id="deep_item", b=0))
|
|
graph.add_node(CollectInvocation(id="deep_collect"))
|
|
graph.add_node(TwoIntegerCollectionsTestInvocation(id="consumer"))
|
|
|
|
graph.add_edge(create_edge("outer_range", "collection", "outer_iter", "collection"))
|
|
graph.add_edge(create_edge("outer_iter", "item", "middle_collection", "value"))
|
|
graph.add_edge(create_edge("middle_collection", "collection", "middle_iter", "collection"))
|
|
graph.add_edge(create_edge("middle_iter", "item", "shallow_item", "a"))
|
|
graph.add_edge(create_edge("shallow_item", "value", "shallow_collect", "item"))
|
|
graph.add_edge(create_edge("middle_iter", "item", "inner_collection", "value"))
|
|
graph.add_edge(create_edge("inner_collection", "collection", "inner_iter", "collection"))
|
|
graph.add_edge(create_edge("inner_iter", "item", "deep_item", "a"))
|
|
graph.add_edge(create_edge("deep_item", "value", "deep_collect", "item"))
|
|
graph.add_edge(create_edge("shallow_collect", "collection", "consumer", "first"))
|
|
graph.add_edge(create_edge("deep_collect", "collection", "consumer", "second"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
execute_all_nodes(g)
|
|
|
|
consumer_results = sorted(g.results[node_id].collection for node_id in g.source_prepared_mapping["consumer"])
|
|
|
|
assert consumer_results == [
|
|
[0, 1, 0, 1],
|
|
[0, 1, 10, 11],
|
|
[10, 11, 100, 101],
|
|
[10, 11, 110, 111],
|
|
]
|
|
|
|
|
|
def test_graph_consumer_reuses_global_parent_for_each_nested_collector_iteration():
|
|
graph = Graph()
|
|
|
|
graph.add_node(RangeInvocation(id="global_collection", start=99, stop=100, step=1))
|
|
graph.add_node(RangeInvocation(id="outer_range", start=0, stop=2, step=1))
|
|
graph.add_node(IterateInvocation(id="outer_iter"))
|
|
graph.add_node(IntegerCollectionFromItemTestInvocation(id="inner_collection"))
|
|
graph.add_node(IterateInvocation(id="inner_iter"))
|
|
graph.add_node(AddInvocation(id="inner_item", b=0))
|
|
graph.add_node(CollectInvocation(id="collect"))
|
|
graph.add_node(TwoIntegerCollectionsTestInvocation(id="consumer"))
|
|
|
|
graph.add_edge(create_edge("global_collection", "collection", "consumer", "second"))
|
|
graph.add_edge(create_edge("outer_range", "collection", "outer_iter", "collection"))
|
|
graph.add_edge(create_edge("outer_iter", "item", "inner_collection", "value"))
|
|
graph.add_edge(create_edge("inner_collection", "collection", "inner_iter", "collection"))
|
|
graph.add_edge(create_edge("inner_iter", "item", "inner_item", "a"))
|
|
graph.add_edge(create_edge("inner_item", "value", "collect", "item"))
|
|
graph.add_edge(create_edge("collect", "collection", "consumer", "first"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
execute_all_nodes(g)
|
|
|
|
consumer_results = sorted(g.results[node_id].collection for node_id in g.source_prepared_mapping["consumer"])
|
|
|
|
assert consumer_results == [[0, 1, 99], [10, 11, 99]]
|
|
|
|
|
|
def test_graph_validate_self_iterator_without_collection_input_raises_invalid_edge_error():
|
|
"""Iterator nodes with no collection input should fail validation cleanly.
|
|
|
|
This test exposes the bug where validation crashes with IndexError instead of raising InvalidEdgeError.
|
|
"""
|
|
from invokeai.app.services.shared.graph import InvalidEdgeError
|
|
|
|
graph = Graph()
|
|
graph.add_node(IterateInvocation(id="iterate"))
|
|
|
|
with pytest.raises(InvalidEdgeError):
|
|
graph.validate_self()
|
|
|
|
|
|
def test_graph_validate_self_collector_without_item_inputs_raises_invalid_edge_error():
|
|
"""Collector nodes with no item inputs should fail validation cleanly.
|
|
|
|
This test exposes the bug where validation can crash (e.g. StopIteration) instead of raising InvalidEdgeError.
|
|
"""
|
|
from invokeai.app.services.shared.graph import InvalidEdgeError
|
|
|
|
graph = Graph()
|
|
graph.add_node(CollectInvocation(id="collect"))
|
|
|
|
with pytest.raises(InvalidEdgeError):
|
|
graph.validate_self()
|
|
|
|
|
|
def test_if_invocation_selects_true_input_value():
|
|
invocation = IfInvocation(id="if", condition=True, true_input="true", false_input="false")
|
|
|
|
output = invocation.invoke(Mock(InvocationContext))
|
|
|
|
assert output.value == "true"
|
|
|
|
|
|
def test_if_invocation_outputs_none_when_selected_input_is_missing():
|
|
invocation = IfInvocation(id="if", condition=False, true_input="true")
|
|
|
|
output = invocation.invoke(Mock(InvocationContext))
|
|
|
|
assert output.value is None
|
|
|
|
|
|
def test_if_invocation_output_allows_missing_value_on_deserialization():
|
|
output = IfInvocationOutput.model_validate({"type": "if_output"})
|
|
|
|
assert output.value is None
|
|
|
|
|
|
def test_if_invocation_output_connects_to_downstream_input():
|
|
graph = Graph()
|
|
graph.add_node(IfInvocation(id="if", condition=True, true_input="connected value", false_input="unused"))
|
|
graph.add_node(PromptTestInvocation(id="prompt"))
|
|
graph.add_edge(create_edge("if", "value", "prompt", "prompt"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
while not g.is_complete():
|
|
invoke_next(g)
|
|
|
|
prepared_prompt_nodes = g.source_prepared_mapping["prompt"]
|
|
assert len(prepared_prompt_nodes) == 1
|
|
prepared_prompt_node_id = next(iter(prepared_prompt_nodes))
|
|
assert g.results[prepared_prompt_node_id].prompt == "connected value"
|
|
|
|
|
|
@pytest.mark.xfail(strict=True, reason="Legacy eager If-node execution should no longer occur")
|
|
def test_if_graph_current_behavior_executes_both_branches_and_shared_ancestors():
|
|
graph = Graph()
|
|
graph.add_node(BooleanInvocation(id="condition", value=True))
|
|
graph.add_node(PromptTestInvocation(id="shared", prompt="shared value"))
|
|
graph.add_node(PromptTestInvocation(id="true_mid"))
|
|
graph.add_node(PromptTestInvocation(id="true_leaf"))
|
|
graph.add_node(PromptTestInvocation(id="false_mid"))
|
|
graph.add_node(PromptTestInvocation(id="false_leaf"))
|
|
graph.add_node(PromptTestInvocation(id="side_consumer"))
|
|
graph.add_node(IfInvocation(id="if"))
|
|
graph.add_node(PromptTestInvocation(id="selected_output"))
|
|
|
|
graph.add_edge(create_edge("condition", "value", "if", "condition"))
|
|
graph.add_edge(create_edge("shared", "prompt", "true_mid", "prompt"))
|
|
graph.add_edge(create_edge("true_mid", "prompt", "true_leaf", "prompt"))
|
|
graph.add_edge(create_edge("true_leaf", "prompt", "if", "true_input"))
|
|
graph.add_edge(create_edge("shared", "prompt", "false_mid", "prompt"))
|
|
graph.add_edge(create_edge("false_mid", "prompt", "false_leaf", "prompt"))
|
|
graph.add_edge(create_edge("false_leaf", "prompt", "if", "false_input"))
|
|
graph.add_edge(create_edge("shared", "prompt", "side_consumer", "prompt"))
|
|
graph.add_edge(create_edge("if", "value", "selected_output", "prompt"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
executed_source_ids = execute_all_nodes(g)
|
|
|
|
assert set(executed_source_ids) == {
|
|
"condition",
|
|
"shared",
|
|
"true_mid",
|
|
"true_leaf",
|
|
"false_mid",
|
|
"false_leaf",
|
|
"side_consumer",
|
|
"if",
|
|
"selected_output",
|
|
}
|
|
assert executed_source_ids.count("false_mid") == 1
|
|
assert executed_source_ids.count("false_leaf") == 1
|
|
|
|
prepared_selected_output_id = next(iter(g.source_prepared_mapping["selected_output"]))
|
|
assert g.results[prepared_selected_output_id].prompt == "shared value"
|
|
|
|
|
|
@pytest.mark.xfail(strict=True, reason="Legacy eager If-node execution should no longer occur")
|
|
def test_if_graph_current_behavior_executes_both_simple_branches():
|
|
graph = Graph()
|
|
graph.add_node(BooleanInvocation(id="condition", value=True))
|
|
graph.add_node(PromptTestInvocation(id="true_value", prompt="true branch"))
|
|
graph.add_node(PromptTestInvocation(id="false_value", prompt="false branch"))
|
|
graph.add_node(IfInvocation(id="if"))
|
|
graph.add_node(PromptTestInvocation(id="selected_output"))
|
|
|
|
graph.add_edge(create_edge("condition", "value", "if", "condition"))
|
|
graph.add_edge(create_edge("true_value", "prompt", "if", "true_input"))
|
|
graph.add_edge(create_edge("false_value", "prompt", "if", "false_input"))
|
|
graph.add_edge(create_edge("if", "value", "selected_output", "prompt"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
executed_source_ids = execute_all_nodes(g)
|
|
|
|
assert set(executed_source_ids) == {"condition", "true_value", "false_value", "if", "selected_output"}
|
|
prepared_selected_output_id = next(iter(g.source_prepared_mapping["selected_output"]))
|
|
assert g.results[prepared_selected_output_id].prompt == "true branch"
|
|
|
|
|
|
def test_if_graph_optimized_behavior_executes_only_selected_simple_branch():
|
|
graph = Graph()
|
|
graph.add_node(BooleanInvocation(id="condition", value=True))
|
|
graph.add_node(PromptTestInvocation(id="true_value", prompt="true branch"))
|
|
graph.add_node(PromptTestInvocation(id="false_value", prompt="false branch"))
|
|
graph.add_node(IfInvocation(id="if"))
|
|
graph.add_node(PromptTestInvocation(id="selected_output"))
|
|
|
|
graph.add_edge(create_edge("condition", "value", "if", "condition"))
|
|
graph.add_edge(create_edge("true_value", "prompt", "if", "true_input"))
|
|
graph.add_edge(create_edge("false_value", "prompt", "if", "false_input"))
|
|
graph.add_edge(create_edge("if", "value", "selected_output", "prompt"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
executed_source_ids = execute_all_nodes(g)
|
|
|
|
assert set(executed_source_ids) == {"condition", "true_value", "if", "selected_output"}
|
|
assert "false_value" not in executed_source_ids
|
|
|
|
|
|
def test_if_graph_optimized_behavior_records_skipped_branch_in_execution_history():
|
|
graph = Graph()
|
|
graph.add_node(BooleanInvocation(id="condition", value=True))
|
|
graph.add_node(PromptTestInvocation(id="true_value", prompt="true branch"))
|
|
graph.add_node(PromptTestInvocation(id="false_value", prompt="false branch"))
|
|
graph.add_node(IfInvocation(id="if"))
|
|
graph.add_node(PromptTestInvocation(id="selected_output"))
|
|
|
|
graph.add_edge(create_edge("condition", "value", "if", "condition"))
|
|
graph.add_edge(create_edge("true_value", "prompt", "if", "true_input"))
|
|
graph.add_edge(create_edge("false_value", "prompt", "if", "false_input"))
|
|
graph.add_edge(create_edge("if", "value", "selected_output", "prompt"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
execute_all_nodes(g)
|
|
|
|
assert set(g.executed_history) == {"condition", "true_value", "false_value", "if", "selected_output"}
|
|
assert g.executed_history.count("false_value") == 1
|
|
|
|
|
|
def test_if_graph_optimized_behavior_skips_unselected_branch_but_keeps_shared_ancestors():
|
|
graph = Graph()
|
|
graph.add_node(BooleanInvocation(id="condition", value=True))
|
|
graph.add_node(PromptTestInvocation(id="shared", prompt="shared value"))
|
|
graph.add_node(PromptTestInvocation(id="true_mid"))
|
|
graph.add_node(PromptTestInvocation(id="true_leaf"))
|
|
graph.add_node(PromptTestInvocation(id="false_mid"))
|
|
graph.add_node(PromptTestInvocation(id="false_leaf"))
|
|
graph.add_node(PromptTestInvocation(id="side_consumer"))
|
|
graph.add_node(IfInvocation(id="if"))
|
|
graph.add_node(PromptTestInvocation(id="selected_output"))
|
|
|
|
graph.add_edge(create_edge("condition", "value", "if", "condition"))
|
|
graph.add_edge(create_edge("shared", "prompt", "true_mid", "prompt"))
|
|
graph.add_edge(create_edge("true_mid", "prompt", "true_leaf", "prompt"))
|
|
graph.add_edge(create_edge("true_leaf", "prompt", "if", "true_input"))
|
|
graph.add_edge(create_edge("shared", "prompt", "false_mid", "prompt"))
|
|
graph.add_edge(create_edge("false_mid", "prompt", "false_leaf", "prompt"))
|
|
graph.add_edge(create_edge("false_leaf", "prompt", "if", "false_input"))
|
|
graph.add_edge(create_edge("shared", "prompt", "side_consumer", "prompt"))
|
|
graph.add_edge(create_edge("if", "value", "selected_output", "prompt"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
executed_source_ids = execute_all_nodes(g)
|
|
|
|
assert set(executed_source_ids) == {
|
|
"condition",
|
|
"shared",
|
|
"true_mid",
|
|
"true_leaf",
|
|
"side_consumer",
|
|
"if",
|
|
"selected_output",
|
|
}
|
|
assert "false_mid" not in executed_source_ids
|
|
assert "false_leaf" not in executed_source_ids
|
|
|
|
|
|
def test_if_graph_optimized_behavior_skips_distant_unselected_ancestors_only_when_exclusive():
|
|
graph = Graph()
|
|
graph.add_node(BooleanInvocation(id="condition", value=False))
|
|
graph.add_node(PromptTestInvocation(id="shared_root", prompt="shared value"))
|
|
graph.add_node(PromptTestInvocation(id="true_shared_mid"))
|
|
graph.add_node(PromptTestInvocation(id="true_exclusive_leaf"))
|
|
graph.add_node(PromptTestInvocation(id="false_mid"))
|
|
graph.add_node(PromptTestInvocation(id="false_leaf"))
|
|
graph.add_node(PromptTestInvocation(id="shared_observer"))
|
|
graph.add_node(IfInvocation(id="if"))
|
|
graph.add_node(PromptTestInvocation(id="selected_output"))
|
|
|
|
graph.add_edge(create_edge("condition", "value", "if", "condition"))
|
|
graph.add_edge(create_edge("shared_root", "prompt", "true_shared_mid", "prompt"))
|
|
graph.add_edge(create_edge("true_shared_mid", "prompt", "true_exclusive_leaf", "prompt"))
|
|
graph.add_edge(create_edge("true_exclusive_leaf", "prompt", "if", "true_input"))
|
|
graph.add_edge(create_edge("shared_root", "prompt", "false_mid", "prompt"))
|
|
graph.add_edge(create_edge("false_mid", "prompt", "false_leaf", "prompt"))
|
|
graph.add_edge(create_edge("false_leaf", "prompt", "if", "false_input"))
|
|
graph.add_edge(create_edge("true_shared_mid", "prompt", "shared_observer", "prompt"))
|
|
graph.add_edge(create_edge("if", "value", "selected_output", "prompt"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
executed_source_ids = execute_all_nodes(g)
|
|
|
|
assert set(executed_source_ids) == {
|
|
"condition",
|
|
"shared_root",
|
|
"true_shared_mid",
|
|
"false_mid",
|
|
"false_leaf",
|
|
"shared_observer",
|
|
"if",
|
|
"selected_output",
|
|
}
|
|
assert "true_exclusive_leaf" not in executed_source_ids
|
|
|
|
|
|
def test_if_graph_optimized_behavior_allows_selected_missing_branch_input():
|
|
graph = Graph()
|
|
graph.add_node(BooleanInvocation(id="condition", value=False))
|
|
graph.add_node(PromptTestInvocation(id="true_value", prompt="true branch"))
|
|
graph.add_node(IfInvocation(id="if"))
|
|
graph.add_node(AnyTypeTestInvocation(id="selected_output"))
|
|
|
|
graph.add_edge(create_edge("condition", "value", "if", "condition"))
|
|
graph.add_edge(create_edge("true_value", "prompt", "if", "true_input"))
|
|
graph.add_edge(create_edge("if", "value", "selected_output", "value"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
executed_source_ids = execute_all_nodes(g)
|
|
|
|
prepared_selected_output_id = next(iter(g.source_prepared_mapping["selected_output"]))
|
|
assert g.results[prepared_selected_output_id].value is None
|
|
assert set(executed_source_ids) == {"condition", "if", "selected_output"}
|
|
assert "true_value" not in executed_source_ids
|
|
|
|
|
|
def test_if_graph_optimized_behavior_does_not_cross_defer_independent_ifs():
|
|
graph = Graph()
|
|
graph.add_node(BooleanInvocation(id="condition_a", value=True))
|
|
graph.add_node(BooleanInvocation(id="condition_b", value=False))
|
|
graph.add_node(PromptTestInvocation(id="true_a", prompt="true a"))
|
|
graph.add_node(PromptTestInvocation(id="false_a", prompt="false a"))
|
|
graph.add_node(PromptTestInvocation(id="true_b", prompt="true b"))
|
|
graph.add_node(PromptTestInvocation(id="false_b", prompt="false b"))
|
|
graph.add_node(IfInvocation(id="if_a"))
|
|
graph.add_node(IfInvocation(id="if_b"))
|
|
graph.add_node(CollectInvocation(id="collect"))
|
|
|
|
graph.add_edge(create_edge("condition_a", "value", "if_a", "condition"))
|
|
graph.add_edge(create_edge("true_a", "prompt", "if_a", "true_input"))
|
|
graph.add_edge(create_edge("false_a", "prompt", "if_a", "false_input"))
|
|
graph.add_edge(create_edge("condition_b", "value", "if_b", "condition"))
|
|
graph.add_edge(create_edge("true_b", "prompt", "if_b", "true_input"))
|
|
graph.add_edge(create_edge("false_b", "prompt", "if_b", "false_input"))
|
|
graph.add_edge(create_edge("if_a", "value", "collect", "item"))
|
|
graph.add_edge(create_edge("if_b", "value", "collect", "item"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
executed_source_ids = execute_all_nodes(g)
|
|
|
|
prepared_collect_id = next(iter(g.source_prepared_mapping["collect"]))
|
|
assert sorted(g.results[prepared_collect_id].collection) == ["false b", "true a"]
|
|
assert set(executed_source_ids) == {
|
|
"condition_a",
|
|
"condition_b",
|
|
"true_a",
|
|
"false_b",
|
|
"if_a",
|
|
"if_b",
|
|
"collect",
|
|
}
|
|
assert "false_a" not in executed_source_ids
|
|
assert "true_b" not in executed_source_ids
|
|
|
|
|
|
def test_if_graph_optimized_behavior_supports_nested_ifs():
|
|
graph = Graph()
|
|
graph.add_node(BooleanInvocation(id="outer_condition", value=True))
|
|
graph.add_node(BooleanInvocation(id="inner_condition", value=False))
|
|
graph.add_node(PromptTestInvocation(id="outer_false", prompt="outer false"))
|
|
graph.add_node(PromptTestInvocation(id="inner_true", prompt="inner true"))
|
|
graph.add_node(PromptTestInvocation(id="inner_false", prompt="inner false"))
|
|
graph.add_node(IfInvocation(id="inner_if"))
|
|
graph.add_node(IfInvocation(id="outer_if"))
|
|
graph.add_node(PromptTestInvocation(id="selected_output"))
|
|
|
|
graph.add_edge(create_edge("inner_condition", "value", "inner_if", "condition"))
|
|
graph.add_edge(create_edge("inner_true", "prompt", "inner_if", "true_input"))
|
|
graph.add_edge(create_edge("inner_false", "prompt", "inner_if", "false_input"))
|
|
graph.add_edge(create_edge("outer_condition", "value", "outer_if", "condition"))
|
|
graph.add_edge(create_edge("inner_if", "value", "outer_if", "true_input"))
|
|
graph.add_edge(create_edge("outer_false", "prompt", "outer_if", "false_input"))
|
|
graph.add_edge(create_edge("outer_if", "value", "selected_output", "prompt"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
executed_source_ids = execute_all_nodes(g)
|
|
|
|
prepared_selected_output_id = next(iter(g.source_prepared_mapping["selected_output"]))
|
|
assert g.results[prepared_selected_output_id].prompt == "inner false"
|
|
assert set(executed_source_ids) == {
|
|
"outer_condition",
|
|
"inner_condition",
|
|
"inner_false",
|
|
"inner_if",
|
|
"outer_if",
|
|
"selected_output",
|
|
}
|
|
assert "inner_true" not in executed_source_ids
|
|
assert "outer_false" not in executed_source_ids
|
|
|
|
|
|
def test_if_graph_optimized_behavior_prunes_branches_per_iteration():
|
|
graph = Graph()
|
|
graph.add_node(BooleanCollectionInvocation(id="conditions", collection=[True, False, True]))
|
|
graph.add_node(IterateInvocation(id="condition_iter"))
|
|
graph.add_node(AnyTypeTestInvocation(id="true_branch"))
|
|
graph.add_node(AnyTypeTestInvocation(id="false_branch"))
|
|
graph.add_node(IfInvocation(id="if"))
|
|
graph.add_node(CollectInvocation(id="collect"))
|
|
|
|
graph.add_edge(create_edge("conditions", "collection", "condition_iter", "collection"))
|
|
graph.add_edge(create_edge("condition_iter", "item", "if", "condition"))
|
|
graph.add_edge(create_edge("condition_iter", "item", "true_branch", "value"))
|
|
graph.add_edge(create_edge("true_branch", "value", "if", "true_input"))
|
|
graph.add_edge(create_edge("condition_iter", "item", "false_branch", "value"))
|
|
graph.add_edge(create_edge("false_branch", "value", "if", "false_input"))
|
|
graph.add_edge(create_edge("if", "value", "collect", "item"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
executed_source_ids = execute_all_nodes(g)
|
|
|
|
prepared_collect_id = next(iter(g.source_prepared_mapping["collect"]))
|
|
assert g.results[prepared_collect_id].collection == [True, False, True]
|
|
assert executed_source_ids.count("condition_iter") == 3
|
|
assert executed_source_ids.count("true_branch") == 2
|
|
assert executed_source_ids.count("false_branch") == 1
|
|
assert executed_source_ids.count("if") == 3
|
|
|
|
|
|
def test_if_graph_optimized_behavior_keeps_shared_live_consumers_per_iteration():
|
|
graph = Graph()
|
|
graph.add_node(BooleanCollectionInvocation(id="conditions", collection=[True, False, False]))
|
|
graph.add_node(IterateInvocation(id="condition_iter"))
|
|
graph.add_node(AnyTypeTestInvocation(id="shared_branch"))
|
|
graph.add_node(AnyTypeTestInvocation(id="true_leaf"))
|
|
graph.add_node(AnyTypeTestInvocation(id="false_branch"))
|
|
graph.add_node(AnyTypeTestInvocation(id="observer"))
|
|
graph.add_node(IfInvocation(id="if"))
|
|
graph.add_node(CollectInvocation(id="selected_collect"))
|
|
graph.add_node(CollectInvocation(id="observer_collect"))
|
|
|
|
graph.add_edge(create_edge("conditions", "collection", "condition_iter", "collection"))
|
|
graph.add_edge(create_edge("condition_iter", "item", "if", "condition"))
|
|
graph.add_edge(create_edge("condition_iter", "item", "shared_branch", "value"))
|
|
graph.add_edge(create_edge("shared_branch", "value", "true_leaf", "value"))
|
|
graph.add_edge(create_edge("true_leaf", "value", "if", "true_input"))
|
|
graph.add_edge(create_edge("condition_iter", "item", "false_branch", "value"))
|
|
graph.add_edge(create_edge("false_branch", "value", "if", "false_input"))
|
|
graph.add_edge(create_edge("shared_branch", "value", "observer", "value"))
|
|
graph.add_edge(create_edge("if", "value", "selected_collect", "item"))
|
|
graph.add_edge(create_edge("observer", "value", "observer_collect", "item"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
executed_source_ids = execute_all_nodes(g)
|
|
|
|
prepared_selected_collect_id = next(iter(g.source_prepared_mapping["selected_collect"]))
|
|
assert g.results[prepared_selected_collect_id].collection == [True, False, False]
|
|
prepared_observer_collect_id = next(iter(g.source_prepared_mapping["observer_collect"]))
|
|
assert g.results[prepared_observer_collect_id].collection == [True, False, False]
|
|
|
|
assert executed_source_ids.count("condition_iter") == 3
|
|
assert executed_source_ids.count("shared_branch") == 3
|
|
assert executed_source_ids.count("observer") == 3
|
|
assert executed_source_ids.count("true_leaf") == 1
|
|
assert executed_source_ids.count("false_branch") == 2
|
|
|
|
|
|
def test_if_graph_optimized_behavior_handles_selected_true_branch_with_shared_false_input_ancestor():
|
|
graph = Graph()
|
|
graph.add_node(BooleanInvocation(id="condition", value=True))
|
|
graph.add_node(AnyTypeTestInvocation(id="shared_item", value="shared"))
|
|
graph.add_node(AnyTypeTestInvocation(id="true_item", value="true"))
|
|
graph.add_node(CollectInvocation(id="shared_collect"))
|
|
graph.add_node(CollectInvocation(id="true_collect"))
|
|
graph.add_node(IfInvocation(id="if"))
|
|
graph.add_node(AnyTypeTestInvocation(id="selected_output"))
|
|
|
|
graph.add_edge(create_edge("condition", "value", "if", "condition"))
|
|
graph.add_edge(create_edge("shared_item", "value", "shared_collect", "item"))
|
|
graph.add_edge(create_edge("shared_collect", "collection", "true_collect", "collection"))
|
|
graph.add_edge(create_edge("true_item", "value", "true_collect", "item"))
|
|
graph.add_edge(create_edge("shared_collect", "collection", "if", "false_input"))
|
|
graph.add_edge(create_edge("true_collect", "collection", "if", "true_input"))
|
|
graph.add_edge(create_edge("if", "value", "selected_output", "value"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
executed_source_ids = execute_all_nodes(g)
|
|
|
|
prepared_selected_output_id = next(iter(g.source_prepared_mapping["selected_output"]))
|
|
assert g.results[prepared_selected_output_id].value == ["shared", "true"]
|
|
assert set(executed_source_ids) == {
|
|
"condition",
|
|
"shared_item",
|
|
"true_item",
|
|
"shared_collect",
|
|
"true_collect",
|
|
"if",
|
|
"selected_output",
|
|
}
|
|
|
|
|
|
def test_if_graph_optimized_behavior_handles_selected_false_branch_with_shared_true_input_ancestor():
|
|
graph = Graph()
|
|
graph.add_node(BooleanInvocation(id="condition", value=False))
|
|
graph.add_node(AnyTypeTestInvocation(id="shared_item", value="shared"))
|
|
graph.add_node(AnyTypeTestInvocation(id="true_item", value="true"))
|
|
graph.add_node(CollectInvocation(id="shared_collect"))
|
|
graph.add_node(CollectInvocation(id="true_collect"))
|
|
graph.add_node(IfInvocation(id="if"))
|
|
graph.add_node(AnyTypeTestInvocation(id="selected_output"))
|
|
|
|
graph.add_edge(create_edge("condition", "value", "if", "condition"))
|
|
graph.add_edge(create_edge("shared_item", "value", "shared_collect", "item"))
|
|
graph.add_edge(create_edge("shared_collect", "collection", "true_collect", "collection"))
|
|
graph.add_edge(create_edge("true_item", "value", "true_collect", "item"))
|
|
graph.add_edge(create_edge("shared_collect", "collection", "if", "false_input"))
|
|
graph.add_edge(create_edge("true_collect", "collection", "if", "true_input"))
|
|
graph.add_edge(create_edge("if", "value", "selected_output", "value"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
executed_source_ids = execute_all_nodes(g)
|
|
|
|
prepared_selected_output_id = next(iter(g.source_prepared_mapping["selected_output"]))
|
|
assert g.results[prepared_selected_output_id].value == ["shared"]
|
|
assert set(executed_source_ids) == {
|
|
"condition",
|
|
"shared_item",
|
|
"shared_collect",
|
|
"if",
|
|
"selected_output",
|
|
}
|
|
assert "true_item" not in executed_source_ids
|
|
assert "true_collect" not in executed_source_ids
|
|
|
|
|
|
def test_prepare_if_inputs_raises_when_selected_branch_source_has_no_result():
|
|
graph = Graph()
|
|
graph.add_node(BooleanInvocation(id="condition", value=True))
|
|
graph.add_node(PromptTestInvocation(id="true_value", prompt="true branch"))
|
|
graph.add_node(IfInvocation(id="if"))
|
|
|
|
graph.add_edge(create_edge("condition", "value", "if", "condition"))
|
|
graph.add_edge(create_edge("true_value", "prompt", "if", "true_input"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
|
|
condition_exec_id = g._create_execution_node("condition", [])[0]
|
|
true_value_exec_id = g._create_execution_node("true_value", [])[0]
|
|
if_exec_id = g._create_execution_node(
|
|
"if",
|
|
[("condition", condition_exec_id), ("true_value", true_value_exec_id)],
|
|
)[0]
|
|
|
|
g.executed.add(condition_exec_id)
|
|
g.results[condition_exec_id] = BooleanOutput(value=True)
|
|
g.executed.add(true_value_exec_id)
|
|
g._resolved_if_exec_branches[if_exec_id] = "true_input"
|
|
|
|
if_node = g.execution_graph.get_node(if_exec_id)
|
|
with pytest.raises(RuntimeError) as exc_info:
|
|
g._prepare_inputs(if_node)
|
|
|
|
message = str(exc_info.value)
|
|
assert if_exec_id in message
|
|
assert true_value_exec_id in message
|
|
assert "iteration_path=()" in message
|
|
|
|
|
|
def test_get_collect_iteration_mapping_groups_ignores_skipped_prepared_exec_nodes():
|
|
graph = Graph()
|
|
graph.add_node(AnyTypeTestInvocation(id="parent", value="value"))
|
|
graph.add_node(CollectInvocation(id="collect"))
|
|
graph.add_edge(create_edge("parent", "value", "collect", "item"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
|
|
skipped_exec_id = g._create_execution_node("parent", [])[0]
|
|
active_exec_id = g._create_execution_node("parent", [])[0]
|
|
g._set_prepared_exec_state(skipped_exec_id, "skipped")
|
|
|
|
mappings = g._materializer()._get_collect_iteration_mapping_groups(graph._get_input_edges("collect"))
|
|
|
|
assert mappings == [((), [("parent", active_exec_id)])]
|
|
|
|
|
|
def test_get_iteration_node_ignores_skipped_prepared_exec_nodes():
|
|
graph = Graph()
|
|
graph.add_node(PromptTestInvocation(id="value", prompt="branch value"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
|
|
skipped_exec_id = g._create_execution_node("value", [])[0]
|
|
active_exec_id = g._create_execution_node("value", [])[0]
|
|
g._set_prepared_exec_state(skipped_exec_id, "skipped")
|
|
|
|
selected_exec_id = g._get_iteration_node("value", graph.nx_graph_flat(), g.execution_graph.nx_graph_flat(), [])
|
|
|
|
assert selected_exec_id == active_exec_id
|
|
|
|
|
|
def test_get_iteration_node_returns_single_active_prepared_exec_node():
|
|
graph = Graph()
|
|
graph.add_node(PromptTestInvocation(id="value", prompt="branch value"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
|
|
active_exec_id = g._create_execution_node("value", [])[0]
|
|
|
|
selected_exec_id = g._get_iteration_node("value", graph.nx_graph_flat(), g.execution_graph.nx_graph_flat(), [])
|
|
|
|
assert selected_exec_id == active_exec_id
|
|
|
|
|
|
def test_get_iteration_node_returns_none_when_only_skipped_prepared_exec_nodes_exist():
|
|
graph = Graph()
|
|
graph.add_node(PromptTestInvocation(id="value", prompt="branch value"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
|
|
skipped_exec_id = g._create_execution_node("value", [])[0]
|
|
g._set_prepared_exec_state(skipped_exec_id, "skipped")
|
|
|
|
selected_exec_id = g._get_iteration_node("value", graph.nx_graph_flat(), g.execution_graph.nx_graph_flat(), [])
|
|
|
|
assert selected_exec_id is None
|
|
|
|
|
|
def test_get_iteration_node_does_not_reuse_wrong_iterator_when_only_other_iteration_is_live():
|
|
graph = Graph()
|
|
graph.add_node(BooleanCollectionInvocation(id="conditions", collection=[True, False]))
|
|
graph.add_node(IterateInvocation(id="condition_iter"))
|
|
graph.add_node(AnyTypeTestInvocation(id="value"))
|
|
|
|
graph.add_edge(create_edge("conditions", "collection", "condition_iter", "collection"))
|
|
graph.add_edge(create_edge("condition_iter", "item", "value", "value"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
|
|
conditions_exec_id = g._create_execution_node("conditions", [])[0]
|
|
g.executed.add(conditions_exec_id)
|
|
g.results[conditions_exec_id] = BooleanCollectionOutput(collection=[True, False])
|
|
|
|
iterator_exec_ids = g._create_execution_node("condition_iter", [("conditions", conditions_exec_id)])
|
|
assert len(iterator_exec_ids) == 2
|
|
iterator_exec_ids_by_index = {g.execution_graph.get_node(exec_id).index: exec_id for exec_id in iterator_exec_ids}
|
|
first_iter_exec_id = iterator_exec_ids_by_index[0]
|
|
second_iter_exec_id = iterator_exec_ids_by_index[1]
|
|
|
|
value_exec_ids = []
|
|
value_exec_ids.extend(g._create_execution_node("value", [("condition_iter", first_iter_exec_id)]))
|
|
value_exec_ids.extend(g._create_execution_node("value", [("condition_iter", second_iter_exec_id)]))
|
|
assert len(value_exec_ids) == 2
|
|
|
|
for exec_id in value_exec_ids:
|
|
if g._get_iteration_path(exec_id) == (1,):
|
|
active_value_exec_id = exec_id
|
|
else:
|
|
skipped_value_exec_id = exec_id
|
|
|
|
g._set_prepared_exec_state(skipped_value_exec_id, "skipped")
|
|
|
|
selected_exec_id = g._get_iteration_node(
|
|
"value", graph.nx_graph_flat(), g.execution_graph.nx_graph_flat(), [first_iter_exec_id]
|
|
)
|
|
|
|
assert selected_exec_id is None
|
|
assert active_value_exec_id != skipped_value_exec_id
|
|
|
|
|
|
def test_mark_exec_node_skipped_does_not_hide_already_executed_results():
|
|
graph = Graph()
|
|
graph.add_node(AnyTypeTestInvocation(id="value", value="value"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
|
|
exec_id = g._create_execution_node("value", [])[0]
|
|
g.results[exec_id] = AnyTypeTestInvocationOutput(value="value")
|
|
g.executed.add(exec_id)
|
|
g._set_prepared_exec_state(exec_id, "executed")
|
|
|
|
g._if_scheduler().mark_exec_node_skipped(exec_id)
|
|
|
|
assert g._get_prepared_exec_metadata(exec_id).state == "executed"
|
|
assert g.results[exec_id].value == "value"
|
|
|
|
|
|
def test_mark_exec_node_skipped_is_idempotent_for_skipped_state():
|
|
graph = Graph()
|
|
graph.add_node(AnyTypeTestInvocation(id="value", value="value"))
|
|
|
|
g = GraphExecutionState(graph=graph)
|
|
|
|
exec_id = g._create_execution_node("value", [])[0]
|
|
|
|
g._if_scheduler().mark_exec_node_skipped(exec_id)
|
|
g._if_scheduler().mark_exec_node_skipped(exec_id)
|
|
|
|
assert g._get_prepared_exec_metadata(exec_id).state == "skipped"
|
|
assert g.executed_history.count("value") == 1
|
|
|
|
|
|
def test_are_connection_types_compatible_accepts_subclass_to_base():
|
|
"""A subclass output should be connectable to a base-class input.
|
|
|
|
This test exposes the bug where non-Union targets reject valid subclass connections.
|
|
"""
|
|
from invokeai.app.services.shared.graph import are_connection_types_compatible
|
|
|
|
class Base:
|
|
pass
|
|
|
|
class Child(Base):
|
|
pass
|
|
|
|
assert are_connection_types_compatible(Child, Base) is True
|