From 99bf3ff565e45e4328c29a53608e2375ad8cb6b1 Mon Sep 17 00:00:00 2001
From: Joel <iamjoel007@gmail.com>
Date: Thu, 8 Aug 2024 14:22:54 +0800
Subject: [PATCH] feat: params support vison

---
 .../nodes/parameter-extractor/default.ts      |  3 ++
 .../nodes/parameter-extractor/panel.tsx       | 42 ++++++++++++-------
 .../nodes/parameter-extractor/types.ts        |  6 ++-
 .../nodes/parameter-extractor/use-config.ts   | 33 +++++++++++++--
 4 files changed, 66 insertions(+), 18 deletions(-)

diff --git a/web/app/components/workflow/nodes/parameter-extractor/default.ts b/web/app/components/workflow/nodes/parameter-extractor/default.ts
index c1f5776ad5..125d469aea 100644
--- a/web/app/components/workflow/nodes/parameter-extractor/default.ts
+++ b/web/app/components/workflow/nodes/parameter-extractor/default.ts
@@ -16,6 +16,9 @@ const nodeDefault: NodeDefault<ParameterExtractorNodeType> = {
       },
     },
     reasoning_mode: ReasoningModeType.prompt,
+    vision: {
+      enabled: false,
+    },
   },
   getAvailablePrevNodes(isChatMode: boolean) {
     const nodes = isChatMode
diff --git a/web/app/components/workflow/nodes/parameter-extractor/panel.tsx b/web/app/components/workflow/nodes/parameter-extractor/panel.tsx
index 7c5686fadc..f3ccf14277 100644
--- a/web/app/components/workflow/nodes/parameter-extractor/panel.tsx
+++ b/web/app/components/workflow/nodes/parameter-extractor/panel.tsx
@@ -8,6 +8,7 @@ import MemoryConfig from '../_base/components/memory-config'
 import VarReferencePicker from '../_base/components/variable/var-reference-picker'
 import Editor from '../_base/components/prompt/editor'
 import ResultPanel from '../../run/result-panel'
+import ConfigVision from '../_base/components/config-vision'
 import useConfig from './use-config'
 import type { ParameterExtractorNodeType } from './types'
 import ExtractParameter from './components/extract-parameter/list'
@@ -54,6 +55,9 @@ const Panel: FC<NodePanelProps<ParameterExtractorNodeType>> = ({
     availableNodesWithParent,
     inputVarValues,
     varInputs,
+    isVisionModel,
+    handleVisionResolutionChange,
+    handleVisionResolutionEnabledChange,
     isShowSingleRun,
     hideSingleRun,
     runningStatus,
@@ -68,20 +72,6 @@ const Panel: FC<NodePanelProps<ParameterExtractorNodeType>> = ({
   return (
     <div className='mt-2'>
       <div className='px-4 pb-4 space-y-4'>
-        <Field
-          title={t(`${i18nPrefix}.inputVar`)}
-        >
-          <>
-            <VarReferencePicker
-              readonly={readOnly}
-              nodeId={id}
-              isShowNodeName
-              value={inputs.query || []}
-              onChange={handleInputVarChange}
-              filterVar={filterVar}
-            />
-          </>
-        </Field>
         <Field
           title={t(`${i18nCommonPrefix}.model`)}
         >
@@ -100,6 +90,30 @@ const Panel: FC<NodePanelProps<ParameterExtractorNodeType>> = ({
             readonly={readOnly}
           />
         </Field>
+        <Field
+          title={t(`${i18nPrefix}.inputVar`)}
+        >
+          <>
+            <VarReferencePicker
+              readonly={readOnly}
+              nodeId={id}
+              isShowNodeName
+              value={inputs.query || []}
+              onChange={handleInputVarChange}
+              filterVar={filterVar}
+            />
+          </>
+        </Field>
+        <Split />
+        <ConfigVision
+          nodeId={id}
+          readOnly={readOnly}
+          isVisionModel={isVisionModel}
+          enabled={inputs.vision.enabled}
+          onEnabledChange={handleVisionResolutionEnabledChange}
+          config={inputs.vision.configs}
+          onConfigChange={handleVisionResolutionChange}
+        />
         <Field
           title={t(`${i18nPrefix}.extractParameters`)}
           operations={
diff --git a/web/app/components/workflow/nodes/parameter-extractor/types.ts b/web/app/components/workflow/nodes/parameter-extractor/types.ts
index 5038bd8a9c..f5ba717be8 100644
--- a/web/app/components/workflow/nodes/parameter-extractor/types.ts
+++ b/web/app/components/workflow/nodes/parameter-extractor/types.ts
@@ -1,4 +1,4 @@
-import type { CommonNodeType, Memory, ModelConfig, ValueSelector } from '@/app/components/workflow/types'
+import type { CommonNodeType, Memory, ModelConfig, ValueSelector, VisionSetting } from '@/app/components/workflow/types'
 
 export enum ParamType {
   string = 'string',
@@ -30,4 +30,8 @@ export type ParameterExtractorNodeType = CommonNodeType & {
   parameters: Param[]
   instruction: string
   memory?: Memory
+  vision: {
+    enabled: boolean
+    configs?: VisionSetting
+  }
 }
diff --git a/web/app/components/workflow/nodes/parameter-extractor/use-config.ts b/web/app/components/workflow/nodes/parameter-extractor/use-config.ts
index ddfaed507e..b6c864a7bd 100644
--- a/web/app/components/workflow/nodes/parameter-extractor/use-config.ts
+++ b/web/app/components/workflow/nodes/parameter-extractor/use-config.ts
@@ -9,6 +9,7 @@ import {
   useWorkflow,
 } from '../../hooks'
 import useOneStepRun from '../_base/hooks/use-one-step-run'
+import useConfigVision from '../../hooks/use-config-vision'
 import type { Param, ParameterExtractorNodeType, ReasoningModeType } from './types'
 import { useModelListAndDefaultModelAndCurrentProviderAndModel, useTextGenerationCurrentProviderAndModelAndModelList } from '@/app/components/header/account-setting/model-provider-page/hooks'
 import {
@@ -84,9 +85,23 @@ const useConfig = (id: string, payload: ParameterExtractorNodeType) => {
   }
   const modelMode = inputs.model?.mode
   const isChatModel = modelMode === 'chat'
-
   const isCompletionModel = !isChatModel
 
+  const {
+    isVisionModel,
+    handleVisionResolutionEnabledChange,
+    handleVisionResolutionChange,
+    handleModelChanged: handleVisionConfigAfterModelChanged,
+  } = useConfigVision(model, {
+    payload: inputs.vision,
+    onChange: (newPayload) => {
+      const newInputs = produce(inputs, (draft) => {
+        draft.vision = newPayload
+      })
+      setInputs(newInputs)
+    },
+  })
+
   const appendDefaultPromptConfig = useCallback((draft: ParameterExtractorNodeType, defaultConfig: any, _passInIsChatMode?: boolean) => {
     const promptTemplates = defaultConfig.prompt_templates
     if (!isChatModel) {
@@ -97,7 +112,7 @@ const useConfig = (id: string, payload: ParameterExtractorNodeType) => {
     }
   }, [isChatModel])
 
-  // const [modelChanged, setModelChanged] = useState(false)
+  const [modelChanged, setModelChanged] = useState(false)
   const {
     currentProvider,
     currentModel,
@@ -113,7 +128,7 @@ const useConfig = (id: string, payload: ParameterExtractorNodeType) => {
         appendDefaultPromptConfig(draft, defaultConfig, model.mode === 'chat')
     })
     setInputs(newInputs)
-    // setModelChanged(true)
+    setModelChanged(true)
   }, [setInputs, defaultConfig, appendDefaultPromptConfig])
 
   useEffect(() => {
@@ -126,6 +141,15 @@ const useConfig = (id: string, payload: ParameterExtractorNodeType) => {
     }
   }, [model?.provider, currentProvider, currentModel, handleModelChanged])
 
+  // change to vision model to set vision enabled, else disabled
+  useEffect(() => {
+    if (!modelChanged)
+      return
+    setModelChanged(false)
+    handleVisionConfigAfterModelChanged()
+  // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [isVisionModel, modelChanged])
+
   const {
     currentModel: currModel,
   } = useTextGenerationCurrentProviderAndModelAndModelList(
@@ -245,6 +269,9 @@ const useConfig = (id: string, payload: ParameterExtractorNodeType) => {
     handleMemoryChange,
     varInputs,
     inputVarValues,
+    isVisionModel,
+    handleVisionResolutionEnabledChange,
+    handleVisionResolutionChange,
     isShowSingleRun,
     hideSingleRun,
     runningStatus,