feat: Support variable assign

Author: liuruibin

apps/application/flow/step_node/__init__.py

@@ -24,11 +24,14 @@
 from .speech_to_text_step_node import BaseSpeechToTextNode
 from .start_node import *
 from .text_to_speech_step_node.impl.base_text_to_speech_node import BaseTextToSpeechNode
+from .variable_assign import BaseVariableAssignNode
 
-node_list = [BaseStartStepNode, BaseChatNode, BaseSearchDatasetNode, BaseQuestionNode, BaseConditionNode, BaseReplyNode,
+node_list = [BaseStartStepNode, BaseChatNode, BaseSearchDatasetNode, BaseQuestionNode,
+             BaseConditionNode, BaseReplyNode,
              BaseFunctionNodeNode, BaseFunctionLibNodeNode, BaseRerankerNode, BaseApplicationNode,
              BaseDocumentExtractNode,
-             BaseImageUnderstandNode, BaseFormNode, BaseSpeechToTextNode, BaseTextToSpeechNode,BaseImageGenerateNode]
+             BaseImageUnderstandNode, BaseFormNode, BaseSpeechToTextNode, BaseTextToSpeechNode,
+             BaseImageGenerateNode, BaseVariableAssignNode]
 
 
 def get_node(node_type):

apps/application/flow/step_node/variable_assign/__init__.py

@@ -0,0 +1,3 @@
+# coding=utf-8
+
+from .impl import *

apps/application/flow/step_node/variable_assign/i_variable_assign_node.py

@@ -0,0 +1,27 @@
+# coding=utf-8
+
+from typing import Type
+
+from django.utils.translation import gettext_lazy as _
+from rest_framework import serializers
+
+from application.flow.i_step_node import INode, NodeResult
+from common.util.field_message import ErrMessage
+
+
+class VariableAssignNodeParamsSerializer(serializers.Serializer):
+    variable_list = serializers.ListField(required=True,
+                                          error_messages=ErrMessage.list(_("Reference Field")))
+
+
+class IVariableAssignNode(INode):
+    type = 'variable-assign-node'
+
+    def get_node_params_serializer_class(self) -> Type[serializers.Serializer]:
+        return VariableAssignNodeParamsSerializer
+
+    def _run(self):
+        return self.execute(**self.node_params_serializer.data, **self.flow_params_serializer.data)
+
+    def execute(self, variable_list, stream, **kwargs) -> NodeResult:
+        pass

apps/application/flow/step_node/variable_assign/impl/__init__.py

@@ -0,0 +1,9 @@
+# coding=utf-8
+"""
+    @project: maxkb
+    @Author：虎
+    @file： __init__.py
+    @date：2024/6/11 17:49
+    @desc:
+"""
+from .base_variable_assign_node import *

apps/application/flow/step_node/variable_assign/impl/base_variable_assign_node.py

@@ -0,0 +1,41 @@
+# coding=utf-8
+from typing import List
+
+from application.flow.i_step_node import NodeResult
+from application.flow.step_node.variable_assign.i_variable_assign_node import IVariableAssignNode
+
+
+class BaseVariableAssignNode(IVariableAssignNode):
+    def save_context(self, details, workflow_manage):
+        self.context['variable_list'] = details.get('variable_list')
+
+    def execute(self, variable_list, stream, **kwargs) -> NodeResult:
+        #
+        for variable in variable_list:
+            if 'fields' not in variable:
+                continue
+            if 'global' == variable['fields'][0]:
+                if variable['source'] == 'custom':
+                    self.workflow_manage.context[variable['fields'][1]] = variable['value']
+                else:
+                    reference = self.get_reference_content(variable['reference'])
+                    self.workflow_manage.context[variable['fields'][1]] = reference
+        # print('variable_list:', variable_list)
+
+        return NodeResult({'variable_list': variable_list}, {})
+
+    def get_reference_content(self, fields: List[str]):
+        return str(self.workflow_manage.get_reference_field(
+            fields[0],
+            fields[1:]))
+
+    def get_details(self, index: int, **kwargs):
+        return {
+            'name': self.node.properties.get('stepName'),
+            "index": index,
+            'run_time': self.context.get('run_time'),
+            'type': self.node.type,
+            'variable_list': self.context.get('variable_list'),
+            'status': self.status,
+            'err_message': self.err_message
+        }

