From 5c246285da6c97f015354ae535b0875d9f4d2085 Mon Sep 17 00:00:00 2001
From: Joel <iamjoel007@gmail.com>
Date: Thu, 14 Mar 2024 17:55:12 +0800
Subject: [PATCH] feat: support node type filter

---
 .../nodes/_base/components/variable/utils.ts  | 12 ++++--
 .../_base/components/variable/var-list.tsx    |  8 +++-
 .../variable/var-reference-picker.tsx         | 11 +++--
 .../variable/var-reference-popup.tsx          | 40 +++++++++++--------
 .../components/workflow/nodes/end/panel.tsx   |  2 +-
 .../workflow/nodes/end/use-config.ts          |  4 +-
 .../components/var-list/index.tsx             | 13 ++++--
 .../workflow/nodes/variable-assigner/node.tsx |  2 +-
 .../nodes/variable-assigner/panel.tsx         | 11 +++--
 web/i18n/en-US/workflow.ts                    |  1 +
 web/i18n/zh-Hans/workflow.ts                  |  1 +
 11 files changed, 70 insertions(+), 35 deletions(-)

diff --git a/web/app/components/workflow/nodes/_base/components/variable/utils.ts b/web/app/components/workflow/nodes/_base/components/variable/utils.ts
index b4e84eb176..bc0249f22f 100644
--- a/web/app/components/workflow/nodes/_base/components/variable/utils.ts
+++ b/web/app/components/workflow/nodes/_base/components/variable/utils.ts
@@ -21,7 +21,7 @@ const inputVarTypeToVarType = (type: InputVarType): VarType => {
   return VarType.string
 }
 
