From d673b4c2199a0bfbfc01cbd65da9557922725720 Mon Sep 17 00:00:00 2001
From: StyleZhang <jasonapring2015@outlook.com>
Date: Fri, 29 Mar 2024 17:20:14 +0800
Subject: [PATCH] fix: prompt editor

---
 .../plugins/component-picker-block/hooks.tsx   | 13 +++++++++++--
 .../plugins/component-picker-block/index.tsx   | 18 ++++++++++++------
 ...rkflow-variable-block-replacement-block.tsx |  1 -
 .../components/variable/var-reference-vars.tsx |  8 ++++----
 4 files changed, 27 insertions(+), 13 deletions(-)

diff --git a/web/app/components/base/prompt-editor/plugins/component-picker-block/hooks.tsx b/web/app/components/base/prompt-editor/plugins/component-picker-block/hooks.tsx
index 2e2e5749fd..0b1cc3ce7b 100644
--- a/web/app/components/base/prompt-editor/plugins/component-picker-block/hooks.tsx
+++ b/web/app/components/base/prompt-editor/plugins/component-picker-block/hooks.tsx
@@ -8,6 +8,7 @@ import type {
   HistoryBlockType,
   QueryBlockType,
   VariableBlockType,
+  WorkflowVariableBlockType,
 } from '../../types'
 import { INSERT_CONTEXT_BLOCK_COMMAND } from '../context-block'
 import { INSERT_HISTORY_BLOCK_COMMAND } from '../history-block'
@@ -175,18 +176,26 @@ export const useOptions = (
   historyBlock?: HistoryBlockType,
   variableBlock?: VariableBlockType,
   externalToolBlockType?: ExternalToolBlockType,
+  workflowVariableBlockType?: WorkflowVariableBlockType,
   queryString?: string,
 ) => {
   const promptOptions = usePromptOptions(contextBlock, queryBlock, historyBlock)
   const variableOptions = useVariableOptions(variableBlock, queryString)
   const externalToolOptions = useExternalToolOptions(externalToolBlockType, queryString)
+  const workflowVariableOptions = useMemo(() => {
+    if (!workflowVariableBlockType?.show)
+      return []
+
+    return workflowVariableBlockType.variables || []
+  }, [workflowVariableBlockType])
 
   return useMemo(() => {
     return {
       promptOptions,
       variableOptions,
       externalToolOptions,
-      allOptions: [...promptOptions, ...variableOptions, ...externalToolOptions],
+      workflowVariableOptions,
+      allOptions: [...promptOptions, ...variableOptions, ...externalToolOptions, ...workflowVariableOptions],
     }
-  }, [promptOptions, variableOptions, externalToolOptions])
+  }, [promptOptions, variableOptions, externalToolOptions, workflowVariableOptions])
 }
