From 9c0bad3883a8648b45c6f8a18690ffe40e209bc2 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Tue, 26 Mar 2024 00:33:26 -0300
Subject: [PATCH] Update AstraDB and AstraDBSearch components

---
 .../components/vectorstores/AstraDB.py        |  2 +-
 .../components/vectorstores/AstraDBSearch.py  | 52 ++++++++++++++++++-
 2 files changed, 51 insertions(+), 3 deletions(-)

diff --git a/src/backend/langflow/components/vectorstores/AstraDB.py b/src/backend/langflow/components/vectorstores/AstraDB.py
index d75ac43f5..4494d950c 100644
--- a/src/backend/langflow/components/vectorstores/AstraDB.py
+++ b/src/backend/langflow/components/vectorstores/AstraDB.py
@@ -16,7 +16,7 @@ class AstraDBVectorStoreComponent(CustomComponent):
             "inputs": {"display_name": "Inputs", "info": "Inputs to AstraDB"},
             "embedding": {"display_name": "Embedding", "info": "Embedding to use"},
             "collection_name": {"display_name": "Collection Name", "info": "Collection name"},
-            "token": {"display_name": "Token", "info": "Token to use"},
+            "token": {"display_name": "Token", "info": "Token to use", "password": True},
             "api_endpoint": {"display_name": "API Endpoint", "info": "API Endpoint to use"},
             "namespace": {"display_name": "Namespace", "info": "Namespace to use", "advanced": True},
             "metric": {"display_name": "Metric", "info": "Metric to use", "advanced": True},
diff --git a/src/backend/langflow/components/vectorstores/AstraDBSearch.py b/src/backend/langflow/components/vectorstores/AstraDBSearch.py
index 9058ef4e7..8432f01f1 100644
--- a/src/backend/langflow/components/vectorstores/AstraDBSearch.py
+++ b/src/backend/langflow/components/vectorstores/AstraDBSearch.py
@@ -2,7 +2,7 @@ from typing import List, Optional
 
 from langflow.components.vectorstores.AstraDB import AstraDBVectorStoreComponent
 from langflow.components.vectorstores.base.model import LCVectorStoreComponent
-from langflow.field_typing import Embeddings
+from langflow.field_typing import Embeddings, Text
 from langflow.schema import Record
 
 
@@ -20,13 +20,61 @@ class AstraDBSearchComponent(AstraDBVectorStoreComponent, LCVectorStoreComponent
                 "display_name": "Input Value",
                 "info": "Input value to search",
             },
+            "embedding": {"display_name": "Embedding", "info": "Embedding to use"},
+            "collection_name": {"display_name": "Collection Name", "info": "Collection name"},
+            "token": {"display_name": "Token", "info": "Token to use", "password": True},
+            "api_endpoint": {"display_name": "API Endpoint", "info": "API Endpoint to use"},
+            "namespace": {"display_name": "Namespace", "info": "Namespace to use", "advanced": True},
+            "metric": {"display_name": "Metric", "info": "Metric to use", "advanced": True},
+            "batch_size": {"display_name": "Batch Size", "info": "Batch size to use", "advanced": True},
+            "bulk_insert_batch_concurrency": {
+                "display_name": "Bulk Insert Batch Concurrency",
+                "info": "Bulk Insert Batch Concurrency to use",
+                "advanced": True,
+            },
+            "bulk_insert_overwrite_concurrency": {
+                "display_name": "Bulk Insert Overwrite Concurrency",
+                "info": "Bulk Insert Overwrite Concurrency to use",
+                "advanced": True,
+            },
+            "bulk_delete_concurrency": {
+                "display_name": "Bulk Delete Concurrency",
+                "info": "Bulk Delete Concurrency to use",
+                "advanced": True,
+            },
+            "setup_mode": {
+                "display_name": "Setup Mode",
+                "info": "Setup mode for the vector store",
+                "options": ["Sync", "Async", "Off"],
+                "advanced": True,
+            },
+            "pre_delete_collection": {
+                "display_name": "Pre Delete Collection",
+                "info": "Pre delete collection",
+                "advanced": True,
+            },
+            "metadata_indexing_include": {
+                "display_name": "Metadata Indexing Include",
+                "info": "Metadata Indexing Include",
+                "advanced": True,
+            },
+            "metadata_indexing_exclude": {
+                "display_name": "Metadata Indexing Exclude",
+                "info": "Metadata Indexing Exclude",
+                "advanced": True,
+            },
+            "collection_indexing_policy": {
+                "display_name": "Collection Indexing Policy",
+                "info": "Collection Indexing Policy",
+                "advanced": True,
+            },
         }
 
     def build(
         self,
         embedding: Embeddings,
         collection_name: str,
-        input_value: Optional[List[Record]] = None,
+        input_value: Optional[Text] = None,
         search_type: str = "Similarity",
         token: Optional[str] = None,
         api_endpoint: Optional[str] = None,