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

import type { ValueSelector } from '../../types'
import type { KnowledgeRetrievalNodeType } from './types'
import { produce } from 'immer'
import {
  useEffect,
  useMemo,
} from 'react'
import { ModelTypeEnum } from '@/app/components/header/account-setting/model-provider-page/declarations'
import { useCurrentProviderAndModel, useModelListAndDefaultModelAndCurrentProviderAndModel } from '@/app/components/header/account-setting/model-provider-page/hooks'
import useNodeCrud from '@/app/components/workflow/nodes/_base/hooks/use-node-crud'
import { useDatasetsDetailStore } from '../../datasets-detail-store/store'
import {
  useIsChatMode,
  useNodesReadOnly,
  useWorkflow,
} from '../../hooks'
import { BlockEnum } from '../../types'
import useKnowledgeDatasetSelection from './hooks/use-knowledge-dataset-selection'
import useKnowledgeInputManager from './hooks/use-knowledge-input-manager'
import useKnowledgeMetadataConfig from './hooks/use-knowledge-metadata-config'
import useKnowledgeModelConfig from './hooks/use-knowledge-model-config'

const useConfig = (id: string, payload: KnowledgeRetrievalNodeType) => {
  const { nodesReadOnly: readOnly } = useNodesReadOnly()
  const isChatMode = useIsChatMode()
  const { getBeforeNodesInSameBranch } = useWorkflow()
  const startNode = getBeforeNodesInSameBranch(id).find(node => node.data.type === BlockEnum.Start)
  const startNodeId = startNode?.id
  const { inputs, setInputs: doSetInputs } = useNodeCrud<KnowledgeRetrievalNodeType>(id, payload)
  const updateDatasetsDetail = useDatasetsDetailStore(s => s.updateDatasetsDetail)
  const {
    inputRef,
    setInputs,
    handleQueryVarChange,
    handleQueryAttachmentChange,
  } = useKnowledgeInputManager({
    inputs,
    doSetInputs,
  })

  const {
    currentProvider,
    currentModel,
  } = useModelListAndDefaultModelAndCurrentProviderAndModel(ModelTypeEnum.textGeneration)

  const {
    modelList: rerankModelList,
    defaultModel: rerankDefaultModel,
  } = useModelListAndDefaultModelAndCurrentProviderAndModel(ModelTypeEnum.rerank)

  const {
    currentModel: currentRerankModel,
    currentProvider: currentRerankProvider,
  } = useCurrentProviderAndModel(
    rerankModelList,
    rerankDefaultModel
      ? {
          ...rerankDefaultModel,
          provider: rerankDefaultModel.provider.provider,
        }
      : undefined,
  )

  const fallbackRerankModel = useMemo(() => ({
    provider: currentRerankProvider?.provider,
    model: currentRerankModel?.model,
  }), [currentRerankModel?.model, currentRerankProvider?.provider])

  const {
    selectedDatasets,
    selectedDatasetsLoaded,
    rerankModelOpen,
    setRerankModelOpen,
    handleOnDatasetsChange,
    showImageQueryVarSelector,
  } = useKnowledgeDatasetSelection({
    inputs,
    inputRef,
    setInputs,
    payloadRetrievalMode: payload.retrieval_mode,
    updateDatasetsDetail,
    fallbackRerankModel,
  })

  const {
    handleModelChanged,
    handleCompletionParamsChange,
    handleRetrievalModeChange,
    handleMultipleRetrievalConfigChange,
  } = useKnowledgeModelConfig({
    inputs,
    inputRef,
    setInputs,
    selectedDatasets,
    currentProvider,
    currentModel,
    fallbackRerankModel,
    hasRerankDefaultModel: Boolean(currentRerankModel && rerankDefaultModel),
  })

  useEffect(() => {
    const currentInputs = inputRef.current
    let nextQueryVariableSelector: ValueSelector = currentInputs.query_variable_selector
    if (isChatMode && currentInputs.query_variable_selector.length === 0 && startNodeId)
      nextQueryVariableSelector = [startNodeId, 'sys.query']

    setInputs(produce(currentInputs, (draft) => {
      draft.query_variable_selector = nextQueryVariableSelector
    }))
  }, [inputRef, isChatMode, setInputs, startNodeId])

  const metadataConfig = useKnowledgeMetadataConfig({
    id,
    inputRef,
    setInputs,
  })

  return {
    readOnly,
    inputs,
    handleQueryVarChange,
    handleQueryAttachmentChange,
    filterStringVar: metadataConfig.filterStringVar,
    filterFileVar: metadataConfig.filterFileVar,
    handleRetrievalModeChange,
    handleMultipleRetrievalConfigChange,
    handleModelChanged,
    handleCompletionParamsChange,
    selectedDatasets: selectedDatasets.filter(d => d.name),
    selectedDatasetsLoaded,
    handleOnDatasetsChange,
    rerankModelOpen,
    setRerankModelOpen,
    handleMetadataFilterModeChange: metadataConfig.handleMetadataFilterModeChange,
    handleUpdateCondition: metadataConfig.handleUpdateCondition,
    handleAddCondition: metadataConfig.handleAddCondition,
    handleRemoveCondition: metadataConfig.handleRemoveCondition,
    handleToggleConditionLogicalOperator: metadataConfig.handleToggleConditionLogicalOperator,
    handleMetadataModelChange: metadataConfig.handleMetadataModelChange,
    handleMetadataCompletionParamsChange: metadataConfig.handleMetadataCompletionParamsChange,
    availableStringVars: metadataConfig.availableStringVars,
    availableStringNodesWithParent: metadataConfig.availableStringNodesWithParent,
    availableNumberVars: metadataConfig.availableNumberVars,
    availableNumberNodesWithParent: metadataConfig.availableNumberNodesWithParent,
    showImageQueryVarSelector,
  }
}

export default useConfig