Merge branch 'refs/heads/main' into feat/workflow-parallel-support

# Conflicts:
#	api/core/app/apps/advanced_chat/app_generator.py
#	api/core/app/apps/advanced_chat/generate_task_pipeline.py
#	api/core/app/apps/workflow/app_runner.py
#	api/core/app/apps/workflow/generate_task_pipeline.py
#	api/core/app/task_pipeline/workflow_cycle_state_manager.py
#	api/core/workflow/entities/variable_pool.py
#	api/core/workflow/nodes/code/code_node.py
#	api/core/workflow/nodes/llm/llm_node.py
#	api/core/workflow/nodes/start/start_node.py
#	api/core/workflow/nodes/variable_assigner/__init__.py
#	api/tests/integration_tests/workflow/nodes/test_llm.py
#	api/tests/integration_tests/workflow/nodes/test_parameter_extractor.py
#	api/tests/unit_tests/core/workflow/nodes/test_answer.py
#	api/tests/unit_tests/core/workflow/nodes/test_if_else.py
#	api/tests/unit_tests/core/workflow/nodes/test_variable_assigner.py

api/core/app/apps/advanced_chat/app_generator.py

@@ -23,6 +23,8 @@
 from core.file.message_file_parser import MessageFileParser
 from core.model_runtime.errors.invoke import InvokeAuthorizationError, InvokeError
 from core.ops.ops_trace_manager import TraceQueueManager
+from core.workflow.entities.variable_pool import VariablePool
+from core.workflow.enums import SystemVariableKey
 from extensions.ext_database import db
 from models.account import Account
 from models.model import App, Conversation, EndUser, Message
