mirror of
https://github.com/langgenius/dify.git
synced 2026-03-10 09:56:13 +08:00
feat(graph_engine): add ready_queue state persistence to GraphRuntimeState
- Add ReadyQueueState TypedDict for type-safe queue serialization - Add ready_queue attribute to GraphRuntimeState for initializing with pre-existing queue state - Update GraphEngine to load ready_queue from GraphRuntimeState on initialization - Implement proper type hints using ReadyQueueState for better type safety - Add comprehensive tests for ready_queue loading functionality The ready_queue is read-only after initialization and allows resuming workflow execution with a pre-populated queue of nodes ready to execute.
This commit is contained in:
@ -4,6 +4,7 @@ import pytest
|
||||
|
||||
from core.workflow.entities.graph_runtime_state import GraphRuntimeState
|
||||
from core.workflow.entities.variable_pool import VariablePool
|
||||
from core.workflow.graph_engine.ready_queue import ReadyQueueState
|
||||
|
||||
|
||||
class TestGraphRuntimeState:
|
||||
@ -109,3 +110,30 @@ class TestGraphRuntimeState:
|
||||
|
||||
# Original should remain unchanged
|
||||
assert state.get_output("nested")["level1"]["level2"]["value"] == "test"
|
||||
|
||||
def test_ready_queue_property(self):
|
||||
variable_pool = VariablePool()
|
||||
|
||||
# Test default empty ready_queue
|
||||
state = GraphRuntimeState(variable_pool=variable_pool, start_at=time())
|
||||
assert state.ready_queue == {}
|
||||
|
||||
# Test initialization with ready_queue data as ReadyQueueState
|
||||
queue_data = ReadyQueueState(type="InMemoryReadyQueue", version="1.0", items=["node1", "node2"], maxsize=0)
|
||||
state = GraphRuntimeState(variable_pool=variable_pool, start_at=time(), ready_queue=queue_data)
|
||||
assert state.ready_queue == queue_data
|
||||
|
||||
# Test with different ready_queue data at initialization
|
||||
another_queue_data = ReadyQueueState(
|
||||
type="InMemoryReadyQueue",
|
||||
version="1.0",
|
||||
items=["node3", "node4", "node5"],
|
||||
maxsize=0,
|
||||
)
|
||||
another_state = GraphRuntimeState(variable_pool=variable_pool, start_at=time(), ready_queue=another_queue_data)
|
||||
assert another_state.ready_queue == another_queue_data
|
||||
|
||||
# Test immutability - modifying retrieved queue doesn't affect internal state
|
||||
retrieved_queue = state.ready_queue
|
||||
retrieved_queue["items"].append("node6")
|
||||
assert len(state.ready_queue["items"]) == 2 # Should still be 2, not 3
|
||||
|
||||
@ -744,3 +744,78 @@ def test_event_sequence_validation_with_table_tests():
|
||||
else:
|
||||
assert result.event_sequence_match is True
|
||||
assert result.success, f"Test {i + 1} failed: {result.event_mismatch_details or result.error}"
|
||||
|
||||
|
||||
def test_ready_queue_state_loading():
|
||||
"""
|
||||
Test that the ready_queue state is properly loaded from GraphRuntimeState
|
||||
during GraphEngine initialization.
|
||||
"""
|
||||
# Use the TableTestRunner to create a proper workflow instance
|
||||
runner = TableTestRunner()
|
||||
|
||||
# Create a simple workflow
|
||||
test_case = WorkflowTestCase(
|
||||
fixture_path="simple_passthrough_workflow",
|
||||
inputs={"query": "test"},
|
||||
expected_outputs={"query": "test"},
|
||||
description="Test ready_queue loading",
|
||||
)
|
||||
|
||||
# Load the workflow fixture
|
||||
workflow_runner = runner.workflow_runner
|
||||
fixture_data = workflow_runner.load_fixture("simple_passthrough_workflow")
|
||||
|
||||
# Create graph and runtime state with pre-populated ready_queue
|
||||
ready_queue_data = {
|
||||
"type": "InMemoryReadyQueue",
|
||||
"version": "1.0",
|
||||
"items": ["node1", "node2", "node3"],
|
||||
"maxsize": 0,
|
||||
}
|
||||
|
||||
# We need to create the graph first, then create a new GraphRuntimeState with ready_queue
|
||||
graph, original_runtime_state = workflow_runner.create_graph_from_fixture(fixture_data, query="test")
|
||||
|
||||
# Create a new GraphRuntimeState with the ready_queue data
|
||||
from core.workflow.entities import GraphRuntimeState
|
||||
from core.workflow.graph_engine.ready_queue import ReadyQueueState
|
||||
|
||||
# Convert ready_queue_data to ReadyQueueState
|
||||
ready_queue_state = ReadyQueueState(**ready_queue_data)
|
||||
|
||||
graph_runtime_state = GraphRuntimeState(
|
||||
variable_pool=original_runtime_state.variable_pool,
|
||||
start_at=original_runtime_state.start_at,
|
||||
ready_queue=ready_queue_state,
|
||||
)
|
||||
|
||||
# Update all nodes to use the new GraphRuntimeState
|
||||
for node in graph.nodes.values():
|
||||
node.graph_runtime_state = graph_runtime_state
|
||||
|
||||
# Create GraphEngine
|
||||
command_channel = InMemoryChannel()
|
||||
engine = GraphEngine(
|
||||
tenant_id="test-tenant",
|
||||
app_id="test-app",
|
||||
workflow_id="test-workflow",
|
||||
user_id="test-user",
|
||||
user_from=UserFrom.ACCOUNT,
|
||||
invoke_from=InvokeFrom.DEBUGGER,
|
||||
call_depth=0,
|
||||
graph=graph,
|
||||
graph_config={},
|
||||
graph_runtime_state=graph_runtime_state,
|
||||
command_channel=command_channel,
|
||||
)
|
||||
|
||||
# Verify that the ready_queue was loaded from GraphRuntimeState
|
||||
assert engine._ready_queue.qsize() == 3
|
||||
|
||||
# Verify the initial state matches what was provided
|
||||
initial_queue_state = engine.graph_runtime_state.ready_queue
|
||||
assert initial_queue_state["type"] == "InMemoryReadyQueue"
|
||||
assert initial_queue_state["version"] == "1.0"
|
||||
assert len(initial_queue_state["items"]) == 3
|
||||
assert initial_queue_state["items"] == ["node1", "node2", "node3"]
|
||||
|
||||
Reference in New Issue
Block a user