import { useCallback } from 'react' import produce from 'immer' import type { Memory, ValueSelector } from '../../types' import type { QuestionClassifierNodeType } from './types' 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' const useConfig = (id: string, payload: QuestionClassifierNodeType) => { const { inputs, setInputs } = useNodeCrud(id, payload) // 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) => { const newInputs = produce(inputs, (draft) => { draft.model.completion_params = newParams }) setInputs(newInputs) }, [inputs, setInputs]) const handleQueryVarChange = useCallback((newVar: ValueSelector) => { const newInputs = produce(inputs, (draft) => { draft.query_variable_selector = newVar }) setInputs(newInputs) }, [inputs, setInputs]) const handleClassesChange = useCallback((newClasses: any) => { const newInputs = produce(inputs, (draft) => { draft.classes = newClasses draft._targetBranches = newClasses }) setInputs(newInputs) }, [inputs, setInputs]) const handleInstructionChange = useCallback((instruction: string) => { const newInputs = produce(inputs, (draft) => { draft.instruction = instruction }) setInputs(newInputs) }, [inputs, setInputs]) const handleMemoryChange = useCallback((memory: Memory) => { const newInputs = produce(inputs, (draft) => { draft.memory = memory }) setInputs(newInputs) }, [inputs, setInputs]) // single run const { isShowSingleRun, hideSingleRun, runningStatus, handleRun, handleStop, runInputData, setRunInputData, runResult, } = useOneStepRun({ id, data: inputs, defaultRunInputData: { query: '', }, }) const query = runInputData.query const setQuery = useCallback((newQuery: string) => { setRunInputData({ ...runInputData, query: newQuery, }) }, [runInputData, setRunInputData]) return { inputs, handleModelChanged, handleCompletionParamsChange, handleQueryVarChange, handleTopicsChange: handleClassesChange, handleInstructionChange, handleMemoryChange, isShowSingleRun, hideSingleRun, runningStatus, handleRun, handleStop, query, setQuery, runResult, } } export default useConfig