From ab6a01b476b4712e9af20f71d0d40dad7508c703 Mon Sep 17 00:00:00 2001
From: Joel <iamjoel007@gmail.com>
Date: Sun, 18 Feb 2024 18:29:54 +0800
Subject: [PATCH] chore: handle llm model type

---
 .../workflow/nodes/_base/components/field.tsx |  4 +--
 web/app/components/workflow/nodes/llm/mock.ts |  2 +-
 .../components/workflow/nodes/llm/panel.tsx   | 27 +++++++++++++++----
 .../workflow/nodes/llm/use-input.ts           | 27 ++++++++++++++-----
 4 files changed, 45 insertions(+), 15 deletions(-)

diff --git a/web/app/components/workflow/nodes/_base/components/field.tsx b/web/app/components/workflow/nodes/_base/components/field.tsx
index b5458f230a..24c7187d2e 100644
--- a/web/app/components/workflow/nodes/_base/components/field.tsx
+++ b/web/app/components/workflow/nodes/_base/components/field.tsx
@@ -7,7 +7,7 @@ import TooltipPlus from '@/app/components/base/tooltip-plus'
 type Props = {
   title: string
   tooltip?: string
-  children: JSX.Element | string
+  children?: JSX.Element | string | null
   operations?: JSX.Element
   inline?: boolean
 }
@@ -33,7 +33,7 @@ const Filed: FC<Props> = ({
         </div>
         {operations && <div>{operations}</div>}
       </div>
-      <div>{children}</div>
+      {children && <div className={cn(!inline && 'mt-1')}>{children}</div>}
     </div>
   )
 }
diff --git a/web/app/components/workflow/nodes/llm/mock.ts b/web/app/components/workflow/nodes/llm/mock.ts
index 1a2df388da..15a6262c58 100644
--- a/web/app/components/workflow/nodes/llm/mock.ts
+++ b/web/app/components/workflow/nodes/llm/mock.ts
@@ -9,7 +9,7 @@ export const mockLLMNodeData: LLMNodeData = {
   model: {
     provider: 'openai',
     name: 'gpt-4',
-    mode: 'completion',
+    mode: 'chat',
     completion_params: {
       temperature: 0.7,
     },
diff --git a/web/app/components/workflow/nodes/llm/panel.tsx b/web/app/components/workflow/nodes/llm/panel.tsx
index 12416bf299..de815ba455 100644
--- a/web/app/components/workflow/nodes/llm/panel.tsx
+++ b/web/app/components/workflow/nodes/llm/panel.tsx
@@ -7,15 +7,21 @@ 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 ModelSelector from '@/app/components/header/account-setting/model-provider-page/model-selector'
-import { useTextGenerationCurrentProviderAndModelAndModelList } from '@/app/components/header/account-setting/model-provider-page/hooks'
+import Switch from '@/app/components/base/switch'
 const i18nPrefix = 'workflow.nodes.llm'
 
 const Panel: FC = () => {
   const { t } = useTranslation()
-  const { inputs, handleModelChanged } = useInput(mockLLMNodeData)
   const {
     textGenerationModelList,
-  } = useTextGenerationCurrentProviderAndModelAndModelList()
+    inputs,
+    handleModelChanged,
+    toggleContextEnabled,
+  } = useInput(mockLLMNodeData)
+
+  const modelMode = inputs.model.mode
+  const isChatMode = modelMode === 'chat'
+
   const handleAddVariable = () => {
     console.log('add variable')
   }
@@ -49,11 +55,22 @@ const Panel: FC = () => {
 
           <Field
             title={t(`${i18nPrefix}.context`)}
+            operations={
+              <Switch
+                defaultValue={inputs.context.enabled}
+                onChange={toggleContextEnabled}
+                size='md'
+              />
+            }
           >
-            Context
+            {inputs.context.enabled
+              ? (
+                <div>Context</div>
+              )
+              : null}
           </Field>
           <Field
-            title={t(`${i18nPrefix}.context`)}
+            title={t(`${i18nPrefix}.prompt`)}
           >
             Prompt
           </Field>
diff --git a/web/app/components/workflow/nodes/llm/use-input.ts b/web/app/components/workflow/nodes/llm/use-input.ts
index 44380a4122..7f864f5d80 100644
--- a/web/app/components/workflow/nodes/llm/use-input.ts
+++ b/web/app/components/workflow/nodes/llm/use-input.ts
@@ -1,25 +1,38 @@
 import { useCallback, useState } from 'react'
 import produce from 'immer'
 import type { LLMNodeData } from '../../types'
+import { useTextGenerationCurrentProviderAndModelAndModelList } from '@/app/components/header/account-setting/model-provider-page/hooks'
+
 const useInput = (initInputs: LLMNodeData) => {
+  const {
+    textGenerationModelList,
+  } = useTextGenerationCurrentProviderAndModelAndModelList()
+
   const [inputs, setInputs] = useState<LLMNodeData>(initInputs)
 
   const handleModelChanged = useCallback((model: { provider: string; model: string }) => {
+    const targetProvider = textGenerationModelList.find(modelItem => modelItem.provider === model.provider)
+    const targetModelItem = targetProvider?.models.find(modelItem => modelItem.model === model.model)
     const newInputs = produce(inputs, (draft) => {
       draft.model.provider = model.provider
       draft.model.name = model.model
+      draft.model.mode = targetModelItem?.model_properties.mode as string
     })
     setInputs(newInputs)
-  }, [inputs.model])
+  }, [inputs.model, textGenerationModelList])
+
+  const toggleContextEnabled = useCallback(() => {
+    const newInputs = produce(inputs, (draft) => {
+      draft.context.enabled = !draft.context.enabled
+    })
+    setInputs(newInputs)
+  }, [inputs.context.enabled])
+
   return {
+    textGenerationModelList,
     inputs,
-    setInputs: (key: string, payload: any) => {
-      setInputs({
-        ...inputs,
-        [key]: payload,
-      } as LLMNodeData)
-    },
     handleModelChanged,
+    toggleContextEnabled,
   }
 }