dify/web/app/components/workflow/nodes/knowledge-retrieval/hooks/use-knowledge-dataset-selection.ts

import type { MutableRefObject } from 'react'
import type { KnowledgeRetrievalNodeType } from '../types'
import type { DataSet } from '@/models/datasets'
import { produce } from 'immer'
import { useCallback, useEffect, useMemo, useState } from 'react'
import { fetchDatasets } from '@/service/datasets'
import { RETRIEVE_TYPE } from '@/types/app'
import { getMultipleRetrievalConfig, getSelectedDatasetsMode } from '../utils'

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

type Params = {
  inputs: KnowledgeRetrievalNodeType
  inputRef: MutableRefObject<KnowledgeRetrievalNodeType>
  setInputs: (inputs: KnowledgeRetrievalNodeType) => void
  payloadRetrievalMode: RETRIEVE_TYPE
  updateDatasetsDetail: (datasets: DataSet[]) => void
  fallbackRerankModel: ModelIdentity
}

const useKnowledgeDatasetSelection = ({
  inputs,
  inputRef,
  setInputs,
  payloadRetrievalMode,
  updateDatasetsDetail,
  fallbackRerankModel,
}: Params) => {
  const [selectedDatasets, setSelectedDatasets] = useState<DataSet[]>([])
  const [selectedDatasetsLoaded, setSelectedDatasetsLoaded] = useState(false)
  const [rerankModelOpen, setRerankModelOpen] = useState(false)

  useEffect(() => {
    void (async () => {
      const currentInputs = inputRef.current
      const datasetIds = currentInputs.dataset_ids
      if (datasetIds.length > 0) {
        const { data: dataSetsWithDetail } = await fetchDatasets({
          url: '/datasets',
          params: {
            page: 1,
            ids: datasetIds,
          },
        })
        setSelectedDatasets(dataSetsWithDetail)
      }

      const nextInputs = produce(currentInputs, (draft) => {
        draft.dataset_ids = datasetIds
      })
      setInputs(nextInputs)
      setSelectedDatasetsLoaded(true)
    })()
  }, [inputRef, setInputs])

  const handleOnDatasetsChange = useCallback(
    (newDatasets: DataSet[]) => {
      const {
        mixtureHighQualityAndEconomic,
        mixtureInternalAndExternal,
        inconsistentEmbeddingModel,
        allInternal,
        allExternal,
      } = getSelectedDatasetsMode(newDatasets)
      const noMultiModalDatasets = newDatasets.every((dataset) => !dataset.is_multimodal)

      const nextInputs = produce(inputs, (draft) => {
        draft.dataset_ids = newDatasets.map((dataset) => dataset.id)

        if (payloadRetrievalMode === RETRIEVE_TYPE.multiWay && newDatasets.length > 0) {
          draft.multiple_retrieval_config = getMultipleRetrievalConfig(
            draft.multiple_retrieval_config!,
            newDatasets,
            selectedDatasets,
            fallbackRerankModel,
          )
        }

        if (noMultiModalDatasets) draft.query_attachment_selector = []
      })

      updateDatasetsDetail(newDatasets)
      setInputs(nextInputs)
      setSelectedDatasets(newDatasets)

      if (
        (allInternal && (mixtureHighQualityAndEconomic || inconsistentEmbeddingModel)) ||
        mixtureInternalAndExternal ||
        allExternal
      ) {
        setRerankModelOpen(true)
      }
    },
    [
      fallbackRerankModel,
      inputs,
      payloadRetrievalMode,
      selectedDatasets,
      setInputs,
      updateDatasetsDetail,
    ],
  )

  const showImageQueryVarSelector = useMemo(() => {
    return selectedDatasets.some((dataset) => dataset.is_multimodal)
  }, [selectedDatasets])

  return {
    selectedDatasets,
    selectedDatasetsLoaded,
    rerankModelOpen,
    setRerankModelOpen,
    handleOnDatasetsChange,
    showImageQueryVarSelector,
  }
}

export default useKnowledgeDatasetSelection