-const formatItem = (item: any, isChatMode: boolean): NodeOutPutVar => {
+const formatItem = (item: any, isChatMode: boolean, varType?: VarType): NodeOutPutVar => {
   const { id, data } = item
   const res: NodeOutPutVar = {
     nodeId: id,
@@ -109,9 +109,15 @@ const formatItem = (item: any, isChatMode: boolean): NodeOutPutVar => {
       break
     }
   }
+  if (varType)
+    res.vars = res.vars.filter(v => v.type === varType)
 
   return res
 }
-export const toNodeOutputVars = (nodes: any[], isChatMode: boolean): NodeOutPutVar[] => {
-  return nodes.filter(node => SUPPORT_OUTPUT_VARS_NODE.includes(node.data.type)).map(node => formatItem(node, isChatMode))
+export const toNodeOutputVars = (nodes: any[], isChatMode: boolean, varType?: VarType): NodeOutPutVar[] => {
+  const res = nodes
+    .filter(node => SUPPORT_OUTPUT_VARS_NODE.includes(node.data.type))
+    .map(node => formatItem(node, isChatMode, varType))
+    .filter(item => item.vars.length > 0)
+  return res
 }
diff --git a/web/app/components/workflow/nodes/_base/components/variable/var-list.tsx b/web/app/components/workflow/nodes/_base/components/variable/var-list.tsx
index a96090133a..71a360c4ca 100644
--- a/web/app/components/workflow/nodes/_base/components/variable/var-list.tsx
+++ b/web/app/components/workflow/nodes/_base/components/variable/var-list.tsx
@@ -4,7 +4,7 @@ import React, { useCallback } from 'react'
 import produce from 'immer'
 import RemoveButton from '../remove-button'
 import VarReferencePicker from './var-reference-picker'
-import { type ValueSelector, type Variable } from '@/app/components/workflow/types'
+import { type ValueSelector, type VarType, type Variable } from '@/app/components/workflow/types'
 import { VarType as VarKindType } from '@/app/components/workflow/nodes/tool/types'
 
 type Props = {
@@ -13,6 +13,8 @@ type Props = {
   list: Variable[]
   onChange: (list: Variable[]) => void
   isSupportConstantValue?: boolean
+  onlyLeafNodeVar?: boolean
+  onlyVarType?: VarType
 }
 
 const VarList: FC<Props> = ({
@@ -21,6 +23,8 @@ const VarList: FC<Props> = ({
   list,
   onChange,
   isSupportConstantValue,
+  onlyLeafNodeVar,
+  onlyVarType,
 }) => {
   const handleVarNameChange = useCallback((index: number) => {
     return (e: React.ChangeEvent<HTMLInputElement>) => {
@@ -80,6 +84,8 @@ const VarList: FC<Props> = ({
             isSupportConstantValue={isSupportConstantValue}
             onChange={handleVarReferenceChange(index)}
             defaultVarKindType={item.variable_type}
+            onlyLeafNodeVar={onlyLeafNodeVar}
+            onlyVarType={onlyVarType}
           />
           <RemoveButton
             className='!p-2 !bg-gray-100 hover:!bg-gray-200'
diff --git a/web/app/components/workflow/nodes/_base/components/variable/var-reference-picker.tsx b/web/app/components/workflow/nodes/_base/components/variable/var-reference-picker.tsx
index c0b9236295..958269b1b8 100644
--- a/web/app/components/workflow/nodes/_base/components/variable/var-reference-picker.tsx
+++ b/web/app/components/workflow/nodes/_base/components/variable/var-reference-picker.tsx
@@ -33,6 +33,8 @@ type Props = {
   onChange: (value: ValueSelector | string, varKindType: VarKindType) => void
   isSupportConstantValue?: boolean
   defaultVarKindType?: VarKindType
+  onlyLeafNodeVar?: boolean
+  onlyVarType?: VarType
 }
 
 export const getNodeInfoById = (nodes: any, id: string) => {
@@ -52,13 +54,15 @@ const VarReferencePicker: FC<Props> = ({
   onChange,
   isSupportConstantValue,
   defaultVarKindType = VarKindType.static,
+  onlyLeafNodeVar,
+  onlyVarType,
 }) => {
   const isChatMode = useIsChatMode()
   const [varKindType, setVarKindType] = useState<VarKindType>(defaultVarKindType)
   const isConstant = isSupportConstantValue && varKindType === VarKindType.static
   const { getTreeLeafNodes, getBeforeNodesInSameBranch } = useWorkflow()
-  const availableNodes = getBeforeNodesInSameBranch(nodeId)
-  const outputVars = toNodeOutputVars(availableNodes, isChatMode)
+  const availableNodes = onlyLeafNodeVar ? getTreeLeafNodes() : getBeforeNodesInSameBranch(nodeId)
+  const outputVars = toNodeOutputVars(availableNodes, isChatMode, onlyVarType)
   const [open, setOpen] = useState(false)
   const hasValue = !isConstant && value.length > 0
   const outputVarNodeId = hasValue ? value[0] : ''
@@ -128,12 +132,13 @@ const VarReferencePicker: FC<Props> = ({
         onOpenChange={setOpen}
         placement='bottom-start'
       >
-        <PortalToFollowElemTrigger onClick={() => !isConstant && setOpen(!open)} className='!flex'>
+        <PortalToFollowElemTrigger onClick={() => !isConstant ? setOpen(!open) : setControlFocus(Date.now())} className='!flex'>
           <div className={cn((open || isFocus) && 'border border-gray-300', 'flex items-center w-full h-8 p-1 rounded-lg bg-gray-100')}>
             {isSupportConstantValue
               ? <div onClick={(e) => {
                 e.stopPropagation()
                 setOpen(false)
+                setControlFocus(Date.now())
               }} className='mr-1 flex items-center space-x-1'>
                 <TypeSelector
                   noLeft
diff --git a/web/app/components/workflow/nodes/_base/components/variable/var-reference-popup.tsx b/web/app/components/workflow/nodes/_base/components/variable/var-reference-popup.tsx
index fdbe9e788e..cc995c74dd 100644
--- a/web/app/components/workflow/nodes/_base/components/variable/var-reference-popup.tsx
+++ b/web/app/components/workflow/nodes/_base/components/variable/var-reference-popup.tsx
@@ -3,6 +3,7 @@ import type { FC } from 'react'
 import React, { useRef } from 'react'
 import { useHover } from 'ahooks'
 import cn from 'classnames'
+import { useTranslation } from 'react-i18next'
 import { type NodeOutPutVar, type ValueSelector, type Var, VarType } from '@/app/components/workflow/types'
 import { Variable02 } from '@/app/components/base/icons/src/vender/solid/development'
 import { ChevronRight } from '@/app/components/base/icons/src/vender/line/arrows'
@@ -90,7 +91,8 @@ const ObjectChildren: FC<ObjectChildrenProps> = ({
     }}>
       <div className='flex items-center h-[22px] px-3 text-xs font-normal text-gray-700'><span className='text-gray-500'>{title}.</span>{currObjPath.join('.')}</div>
       {
-        data?.map((v, i) => (
+        (data && data.length > 0)
+        && data.map((v, i) => (
           <Item
             key={i}
             nodeId={nodeId}
@@ -115,26 +117,30 @@ const VarReferencePopup: FC<Props> = ({
   onChange,
   itemWidth,
 }) => {
+  const { t } = useTranslation()
+
   return (
     <div className='p-1 bg-white rounded-lg border border-gray-200 shadow-lg space-y-1' style={{
       width: itemWidth || 228,
     }}>
-      {vars.map((item, i) => (
-        <div key={i}>
-          <div className='flex items-center h-[22px] px-3 text-xs font-medium text-gray-500 uppercase'>{item.title}</div>
-          {item.vars.map((v, j) => (
-            <Item
-              key={j}
-              title={item.title}
-              nodeId={item.nodeId}
-              objPath={[]}
-              itemData={v}
-              onChange={onChange}
-              itemWidth={itemWidth}
-            />
-          ))}
-        </div>
-      ))}
+      {vars.length > 0
+        ? vars.map((item, i) => (
+          <div key={i}>
+            <div className='flex items-center h-[22px] px-3 text-xs font-medium text-gray-500 uppercase'>{item.title}</div>
+            {item.vars.map((v, j) => (
+              <Item
+                key={j}
+                title={item.title}
+                nodeId={item.nodeId}
+                objPath={[]}
+                itemData={v}
+                onChange={onChange}
+                itemWidth={itemWidth}
+              />
+            ))}
+          </div>
+        ))
+        : <div className='pl-3 leading-[18px] text-xs font-medium text-gray-500 uppercase'>{t('workflow.common.noVar')}</div>}
     </div>
   )
 }
diff --git a/web/app/components/workflow/nodes/end/panel.tsx b/web/app/components/workflow/nodes/end/panel.tsx
index a1bbcc9ce7..1855dc8c31 100644
--- a/web/app/components/workflow/nodes/end/panel.tsx
+++ b/web/app/components/workflow/nodes/end/panel.tsx
@@ -9,7 +9,7 @@ import { EndVarType } from './types'
 import VarList from '@/app/components/workflow/nodes/_base/components/variable/var-list'
 import Field from '@/app/components/workflow/nodes/_base/components/field'
 import AddButton from '@/app/components/base/button/add-button'
-import type { NodePanelProps } from '@/app/components/workflow/types'
+import { type NodePanelProps } from '@/app/components/workflow/types'
 
 const i18nPrefix = 'workflow.nodes.end'
 
diff --git a/web/app/components/workflow/nodes/end/use-config.ts b/web/app/components/workflow/nodes/end/use-config.ts
index 7992c619ea..6a5abb800a 100644
--- a/web/app/components/workflow/nodes/end/use-config.ts
+++ b/web/app/components/workflow/nodes/end/use-config.ts
@@ -13,9 +13,9 @@ const useConfig = (id: string, payload: EndNodeType) => {
     setInputs(newInputs)
   }, [inputs, setInputs])
 
-  const handelPlainTextSelectorChange = useCallback((newList: string[]) => {
+  const handelPlainTextSelectorChange = useCallback((newList: string[] | string) => {
     const newInputs = produce(inputs, (draft: any) => {
-      draft.outputs.plain_text_selector = newList
+      draft.outputs.plain_text_selector = newList as string[]
     })
     setInputs(newInputs)
   }
diff --git a/web/app/components/workflow/nodes/variable-assigner/components/var-list/index.tsx b/web/app/components/workflow/nodes/variable-assigner/components/var-list/index.tsx
index ed7bf80178..21b31c8d03 100644
--- a/web/app/components/workflow/nodes/variable-assigner/components/var-list/index.tsx
+++ b/web/app/components/workflow/nodes/variable-assigner/components/var-list/index.tsx
@@ -5,13 +5,15 @@ import React, { useCallback } from 'react'
 import produce from 'immer'
 import RemoveButton from '../../../_base/components/remove-button'
 import VarReferencePicker from '@/app/components/workflow/nodes/_base/components/variable/var-reference-picker'
-import type { ValueSelector } from '@/app/components/workflow/types'
+import type { ValueSelector, VarType } from '@/app/components/workflow/types'
 
 type Props = {
   readonly: boolean
   nodeId: string
   list: ValueSelector[]
   onChange: (list: ValueSelector[]) => void
+  onlyLeafNodeVar?: boolean
+  onlyVarType?: VarType
 }
 
 const VarList: FC<Props> = ({
@@ -19,12 +21,14 @@ const VarList: FC<Props> = ({
   nodeId,
   list,
   onChange,
+  onlyLeafNodeVar,
+  onlyVarType,
 }) => {
   const { t } = useTranslation()
   const handleVarReferenceChange = useCallback((index: number) => {
-    return (value: ValueSelector) => {
+    return (value: ValueSelector | string) => {
       const newList = produce(list, (draft) => {
-        draft[index] = value
+        draft[index] = value as ValueSelector
       })
       onChange(newList)
     }
@@ -58,6 +62,9 @@ const VarList: FC<Props> = ({
             className='grow'
             value={item}
             onChange={handleVarReferenceChange(index)}
+            onlyLeafNodeVar={onlyLeafNodeVar}
+            onlyVarType={onlyVarType}
+            width={350}
           />
           <RemoveButton
             className='!p-2 !bg-gray-100 hover:!bg-gray-200'
diff --git a/web/app/components/workflow/nodes/variable-assigner/node.tsx b/web/app/components/workflow/nodes/variable-assigner/node.tsx
index d6687ab08c..4166c81b4a 100644
--- a/web/app/components/workflow/nodes/variable-assigner/node.tsx
+++ b/web/app/components/workflow/nodes/variable-assigner/node.tsx
@@ -34,7 +34,7 @@ const Node: FC<NodeProps<VariableAssignerNodeType>> = (props) => {
         <>
           <div className='space-y-0.5'>
             {variables.map((item, index) => {
-              const node = getNodeInfoById(item[0])
+              const node = getNodeInfoById([], item[0]) // TODO: can not get all nodes
               const varName = item[item.length - 1]
               return (
                 <div key={index} className='relative flex items-center h-6 bg-gray-100 rounded-md  px-1 text-xs font-normal text-gray-700' >
diff --git a/web/app/components/workflow/nodes/variable-assigner/panel.tsx b/web/app/components/workflow/nodes/variable-assigner/panel.tsx
index a7bc4c3456..cd831ab9a6 100644
--- a/web/app/components/workflow/nodes/variable-assigner/panel.tsx
+++ b/web/app/components/workflow/nodes/variable-assigner/panel.tsx
@@ -9,6 +9,7 @@ import Selector from '@/app/components/workflow/nodes/_base/components/selector'
 import AddButton from '@/app/components/base/button/add-button'
 import { ChevronDown } from '@/app/components/base/icons/src/vender/line/arrows'
 import type { NodePanelProps } from '@/app/components/workflow/types'
+import { VarType } from '@/app/components/workflow/types'
 
 const i18nPrefix = 'workflow.nodes.variableAssigner'
 const Panel: FC<NodePanelProps<VariableAssignerNodeType>> = ({
@@ -26,10 +27,10 @@ const Panel: FC<NodePanelProps<VariableAssignerNodeType>> = ({
   } = useConfig(id, data)
 
   const typeOptions = [
-    { label: t(`${i18nPrefix}.type.string`), value: 'string' },
-    { label: t(`${i18nPrefix}.type.number`), value: 'number' },
-    { label: t(`${i18nPrefix}.type.object`), value: 'object' },
-    { label: t(`${i18nPrefix}.type.array`), value: 'array' },
+    { label: t(`${i18nPrefix}.type.string`), value: VarType.string },
+    { label: t(`${i18nPrefix}.type.number`), value: VarType.number },
+    { label: t(`${i18nPrefix}.type.object`), value: VarType.object },
+    { label: t(`${i18nPrefix}.type.array`), value: VarType.array },
   ]
 
   return (
@@ -64,6 +65,8 @@ const Panel: FC<NodePanelProps<VariableAssignerNodeType>> = ({
             nodeId={id}
             list={inputs.variables}
             onChange={handleVarListChange}
+            onlyLeafNodeVar
+            onlyVarType={inputs.output_type}
           />
         </Field>
       </div>
diff --git a/web/i18n/en-US/workflow.ts b/web/i18n/en-US/workflow.ts
index 5a4caa2f3c..b50df284a1 100644
--- a/web/i18n/en-US/workflow.ts
+++ b/web/i18n/en-US/workflow.ts
@@ -23,6 +23,7 @@ const translation = {
     restore: 'Restore',
     addTitle: 'Add title...',
     addDescription: 'Add description...',
+    noVar: 'No variable',
   },
   singleRun: {
     testRun: 'Test Run ',
diff --git a/web/i18n/zh-Hans/workflow.ts b/web/i18n/zh-Hans/workflow.ts
index 0c8f748f66..20143c85f7 100644
--- a/web/i18n/zh-Hans/workflow.ts
+++ b/web/i18n/zh-Hans/workflow.ts
@@ -23,6 +23,7 @@ const translation = {
     restore: '恢复',
     addTitle: '添加标题...',
     addDescription: '添加描述...',
+    noVar: '没有变量',
   },
   singleRun: {
     testRun: '测试运行 ',