diff --git a/web/app/components/base/prompt-editor/plugins/component-picker-block/index.tsx b/web/app/components/base/prompt-editor/plugins/component-picker-block/index.tsx
index a27fb5e44e..2dd1129701 100644
--- a/web/app/components/base/prompt-editor/plugins/component-picker-block/index.tsx
+++ b/web/app/components/base/prompt-editor/plugins/component-picker-block/index.tsx
@@ -64,12 +64,14 @@ const ComponentPicker = ({
     promptOptions,
     variableOptions,
     externalToolOptions,
+    workflowVariableOptions,
   } = useOptions(
     contextBlock,
     queryBlock,
     historyBlock,
     variableBlock,
     externalToolBlock,
+    workflowVariableBlock,
   )
 
   const onSelectOption = useCallback(
@@ -83,7 +85,8 @@ const ComponentPicker = ({
         if (nodeToRemove)
           nodeToRemove.remove()
 
-        selectedOption.onSelect(matchingString)
+        if (selectedOption?.onSelect)
+          selectedOption.onSelect(matchingString)
         closeMenu()
       })
     },
@@ -180,7 +183,7 @@ const ComponentPicker = ({
             )
           }
           {
-            workflowVariableBlock?.show && !!workflowVariableBlock?.variables?.length && (
+            !!workflowVariableOptions.length && (
               <>
                 {
                   (!!promptOptions.length || !!variableOptions.length || !!externalToolOptions.length) && (
@@ -189,8 +192,11 @@ const ComponentPicker = ({
                 }
                 <VarReferenceVars
                   hideSearch
-                  vars={workflowVariableBlock?.variables}
-                  onChange={handleSelectWorkflowVariable}
+                  vars={workflowVariableOptions}
+                  onChange={(variables: string[], item: any) => {
+                    selectOptionAndCleanUp(item)
+                    handleSelectWorkflowVariable(variables)
+                  }}
                 />
               </>
             )
@@ -206,14 +212,14 @@ const ComponentPicker = ({
     promptOptions,
     variableOptions,
     externalToolOptions,
+    workflowVariableOptions,
     queryString,
-    workflowVariableBlock,
     handleSelectWorkflowVariable,
   ])
 
   return (
     <LexicalTypeaheadMenuPlugin
-      options={allOptions}
+      options={allOptions as any}
       onQueryChange={setQueryString}
       onSelectOption={onSelectOption}
       anchorClassName='z-[999999]'
diff --git a/web/app/components/base/prompt-editor/plugins/workflow-variable-block/workflow-variable-block-replacement-block.tsx b/web/app/components/base/prompt-editor/plugins/workflow-variable-block/workflow-variable-block-replacement-block.tsx
index d2ed697d2e..d0589a9485 100644
--- a/web/app/components/base/prompt-editor/plugins/workflow-variable-block/workflow-variable-block-replacement-block.tsx
+++ b/web/app/components/base/prompt-editor/plugins/workflow-variable-block/workflow-variable-block-replacement-block.tsx
@@ -40,7 +40,6 @@ const WorkflowVariableBlockReplacementBlock = ({
     if (matchArr === null)
       return null
 
-    console.log(matchArr, 'ma')
     const startOffset = matchArr.index
     const endOffset = startOffset + matchArr[0].length
     return {
diff --git a/web/app/components/workflow/nodes/_base/components/variable/var-reference-vars.tsx b/web/app/components/workflow/nodes/_base/components/variable/var-reference-vars.tsx
index 637a607454..62c7227e99 100644
--- a/web/app/components/workflow/nodes/_base/components/variable/var-reference-vars.tsx
+++ b/web/app/components/workflow/nodes/_base/components/variable/var-reference-vars.tsx
@@ -23,7 +23,7 @@ type ObjectChildrenProps = {
   title: string
   data: Var[]
   objPath: string[]
-  onChange: (value: ValueSelector) => void
+  onChange: (value: ValueSelector, item: Var) => void
   onHovering?: (value: boolean) => void
   itemWidth?: number
 }
@@ -33,7 +33,7 @@ type ItemProps = {
   title: string
   objPath: string[]
   itemData: Var
-  onChange: (value: ValueSelector) => void
+  onChange: (value: ValueSelector, item: Var) => void
   onHovering?: (value: boolean) => void
   itemWidth?: number
 }
@@ -76,7 +76,7 @@ const Item: FC<ItemProps> = ({
   }, [isHovering])
   const handleChosen = (e: React.MouseEvent) => {
     e.stopPropagation()
-    onChange([nodeId, ...objPath, itemData.variable])
+    onChange([nodeId, ...objPath, itemData.variable], itemData)
   }
   return (
     <PortalToFollowElem
@@ -188,7 +188,7 @@ type Props = {
   hideSearch?: boolean
   searchBoxClassName?: string
   vars: NodeOutPutVar[]
-  onChange: (value: ValueSelector) => void
+  onChange: (value: ValueSelector, item: Var) => void
   itemWidth?: number
 }
 const VarReferenceVars: FC<Props> = ({