diff --git a/src/main/reranker/BaseReranker.ts b/src/main/reranker/BaseReranker.ts index c129e19115..16d36f87fc 100644 --- a/src/main/reranker/BaseReranker.ts +++ b/src/main/reranker/BaseReranker.ts @@ -61,6 +61,12 @@ export default abstract class BaseReranker { top_n: topN } } + } else if (provider?.includes('tei')) { + return { + query, + texts: documents, + return_text: true + } } else { return { model: this.base.rerankModel, @@ -80,6 +86,13 @@ export default abstract class BaseReranker { return data.output.results } else if (provider === 'voyageai') { return data.data + } else if (provider === 'mis-tei') { + return data.map((item: any) => { + return { + index: item.index, + relevance_score: item.score + } + }) } else { return data.results } diff --git a/src/renderer/src/services/KnowledgeService.ts b/src/renderer/src/services/KnowledgeService.ts index 4ddc4360b1..03405b13cc 100644 --- a/src/renderer/src/services/KnowledgeService.ts +++ b/src/renderer/src/services/KnowledgeService.ts @@ -48,7 +48,7 @@ export const getKnowledgeBaseParams = (base: KnowledgeBase): KnowledgeBaseParams rerankBaseURL: rerankHost, rerankApiKey: rerankAiProvider.getApiKey() || 'secret', rerankModel: base.rerankModel?.id, - rerankModelProvider: base.rerankModel?.provider + rerankModelProvider: rerankProvider.name.toLowerCase() // topN: base.topN } }