ui/src/enums/workflow.ts

@@ -12,6 +12,7 @@ export enum WorkflowType {
   Application = 'application-node',
   DocumentExtractNode = 'document-extract-node',
   ImageUnderstandNode = 'image-understand-node',
+  VariableAssignNode = 'variable-assign-node',
   FormNode = 'form-node',
   TextToSpeechNode = 'text-to-speech-node',
   SpeechToTextNode = 'speech-to-text-node',

ui/src/locales/lang/en-US/views/application-workflow.ts

@@ -212,6 +212,22 @@ export default {
         requiredMessage: 'Please select an image'
       }
     },
+    variableAssignNode: {
+      label: 'Variable assign',
+      text: 'Update the value of the global variable',
+      variable: {
+        label: 'select variable',
+        requiredMessage: 'Please select a variable'
+      },
+      value: {
+        label: 'Set value',
+      },
+      source: {
+        reference: 'Reference Variable',
+        custom: 'Custom'
+      },
+      addVariable: 'Add Variable',
+    },
     imageGenerateNode: {
       label: 'Image Generation',
       text: 'Generate images based on provided text content',

ui/src/locales/lang/zh-CN/views/application-workflow.ts

@@ -211,6 +211,22 @@ export default {
         requiredMessage: '请选择图片'
       }
     },
+    variableAssignNode: {
+      label: '变量赋值',
+      text: '更新全局变量的值',
+      variable: {
+        label: '选择变量',
+        requiredMessage: '请选择变量'
+      },
+      value: {
+        label: '赋值',
+      },
+      source: {
+        reference: '引用变量',
+        custom: '自定义'
+      },
+      addVariable: '添加变量'
+    },
     imageGenerateNode: {
       label: '图片生成',
       text: '根据提供的文本内容生成图片',

ui/src/locales/lang/zh-Hant/views/application-workflow.ts

@@ -210,6 +210,22 @@ export default {
         requiredMessage: '請選擇圖片'
       }
     },
+    variableAssignNode: {
+      label: '變數賦值',
+      text: '更新全域變數的值',
+      variable: {
+        label: '選擇變數',
+        requiredMessage: '請選擇變數'
+      },
+      value: {
+        label: '設置值',
+      },
+      source: {
+        reference: '引用變數',
+        custom: '自定義'
+      },
+      addVariable: '添加變數',
+    },
     imageGenerateNode: {
       label: '圖片生成',
       text: '根據提供的文本內容生成圖片',

ui/src/workflow/common/data.ts

@@ -239,6 +239,19 @@ export const imageUnderstandNode = {
   }
 }
 
+export const variableAssignNode = {
+  type: WorkflowType.VariableAssignNode,
+  text: t('views.applicationWorkflow.nodes.variableAssignNode.text'),
+  label: t('views.applicationWorkflow.nodes.variableAssignNode.label'),
+  height: 252,
+  properties: {
+    stepName: t('views.applicationWorkflow.nodes.variableAssignNode.label'),
+    config: {
+
+    }
+  }
+}
+
 export const imageGenerateNode = {
   type: WorkflowType.ImageGenerateNode,
   text: t('views.applicationWorkflow.nodes.imageGenerateNode.text'),
@@ -307,7 +320,8 @@ export const menuNodes = [
   questionNode,
   documentExtractNode,
   speechToTextNode,
-  textToSpeechNode
+  textToSpeechNode,
+  variableAssignNode
 ]
 
 /**
@@ -400,7 +414,8 @@ export const nodeDict: any = {
   [WorkflowType.ImageUnderstandNode]: imageUnderstandNode,
   [WorkflowType.TextToSpeechNode]: textToSpeechNode,
   [WorkflowType.SpeechToTextNode]: speechToTextNode,
-  [WorkflowType.ImageGenerateNode]: imageGenerateNode
+  [WorkflowType.ImageGenerateNode]: imageGenerateNode,
+  [WorkflowType.VariableAssignNode]: variableAssignNode,
 }
 export function isWorkFlow(type: string | undefined) {
   return type === 'WORK_FLOW'