61 lines
1.7 KiB
TypeScript
61 lines
1.7 KiB
TypeScript
import { ExtractChunkData } from '@cherrystudio/embedjs-interfaces'
|
|
import { KnowledgeBaseParams } from '@types'
|
|
import axios from 'axios'
|
|
|
|
import BaseReranker from './BaseReranker'
|
|
|
|
export default class VoyageReranker extends BaseReranker {
|
|
constructor(base: KnowledgeBaseParams) {
|
|
super(base)
|
|
}
|
|
|
|
public rerank = async (query: string, searchResults: ExtractChunkData[]): Promise<ExtractChunkData[]> => {
|
|
let baseURL = this.base?.rerankBaseURL?.endsWith('/')
|
|
? this.base.rerankBaseURL.slice(0, -1)
|
|
: this.base.rerankBaseURL
|
|
|
|
if (baseURL && !baseURL.endsWith('/v1')) {
|
|
baseURL = `${baseURL}/v1`
|
|
}
|
|
|
|
const url = `${baseURL}/rerank`
|
|
|
|
const requestBody = {
|
|
model: this.base.rerankModel,
|
|
query,
|
|
documents: searchResults.map((doc) => doc.pageContent),
|
|
top_k: this.base.topN,
|
|
return_documents: false,
|
|
truncation: true
|
|
}
|
|
|
|
try {
|
|
const { data } = await axios.post(url, requestBody, {
|
|
headers: {
|
|
...this.defaultHeaders()
|
|
}
|
|
})
|
|
|
|
const rerankResults = data.data
|
|
|
|
const resultMap = new Map(rerankResults.map((result: any) => [result.index, result.relevance_score || 0]))
|
|
|
|
return searchResults
|
|
.map((doc: ExtractChunkData, index: number) => {
|
|
const score = resultMap.get(index)
|
|
if (score === undefined) return undefined
|
|
|
|
return {
|
|
...doc,
|
|
score
|
|
}
|
|
})
|
|
.filter((doc): doc is ExtractChunkData => doc !== undefined)
|
|
.sort((a, b) => b.score - a.score)
|
|
} catch (error: any) {
|
|
console.error('Voyage Reranker API 错误:', error.message || error)
|
|
throw new Error(`${error} - BaseUrl: ${baseURL}`)
|
|
}
|
|
}
|
|
}
|