From 70706a7183e91829f9086abc72876e27b1a3233c Mon Sep 17 00:00:00 2001
From: Henry Heng <henryheng@flowiseai.com>
Date: Sun, 31 Mar 2024 00:07:59 +0800
Subject: [PATCH] Bugfix/Add astradb collection name and namespace (#2071)

add astradb collection name and namespace
---
 .../credentials/AstraApi.credential.ts        |  7 +-----
 .../nodes/vectorstores/Astra/Astra.ts         | 24 +++++++++++++++----
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/packages/components/credentials/AstraApi.credential.ts b/packages/components/credentials/AstraApi.credential.ts
index a89a259f5..abe1457c8 100644
--- a/packages/components/credentials/AstraApi.credential.ts
+++ b/packages/components/credentials/AstraApi.credential.ts
@@ -10,13 +10,8 @@ class AstraDBApi implements INodeCredential {
     constructor() {
         this.label = 'Astra DB API'
         this.name = 'AstraDBApi'
-        this.version = 1.0
+        this.version = 2.0
         this.inputs = [
-            {
-                label: 'Astra DB Collection Name',
-                name: 'collectionName',
-                type: 'string'
-            },
             {
                 label: 'Astra DB Application Token',
                 name: 'applicationToken',
diff --git a/packages/components/nodes/vectorstores/Astra/Astra.ts b/packages/components/nodes/vectorstores/Astra/Astra.ts
index b8afe3a16..f6a583e09 100644
--- a/packages/components/nodes/vectorstores/Astra/Astra.ts
+++ b/packages/components/nodes/vectorstores/Astra/Astra.ts
@@ -23,13 +23,12 @@ class Astra_VectorStores implements INode {
     constructor() {
         this.label = 'Astra'
         this.name = 'Astra'
-        this.version = 1.0
+        this.version = 2.0
         this.type = 'Astra'
         this.icon = 'astra.svg'
         this.category = 'Vector Stores'
         this.description = `Upsert embedded data and perform similarity or mmr search upon query using DataStax Astra DB, a serverless vector database that’s perfect for managing mission-critical AI workloads`
         this.baseClasses = [this.type, 'VectorStoreRetriever', 'BaseRetriever']
-        this.badge = 'NEW'
         this.credential = {
             label: 'Connect Credential',
             name: 'credential',
@@ -49,6 +48,16 @@ class Astra_VectorStores implements INode {
                 name: 'embeddings',
                 type: 'Embeddings'
             },
+            {
+                label: 'Namespace',
+                name: 'astraNamespace',
+                type: 'string'
+            },
+            {
+                label: 'Collection',
+                name: 'astraCollection',
+                type: 'string'
+            },
             {
                 label: 'Vector Dimension',
                 name: 'vectorDimension',
@@ -96,6 +105,8 @@ class Astra_VectorStores implements INode {
             const docs = nodeData.inputs?.document as Document[]
             const embeddings = nodeData.inputs?.embeddings as Embeddings
             const vectorDimension = nodeData.inputs?.vectorDimension as number
+            const astraNamespace = nodeData.inputs?.astraNamespace as string
+            const astraCollection = nodeData.inputs?.astraCollection as string
             const similarityMetric = nodeData.inputs?.similarityMetric as 'cosine' | 'euclidean' | 'dot_product' | undefined
             const credentialData = await getCredentialData(nodeData.credential ?? '', options)
 
@@ -111,7 +122,8 @@ class Astra_VectorStores implements INode {
 
             const astraConfig: AstraLibArgs = {
                 ...clientConfig,
-                collection: credentialData.collectionName ?? 'flowise_test',
+                namespace: astraNamespace ?? 'default_keyspace',
+                collection: astraCollection ?? credentialData.collectionName ?? 'flowise_test',
                 collectionOptions: {
                     vector: {
                         dimension: vectorDimension ?? 1536,
@@ -141,7 +153,8 @@ class Astra_VectorStores implements INode {
         const embeddings = nodeData.inputs?.embeddings as Embeddings
         const vectorDimension = nodeData.inputs?.vectorDimension as number
         const similarityMetric = nodeData.inputs?.similarityMetric as 'cosine' | 'euclidean' | 'dot_product' | undefined
-
+        const astraNamespace = nodeData.inputs?.astraNamespace as string
+        const astraCollection = nodeData.inputs?.astraCollection as string
         const credentialData = await getCredentialData(nodeData.credential ?? '', options)
 
         const expectedSimilarityMetric = ['cosine', 'euclidean', 'dot_product']
@@ -156,7 +169,8 @@ class Astra_VectorStores implements INode {
 
         const astraConfig: AstraLibArgs = {
             ...clientConfig,
-            collection: credentialData.collectionName ?? 'flowise_test',
+            namespace: astraNamespace ?? 'default_keyspace',
+            collection: astraCollection ?? credentialData.collectionName ?? 'flowise_test',
             collectionOptions: {
                 vector: {
                     dimension: vectorDimension ?? 1536,