admin/dify
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

feat: Introduce WorkflowExecution Domain Entity and Repository, Replace WorkflowRun Direct Usage, and Unify Stream Response Logic (#20067)

Browse Source 
Signed-off-by: -LAN- <laipz8200@outlook.com>
This commit is contained in:
-LAN-
2025-05-21 22:01:53 +08:00
committed by GitHub
parent 7d230acf40
commit d31235ca13
13 changed files with 1710 additions and 644 deletions
Download Patch File Download Diff File Expand all files Collapse all files

423
api/tests/unit_tests/core/workflow/test_workflow_cycle_manager.py
View File

@@ -1,45 +1,73 @@
import json
import time
from datetime import UTC, datetime
from unittest.mock import MagicMock, patch
from unittest.mock import MagicMock
import pytest
from sqlalchemy.orm import Session
from core.app.app_config.entities import AppAdditionalFeatures, WorkflowUIBasedAppConfig
from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, InvokeFrom
from core.app.entities.queue_entities import (
QueueNodeFailedEvent,
QueueNodeStartedEvent,
QueueNodeSucceededEvent,
)
from core.workflow.entities.node_entities import NodeRunMetadataKey
from core.workflow.entities.node_execution_entities import NodeExecution, NodeExecutionStatus
from core.workflow.entities.workflow_execution_entities import WorkflowExecution, WorkflowExecutionStatus, WorkflowType
from core.workflow.enums import SystemVariableKey
from core.workflow.nodes import NodeType
from core.workflow.repository.workflow_execution_repository import WorkflowExecutionRepository
from core.workflow.repository.workflow_node_execution_repository import WorkflowNodeExecutionRepository
from core.workflow.workflow_cycle_manager import WorkflowCycleManager
from models.enums import CreatorUserRole
from models.model import AppMode
from models.workflow import (
Workflow,
WorkflowNodeExecutionStatus,
WorkflowRun,
WorkflowRunStatus,
)
@pytest.fixture
def mock_app_generate_entity():
entity = MagicMock(spec=AdvancedChatAppGenerateEntity)
entity.inputs = {"query": "test query"}
entity.invoke_from = InvokeFrom.WEB_APP
# Create app_config as a separate mock
app_config = MagicMock()
app_config.tenant_id = "test-tenant-id"
app_config.app_id = "test-app-id"
entity.app_config = app_config
def real_app_generate_entity():
additional_features = AppAdditionalFeatures(
file_upload=None,
opening_statement=None,
suggested_questions=[],
suggested_questions_after_answer=False,
show_retrieve_source=False,
more_like_this=False,
speech_to_text=False,
text_to_speech=None,
trace_config=None,
)
app_config = WorkflowUIBasedAppConfig(
tenant_id="test-tenant-id",
app_id="test-app-id",
app_mode=AppMode.WORKFLOW,
additional_features=additional_features,
workflow_id="test-workflow-id",
)
entity = AdvancedChatAppGenerateEntity(
task_id="test-task-id",
app_config=app_config,
inputs={"query": "test query"},
files=[],
user_id="test-user-id",
stream=False,
invoke_from=InvokeFrom.WEB_APP,
query="test query",
conversation_id="test-conversation-id",
)
return entity
@pytest.fixture
def mock_workflow_system_variables():
def real_workflow_system_variables():
return {
SystemVariableKey.QUERY: "test query",
SystemVariableKey.CONVERSATION_ID: "test-conversation-id",
@@ -59,10 +87,23 @@ def mock_node_execution_repository():
@pytest.fixture
def workflow_cycle_manager(mock_app_generate_entity, mock_workflow_system_variables, mock_node_execution_repository):
def mock_workflow_execution_repository():
repo = MagicMock(spec=WorkflowExecutionRepository)
repo.get.return_value = None
return repo
@pytest.fixture
def workflow_cycle_manager(
real_app_generate_entity,
real_workflow_system_variables,
mock_workflow_execution_repository,
mock_node_execution_repository,
):
return WorkflowCycleManager(
application_generate_entity=mock_app_generate_entity,
workflow_system_variables=mock_workflow_system_variables,
application_generate_entity=real_app_generate_entity,
workflow_system_variables=real_workflow_system_variables,
workflow_execution_repository=mock_workflow_execution_repository,
workflow_node_execution_repository=mock_node_execution_repository,
)
@@ -74,121 +115,173 @@ def mock_session():
@pytest.fixture
def mock_workflow():
workflow = MagicMock(spec=Workflow)
def real_workflow():
workflow = Workflow()
workflow.id = "test-workflow-id"
workflow.tenant_id = "test-tenant-id"
workflow.app_id = "test-app-id"
workflow.type = "chat"
workflow.version = "1.0"
workflow.graph = json.dumps({"nodes": [], "edges": []})
graph_data = {"nodes": [], "edges": []}
workflow.graph = json.dumps(graph_data)
workflow.features = json.dumps({"file_upload": {"enabled": False}})
workflow.created_by = "test-user-id"
workflow.created_at = datetime.now(UTC).replace(tzinfo=None)
workflow.updated_at = datetime.now(UTC).replace(tzinfo=None)
workflow._environment_variables = "{}"
workflow._conversation_variables = "{}"
return workflow
@pytest.fixture
def mock_workflow_run():
workflow_run = MagicMock(spec=WorkflowRun)
def real_workflow_run():
workflow_run = WorkflowRun()
workflow_run.id = "test-workflow-run-id"
workflow_run.tenant_id = "test-tenant-id"
workflow_run.app_id = "test-app-id"
workflow_run.workflow_id = "test-workflow-id"
workflow_run.sequence_number = 1
workflow_run.type = "chat"
workflow_run.triggered_from = "app-run"
workflow_run.version = "1.0"
workflow_run.graph = json.dumps({"nodes": [], "edges": []})
workflow_run.inputs = json.dumps({"query": "test query"})
workflow_run.status = WorkflowRunStatus.RUNNING
workflow_run.outputs = json.dumps({"answer": "test answer"})
workflow_run.created_by_role = CreatorUserRole.ACCOUNT
workflow_run.created_by = "test-user-id"
workflow_run.created_at = datetime.now(UTC).replace(tzinfo=None)
workflow_run.inputs_dict = {"query": "test query"}
workflow_run.outputs_dict = {"answer": "test answer"}
return workflow_run
def test_init(
workflow_cycle_manager, mock_app_generate_entity, mock_workflow_system_variables, mock_node_execution_repository
workflow_cycle_manager,
real_app_generate_entity,
real_workflow_system_variables,
mock_workflow_execution_repository,
mock_node_execution_repository,
):
"""Test initialization of WorkflowCycleManager"""
assert workflow_cycle_manager._workflow_run is None
assert workflow_cycle_manager._application_generate_entity == mock_app_generate_entity
assert workflow_cycle_manager._workflow_system_variables == mock_workflow_system_variables
assert workflow_cycle_manager._application_generate_entity == real_app_generate_entity
assert workflow_cycle_manager._workflow_system_variables == real_workflow_system_variables
assert workflow_cycle_manager._workflow_execution_repository == mock_workflow_execution_repository
assert workflow_cycle_manager._workflow_node_execution_repository == mock_node_execution_repository
def test_handle_workflow_run_start(workflow_cycle_manager, mock_session, mock_workflow):
"""Test _handle_workflow_run_start method"""
def test_handle_workflow_run_start(workflow_cycle_manager, mock_session, real_workflow):
"""Test handle_workflow_run_start method"""
# Mock session.scalar to return the workflow and max sequence
mock_session.scalar.side_effect = [mock_workflow, 5]
mock_session.scalar.side_effect = [real_workflow, 5]
# Call the method
workflow_run = workflow_cycle_manager._handle_workflow_run_start(
workflow_execution = workflow_cycle_manager.handle_workflow_run_start(
session=mock_session,
workflow_id="test-workflow-id",
user_id="test-user-id",
created_by_role=CreatorUserRole.ACCOUNT,
)
# Verify the result
assert workflow_run.tenant_id == mock_workflow.tenant_id
assert workflow_run.app_id == mock_workflow.app_id
assert workflow_run.workflow_id == mock_workflow.id
assert workflow_run.sequence_number == 6 # max_sequence + 1
assert workflow_run.status == WorkflowRunStatus.RUNNING
assert workflow_run.created_by_role == CreatorUserRole.ACCOUNT
assert workflow_run.created_by == "test-user-id"
assert workflow_execution.workflow_id == real_workflow.id
assert workflow_execution.sequence_number == 6 # max_sequence + 1
# Verify session.add was called
mock_session.add.assert_called_once_with(workflow_run)
# Verify the workflow_execution_repository.save was called
workflow_cycle_manager._workflow_execution_repository.save.assert_called_once_with(workflow_execution)
def test_handle_workflow_run_success(workflow_cycle_manager, mock_session, mock_workflow_run):
"""Test _handle_workflow_run_success method"""
# Mock _get_workflow_run to return the mock_workflow_run
with patch.object(workflow_cycle_manager, "_get_workflow_run", return_value=mock_workflow_run):
# Call the method
result = workflow_cycle_manager._handle_workflow_run_success(
session=mock_session,
workflow_run_id="test-workflow-run-id",
start_at=time.perf_counter() - 10, # 10 seconds ago
total_tokens=100,
total_steps=5,
outputs={"answer": "test answer"},
)
def test_handle_workflow_run_success(workflow_cycle_manager, mock_workflow_execution_repository):
"""Test handle_workflow_run_success method"""
# Create a real WorkflowExecution
# Verify the result
assert result == mock_workflow_run
assert result.status == WorkflowRunStatus.SUCCEEDED
assert result.outputs == json.dumps({"answer": "test answer"})
assert result.total_tokens == 100
assert result.total_steps == 5
assert result.finished_at is not None
workflow_execution = WorkflowExecution(
id="test-workflow-run-id",
workflow_id="test-workflow-id",
workflow_version="1.0",
sequence_number=1,
type=WorkflowType.CHAT,
graph={"nodes": [], "edges": []},
inputs={"query": "test query"},
started_at=datetime.now(UTC).replace(tzinfo=None),
)
# Mock _get_workflow_execution_or_raise_error to return the real workflow_execution
workflow_cycle_manager._workflow_execution_repository.get.return_value = workflow_execution
# Call the method
result = workflow_cycle_manager.handle_workflow_run_success(
workflow_run_id="test-workflow-run-id",
total_tokens=100,
total_steps=5,
outputs={"answer": "test answer"},
)
# Verify the result
assert result == workflow_execution
assert result.status == WorkflowExecutionStatus.SUCCEEDED
assert result.outputs == {"answer": "test answer"}
assert result.total_tokens == 100
assert result.total_steps == 5
assert result.finished_at is not None
def test_handle_workflow_run_failed(workflow_cycle_manager, mock_session, mock_workflow_run):
"""Test _handle_workflow_run_failed method"""
# Mock _get_workflow_run to return the mock_workflow_run
with patch.object(workflow_cycle_manager, "_get_workflow_run", return_value=mock_workflow_run):
# Mock get_running_executions to return an empty list
workflow_cycle_manager._workflow_node_execution_repository.get_running_executions.return_value = []
def test_handle_workflow_run_failed(workflow_cycle_manager, mock_workflow_execution_repository):
"""Test handle_workflow_run_failed method"""
# Create a real WorkflowExecution
# Call the method
result = workflow_cycle_manager._handle_workflow_run_failed(
session=mock_session,
workflow_run_id="test-workflow-run-id",
start_at=time.perf_counter() - 10, # 10 seconds ago
total_tokens=50,
total_steps=3,
status=WorkflowRunStatus.FAILED,
error="Test error message",
)
workflow_execution = WorkflowExecution(
id="test-workflow-run-id",
workflow_id="test-workflow-id",
workflow_version="1.0",
sequence_number=1,
type=WorkflowType.CHAT,
graph={"nodes": [], "edges": []},
inputs={"query": "test query"},
started_at=datetime.now(UTC).replace(tzinfo=None),
)
# Verify the result
assert result == mock_workflow_run
assert result.status == WorkflowRunStatus.FAILED.value
assert result.error == "Test error message"
assert result.total_tokens == 50
assert result.total_steps == 3
assert result.finished_at is not None
# Mock _get_workflow_execution_or_raise_error to return the real workflow_execution
workflow_cycle_manager._workflow_execution_repository.get.return_value = workflow_execution
# Mock get_running_executions to return an empty list
workflow_cycle_manager._workflow_node_execution_repository.get_running_executions.return_value = []
# Call the method
result = workflow_cycle_manager.handle_workflow_run_failed(
workflow_run_id="test-workflow-run-id",
total_tokens=50,
total_steps=3,
status=WorkflowRunStatus.FAILED,
error_message="Test error message",
)
# Verify the result
assert result == workflow_execution
assert result.status == WorkflowExecutionStatus(WorkflowRunStatus.FAILED.value)
assert result.error_message == "Test error message"
assert result.total_tokens == 50
assert result.total_steps == 3
assert result.finished_at is not None
def test_handle_node_execution_start(workflow_cycle_manager, mock_workflow_run):
"""Test _handle_node_execution_start method"""
def test_handle_node_execution_start(workflow_cycle_manager, mock_workflow_execution_repository):
"""Test handle_node_execution_start method"""
# Create a real WorkflowExecution
workflow_execution = WorkflowExecution(
id="test-workflow-execution-id",
workflow_id="test-workflow-id",
workflow_version="1.0",
sequence_number=1,
type=WorkflowType.CHAT,
graph={"nodes": [], "edges": []},
inputs={"query": "test query"},
started_at=datetime.now(UTC).replace(tzinfo=None),
)
# Mock _get_workflow_execution_or_raise_error to return the real workflow_execution
workflow_cycle_manager._workflow_execution_repository.get.return_value = workflow_execution
# Create a mock event
event = MagicMock(spec=QueueNodeStartedEvent)
event.node_execution_id = "test-node-execution-id"
@@ -207,129 +300,171 @@ def test_handle_node_execution_start(workflow_cycle_manager, mock_workflow_run):
event.in_loop_id = "test-loop-id"
# Call the method
result = workflow_cycle_manager._handle_node_execution_start(
workflow_run=mock_workflow_run,
result = workflow_cycle_manager.handle_node_execution_start(
workflow_execution_id=workflow_execution.id,
event=event,
)
# Verify the result
# NodeExecution doesn't have tenant_id attribute, it's handled at repository level
# assert result.tenant_id == mock_workflow_run.tenant_id
# assert result.app_id == mock_workflow_run.app_id
assert result.workflow_id == mock_workflow_run.workflow_id
assert result.workflow_run_id == mock_workflow_run.id
assert result.workflow_id == workflow_execution.workflow_id
assert result.workflow_run_id == workflow_execution.id
assert result.node_execution_id == event.node_execution_id
assert result.node_id == event.node_id
assert result.node_type == event.node_type
assert result.title == event.node_data.title
assert result.status == WorkflowNodeExecutionStatus.RUNNING.value
# NodeExecution doesn't have created_by_role and created_by attributes, they're handled at repository level
# assert result.created_by_role == mock_workflow_run.created_by_role
# assert result.created_by == mock_workflow_run.created_by
assert result.status == NodeExecutionStatus.RUNNING
# Verify save was called
workflow_cycle_manager._workflow_node_execution_repository.save.assert_called_once_with(result)
def test_get_workflow_run(workflow_cycle_manager, mock_session, mock_workflow_run):
"""Test _get_workflow_run method"""
# Mock session.scalar to return the workflow run
mock_session.scalar.return_value = mock_workflow_run
def test_get_workflow_execution_or_raise_error(workflow_cycle_manager, mock_workflow_execution_repository):
"""Test _get_workflow_execution_or_raise_error method"""
# Create a real WorkflowExecution
# Call the method
result = workflow_cycle_manager._get_workflow_run(
session=mock_session,
workflow_run_id="test-workflow-run-id",
workflow_execution = WorkflowExecution(
id="test-workflow-run-id",
workflow_id="test-workflow-id",
workflow_version="1.0",
sequence_number=1,
type=WorkflowType.CHAT,
graph={"nodes": [], "edges": []},
inputs={"query": "test query"},
started_at=datetime.now(UTC).replace(tzinfo=None),
)
# Mock the repository get method to return the real execution
workflow_cycle_manager._workflow_execution_repository.get.return_value = workflow_execution
# Call the method
result = workflow_cycle_manager._get_workflow_execution_or_raise_error("test-workflow-run-id")
# Verify the result
assert result == mock_workflow_run
assert workflow_cycle_manager._workflow_run == mock_workflow_run
assert result == workflow_execution
# Test error case
workflow_cycle_manager._workflow_execution_repository.get.return_value = None
# Expect an error when execution is not found
with pytest.raises(ValueError):
workflow_cycle_manager._get_workflow_execution_or_raise_error("non-existent-id")
def test_handle_workflow_node_execution_success(workflow_cycle_manager):
"""Test _handle_workflow_node_execution_success method"""
"""Test handle_workflow_node_execution_success method"""
# Create a mock event
event = MagicMock(spec=QueueNodeSucceededEvent)
event.node_execution_id = "test-node-execution-id"
event.inputs = {"input": "test input"}
event.process_data = {"process": "test process"}
event.outputs = {"output": "test output"}
event.execution_metadata = {"metadata": "test metadata"}
event.execution_metadata = {NodeRunMetadataKey.TOTAL_TOKENS: 100}
event.start_at = datetime.now(UTC).replace(tzinfo=None)
# Create a mock node execution
node_execution = MagicMock()
node_execution.node_execution_id = "test-node-execution-id"
# Create a real node execution
node_execution = NodeExecution(
id="test-node-execution-record-id",
node_execution_id="test-node-execution-id",
workflow_id="test-workflow-id",
workflow_run_id="test-workflow-run-id",
index=1,
node_id="test-node-id",
node_type=NodeType.LLM,
title="Test Node",
created_at=datetime.now(UTC).replace(tzinfo=None),
)
# Mock the repository to return the node execution
workflow_cycle_manager._workflow_node_execution_repository.get_by_node_execution_id.return_value = node_execution
# Call the method
result = workflow_cycle_manager._handle_workflow_node_execution_success(
result = workflow_cycle_manager.handle_workflow_node_execution_success(
event=event,
)
# Verify the result
assert result == node_execution
assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED.value
assert result.status == NodeExecutionStatus.SUCCEEDED
# Verify save was called
workflow_cycle_manager._workflow_node_execution_repository.save.assert_called_once_with(node_execution)
def test_handle_workflow_run_partial_success(workflow_cycle_manager, mock_session, mock_workflow_run):
"""Test _handle_workflow_run_partial_success method"""
# Mock _get_workflow_run to return the mock_workflow_run
with patch.object(workflow_cycle_manager, "_get_workflow_run", return_value=mock_workflow_run):
# Call the method
result = workflow_cycle_manager._handle_workflow_run_partial_success(
session=mock_session,
workflow_run_id="test-workflow-run-id",
start_at=time.perf_counter() - 10, # 10 seconds ago
total_tokens=75,
total_steps=4,
outputs={"partial_answer": "test partial answer"},
exceptions_count=2,
)
def test_handle_workflow_run_partial_success(workflow_cycle_manager, mock_workflow_execution_repository):
"""Test handle_workflow_run_partial_success method"""
# Create a real WorkflowExecution
# Verify the result
assert result == mock_workflow_run
assert result.status == WorkflowRunStatus.PARTIAL_SUCCEEDED.value
assert result.outputs == json.dumps({"partial_answer": "test partial answer"})
assert result.total_tokens == 75
assert result.total_steps == 4
assert result.exceptions_count == 2
assert result.finished_at is not None
workflow_execution = WorkflowExecution(
id="test-workflow-run-id",
workflow_id="test-workflow-id",
workflow_version="1.0",
sequence_number=1,
type=WorkflowType.CHAT,
graph={"nodes": [], "edges": []},
inputs={"query": "test query"},
started_at=datetime.now(UTC).replace(tzinfo=None),
)
# Mock _get_workflow_execution_or_raise_error to return the real workflow_execution
workflow_cycle_manager._workflow_execution_repository.get.return_value = workflow_execution
# Call the method
result = workflow_cycle_manager.handle_workflow_run_partial_success(
workflow_run_id="test-workflow-run-id",
total_tokens=75,
total_steps=4,
outputs={"partial_answer": "test partial answer"},
exceptions_count=2,
)
# Verify the result
assert result == workflow_execution
assert result.status == WorkflowExecutionStatus.PARTIAL_SUCCEEDED
assert result.outputs == {"partial_answer": "test partial answer"}
assert result.total_tokens == 75
assert result.total_steps == 4
assert result.exceptions_count == 2
assert result.finished_at is not None
def test_handle_workflow_node_execution_failed(workflow_cycle_manager):
"""Test _handle_workflow_node_execution_failed method"""
"""Test handle_workflow_node_execution_failed method"""
# Create a mock event
event = MagicMock(spec=QueueNodeFailedEvent)
event.node_execution_id = "test-node-execution-id"
event.inputs = {"input": "test input"}
event.process_data = {"process": "test process"}
event.outputs = {"output": "test output"}
event.execution_metadata = {"metadata": "test metadata"}
event.execution_metadata = {NodeRunMetadataKey.TOTAL_TOKENS: 100}
event.start_at = datetime.now(UTC).replace(tzinfo=None)
event.error = "Test error message"
# Create a mock node execution
node_execution = MagicMock()
node_execution.node_execution_id = "test-node-execution-id"
# Create a real node execution
node_execution = NodeExecution(
id="test-node-execution-record-id",
node_execution_id="test-node-execution-id",
workflow_id="test-workflow-id",
workflow_run_id="test-workflow-run-id",
index=1,
node_id="test-node-id",
node_type=NodeType.LLM,
title="Test Node",
created_at=datetime.now(UTC).replace(tzinfo=None),
)
# Mock the repository to return the node execution
workflow_cycle_manager._workflow_node_execution_repository.get_by_node_execution_id.return_value = node_execution
# Call the method
result = workflow_cycle_manager._handle_workflow_node_execution_failed(
result = workflow_cycle_manager.handle_workflow_node_execution_failed(
event=event,
)
# Verify the result
assert result == node_execution
assert result.status == WorkflowNodeExecutionStatus.FAILED.value
assert result.status == NodeExecutionStatus.FAILED
assert result.error == "Test error message"
# Verify save was called