dify/web/app/components/workflow/nodes/knowledge-retrieval/hooks/use-knowledge-model-config.ts

import type { MutableRefObject } from 'react'
import type { KnowledgeRetrievalNodeType, MultipleRetrievalConfig } from '../types'
import type { ModelConfig } from '@/app/components/workflow/types'
import type { DataSet } from '@/models/datasets'
import { isEqual } from 'es-toolkit/predicate'
import { produce } from 'immer'
import { useCallback, useEffect } from 'react'
import { DATASET_DEFAULT } from '@/config'
import { AppModeEnum, RETRIEVE_TYPE } from '@/types/app'
import { getMultipleRetrievalConfig } from '../utils'

type ModelIdentity = {
  provider?: string
  model?: string
}

type TextProvider = {
  provider: string
}

type TextModel = {
  model: string
  model_properties?: {
    mode?: string
  }
}

type Params = {
  inputs: KnowledgeRetrievalNodeType
  inputRef: MutableRefObject<KnowledgeRetrievalNodeType>
  setInputs: (inputs: KnowledgeRetrievalNodeType) => void
  selectedDatasets: DataSet[]
  currentProvider?: TextProvider
  currentModel?: TextModel
  fallbackRerankModel: ModelIdentity
  hasRerankDefaultModel: boolean
}

const createSingleRetrievalConfig = (
  model: ModelConfig,
): KnowledgeRetrievalNodeType['single_retrieval_config'] => ({
  model,
})

const useKnowledgeModelConfig = ({
  inputs,
  inputRef,
  setInputs,
  selectedDatasets,
  currentProvider,
  currentModel,
  fallbackRerankModel,
  hasRerankDefaultModel,
}: Params) => {
  const handleModelChanged = useCallback(
    (model: { provider: string; modelId: string; mode?: string }) => {
      const nextInputs = produce(inputRef.current, (draft) => {
        if (!draft.single_retrieval_config) {
          draft.single_retrieval_config = createSingleRetrievalConfig({
            provider: '',
            name: '',
            mode: '',
            completion_params: {},
          })
        }

        const draftModel = draft.single_retrieval_config!.model
        draftModel.provider = model.provider
        draftModel.name = model.modelId
        draftModel.mode = model.mode || AppModeEnum.CHAT
      })
      setInputs(nextInputs)
    },
    [inputRef, setInputs],
  )

  const handleCompletionParamsChange = useCallback(
    (newParams: Record<string, unknown>) => {
      if (isEqual(newParams, inputRef.current.single_retrieval_config?.model.completion_params))
        return

      const nextInputs = produce(inputRef.current, (draft) => {
        if (!draft.single_retrieval_config) {
          draft.single_retrieval_config = createSingleRetrievalConfig({
            provider: '',
            name: '',
            mode: '',
            completion_params: {},
          })
        }

        draft.single_retrieval_config!.model.completion_params = newParams
      })
      setInputs(nextInputs)
    },
    [inputRef, setInputs],
  )

  useEffect(() => {
    const currentInputs = inputRef.current
    if (
      currentInputs.retrieval_mode === RETRIEVE_TYPE.multiWay &&
      currentInputs.multiple_retrieval_config?.reranking_model?.provider &&
      fallbackRerankModel.model &&
      hasRerankDefaultModel
    ) {
      return
    }

    if (
      currentInputs.retrieval_mode === RETRIEVE_TYPE.oneWay &&
      currentInputs.single_retrieval_config?.model?.provider
    )
      return

    const nextInputs = produce(currentInputs, (draft) => {
      if (currentProvider?.provider && currentModel?.model) {
        const hasSetModel = draft.single_retrieval_config?.model?.provider
        if (!hasSetModel) {
          draft.single_retrieval_config = createSingleRetrievalConfig({
            provider: currentProvider.provider,
            name: currentModel.model,
            mode: currentModel.model_properties?.mode || AppModeEnum.CHAT,
            completion_params: {},
          })
        }
      }

      const multipleRetrievalConfig = draft.multiple_retrieval_config
      draft.multiple_retrieval_config = {
        top_k: multipleRetrievalConfig?.top_k || DATASET_DEFAULT.top_k,
        score_threshold: multipleRetrievalConfig?.score_threshold,
        reranking_model: multipleRetrievalConfig?.reranking_model,
        reranking_mode: multipleRetrievalConfig?.reranking_mode,
        weights: multipleRetrievalConfig?.weights,
        reranking_enable:
          multipleRetrievalConfig?.reranking_enable !== undefined
            ? multipleRetrievalConfig.reranking_enable
            : Boolean(fallbackRerankModel.model && hasRerankDefaultModel),
      }
    })
    setInputs(nextInputs)
  }, [
    currentModel,
    currentProvider?.provider,
    fallbackRerankModel.model,
    hasRerankDefaultModel,
    inputRef,
    setInputs,
  ])

  const handleRetrievalModeChange = useCallback(
    (newMode: RETRIEVE_TYPE) => {
      const nextInputs = produce(inputs, (draft) => {
        draft.retrieval_mode = newMode
        if (newMode === RETRIEVE_TYPE.multiWay) {
          draft.multiple_retrieval_config = getMultipleRetrievalConfig(
            draft.multiple_retrieval_config as MultipleRetrievalConfig,
            selectedDatasets,
            selectedDatasets,
            fallbackRerankModel,
          )
          return
        }

        const hasSetModel = draft.single_retrieval_config?.model?.provider
        if (!hasSetModel) {
          draft.single_retrieval_config = createSingleRetrievalConfig({
            provider: currentProvider?.provider || '',
            name: currentModel?.model || '',
            mode: currentModel?.model_properties?.mode || AppModeEnum.CHAT,
            completion_params: {},
          })
        }
      })
      setInputs(nextInputs)
    },
    [
      currentModel?.model,
      currentModel?.model_properties?.mode,
      currentProvider?.provider,
      fallbackRerankModel,
      inputs,
      selectedDatasets,
      setInputs,
    ],
  )

  const handleMultipleRetrievalConfigChange = useCallback(
    (newConfig: MultipleRetrievalConfig) => {
      const nextInputs = produce(inputs, (draft) => {
        draft.multiple_retrieval_config = getMultipleRetrievalConfig(
          newConfig,
          selectedDatasets,
          selectedDatasets,
          fallbackRerankModel,
        )
      })
      setInputs(nextInputs)
    },
    [fallbackRerankModel, inputs, selectedDatasets, setInputs],
  )

  return {
    handleModelChanged,
    handleCompletionParamsChange,
    handleRetrievalModeChange,
    handleMultipleRetrievalConfigChange,
  }
}

export default useKnowledgeModelConfig