29 lines
651 B
TypeScript
Raw Normal View History

2024-10-28 16:02:07 -03:00
import { CohereClient } from "cohere-ai";
import { MapDocument } from "../../controllers/v1/types";
const cohere = new CohereClient({
2024-12-11 19:46:11 -03:00
token: process.env.COHERE_API_KEY
2024-10-28 16:02:07 -03:00
});
export async function rerankDocuments(
documents: (string | Record<string, string>)[],
query: string,
topN = 3,
model = "rerank-english-v3.0"
) {
const rerank = await cohere.v2.rerank({
documents,
query,
topN,
model,
2024-12-11 19:46:11 -03:00
returnDocuments: true
2024-10-28 16:02:07 -03:00
});
2024-12-11 19:46:11 -03:00
return rerank.results
.sort((a, b) => b.relevanceScore - a.relevanceScore)
.map((x) => ({
document: x.document,
index: x.index,
relevanceScore: x.relevanceScore
}));
2024-10-28 16:02:07 -03:00
}