dify/web/app/components/app/configuration/dataset-config/params-config/config-content.tsx

'use client'

import type { FC } from 'react'
import type { ModelParameterModalProps } from '@/app/components/header/account-setting/model-provider-page/model-parameter-modal'
import type { ModelConfig } from '@/app/components/workflow/types'
import type {
  DataSet,
} from '@/models/datasets'
import type {
  DatasetConfigs,
} from '@/models/debug'
import { memo, useCallback, useEffect, useMemo } from 'react'
import { useTranslation } from 'react-i18next'
import Divider from '@/app/components/base/divider'
import ScoreThresholdItem from '@/app/components/base/param-item/score-threshold-item'
import TopKItem from '@/app/components/base/param-item/top-k-item'
import Switch from '@/app/components/base/switch'
import Tooltip from '@/app/components/base/tooltip'
import { toast } from '@/app/components/base/ui/toast'
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 ModelParameterModal from '@/app/components/header/account-setting/model-provider-page/model-parameter-modal'
import ModelSelector from '@/app/components/header/account-setting/model-provider-page/model-selector'
import { useSelectedDatasetsMode } from '@/app/components/workflow/nodes/knowledge-retrieval/hooks'
import { RerankingModeEnum } from '@/models/datasets'
import { RETRIEVE_TYPE } from '@/types/app'
import { cn } from '@/utils/classnames'
import WeightedScore from './weighted-score'

type Props = {
  datasetConfigs: DatasetConfigs
  onChange: (configs: DatasetConfigs, isRetrievalModeChange?: boolean) => void
  selectedDatasets?: DataSet[]
  isInWorkflow?: boolean
  singleRetrievalModelConfig?: ModelConfig
  onSingleRetrievalModelChange?: ModelParameterModalProps['setModel']
  onSingleRetrievalModelParamsChange?: ModelParameterModalProps['onCompletionParamsChange']
}

const noopModelChange: ModelParameterModalProps['setModel'] = () => {}
const noopParamsChange: ModelParameterModalProps['onCompletionParamsChange'] = () => {}

