From 21db8e3be48a55ad834a61766e67128bd29fc30d Mon Sep 17 00:00:00 2001
From: Joel <iamjoel007@gmail.com>
Date: Mon, 19 Feb 2024 11:26:25 +0800
Subject: [PATCH] feat: add var struct

---
 web/app/components/workflow/nodes/llm/mock.ts | 11 +++++++-
 .../components/workflow/nodes/llm/panel.tsx   | 15 ++++++-----
 .../workflow/nodes/llm/use-input.ts           | 26 +++++++++++++++++--
 3 files changed, 43 insertions(+), 9 deletions(-)

diff --git a/web/app/components/workflow/nodes/llm/mock.ts b/web/app/components/workflow/nodes/llm/mock.ts
index 15a6262c58..ef7403ac87 100644
--- a/web/app/components/workflow/nodes/llm/mock.ts
+++ b/web/app/components/workflow/nodes/llm/mock.ts
@@ -14,7 +14,16 @@ export const mockLLMNodeData: LLMNodeData = {
       temperature: 0.7,
     },
   },
-  variables: [],
+  variables: [
+    {
+      variable: 'name',
+      value_selector: ['start', 'name'],
+    },
+    {
+      variable: 'age',
+      value_selector: ['a', 'b', 'c'],
+    },
+  ],
   prompt: [],
   memory: {
     role_prefix: MemoryRole.assistant,
diff --git a/web/app/components/workflow/nodes/llm/panel.tsx b/web/app/components/workflow/nodes/llm/panel.tsx
index 8dc6eb93fc..cbcb7fac30 100644
--- a/web/app/components/workflow/nodes/llm/panel.tsx
+++ b/web/app/components/workflow/nodes/llm/panel.tsx
@@ -1,14 +1,15 @@
 import type { FC } from 'react'
 import { useTranslation } from 'react-i18next'
 import BasePanel from '../_base/panel'
+import VarList from '../_base/components/var/var-list'
 import useInput from './use-input'
 import { mockLLMNodeData } from './mock'
 import Field from '@/app/components/workflow/nodes/_base/components/field'
 import AddButton from '@/app/components/base/button/add-button'
 import Split from '@/app/components/workflow/nodes/_base/components/split'
 import ModelParameterModal from '@/app/components/header/account-setting/model-provider-page/model-parameter-modal'
-
 import Switch from '@/app/components/base/switch'
+
 const i18nPrefix = 'workflow.nodes.llm'
 
 const Panel: FC = () => {
@@ -16,16 +17,15 @@ const Panel: FC = () => {
   const {
     inputs,
     handleModelChanged,
-    toggleContextEnabled,
     handleCompletionParamsChange,
+    handleVarListChange,
+    handleAddVariable,
+    toggleContextEnabled,
   } = useInput(mockLLMNodeData)
   const model = inputs.model
   const modelMode = inputs.model?.mode
   const isChatMode = modelMode === 'chat'
 
-  const handleAddVariable = () => {
-    console.log('add variable')
-  }
   return (
     <BasePanel
       inputsElement={
@@ -53,7 +53,10 @@ const Panel: FC = () => {
               <AddButton onClick={handleAddVariable} />
             }
           >
-            Var Selector
+            <VarList
+              list={inputs.variables}
+              onChange={handleVarListChange}
+            />
           </Field>
 
           <Field
diff --git a/web/app/components/workflow/nodes/llm/use-input.ts b/web/app/components/workflow/nodes/llm/use-input.ts
index f7a32ab39a..263681a577 100644
--- a/web/app/components/workflow/nodes/llm/use-input.ts
+++ b/web/app/components/workflow/nodes/llm/use-input.ts
@@ -1,11 +1,12 @@
 /* eslint-disable react-hooks/exhaustive-deps */
 import { useCallback, useState } from 'react'
 import produce from 'immer'
-import type { LLMNodeData } from '../../types'
+import type { LLMNodeData, Variable } from '../../types'
 
 const useInput = (initInputs: LLMNodeData) => {
   const [inputs, setInputs] = useState<LLMNodeData>(initInputs)
 
+  // model
   const handleModelChanged = useCallback((model: { provider: string; modelId: string; mode?: string }) => {
     const newInputs = produce(inputs, (draft) => {
       draft.model.provider = model.provider
@@ -22,6 +23,25 @@ const useInput = (initInputs: LLMNodeData) => {
     setInputs(newInputs)
   }, [inputs.model])
 
+  // variables
+  const handleVarListChange = useCallback((newList: Variable[]) => {
+    const newInputs = produce(inputs, (draft) => {
+      draft.variables = newList
+    })
+    setInputs(newInputs)
+  }, [inputs.variables])
+
+  const handleAddVariable = useCallback(() => {
+    const newInputs = produce(inputs, (draft) => {
+      draft.variables.push({
+        variable: '',
+        value_selector: [],
+      })
+    })
+    setInputs(newInputs)
+  }, [inputs.variables])
+
+  // context
   const toggleContextEnabled = useCallback(() => {
     const newInputs = produce(inputs, (draft) => {
       draft.context.enabled = !draft.context.enabled
@@ -30,9 +50,11 @@ const useInput = (initInputs: LLMNodeData) => {
   }, [inputs.context.enabled])
 
   return {
-    handleCompletionParamsChange,
     inputs,
     handleModelChanged,
+    handleCompletionParamsChange,
+    handleVarListChange,
+    handleAddVariable,
     toggleContextEnabled,
   }
 }