From 1ffa56897446cb36bd1e0a5699b6bc9c89e844a4 Mon Sep 17 00:00:00 2001
From: vinodkiran <vinodkiran@usa.net>
Date: Fri, 22 Dec 2023 09:08:53 +0530
Subject: [PATCH] Addition of Pinecone MMR search

---
 .../nodes/vectorstores/Pinecone/Pinecone.ts   | 63 +++++++++++++++++--
 1 file changed, 59 insertions(+), 4 deletions(-)

diff --git a/packages/components/nodes/vectorstores/Pinecone/Pinecone.ts b/packages/components/nodes/vectorstores/Pinecone/Pinecone.ts
index 4e6967bc..3f2e3ef1 100644
--- a/packages/components/nodes/vectorstores/Pinecone/Pinecone.ts
+++ b/packages/components/nodes/vectorstores/Pinecone/Pinecone.ts
@@ -23,11 +23,11 @@ class Pinecone_VectorStores implements INode {
     constructor() {
         this.label = 'Pinecone'
         this.name = 'pinecone'
-        this.version = 1.0
+        this.version = 2.0
         this.type = 'Pinecone'
         this.icon = 'pinecone.svg'
         this.category = 'Vector Stores'
-        this.description = `Upsert embedded data and perform similarity search upon query using Pinecone, a leading fully managed hosted vector database`
+        this.description = `Upsert embedded data and perform search upon query using Pinecone, a leading fully managed hosted vector database`
         this.baseClasses = [this.type, 'VectorStoreRetriever', 'BaseRetriever']
         this.badge = 'NEW'
         this.credential = {
@@ -77,6 +77,43 @@ class Pinecone_VectorStores implements INode {
                 type: 'number',
                 additionalParams: true,
                 optional: true
+            },
+            {
+                label: 'Search Type',
+                name: 'searchType',
+                type: 'options',
+                default: 'similarity',
+                options: [
+                    {
+                        label: 'Similarity',
+                        name: 'similarity'
+                    },
+                    {
+                        label: 'Max Marginal Relevance',
+                        name: 'mmr'
+                    }
+                ],
+                additionalParams: true,
+                optional: true
+            },
+            {
+                label: 'Fetch K (for MMR Search)',
+                name: 'fetchK',
+                description: 'Number of initial documents to fetch for MMR reranking. Default to 20. Used only when the search type is MMR',
+                placeholder: '20',
+                type: 'number',
+                additionalParams: true,
+                optional: true
+            },
+            {
+                label: 'Lambda (for MMR Search)',
+                name: 'lambda',
+                description:
+                    'Number between 0 and 1 that determines the degree of diversity among the results, where 0 corresponds to maximum diversity and 1 to minimum diversity. Used only when the search type is MMR',
+                placeholder: '0.5',
+                type: 'number',
+                additionalParams: true,
+                optional: true
             }
         ]
         this.outputs = [
@@ -141,6 +178,7 @@ class Pinecone_VectorStores implements INode {
         const docs = nodeData.inputs?.document as Document[]
         const embeddings = nodeData.inputs?.embeddings as Embeddings
         const output = nodeData.outputs?.output as string
+        const searchType = nodeData.outputs?.searchType as string
         const topK = nodeData.inputs?.topK as string
         const k = topK ? parseFloat(topK) : 4
 
@@ -176,8 +214,25 @@ class Pinecone_VectorStores implements INode {
         const vectorStore = await PineconeStore.fromExistingIndex(embeddings, obj)
 
         if (output === 'retriever') {
-            const retriever = vectorStore.asRetriever(k)
-            return retriever
+            if ('mmr' === searchType) {
+                const fetchK = nodeData.inputs?.fetchK as string
+                const lambda = nodeData.inputs?.lambda as string
+                const f = fetchK ? parseInt(fetchK) : 20
+                const l = lambda ? parseFloat(lambda) : 0.5
+                const retriever = vectorStore.asRetriever({
+                    searchType: 'mmr',
+                    k: 5,
+                    searchKwargs: {
+                        fetchK: f,
+                        lambda: l
+                    }
+                })
+                return retriever
+            } else {
+                // "searchType" is "similarity"
+                const retriever = vectorStore.asRetriever(k)
+                return retriever
+            }
         } else if (output === 'vectorStore') {
             ;(vectorStore as any).k = k
             return vectorStore