dify/web/app/components/workflow/nodes/llm/use-config.ts
2024-03-04 20:14:18 +08:00

84 lines
2.6 KiB
TypeScript

import { useCallback } from 'react'
import produce from 'immer'
import useVarList from '../_base/hooks/use-var-list'
import type { Memory, ValueSelector } from '../../types'
import type { LLMNodeType } from './types'
import type { Resolution } from '@/types/app'
import { 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'
const useConfig = (id: string, payload: LLMNodeType) => {
const { inputs, setInputs } = useNodeCrud<LLMNodeType>(id, payload)
// model
const model = inputs.model
const handleModelChanged = useCallback((model: { provider: string; modelId: string; mode?: string }) => {
const newInputs = produce(inputs, (draft) => {
draft.model.provider = model.provider
draft.model.name = model.modelId
draft.model.mode = model.mode!
})
setInputs(newInputs)
}, [inputs, setInputs])
const handleCompletionParamsChange = useCallback((newParams: Record<string, any>) => {
const newInputs = produce(inputs, (draft) => {
draft.model.completion_params = newParams
})
setInputs(newInputs)
}, [inputs, setInputs])
const {
currentModel: currModel,
} = useTextGenerationCurrentProviderAndModelAndModelList(
{
provider: model.provider,
model: model.name,
},
)
const isShowVisionConfig = !!currModel?.features?.includes(ModelFeatureEnum.vision)
// variables
const { handleVarListChange, handleAddVariable } = useVarList<LLMNodeType>({
inputs,
setInputs,
})
// context
const handleContextVarChange = useCallback((newVar: ValueSelector) => {
const newInputs = produce(inputs, (draft) => {
draft.context.variable_selector = newVar
})
setInputs(newInputs)
}, [inputs, setInputs])
const handleMemoryChange = useCallback((newMemory: Memory) => {
const newInputs = produce(inputs, (draft) => {
draft.memory = newMemory
})
setInputs(newInputs)
}, [inputs, setInputs])
const handleVisionResolutionChange = useCallback((newResolution: Resolution) => {
const newInputs = produce(inputs, (draft) => {
draft.vision.configs.detail = newResolution
})
setInputs(newInputs)
}, [inputs, setInputs])
return {
inputs,
isShowVisionConfig,
handleModelChanged,
handleCompletionParamsChange,
handleVarListChange,
handleAddVariable,
handleContextVarChange,
handleMemoryChange,
handleVisionResolutionChange,
}
}
export default useConfig