Compression Retriever - Cohere Rerank - Add max chunks per document a…

…s optional parameter
FlowiseAI · Jan 17, 2024 · 3407fa9 · 3407fa9
1 parent 0d19dc5
commit 3407fa9
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 3 deletions.
diff --git 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<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)

diff --git 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