From 8f1528b21c98a5a8b721a50f0ba379404e1121e8 Mon Sep 17 00:00:00 2001
From: SuYao <sy20010504@gmail.com>
Date: Thu, 20 Mar 2025 14:50:09 +0800
Subject: [PATCH] fix(reranker): fix reranking API integration with own
 parameters (#3629)

---
 src/main/reranker/BaseReranker.ts             | 2 +-
 src/main/reranker/SiliconFlowReranker.ts      | 5 ++++-
 src/renderer/src/services/KnowledgeService.ts | 6 +++++-
 src/renderer/src/types/index.ts               | 2 ++
 4 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/src/main/reranker/BaseReranker.ts b/src/main/reranker/BaseReranker.ts
index 469ed2e6..d31bb40b 100644
--- a/src/main/reranker/BaseReranker.ts
+++ b/src/main/reranker/BaseReranker.ts
@@ -13,7 +13,7 @@ export default abstract class BaseReranker {
 
   public defaultHeaders() {
     return {
-      Authorization: `Bearer ${this.base.apiKey}`,
+      Authorization: `Bearer ${this.base.rerankApiKey}`,
       'Content-Type': 'application/json'
     }
   }
diff --git a/src/main/reranker/SiliconFlowReranker.ts b/src/main/reranker/SiliconFlowReranker.ts
index 8fa3de35..1730efda 100644
--- a/src/main/reranker/SiliconFlowReranker.ts
+++ b/src/main/reranker/SiliconFlowReranker.ts
@@ -10,7 +10,10 @@ export default class SiliconFlowReranker extends BaseReranker {
   }
 
   public rerank = async (query: string, searchResults: ExtractChunkData[]): Promise<ExtractChunkData[]> => {
-    const url = `${this.base.baseURL}/rerank`
+    const baseURL = this.base?.rerankBaseURL?.endsWith('/')
+      ? this.base.rerankBaseURL.slice(0, -1)
+      : this.base.rerankBaseURL
+    const url = `${baseURL}/rerank`
 
     const { data } = await axios.post(
       url,
diff --git a/src/renderer/src/services/KnowledgeService.ts b/src/renderer/src/services/KnowledgeService.ts
index 12542cef..3ebe542b 100644
--- a/src/renderer/src/services/KnowledgeService.ts
+++ b/src/renderer/src/services/KnowledgeService.ts
@@ -11,10 +11,12 @@ import FileManager from './FileManager'
 
 export const getKnowledgeBaseParams = (base: KnowledgeBase): KnowledgeBaseParams => {
   const provider = getProviderByModel(base.model)
+  const rerankProvider = getProviderByModel(base.rerankModel)
   const aiProvider = new AiProvider(provider)
+  const rerankAiProvider = new AiProvider(rerankProvider)
 
   let host = aiProvider.getBaseURL()
-
+  const rerankHost = rerankAiProvider.getBaseURL()
   if (provider.type === 'gemini') {
     host = host + '/v1beta/openai/'
   }
@@ -40,6 +42,8 @@ export const getKnowledgeBaseParams = (base: KnowledgeBase): KnowledgeBaseParams
     baseURL: host,
     chunkSize,
     chunkOverlap: base.chunkOverlap,
+    rerankBaseURL: rerankHost,
+    rerankApiKey: rerankAiProvider.getApiKey() || 'secret',
     rerankModel: base.rerankModel?.id,
     rerankModelProvider: base.rerankModel?.provider,
     topN: base.topN
diff --git a/src/renderer/src/types/index.ts b/src/renderer/src/types/index.ts
index f520895f..02351c11 100644
--- a/src/renderer/src/types/index.ts
+++ b/src/renderer/src/types/index.ts
@@ -270,6 +270,8 @@ export type KnowledgeBaseParams = {
   baseURL: string
   chunkSize?: number
   chunkOverlap?: number
+  rerankApiKey?: string
+  rerankBaseURL?: string
   rerankModel?: string
   rerankModelProvider?: string
   topN?: number