Compression Retriever - Cohere Rerank - Add max chunks per document as optional parameter

This commit is contained in:
vinodkiran
2024-01-17 18:29:37 +05:30
parent 0d19dc5be4
commit 3407fa92f4
2 changed files with 16 additions and 3 deletions
@@ -7,11 +7,13 @@ export class CohereRerank extends BaseDocumentCompressor {
private COHERE_API_URL = 'https://api.cohere.ai/v1/rerank'
private readonly model: string
private readonly k: number
constructor(cohereAPIKey: string, model: string, k: number) {
private readonly maxChunksPerDoc: number
constructor(cohereAPIKey: string, model: string, k: number, maxChunksPerDoc: number) {
super()
this.cohereAPIKey = cohereAPIKey
this.model = model
this.k = k
this.maxChunksPerDoc = maxChunksPerDoc
}
async compressDocuments(
documents: Document<Record<string, any>>[],
@@ -32,7 +34,7 @@ export class CohereRerank extends BaseDocumentCompressor {
const data = {
model: this.model,
topN: this.k,
max_chunks_per_doc: 10,
max_chunks_per_doc: this.maxChunksPerDoc,
query: query,
return_documents: false,
documents: documents.map((doc) => doc.pageContent)