dify

dify/web/app/components/datasets/common/check-rerank-model.ts

@@ -0,0 +1,70 @@
+import { RETRIEVE_METHOD, type RetrievalConfig } from '@/types/app'
+import type {
+  DefaultModelResponse,
+  Model,
+} from '@/app/components/header/account-setting/model-provider-page/declarations'
+import { RerankingModeEnum } from '@/models/datasets'
+
+export const isReRankModelSelected = ({
+  retrievalConfig,
+  rerankModelList,
+  indexMethod,
+}: {
+  retrievalConfig: RetrievalConfig
+  rerankModelList: Model[]
+  indexMethod?: string
+}) => {
+  const rerankModelSelected = (() => {
+    if (retrievalConfig.reranking_model?.reranking_model_name) {
+      const provider = rerankModelList.find(({ provider }) => provider === retrievalConfig.reranking_model?.reranking_provider_name)
+
+      return provider?.models.find(({ model }) => model === retrievalConfig.reranking_model?.reranking_model_name)
+    }
+
+    return false
+  })()
+
+  if (
+    indexMethod === 'high_quality'
+    && ([RETRIEVE_METHOD.semantic, RETRIEVE_METHOD.fullText].includes(retrievalConfig.search_method))
+    && retrievalConfig.reranking_enable
+    && !rerankModelSelected
+  )
+    return false
+
+  if (
+    indexMethod === 'high_quality'
+    && (retrievalConfig.search_method === RETRIEVE_METHOD.hybrid && retrievalConfig.reranking_mode !== RerankingModeEnum.WeightedScore)
+    && !rerankModelSelected
+  )
+    return false
+
+  return true
+}
+
+export const ensureRerankModelSelected = ({
+  rerankDefaultModel,
+  indexMethod,
+  retrievalConfig,
+}: {
+  rerankDefaultModel: DefaultModelResponse
+  retrievalConfig: RetrievalConfig
+  indexMethod?: string
+}) => {
+  const rerankModel = retrievalConfig.reranking_model?.reranking_model_name ? retrievalConfig.reranking_model : undefined
+  if (
+    indexMethod === 'high_quality'
+    && (retrievalConfig.reranking_enable || retrievalConfig.search_method === RETRIEVE_METHOD.hybrid)
+    && !rerankModel
+    && rerankDefaultModel
+  ) {
+    return {
+      ...retrievalConfig,
+      reranking_model: {
+        reranking_provider_name: rerankDefaultModel.provider.provider,
+        reranking_model_name: rerankDefaultModel.model,
+      },
+    }
+  }
+  return retrievalConfig
+}