Skip to content

D
dify
Commit 2db67c41 authored by takatost's avatar takatost
Browse files
Options

refactor pipeline and remove node run run_args

parent 80b4db08
Show whitespace changes
Inline Side-by-side
Showing with 201 additions and 114 deletions
api/core/app/apps/advanced_chat/generate_task_pipeline.py
View file @ 2db67c41
......@@ -55,6 +55,19 @@ class TaskState(BaseModel):
"""
TaskState entity
"""
class NodeExecutionInfo(BaseModel):
"""
NodeExecutionInfo entity
"""
workflow_node_execution: WorkflowNodeExecution
start_at: float
class Config:
"""Configuration for this pydantic object."""
extra = Extra.forbid
arbitrary_types_allowed = True
answer: str = ""
metadata: dict = {}
usage: LLMUsage
......@@ -64,8 +77,8 @@ class TaskState(BaseModel):
total_tokens: int = 0
total_steps: int = 0
current_node_execution: Optional[WorkflowNodeExecution] = None
current_node_execution_start_at: Optional[float] = None
running_node_execution_infos: dict[str, NodeExecutionInfo] = {}
latest_node_execution_info: Optional[NodeExecutionInfo] = None
class Config:
"""Configuration for this pydantic object."""
......@@ -218,7 +231,7 @@ class AdvancedChatAppGenerateTaskPipeline(WorkflowBasedGenerateTaskPipeline):
yield self._yield_response(response)
elif isinstance(event, QueueNodeStartedEvent):
self._on_node_start(event)
workflow_node_execution = self._task_state.current_node_execution
workflow_node_execution = self._task_state.latest_node_execution_info.workflow_node_execution
response = {
'event': 'node_started',
......@@ -237,7 +250,7 @@ class AdvancedChatAppGenerateTaskPipeline(WorkflowBasedGenerateTaskPipeline):
yield self._yield_response(response)
elif isinstance(event, QueueNodeSucceededEvent | QueueNodeFailedEvent):
self._on_node_finished(event)
workflow_node_execution = self._task_state.current_node_execution
workflow_node_execution = self._task_state.latest_node_execution_info.workflow_node_execution
if workflow_node_execution.status == WorkflowNodeExecutionStatus.SUCCEEDED.value:
if workflow_node_execution.node_type == NodeType.LLM.value:
......@@ -447,15 +460,21 @@ class AdvancedChatAppGenerateTaskPipeline(WorkflowBasedGenerateTaskPipeline):
predecessor_node_id=event.predecessor_node_id
)
self._task_state.current_node_execution = workflow_node_execution
self._task_state.current_node_execution_start_at = time.perf_counter()
latest_node_execution_info = TaskState.NodeExecutionInfo(
workflow_node_execution=workflow_node_execution,
start_at=time.perf_counter()
)
self._task_state.running_node_execution_infos[event.node_id] = latest_node_execution_info
self._task_state.latest_node_execution_info = latest_node_execution_info
self._task_state.total_steps += 1
def _on_node_finished(self, event: QueueNodeSucceededEvent | QueueNodeFailedEvent) -> None:
current_node_execution = self._task_state.running_node_execution_infos[event.node_id]
if isinstance(event, QueueNodeSucceededEvent):
workflow_node_execution = self._workflow_node_execution_success(
workflow_node_execution=self._task_state.current_node_execution,
start_at=self._task_state.current_node_execution_start_at,
workflow_node_execution=current_node_execution.workflow_node_execution,
start_at=current_node_execution.start_at,
inputs=event.inputs,
process_data=event.process_data,
outputs=event.outputs,
......@@ -472,12 +491,14 @@ class AdvancedChatAppGenerateTaskPipeline(WorkflowBasedGenerateTaskPipeline):
self._task_state.metadata['usage'] = usage_dict
else:
workflow_node_execution = self._workflow_node_execution_failed(
workflow_node_execution=self._task_state.current_node_execution,
start_at=self._task_state.current_node_execution_start_at,
workflow_node_execution=current_node_execution.workflow_node_execution,
start_at=current_node_execution.start_at,
error=event.error
)
self._task_state.current_node_execution = workflow_node_execution
# remove running node execution info
del self._task_state.running_node_execution_infos[event.node_id]
self._task_state.latest_node_execution_info.workflow_node_execution = workflow_node_execution
def _on_workflow_finished(self, event: QueueStopEvent | QueueWorkflowSucceededEvent | QueueWorkflowFailedEvent) -> None:
if isinstance(event, QueueStopEvent):
......@@ -504,8 +525,8 @@ class AdvancedChatAppGenerateTaskPipeline(WorkflowBasedGenerateTaskPipeline):
start_at=self._task_state.start_at,
total_tokens=self._task_state.total_tokens,
total_steps=self._task_state.total_steps,
outputs=self._task_state.current_node_execution.outputs
if self._task_state.current_node_execution else None
outputs=self._task_state.latest_node_execution_info.workflow_node_execution.outputs
if self._task_state.latest_node_execution_info else None
)
self._task_state.workflow_run = workflow_run
......
api/core/app/apps/workflow/generate_task_pipeline.py
View file @ 2db67c41
......@@ -41,6 +41,19 @@ class TaskState(BaseModel):
"""
TaskState entity
"""
class NodeExecutionInfo(BaseModel):
"""
NodeExecutionInfo entity
"""
workflow_node_execution: WorkflowNodeExecution
start_at: float
class Config:
"""Configuration for this pydantic object."""
extra = Extra.forbid
arbitrary_types_allowed = True
answer: str = ""
metadata: dict = {}
......@@ -49,8 +62,8 @@ class TaskState(BaseModel):
total_tokens: int = 0
total_steps: int = 0
current_node_execution: Optional[WorkflowNodeExecution] = None
current_node_execution_start_at: Optional[float] = None
running_node_execution_infos: dict[str, NodeExecutionInfo] = {}
latest_node_execution_info: Optional[NodeExecutionInfo] = None
class Config:
"""Configuration for this pydantic object."""
......@@ -179,7 +192,7 @@ class WorkflowAppGenerateTaskPipeline(WorkflowBasedGenerateTaskPipeline):
yield self._yield_response(response)
elif isinstance(event, QueueNodeStartedEvent):
self._on_node_start(event)
workflow_node_execution = self._task_state.current_node_execution
workflow_node_execution = self._task_state.latest_node_execution_info.workflow_node_execution
response = {
'event': 'node_started',
......@@ -198,7 +211,7 @@ class WorkflowAppGenerateTaskPipeline(WorkflowBasedGenerateTaskPipeline):
yield self._yield_response(response)
elif isinstance(event, QueueNodeSucceededEvent | QueueNodeFailedEvent):
self._on_node_finished(event)
workflow_node_execution = self._task_state.current_node_execution
workflow_node_execution = self._task_state.latest_node_execution_info.workflow_node_execution
response = {
'event': 'node_finished',
......@@ -339,15 +352,22 @@ class WorkflowAppGenerateTaskPipeline(WorkflowBasedGenerateTaskPipeline):
predecessor_node_id=event.predecessor_node_id
)
self._task_state.current_node_execution = workflow_node_execution
self._task_state.current_node_execution_start_at = time.perf_counter()
latest_node_execution_info = TaskState.NodeExecutionInfo(
workflow_node_execution=workflow_node_execution,
start_at=time.perf_counter()
)
self._task_state.running_node_execution_infos[event.node_id] = latest_node_execution_info
self._task_state.latest_node_execution_info = latest_node_execution_info
self._task_state.total_steps += 1
def _on_node_finished(self, event: QueueNodeSucceededEvent | QueueNodeFailedEvent) -> None:
current_node_execution = self._task_state.running_node_execution_infos[event.node_id]
if isinstance(event, QueueNodeSucceededEvent):
workflow_node_execution = self._workflow_node_execution_success(
workflow_node_execution=self._task_state.current_node_execution,
start_at=self._task_state.current_node_execution_start_at,
workflow_node_execution=current_node_execution.workflow_node_execution,
start_at=current_node_execution.start_at,
inputs=event.inputs,
process_data=event.process_data,
outputs=event.outputs,
......@@ -359,12 +379,14 @@ class WorkflowAppGenerateTaskPipeline(WorkflowBasedGenerateTaskPipeline):
int(event.execution_metadata.get(NodeRunMetadataKey.TOTAL_TOKENS)))
else:
workflow_node_execution = self._workflow_node_execution_failed(
workflow_node_execution=self._task_state.current_node_execution,
start_at=self._task_state.current_node_execution_start_at,
workflow_node_execution=current_node_execution.workflow_node_execution,
start_at=current_node_execution.start_at,
error=event.error
)
self._task_state.current_node_execution = workflow_node_execution
# remove running node execution info
del self._task_state.running_node_execution_infos[event.node_id]
self._task_state.latest_node_execution_info.workflow_node_execution = workflow_node_execution
def _on_workflow_finished(self, event: QueueStopEvent | QueueWorkflowSucceededEvent | QueueWorkflowFailedEvent) -> None:
if isinstance(event, QueueStopEvent):
......@@ -391,8 +413,8 @@ class WorkflowAppGenerateTaskPipeline(WorkflowBasedGenerateTaskPipeline):
start_at=self._task_state.start_at,
total_tokens=self._task_state.total_tokens,
total_steps=self._task_state.total_steps,
outputs=self._task_state.current_node_execution.outputs
if self._task_state.current_node_execution else None
outputs=self._task_state.latest_node_execution_info.workflow_node_execution.outputs
if self._task_state.latest_node_execution_info else None
)
self._task_state.workflow_run = workflow_run
......
api/core/workflow/entities/variable_pool.py
View file @ 2db67c41
......@@ -19,14 +19,17 @@ class ValueType(Enum):
class VariablePool:
variables_mapping = {}
user_inputs: dict
def __init__(self, system_variables: dict[SystemVariable, Any]) -> None:
def __init__(self, system_variables: dict[SystemVariable, Any],
user_inputs: dict) -> None:
# system variables
# for example:
# {
# 'query': 'abc',
# 'files': []
# }
self.user_inputs = user_inputs
for system_variable, value in system_variables.items():
self.append_variable('sys', [system_variable.value], value)
......
api/core/workflow/entities/workflow_entities.py
View file @ 2db67c41
......@@ -18,15 +18,13 @@ class WorkflowNodeAndResult:
class WorkflowRunState:
workflow: Workflow
start_at: float
user_inputs: dict
variable_pool: VariablePool
total_tokens: int = 0
workflow_nodes_and_results: list[WorkflowNodeAndResult] = []
def __init__(self, workflow: Workflow, start_at: float, user_inputs: dict, variable_pool: VariablePool):
def __init__(self, workflow: Workflow, start_at: float, variable_pool: VariablePool):
self.workflow = workflow
self.start_at = start_at
self.user_inputs = user_inputs
self.variable_pool = variable_pool
api/core/workflow/nodes/base_node.py
View file @ 2db67c41
......@@ -28,31 +28,23 @@ class BaseNode(ABC):
self.callbacks = callbacks or []
@abstractmethod
def _run(self, variable_pool: Optional[VariablePool] = None,
run_args: Optional[dict] = None) -> NodeRunResult:
def _run(self, variable_pool: VariablePool) -> NodeRunResult:
"""
Run node
:param variable_pool: variable pool
:param run_args: run args
:return:
"""
raise NotImplementedError
def run(self, variable_pool: Optional[VariablePool] = None,
run_args: Optional[dict] = None) -> NodeRunResult:
def run(self, variable_pool: VariablePool) -> NodeRunResult:
"""
Run node entry
:param variable_pool: variable pool
:param run_args: run args
:return:
"""
if variable_pool is None and run_args is None:
raise ValueError("At least one of `variable_pool` or `run_args` must be provided.")
try:
result = self._run(
variable_pool=variable_pool,
run_args=run_args
variable_pool=variable_pool
)
except Exception as e:
# process unhandled exception
......@@ -77,6 +69,26 @@ class BaseNode(ABC):
text=text
)
@classmethod
def extract_variable_selector_to_variable_mapping(cls, config: dict) -> dict:
"""
Extract variable selector to variable mapping
:param config: node config
:return:
"""
node_data = cls._node_data_cls(**config.get("data", {}))
return cls._extract_variable_selector_to_variable_mapping(node_data)
@classmethod
@abstractmethod
def _extract_variable_selector_to_variable_mapping(cls, node_data: BaseNodeData) -> dict[list[str], str]:
"""
Extract variable selector to variable mapping
:param node_data: node data
:return:
"""
raise NotImplementedError
@classmethod
def get_default_config(cls, filters: Optional[dict] = None) -> dict:
"""
......
api/core/workflow/nodes/code/code_node.py
View file @ 2db67c41
from typing import Optional, Union, cast
from core.workflow.entities.base_node_data_entities import BaseNodeData
from core.workflow.entities.node_entities import NodeRunResult, NodeType
from core.workflow.entities.variable_pool import VariablePool
from core.workflow.nodes.base_node import BaseNode
......@@ -15,6 +16,7 @@ MAX_STRING_LENGTH = 1000
MAX_STRING_ARRAY_LENGTH = 30
MAX_NUMBER_ARRAY_LENGTH = 1000
class CodeNode(BaseNode):
_node_data_cls = CodeNodeData
node_type = NodeType.CODE
......@@ -78,20 +80,14 @@ class CodeNode(BaseNode):
}
}
def _run(self, variable_pool: Optional[VariablePool] = None,
run_args: Optional[dict] = None) -> NodeRunResult:
def _run(self, variable_pool: VariablePool) -> NodeRunResult:
"""
Run code
:param variable_pool: variable pool
:param run_args: run args
:return:
"""
node_data = self.node_data
node_data: CodeNodeData = cast(self._node_data_cls, node_data)
# SINGLE DEBUG NOT IMPLEMENTED YET
if variable_pool is None and run_args:
raise ValueError("Not support single step debug.")
node_data = cast(self._node_data_cls, node_data)
# Get code language
code_language = node_data.code_language
......@@ -134,7 +130,6 @@ class CodeNode(BaseNode):
Check string
:param value: value
:param variable: variable
:param max_length: max length
:return:
"""
if not isinstance(value, str):
......@@ -244,3 +239,13 @@ class CodeNode(BaseNode):
raise ValueError(f'Output type {output_config.type} is not supported.')
return transformed_result
@classmethod
def _extract_variable_selector_to_variable_mapping(cls, node_data: BaseNodeData) -> dict[list[str], str]:
"""
Extract variable selector to variable mapping
:param node_data: node data
:return:
"""
# TODO extract variable selector to variable mapping for single step debugging
return {}
api/core/workflow/nodes/direct_answer/direct_answer_node.py
View file @ 2db67c41
import time
from typing import Optional, cast
from typing import cast
from core.prompt.utils.prompt_template_parser import PromptTemplateParser
from core.workflow.entities.base_node_data_entities import BaseNodeData
from core.workflow.entities.node_entities import NodeRunResult, NodeType
from core.workflow.entities.variable_pool import ValueType, VariablePool
from core.workflow.nodes.base_node import BaseNode
......@@ -13,20 +14,15 @@ class DirectAnswerNode(BaseNode):
_node_data_cls = DirectAnswerNodeData
node_type = NodeType.DIRECT_ANSWER
def _run(self, variable_pool: Optional[VariablePool] = None,
run_args: Optional[dict] = None) -> NodeRunResult:
def _run(self, variable_pool: VariablePool) -> NodeRunResult:
"""
Run node
:param variable_pool: variable pool
:param run_args: run args
:return:
"""
node_data = self.node_data
node_data = cast(self._node_data_cls, node_data)
if variable_pool is None and run_args:
raise ValueError("Not support single step debug.")
variable_values = {}
for variable_selector in node_data.variables:
value = variable_pool.get_variable_value(
......@@ -43,7 +39,7 @@ class DirectAnswerNode(BaseNode):
# publish answer as stream
for word in answer:
self.publish_text_chunk(word)
time.sleep(0.01) # todo sleep 0.01
time.sleep(0.01)
return NodeRunResult(
status=WorkflowNodeExecutionStatus.SUCCEEDED,
......@@ -52,3 +48,12 @@ class DirectAnswerNode(BaseNode):
"answer": answer
}
)
@classmethod
def _extract_variable_selector_to_variable_mapping(cls, node_data: BaseNodeData) -> dict[list[str], str]:
"""
Extract variable selector to variable mapping
:param node_data: node data
:return:
"""
return {}
api/core/workflow/nodes/end/end_node.py
View file @ 2db67c41
from typing import Optional, cast
from typing import cast
from core.workflow.entities.base_node_data_entities import BaseNodeData
from core.workflow.entities.node_entities import NodeRunResult, NodeType
from core.workflow.entities.variable_pool import ValueType, VariablePool
from core.workflow.nodes.base_node import BaseNode
......@@ -11,19 +12,16 @@ class EndNode(BaseNode):
_node_data_cls = EndNodeData
node_type = NodeType.END
def _run(self, variable_pool: Optional[VariablePool] = None,
run_args: Optional[dict] = None) -> NodeRunResult:
def _run(self, variable_pool: VariablePool) -> NodeRunResult:
"""
Run node
:param variable_pool: variable pool
:param run_args: run args
:return:
"""
node_data = self.node_data
node_data = cast(self._node_data_cls, node_data)
outputs_config = node_data.outputs
if variable_pool is not None:
outputs = None
if outputs_config:
if outputs_config.type == EndNodeDataOutputs.OutputType.PLAIN_TEXT:
......@@ -50,11 +48,18 @@ class EndNode(BaseNode):
outputs[variable_selector.variable] = variable_value
else:
outputs = {}
else:
raise ValueError("Not support single step debug.")
return NodeRunResult(
status=WorkflowNodeExecutionStatus.SUCCEEDED,
inputs=outputs,
outputs=outputs
)
@classmethod
def _extract_variable_selector_to_variable_mapping(cls, node_data: BaseNodeData) -> dict[list[str], str]:
"""
Extract variable selector to variable mapping
:param node_data: node data
:return:
"""
return {}
api/core/workflow/nodes/llm/llm_node.py
View file @ 2db67c41
from typing import Optional, cast
from core.workflow.entities.base_node_data_entities import BaseNodeData
from core.workflow.entities.node_entities import NodeRunResult, NodeType
from core.workflow.entities.variable_pool import VariablePool
from core.workflow.nodes.base_node import BaseNode
......@@ -10,12 +11,10 @@ class LLMNode(BaseNode):
_node_data_cls = LLMNodeData
node_type = NodeType.LLM
def _run(self, variable_pool: Optional[VariablePool] = None,
run_args: Optional[dict] = None) -> NodeRunResult:
def _run(self, variable_pool: VariablePool) -> NodeRunResult:
"""
Run node
:param variable_pool: variable pool
:param run_args: run args
:return:
"""
node_data = self.node_data
......@@ -23,6 +22,17 @@ class LLMNode(BaseNode):
pass
@classmethod
def _extract_variable_selector_to_variable_mapping(cls, node_data: BaseNodeData) -> dict[list[str], str]:
"""
Extract variable selector to variable mapping
:param node_data: node data
:return:
"""
# TODO extract variable selector to variable mapping for single step debugging
return {}
@classmethod
def get_default_config(cls, filters: Optional[dict] = None) -> dict:
"""
......
api/core/workflow/nodes/start/start_node.py
View file @ 2db67c41
from typing import Optional, cast
from typing import cast
from core.app.app_config.entities import VariableEntity
from core.workflow.entities.base_node_data_entities import BaseNodeData
from core.workflow.entities.node_entities import NodeRunResult, NodeType
from core.workflow.entities.variable_pool import VariablePool
from core.workflow.nodes.base_node import BaseNode
......@@ -12,12 +13,10 @@ class StartNode(BaseNode):
_node_data_cls = StartNodeData
node_type = NodeType.START
def _run(self, variable_pool: Optional[VariablePool] = None,
run_args: Optional[dict] = None) -> NodeRunResult:
def _run(self, variable_pool: VariablePool) -> NodeRunResult:
"""
Run node
:param variable_pool: variable pool
:param run_args: run args
:return:
"""
node_data = self.node_data
......@@ -25,7 +24,7 @@ class StartNode(BaseNode):
variables = node_data.variables
# Get cleaned inputs
cleaned_inputs = self._get_cleaned_inputs(variables, run_args)
cleaned_inputs = self._get_cleaned_inputs(variables, variable_pool.user_inputs)
return NodeRunResult(
status=WorkflowNodeExecutionStatus.SUCCEEDED,
......@@ -68,3 +67,12 @@ class StartNode(BaseNode):
filtered_inputs[variable] = value.replace('\x00', '') if value else None
return filtered_inputs
@classmethod
def _extract_variable_selector_to_variable_mapping(cls, node_data: BaseNodeData) -> dict[list[str], str]:
"""
Extract variable selector to variable mapping
:param node_data: node data
:return:
"""
return {}
api/core/workflow/workflow_engine_manager.py
View file @ 2db67c41
......@@ -109,9 +109,9 @@ class WorkflowEngineManager:
workflow_run_state = WorkflowRunState(
workflow=workflow,
start_at=time.perf_counter(),
user_inputs=user_inputs,
variable_pool=VariablePool(
system_variables=system_inputs,
user_inputs=user_inputs
)
)
......@@ -292,9 +292,7 @@ class WorkflowEngineManager:
# run node, result must have inputs, process_data, outputs, execution_metadata
node_run_result = node.run(
variable_pool=workflow_run_state.variable_pool,
run_args=workflow_run_state.user_inputs
if (not predecessor_node and node.node_type == NodeType.START) else None # only on start node
variable_pool=workflow_run_state.variable_pool
)
if node_run_result.status == WorkflowNodeExecutionStatus.FAILED:
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment