diff --git a/packages/components/nodes/retrievers/CohereRerankRetriever/CohereRerank.ts b/packages/components/nodes/retrievers/CohereRerankRetriever/CohereRerank.ts index 55f3c4aadf3..f74b836559e 100644 --- a/packages/components/nodes/retrievers/CohereRerankRetriever/CohereRerank.ts +++ b/packages/components/nodes/retrievers/CohereRerankRetriever/CohereRerank.ts @@ -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>[], @@ -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) diff --git a/packages/components/nodes/retrievers/CohereRerankRetriever/CohereRerankRetriever.ts b/packages/components/nodes/retrievers/CohereRerankRetriever/CohereRerankRetriever.ts index 3c1872b3f67..ca89ca77180 100644 --- a/packages/components/nodes/retrievers/CohereRerankRetriever/CohereRerankRetriever.ts +++ b/packages/components/nodes/retrievers/CohereRerankRetriever/CohereRerankRetriever.ts @@ -67,6 +67,15 @@ class CohereRerankRetriever_Retrievers implements INode { default: 0, additionalParams: true, optional: true + }, + { + label: 'Max Chunks Per Document', + name: 'maxChunksPerDoc', + placeholder: '10', + type: 'number', + default: 10, + additionalParams: true, + optional: true } ] } @@ -78,12 +87,14 @@ class CohereRerankRetriever_Retrievers implements INode { const cohereApiKey = getCredentialParam('cohereApiKey', credentialData, nodeData) const topK = nodeData.inputs?.topK as string let k = topK ? parseFloat(topK) : 4 + const maxChunks = nodeData.inputs?.maxChunksPerDoc as string + let max = maxChunks ? parseInt(maxChunks) : 10 if (k <= 0) { k = (baseRetriever as VectorStoreRetriever).k } - const cohereCompressor = new CohereRerank(cohereApiKey, model, k) + const cohereCompressor = new CohereRerank(cohereApiKey, model, k, max) return new ContextualCompressionRetriever({ baseCompressor: cohereCompressor, baseRetriever: baseRetriever