From 6b9cc927c0faba6845a08f74be228f0879ba82e1 Mon Sep 17 00:00:00 2001
From: Joel <iamjoel007@gmail.com>
Date: Fri, 15 Mar 2024 21:41:18 +0800
Subject: [PATCH] feat: llm default value

---
 .../components/workflow/hooks/use-workflow.ts |  2 +-
 .../nodes/_base/components/memory-config.tsx  |  2 +-
 .../nodes/_base/components/prompt/editor.tsx  |  2 +-
 .../nodes/llm/components/config-prompt.tsx    |  2 +-
 .../components/workflow/nodes/llm/panel.tsx   | 22 +-----
 .../workflow/nodes/llm/use-config.ts          | 79 ++++++++++++++++---
 .../nodes/question-classifier/use-config.ts   |  1 -
 web/app/components/workflow/types.ts          |  1 -
 8 files changed, 77 insertions(+), 34 deletions(-)

diff --git a/web/app/components/workflow/hooks/use-workflow.ts b/web/app/components/workflow/hooks/use-workflow.ts
index 7031829593..f95ede2390 100644
--- a/web/app/components/workflow/hooks/use-workflow.ts
+++ b/web/app/components/workflow/hooks/use-workflow.ts
@@ -236,7 +236,7 @@ export const useWorkflowInit = () => {
       workflowStore.setState({
         nodesDefaultConfigs: nodesDefaultConfigsData.reduce((acc, block) => {
           if (!acc[block.type])
-            acc[block.type] = { ...block.config, _isReady: true }
+            acc[block.type] = { ...block.config }
           return acc
         }, {} as Record<string, any>),
       })
diff --git a/web/app/components/workflow/nodes/_base/components/memory-config.tsx b/web/app/components/workflow/nodes/_base/components/memory-config.tsx
index e6ff4bba8f..31ff34c89d 100644
--- a/web/app/components/workflow/nodes/_base/components/memory-config.tsx
+++ b/web/app/components/workflow/nodes/_base/components/memory-config.tsx
@@ -141,7 +141,7 @@ const MemoryConfig: FC<Props> = ({
                 disabled={readonly}
               />
               <input
-                value={payload.window?.size as number}
+                value={(payload.window?.size || '') as number}
                 className='shrink-0 block ml-4 pl-3 w-12 h-8 appearance-none outline-none rounded-lg bg-gray-100 text-[13px] text-gra-900'
                 type='number'
                 min={WINDOW_SIZE_MIN}
diff --git a/web/app/components/workflow/nodes/_base/components/prompt/editor.tsx b/web/app/components/workflow/nodes/_base/components/prompt/editor.tsx
index 2f97f6f19b..b0e5a0d600 100644
--- a/web/app/components/workflow/nodes/_base/components/prompt/editor.tsx
+++ b/web/app/components/workflow/nodes/_base/components/prompt/editor.tsx
@@ -61,7 +61,7 @@ const Editor: FC<Props> = ({
           <div className='pt-1 pl-3 pr-2 flex justify-between h-6 items-center'>
             <div className='leading-4 text-xs font-semibold text-gray-700 uppercase'>{title}</div>
             <div className='flex items-center'>
-              <div className='leading-[18px] text-xs font-medium text-gray-500'>{value.length}</div>
+              <div className='leading-[18px] text-xs font-medium text-gray-500'>{value?.length || 0}</div>
               <div className='w-px h-3 ml-2 mr-2 bg-gray-200'></div>
               {/* Operations */}
               <div className='flex items-center space-x-2'>
diff --git a/web/app/components/workflow/nodes/llm/components/config-prompt.tsx b/web/app/components/workflow/nodes/llm/components/config-prompt.tsx
index 7593a38a3e..40e5d67a53 100644
--- a/web/app/components/workflow/nodes/llm/components/config-prompt.tsx
+++ b/web/app/components/workflow/nodes/llm/components/config-prompt.tsx
@@ -89,7 +89,7 @@ const ConfigPrompt: FC<Props> = ({
 
   return (
     <div>
-      {isChatModel
+      {(isChatModel && Array.isArray(payload))
         ? (
           <div>
             <div className='space-y-2'>
diff --git a/web/app/components/workflow/nodes/llm/panel.tsx b/web/app/components/workflow/nodes/llm/panel.tsx
index a462c55fc0..c293fd85dd 100644
--- a/web/app/components/workflow/nodes/llm/panel.tsx
+++ b/web/app/components/workflow/nodes/llm/panel.tsx
@@ -1,5 +1,5 @@
 import type { FC } from 'react'
-import React, { useEffect } from 'react'
+import React from 'react'
 import { useTranslation } from 'react-i18next'
 import MemoryConfig from '../_base/components/memory-config'
 import VarReferencePicker from '../_base/components/variable/var-reference-picker'
@@ -18,7 +18,6 @@ import { InputVarType, type NodePanelProps } from '@/app/components/workflow/typ
 import BeforeRunForm from '@/app/components/workflow/nodes/_base/components/before-run-form'
 import type { Props as FormProps } from '@/app/components/workflow/nodes/_base/components/before-run-form/form'
 import ResultPanel from '@/app/components/workflow/run/result-panel'
-import { useModelListAndDefaultModelAndCurrentProviderAndModel } from '@/app/components/header/account-setting/model-provider-page/hooks'
 
 const i18nPrefix = 'workflow.nodes.llm'
 
@@ -28,10 +27,6 @@ const Panel: FC<NodePanelProps<LLMNodeType>> = ({
 }) => {
   const { t } = useTranslation()
   const readOnly = false
-  const {
-    currentProvider,
-    currentModel,
-  } = useModelListAndDefaultModelAndCurrentProviderAndModel(1)
 
   const {
     inputs,
@@ -62,7 +57,6 @@ const Panel: FC<NodePanelProps<LLMNodeType>> = ({
     runResult,
   } = useConfig(id, data)
 
-  const isChatApp = true // TODO: get from app context
   const model = inputs.model
 
   const singleRunForms = (() => {
@@ -114,16 +108,6 @@ const Panel: FC<NodePanelProps<LLMNodeType>> = ({
     return forms
   })()
 
-  useEffect(() => {
-    if (currentProvider?.provider && currentModel?.model && !model.provider) {
-      handleModelChanged({
-        provider: currentProvider?.provider,
-        modelId: currentModel?.model,
-        mode: currentModel?.model_properties?.mode as string,
-      })
-    }
-  }, [model.provider, currentProvider, currentModel, handleModelChanged])
-
   return (
     <div className='mt-2'>
       <div className='px-4 pb-4 space-y-4'>
@@ -187,11 +171,11 @@ const Panel: FC<NodePanelProps<LLMNodeType>> = ({
         )}
 
         {/* Memory examples. Wait for design */}
-        {/* {isChatApp && isChatModel && (
+        {/* {isChatModel && (
           <div className='text-xs text-gray-300'>Memory examples(Designing)</div>
         )} */}
         {/* Memory */}
-        {isChatApp && (
+        {isChatModel && (
           <>
             <MemoryConfig
               readonly={readOnly}
diff --git a/web/app/components/workflow/nodes/llm/use-config.ts b/web/app/components/workflow/nodes/llm/use-config.ts
index 80937d50df..e6e66c6c26 100644
--- a/web/app/components/workflow/nodes/llm/use-config.ts
+++ b/web/app/components/workflow/nodes/llm/use-config.ts
@@ -1,11 +1,13 @@
-import { useCallback } from 'react'
+import { useCallback, useEffect, useRef } from 'react'
 import produce from 'immer'
 import useVarList from '../_base/hooks/use-var-list'
-import { PromptRole, VarType } from '../../types'
+import { VarType } from '../../types'
 import type { Memory, ValueSelector, Var } from '../../types'
+import { useStore } from '../../store'
+import { useIsChatMode } from '../../hooks'
 import type { LLMNodeType } from './types'
 import { Resolution } from '@/types/app'
-import { useTextGenerationCurrentProviderAndModelAndModelList } from '@/app/components/header/account-setting/model-provider-page/hooks'
+import { useModelListAndDefaultModelAndCurrentProviderAndModel, useTextGenerationCurrentProviderAndModelAndModelList } from '@/app/components/header/account-setting/model-provider-page/hooks'
 import { ModelFeatureEnum } from '@/app/components/header/account-setting/model-provider-page/declarations'
 import useNodeCrud from '@/app/components/workflow/nodes/_base/hooks/use-node-crud'
 import useOneStepRun from '@/app/components/workflow/nodes/_base/hooks/use-one-step-run'
@@ -13,25 +15,83 @@ import type { PromptItem } from '@/models/debug'
 import { RETRIEVAL_OUTPUT_STRUCT } from '@/app/components/workflow/constants'
 
 const useConfig = (id: string, payload: LLMNodeType) => {
-  const { inputs, setInputs } = useNodeCrud<LLMNodeType>(id, payload)
+  const isChatMode = useIsChatMode()
 
+  const defaultConfig = useStore(s => s.nodesDefaultConfigs)[payload.type]
+  const { inputs, setInputs } = useNodeCrud<LLMNodeType>(id, payload)
+  const inputRef = useRef(inputs)
+  useEffect(() => {
+    inputRef.current = inputs
+  }, [inputs])
   // model
   const model = inputs.model
   const modelMode = inputs.model?.mode
   const isChatModel = modelMode === 'chat'
   const isCompletionModel = !isChatModel
 
+  const appendDefaultPromptConfig = useCallback((draft: LLMNodeType, defaultConfig: any, passInIsChatMode?: boolean) => {
+    const promptTemplates = defaultConfig.prompt_templates
+    if (passInIsChatMode === undefined ? isChatModel : passInIsChatMode) {
+      draft.prompt_template = promptTemplates.chat_model.prompts
+    }
+    else {
+      draft.prompt_template = promptTemplates.completion_model.prompt
+      if (!draft.memory) {
+        draft.memory = {
+          role_prefix: {
+            user: '',
+            assistant: '',
+          },
+          window: {
+            enabled: false,
+            size: '',
+          },
+        }
+      }
+
+      draft.memory.role_prefix = {
+        user: promptTemplates.completion_model.conversation_histories_role.user_prefix,
+        assistant: promptTemplates.completion_model.conversation_histories_role.assistant_prefix,
+      }
+    }
+  }, [isChatModel])
+  useEffect(() => {
+    const isReady = defaultConfig && Object.keys(defaultConfig).length > 0
+    if (isReady) {
+      const newInputs = produce(inputs, (draft) => {
+        appendDefaultPromptConfig(draft, defaultConfig)
+      })
+      setInputs(newInputs)
+    }
+  // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [defaultConfig, isChatModel])
+
+  const {
+    currentProvider,
+    currentModel,
+  } = useModelListAndDefaultModelAndCurrentProviderAndModel(1)
+
   const handleModelChanged = useCallback((model: { provider: string; modelId: string; mode?: string }) => {
-    const newInputs = produce(inputs, (draft) => {
+    const newInputs = produce(inputRef.current, (draft) => {
       draft.model.provider = model.provider
       draft.model.name = model.modelId
       draft.model.mode = model.mode!
-      const isModeChange = model.mode !== inputs.model.mode
-      if (isModeChange)
-        draft.prompt_template = model.mode === 'chat' ? [{ role: PromptRole.system, text: '' }] : { text: '' }
+      const isModeChange = model.mode !== inputRef.current.model.mode
+      if (isModeChange && defaultConfig && Object.keys(defaultConfig).length > 0)
+        appendDefaultPromptConfig(draft, defaultConfig, model.mode === 'chat')
     })
     setInputs(newInputs)
-  }, [inputs, setInputs])
+  }, [setInputs, defaultConfig, appendDefaultPromptConfig])
+
+  useEffect(() => {
+    if (currentProvider?.provider && currentModel?.model && !model.provider) {
+      handleModelChanged({
+        provider: currentProvider?.provider,
+        modelId: currentModel?.model,
+        mode: currentModel?.model_properties?.mode as string,
+      })
+    }
+  }, [model.provider, currentProvider, currentModel, handleModelChanged])
 
   const handleCompletionParamsChange = useCallback((newParams: Record<string, any>) => {
     const newInputs = produce(inputs, (draft) => {
@@ -152,6 +212,7 @@ const useConfig = (id: string, payload: LLMNodeType) => {
   const varInputs = toVarInputs(inputs.variables)
 
   return {
+    isChatMode,
     inputs,
     isChatModel,
     isCompletionModel,
diff --git a/web/app/components/workflow/nodes/question-classifier/use-config.ts b/web/app/components/workflow/nodes/question-classifier/use-config.ts
index a3db7d99d4..07106e208f 100644
--- a/web/app/components/workflow/nodes/question-classifier/use-config.ts
+++ b/web/app/components/workflow/nodes/question-classifier/use-config.ts
@@ -74,7 +74,6 @@ const useConfig = (id: string, payload: QuestionClassifierNodeType) => {
         ...defaultConfig,
         query_variable_selector: inputs.query_variable_selector.length > 0 ? inputs.query_variable_selector : query_variable_selector,
       })
-      console.log(query_variable_selector)
     }
   // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [defaultConfig])
diff --git a/web/app/components/workflow/types.ts b/web/app/components/workflow/types.ts
index bb8efb3bdc..91d72a1486 100644
--- a/web/app/components/workflow/types.ts
+++ b/web/app/components/workflow/types.ts
@@ -33,7 +33,6 @@ export type CommonNodeType<T = {}> = {
   _isSingleRun?: boolean
   _runningStatus?: NodeRunningStatus
   _singleRunningStatus?: NodeRunningStatus
-  _isReady?: boolean
   selected?: boolean
   title: string
   desc: string