@@ -67,8 +69,9 @@ def generate(
 
         # get conversation
         conversation = None
-        if args.get('conversation_id'):
-            conversation = self._get_conversation_by_user(app_model, args.get('conversation_id', ''), user)
+        conversation_id = args.get('conversation_id')
+        if conversation_id:
+            conversation = self._get_conversation_by_user(app_model=app_model, conversation_id=conversation_id, user=user)
 
         # parse files
         files = args['files'] if args.get('files') else []
@@ -225,6 +228,62 @@ def _generate(self, *,
             message_id=message.id
         )
 
+        # Init conversation variables
+        stmt = select(ConversationVariable).where(
+            ConversationVariable.app_id == conversation.app_id, ConversationVariable.conversation_id == conversation.id
+        )
+        with Session(db.engine) as session:
+            conversation_variables = session.scalars(stmt).all()
+            if not conversation_variables:
+                # Create conversation variables if they don't exist.
+                conversation_variables = [
+                    ConversationVariable.from_variable(
+                        app_id=conversation.app_id, conversation_id=conversation.id, variable=variable
+                    )
+                    for variable in workflow.conversation_variables
+                ]
+                session.add_all(conversation_variables)
+            # Convert database entities to variables.
+            conversation_variables = [item.to_variable() for item in conversation_variables]
+
+            session.commit()
+
+            # Increment dialogue count.
+            conversation.dialogue_count += 1
+
+            conversation_id = conversation.id
+            conversation_dialogue_count = conversation.dialogue_count
+            db.session.commit()
+            db.session.refresh(conversation)
+
+        inputs = application_generate_entity.inputs
+        query = application_generate_entity.query
+        files = application_generate_entity.files
+
+        user_id = None
+        if application_generate_entity.invoke_from in [InvokeFrom.WEB_APP, InvokeFrom.SERVICE_API]:
+            end_user = db.session.query(EndUser).filter(EndUser.id == application_generate_entity.user_id).first()
+            if end_user:
+                user_id = end_user.session_id
+        else:
+            user_id = application_generate_entity.user_id
+
+        # Create a variable pool.
+        system_inputs = {
+            SystemVariableKey.QUERY: query,
+            SystemVariableKey.FILES: files,
+            SystemVariableKey.CONVERSATION_ID: conversation_id,
+            SystemVariableKey.USER_ID: user_id,
+            SystemVariableKey.DIALOGUE_COUNT: conversation_dialogue_count,
+        }
+        variable_pool = VariablePool(
+            system_variables=system_inputs,
+            user_inputs=inputs,
+            environment_variables=workflow.environment_variables,
+            conversation_variables=conversation_variables,
+        )
+        contexts.workflow_variable_pool.set(variable_pool)
+
         # new thread
         worker_thread = threading.Thread(target=self._generate_worker, kwargs={
             'flask_app': current_app._get_current_object(), # type: ignore
@@ -296,7 +355,7 @@ def _generate_worker(self, flask_app: Flask,
                 logger.exception("Validation Error when generating")
                 queue_manager.publish_error(e, PublishFrom.APPLICATION_MANAGER)
             except (ValueError, InvokeError) as e:
-                if os.environ.get("DEBUG") and os.environ.get("DEBUG", "false").lower() == 'true':
+                if os.environ.get("DEBUG", "false").lower() == 'true':
                     logger.exception("Error when generating")
                 queue_manager.publish_error(e, PublishFrom.APPLICATION_MANAGER)
             except Exception as e:

api/core/app/apps/advanced_chat/generate_task_pipeline.py

@@ -47,7 +47,7 @@
 from core.app.task_pipeline.workflow_cycle_manage import WorkflowCycleManage
 from core.model_runtime.utils.encoders import jsonable_encoder
 from core.ops.ops_trace_manager import TraceQueueManager
-from core.workflow.enums import SystemVariable
+from core.workflow.enums import SystemVariableKey
 from core.workflow.graph_engine.entities.graph_runtime_state import GraphRuntimeState
 from events.message_event import message_was_created
 from extensions.ext_database import db
@@ -69,7 +69,7 @@ class AdvancedChatAppGenerateTaskPipeline(BasedGenerateTaskPipeline, WorkflowCyc
     _application_generate_entity: AdvancedChatAppGenerateEntity
     _workflow: Workflow
     _user: Union[Account, EndUser]
-    _workflow_system_variables: dict[SystemVariable, Any]
+    _workflow_system_variables: dict[SystemVariableKey, Any]
 
     def __init__(
             self,
@@ -312,7 +312,7 @@ def _process_stream_response(
             elif isinstance(event, QueueParallelBranchRunStartedEvent):
                 if not workflow_run:
                     raise Exception('Workflow run not initialized.')
-                
+
                 yield self._workflow_parallel_branch_start_to_stream_response(
                     task_id=self._application_generate_entity.task_id,
                     workflow_run=workflow_run,
@@ -321,7 +321,7 @@ def _process_stream_response(
             elif isinstance(event, QueueParallelBranchRunSucceededEvent | QueueParallelBranchRunFailedEvent):
                 if not workflow_run:
                     raise Exception('Workflow run not initialized.')
-                
+
                 yield self._workflow_parallel_branch_finished_to_stream_response(
                     task_id=self._application_generate_entity.task_id,
                     workflow_run=workflow_run,

api/core/app/apps/workflow/app_runner.py

@@ -13,7 +13,7 @@
 from core.workflow.callbacks.base_workflow_callback import WorkflowCallback
 from core.workflow.entities.node_entities import UserFrom
 from core.workflow.entities.variable_pool import VariablePool
-from core.workflow.enums import SystemVariable
+from core.workflow.enums import SystemVariableKey
 from core.workflow.workflow_entry import WorkflowEntry
 from extensions.ext_database import db
 from models.model import App, EndUser
@@ -79,14 +79,14 @@ def run(self) -> None:
                 user_inputs=self.application_generate_entity.single_iteration_run.inputs
             )
         else:
-        
+
             inputs = self.application_generate_entity.inputs
             files = self.application_generate_entity.files
-        
+
             # Create a variable pool.
             system_inputs = {
-                SystemVariable.FILES: files,
-                SystemVariable.USER_ID: user_id,
+                SystemVariableKey.FILES: files,
+                SystemVariableKey.USER_ID: user_id,
             }
 
             variable_pool = VariablePool(
@@ -98,7 +98,7 @@ def run(self) -> None:
 
             # init graph
             graph = self._init_graph(graph_config=workflow.graph_dict)
-        
+
         # RUN WORKFLOW
         workflow_entry = WorkflowEntry(
             tenant_id=workflow.tenant_id,

api/core/app/apps/workflow/generate_task_pipeline.py

@@ -41,7 +41,9 @@
 from core.app.task_pipeline.based_generate_task_pipeline import BasedGenerateTaskPipeline
 from core.app.task_pipeline.workflow_cycle_manage import WorkflowCycleManage
 from core.ops.ops_trace_manager import TraceQueueManager
-from core.workflow.enums import SystemVariable
+from core.workflow.entities.node_entities import NodeType
+from core.workflow.enums import SystemVariableKey
+from core.workflow.nodes.end.end_node import EndNode
 from extensions.ext_database import db
 from models.account import Account
 from models.model import EndUser
@@ -64,7 +66,7 @@ class WorkflowAppGenerateTaskPipeline(BasedGenerateTaskPipeline, WorkflowCycleMa
     _user: Union[Account, EndUser]
     _task_state: WorkflowTaskState
     _application_generate_entity: WorkflowAppGenerateEntity
-    _workflow_system_variables: dict[SystemVariable, Any]
+    _workflow_system_variables: dict[SystemVariableKey, Any]
 
     def __init__(self, application_generate_entity: WorkflowAppGenerateEntity,
                  workflow: Workflow,

api/core/workflow/entities/variable_pool.py

@@ -32,7 +32,7 @@ class VariablePool(BaseModel):
         description='User inputs',
     )
 
-    system_variables: Mapping[SystemVariable, Any] = Field(
+    system_variables: Mapping[SystemVariableKey, Any] = Field(
         description='System variables',
     )
 

api/core/workflow/nodes/code/code_node.py

@@ -316,8 +316,8 @@ def _transform_result(self, result: dict, output_schema: Optional[dict[str, Code
 
     @classmethod
     def _extract_variable_selector_to_variable_mapping(
-        cls, 
-        graph_config: Mapping[str, Any], 
+        cls,
+        graph_config: Mapping[str, Any],
         node_id: str,
         node_data: CodeNodeData
     ) -> Mapping[str, Sequence[str]]:

api/core/workflow/nodes/llm/llm_node.py

@@ -25,7 +25,7 @@
 from core.prompt.utils.prompt_message_util import PromptMessageUtil
 from core.workflow.entities.node_entities import NodeRunMetadataKey, NodeRunResult, NodeType
 from core.workflow.entities.variable_pool import VariablePool
-from core.workflow.enums import SystemVariable
+from core.workflow.enums import SystemVariableKey
 from core.workflow.graph_engine.entities.event import InNodeEvent
 from core.workflow.nodes.base_node import BaseNode
 from core.workflow.nodes.event import RunCompletedEvent, RunEvent, RunRetrieverResourceEvent, RunStreamChunkEvent
@@ -370,7 +370,7 @@ def _fetch_inputs(self, node_data: LLMNodeData, variable_pool: VariablePool) ->
 
                 inputs[variable_selector.variable] = variable_value
 
-        return inputs  # type: ignore
+        return inputs
 
     def _fetch_files(self, node_data: LLMNodeData, variable_pool: VariablePool) -> list["FileVar"]:
         """

api/core/workflow/nodes/start/start_node.py

@@ -3,6 +3,7 @@
 from typing import Any
 
 from core.workflow.entities.node_entities import NodeRunResult, NodeType
+from core.workflow.entities.variable_pool import SYSTEM_VARIABLE_NODE_ID, VariablePool
 from core.workflow.nodes.base_node import BaseNode
 from core.workflow.nodes.start.entities import StartNodeData
 from models.workflow import WorkflowNodeExecutionStatus
@@ -17,11 +18,11 @@ def _run(self) -> NodeRunResult:
         Run node
         :return:
         """
-        # Get cleaned inputs
-        cleaned_inputs = dict(self.graph_runtime_state.variable_pool.user_inputs)
+        node_inputs = dict(self.graph_runtime_state.variable_pool.user_inputs)
+        system_inputs = self.graph_runtime_state.variable_pool.system_variables
 
-        for var in self.graph_runtime_state.variable_pool.system_variables:
-            cleaned_inputs['sys.' + var.value] = self.graph_runtime_state.variable_pool.system_variables[var]
+        for var in system_inputs:
+            node_inputs[SYSTEM_VARIABLE_NODE_ID + '.' + var] = system_inputs[var]
 
         return NodeRunResult(
             status=WorkflowNodeExecutionStatus.SUCCEEDED,
@@ -31,8 +32,8 @@ def _run(self) -> NodeRunResult:
 
     @classmethod
     def _extract_variable_selector_to_variable_mapping(
-        cls, 
-        graph_config: Mapping[str, Any], 
+        cls,
+        graph_config: Mapping[str, Any],
         node_id: str,
         node_data: StartNodeData
     ) -> Mapping[str, Sequence[str]]:

api/core/workflow/nodes/variable_assigner/__init__.py

@@ -1,108 +1,8 @@
 from .node import VariableAssignerNode
 from .node_data import VariableAssignerData, WriteMode
 
-from sqlalchemy import select
-from sqlalchemy.orm import Session
-
-from core.app.segments import SegmentType, Variable, factory
-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
-from extensions.ext_database import db
-from models import ConversationVariable, WorkflowNodeExecutionStatus
-
-
-class VariableAssignerNodeError(Exception):
-    pass
-
-
-class WriteMode(str, Enum):
-    OVER_WRITE = 'over-write'
-    APPEND = 'append'
-    CLEAR = 'clear'
-
-
-class VariableAssignerData(BaseNodeData):
-    title: str = 'Variable Assigner'
-    desc: Optional[str] = 'Assign a value to a variable'
-    assigned_variable_selector: Sequence[str]
-    write_mode: WriteMode
-    input_variable_selector: Sequence[str]
-
-
-class VariableAssignerNode(BaseNode):
-    _node_data_cls: type[BaseNodeData] = VariableAssignerData
-    _node_type: NodeType = NodeType.CONVERSATION_VARIABLE_ASSIGNER
-
-    def _run(self) -> NodeRunResult:
-        data = cast(VariableAssignerData, self.node_data)
-
-        # Should be String, Number, Object, ArrayString, ArrayNumber, ArrayObject
-        original_variable = self.graph_runtime_state.variable_pool.get(data.assigned_variable_selector)
-        if not isinstance(original_variable, Variable):
-            raise VariableAssignerNodeError('assigned variable not found')
-
-        match data.write_mode:
-            case WriteMode.OVER_WRITE:
-                income_value = self.graph_runtime_state.variable_pool.get(data.input_variable_selector)
-                if not income_value:
-                    raise VariableAssignerNodeError('input value not found')
-                updated_variable = original_variable.model_copy(update={'value': income_value.value})
-
-            case WriteMode.APPEND:
-                income_value = self.graph_runtime_state.variable_pool.get(data.input_variable_selector)
-                if not income_value:
-                    raise VariableAssignerNodeError('input value not found')
-                updated_value = original_variable.value + [income_value.value]
-                updated_variable = original_variable.model_copy(update={'value': updated_value})
-
-            case WriteMode.CLEAR:
-                income_value = get_zero_value(original_variable.value_type)
-                updated_variable = original_variable.model_copy(update={'value': income_value.to_object()})
-
-            case _:
-                raise VariableAssignerNodeError(f'unsupported write mode: {data.write_mode}')
-
-        # Over write the variable.
-        self.graph_runtime_state.variable_pool.add(data.assigned_variable_selector, updated_variable)
-
-        # Update conversation variable.
-        # TODO: Find a better way to use the database.
-        conversation_id = self.graph_runtime_state.variable_pool.get(['sys', 'conversation_id'])
-        if not conversation_id:
-            raise VariableAssignerNodeError('conversation_id not found')
-        update_conversation_variable(conversation_id=conversation_id.text, variable=updated_variable)
-
-        return NodeRunResult(
-            status=WorkflowNodeExecutionStatus.SUCCEEDED,
-            inputs={
-                'value': income_value.to_object(),
-            },
-        )
-
-
-def update_conversation_variable(conversation_id: str, variable: Variable):
-    stmt = select(ConversationVariable).where(
-        ConversationVariable.id == variable.id, ConversationVariable.conversation_id == conversation_id
-    )
-    with Session(db.engine) as session:
-        row = session.scalar(stmt)
-        if not row:
-            raise VariableAssignerNodeError('conversation variable not found in the database')
-        row.data = variable.model_dump_json()
-        session.commit()
-
-
-def get_zero_value(t: SegmentType):
-    match t:
-        case SegmentType.ARRAY_OBJECT | SegmentType.ARRAY_STRING | SegmentType.ARRAY_NUMBER:
-            return factory.build_segment([])
-        case SegmentType.OBJECT:
-            return factory.build_segment({})
-        case SegmentType.STRING:
-            return factory.build_segment('')
-        case SegmentType.NUMBER:
-            return factory.build_segment(0)
-        case _:
-            raise VariableAssignerNodeError(f'unsupported variable type: {t}')
+__all__ = [
+    'VariableAssignerNode',
+    'VariableAssignerData',
+    'WriteMode',
+]

api/tests/integration_tests/workflow/nodes/test_llm.py

@@ -12,7 +12,7 @@
 from core.model_runtime.model_providers import ModelProviderFactory
 from core.workflow.entities.node_entities import UserFrom
 from core.workflow.entities.variable_pool import VariablePool
-from core.workflow.enums import SystemVariable
+from core.workflow.enums import SystemVariableKey
 from core.workflow.nodes.llm.llm_node import LLMNode
 from extensions.ext_database import db
 from models.provider import ProviderType

api/tests/integration_tests/workflow/nodes/test_parameter_extractor.py

@@ -13,7 +13,7 @@
 from core.model_runtime.model_providers.model_provider_factory import ModelProviderFactory
 from core.workflow.entities.node_entities import UserFrom
 from core.workflow.entities.variable_pool import VariablePool
-from core.workflow.enums import SystemVariable
+from core.workflow.enums import SystemVariableKey
 from core.workflow.nodes.parameter_extractor.parameter_extractor_node import ParameterExtractorNode
 from extensions.ext_database import db
 from models.provider import ProviderType