const ConfigContent: FC<Props> = ({
  datasetConfigs,
  onChange,
  isInWorkflow,
  singleRetrievalModelConfig: singleRetrievalConfig = {} as ModelConfig,
  onSingleRetrievalModelChange = noopModelChange,
  onSingleRetrievalModelParamsChange = noopParamsChange,
  selectedDatasets = [],
}) => {
  const { t } = useTranslation()
  const selectedDatasetsMode = useSelectedDatasetsMode(selectedDatasets)
  const type = datasetConfigs.retrieval_model

  useEffect(() => {
    if (type === RETRIEVE_TYPE.oneWay) {
      onChange({
        ...datasetConfigs,
        retrieval_model: RETRIEVE_TYPE.multiWay,
      }, isInWorkflow)
    }
  }, [type, datasetConfigs, isInWorkflow, onChange])

  const {
    modelList: rerankModelList,
    currentModel: validDefaultRerankModel,
    currentProvider: validDefaultRerankProvider,
  } = useModelListAndDefaultModelAndCurrentProviderAndModel(ModelTypeEnum.rerank)

  /**
   * If reranking model is set and is valid, use the reranking model
   * Otherwise, check if the default reranking model is valid
   */
  const {
    currentModel: currentRerankModel,
  } = useCurrentProviderAndModel(
    rerankModelList,
    {
      provider: datasetConfigs.reranking_model?.reranking_provider_name || validDefaultRerankProvider?.provider || '',
      model: datasetConfigs.reranking_model?.reranking_model_name || validDefaultRerankModel?.model || '',
    },
  )

  const rerankModel = useMemo(() => {
    return {
      provider_name: datasetConfigs.reranking_model?.reranking_provider_name ?? '',
      model_name: datasetConfigs.reranking_model?.reranking_model_name ?? '',
    }
  }, [datasetConfigs.reranking_model])

  const handleParamChange = (key: string, value: number) => {
    if (key === 'top_k') {
      onChange({
        ...datasetConfigs,
        top_k: value,
      })
    }
    else if (key === 'score_threshold') {
      onChange({
        ...datasetConfigs,
        score_threshold: value,
      })
    }
  }

  const handleSwitch = (key: string, enable: boolean) => {
    if (key === 'top_k')
      return

    onChange({
      ...datasetConfigs,
      score_threshold_enabled: enable,
    })
  }

  const handleWeightedScoreChange = (value: { value: number[] }) => {
    const configs = {
      ...datasetConfigs,
      weights: {
        ...datasetConfigs.weights!,
        vector_setting: {
          ...datasetConfigs.weights!.vector_setting!,
          vector_weight: value.value[0],
        },
        keyword_setting: {
          keyword_weight: value.value[1],
        },
      },
    }
    onChange(configs)
  }

  const handleRerankModeChange = (mode: RerankingModeEnum) => {
    if (mode === datasetConfigs.reranking_mode)
      return

    if (mode === RerankingModeEnum.RerankingModel && !currentRerankModel)
      toast.error(t('errorMsg.rerankModelRequired', { ns: 'workflow' }))

    onChange({
      ...datasetConfigs,
      reranking_mode: mode,
    })
  }

  const model = singleRetrievalConfig // Legacy code, for compatibility, have to keep it

  const rerankingModeOptions = [
    {
      value: RerankingModeEnum.WeightedScore,
      label: t('weightedScore.title', { ns: 'dataset' }),
      tips: t('weightedScore.description', { ns: 'dataset' }),
    },
    {
      value: RerankingModeEnum.RerankingModel,
      label: t('modelProvider.rerankModel.key', { ns: 'common' }),
      tips: t('modelProvider.rerankModel.tip', { ns: 'common' }),
    },
  ]

  const showWeightedScore = selectedDatasetsMode.allHighQuality
    && !selectedDatasetsMode.inconsistentEmbeddingModel

  const showWeightedScorePanel = showWeightedScore && datasetConfigs.reranking_mode === RerankingModeEnum.WeightedScore && datasetConfigs.weights
  const selectedRerankMode = datasetConfigs.reranking_mode || RerankingModeEnum.RerankingModel

  const canManuallyToggleRerank = useMemo(() => {
    return (selectedDatasetsMode.allInternal && selectedDatasetsMode.allEconomic)
      || selectedDatasetsMode.allExternal
  }, [selectedDatasetsMode.allEconomic, selectedDatasetsMode.allExternal, selectedDatasetsMode.allInternal])

  const showRerankModel = useMemo(() => {
    if (!canManuallyToggleRerank)
      return true

    return datasetConfigs.reranking_enable
  }, [datasetConfigs.reranking_enable, canManuallyToggleRerank])

  const handleManuallyToggleRerank = useCallback((enable: boolean) => {
    if (!currentRerankModel && enable)
      toast.error(t('errorMsg.rerankModelRequired', { ns: 'workflow' }))
    onChange({
      ...datasetConfigs,
      reranking_enable: enable,
    })
  }, [currentRerankModel, datasetConfigs, onChange])

  return (
    <div>
      <div className="text-text-primary system-xl-semibold">{t('retrievalSettings', { ns: 'dataset' })}</div>
      <div className="text-text-tertiary system-xs-regular">
        {t('defaultRetrievalTip', { ns: 'dataset' })}
      </div>
      {type === RETRIEVE_TYPE.multiWay && (
        <>
          <div className="my-2 flex flex-col items-center py-1">
            <div className="mb-2 mr-2 shrink-0 text-text-secondary system-xs-semibold-uppercase">
              {t('rerankSettings', { ns: 'dataset' })}
            </div>
            <Divider bgStyle="gradient" className="m-0 !h-px" />
          </div>
          {
            selectedDatasetsMode.inconsistentEmbeddingModel
            && (
              <div className="mt-4 text-text-warning system-xs-medium">
                {t('inconsistentEmbeddingModelTip', { ns: 'dataset' })}
              </div>
            )
          }
          {
            selectedDatasetsMode.mixtureInternalAndExternal && (
              <div className="mt-4 text-text-warning system-xs-medium">
                {t('mixtureInternalAndExternalTip', { ns: 'dataset' })}
              </div>
            )
          }
          {
            selectedDatasetsMode.allExternal && (
              <div className="mt-4 text-text-warning system-xs-medium">
                {t('allExternalTip', { ns: 'dataset' })}
              </div>
            )
          }
          {
            selectedDatasetsMode.mixtureHighQualityAndEconomic
            && (
              <div className="mt-4 text-text-warning system-xs-medium">
                {t('mixtureHighQualityAndEconomicTip', { ns: 'dataset' })}
              </div>
            )
          }
          {
            showWeightedScore && (
              <div className="flex items-center justify-between">
                {
                  rerankingModeOptions.map(option => (
                    <div
                      key={option.value}
                      className={cn(
                        'flex h-8 w-[calc((100%-8px)/2)] cursor-pointer items-center justify-center rounded-lg border border-components-option-card-option-border bg-components-option-card-option-bg text-text-secondary system-sm-medium',
                        selectedRerankMode === option.value && 'border-[1.5px] border-components-option-card-option-selected-border bg-components-option-card-option-selected-bg text-text-primary',
                      )}
                      onClick={() => handleRerankModeChange(option.value)}
                    >
                      <div className="truncate">{option.label}</div>
                      <Tooltip
                        popupContent={(
                          <div className="w-[200px]">
                            {option.tips}
                          </div>
                        )}
                        popupClassName="ml-0.5"
                        triggerClassName="ml-0.5 w-3.5 h-3.5"
                      />
                    </div>
                  ))
                }
              </div>
            )
          }
          {
            !showWeightedScorePanel && (
              <div className="mt-2">
                <div className="flex items-center">
                  {
                    canManuallyToggleRerank && (
                      <Switch
                        size="md"
                        value={showRerankModel ?? false}
                        onChange={handleManuallyToggleRerank}
                      />
                    )
                  }
                  <div className="ml-1 leading-[32px] text-text-secondary system-sm-semibold">{t('modelProvider.rerankModel.key', { ns: 'common' })}</div>
                  <Tooltip
                    popupContent={(
                      <div className="w-[200px]">
                        {t('modelProvider.rerankModel.tip', { ns: 'common' })}
                      </div>
                    )}
                    popupClassName="ml-1"
                    triggerClassName="ml-1 w-4 h-4"
                  />
                </div>
                {
                  showRerankModel && (
                    <div>
                      <ModelSelector
                        defaultModel={rerankModel && { provider: rerankModel?.provider_name, model: rerankModel?.model_name }}
                        onSelect={(v) => {
                          onChange({
                            ...datasetConfigs,
                            reranking_model: {
                              reranking_provider_name: v.provider,
                              reranking_model_name: v.model,
                            },
                          })
                        }}
                        modelList={rerankModelList}
                      />
                    </div>
                  )
                }
              </div>
            )
          }
          {
            showWeightedScorePanel
            && (
              <div className="mt-2 space-y-4">
                <WeightedScore
                  value={{
                    value: [
                      datasetConfigs.weights!.vector_setting.vector_weight,
                      datasetConfigs.weights!.keyword_setting.keyword_weight,
                    ],
                  }}
                  onChange={handleWeightedScoreChange}
                />
                <TopKItem
                  value={datasetConfigs.top_k}
                  onChange={handleParamChange}
                  enable={true}
                />
                <ScoreThresholdItem
                  value={datasetConfigs.score_threshold as number}
                  onChange={handleParamChange}
                  enable={datasetConfigs.score_threshold_enabled}
                  hasSwitch={true}
                  onSwitchChange={handleSwitch}
                />
              </div>
            )
          }
          {
            !showWeightedScorePanel
            && (
              <div className="mt-4 space-y-4">
                <TopKItem
                  value={datasetConfigs.top_k}
                  onChange={handleParamChange}
                  enable={true}
                />
                {
                  showRerankModel && (
                    <ScoreThresholdItem
                      value={datasetConfigs.score_threshold as number}
                      onChange={handleParamChange}
                      enable={datasetConfigs.score_threshold_enabled}
                      hasSwitch={true}
                      onSwitchChange={handleSwitch}
                    />
                  )
                }
              </div>
            )
          }
        </>
      )}

      {isInWorkflow && type === RETRIEVE_TYPE.oneWay && (
        <div className="mt-4">
          <div className="flex items-center space-x-0.5">
            <div className="text-[13px] font-medium leading-[32px] text-text-primary">{t('modelProvider.systemReasoningModel.key', { ns: 'common' })}</div>
            <Tooltip
              popupContent={t('modelProvider.systemReasoningModel.tip', { ns: 'common' })}
            />
          </div>
          <ModelParameterModal
            isInWorkflow={isInWorkflow}
            popupClassName="!w-[387px]"
            isAdvancedMode={true}
            provider={model?.provider}
            completionParams={model?.completion_params}
            modelId={model?.name}
            setModel={onSingleRetrievalModelChange}
            onCompletionParamsChange={onSingleRetrievalModelParamsChange}
            hideDebugWithMultipleModel
            debugWithMultipleModel={false}
          />
        </div>
      )}
    </div>
  )
}
export default memo(ConfigContent)