From a3a3abff2c0deefca5188f9d4f6bcb855fa6d915 Mon Sep 17 00:00:00 2001
From: "codeflash-ai[bot]"
 <148906541+codeflash-ai[bot]@users.noreply.github.com>
Date: Mon, 3 Feb 2025 14:05:13 +0000
Subject: [PATCH] =?UTF-8?q?=E2=9A=A1=EF=B8=8F=20Speed=20up=20method=20`Ast?=
 =?UTF-8?q?raDBVectorStoreComponent.=5Finitialize=5Fcollection=5Foptions`?=
 =?UTF-8?q?=20by=2021%=20in=20PR=20#6085=20(`codeflash/optimize-pr6048-202?=
 =?UTF-8?q?5-02-03T13.54.58`)=20Here's=20the=20optimized=20version.?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../components/vectorstores/astradb.py        | 40 ++++++++-----------
 1 file changed, 17 insertions(+), 23 deletions(-)

diff --git a/src/backend/base/langflow/components/vectorstores/astradb.py b/src/backend/base/langflow/components/vectorstores/astradb.py
index 9ec471dbb9df..7d2b701d4459 100644
--- a/src/backend/base/langflow/components/vectorstores/astradb.py
+++ b/src/backend/base/langflow/components/vectorstores/astradb.py
@@ -376,21 +376,14 @@ def get_api_endpoint(self, *, api_endpoint: str | None = None):
         )
 
     def get_keyspace(self):
-        keyspace = self.keyspace
-
-        if keyspace:
-            return keyspace.strip()
-
-        return None
+        return self._keyspace
 
     def get_database_object(self, api_endpoint: str | None = None):
         try:
-            client = DataAPIClient(token=self.token, environment=self.environment)
-
-            return client.get_database(
+            return self._client.get_database(
                 api_endpoint=self.get_api_endpoint(api_endpoint=api_endpoint),
                 token=self.token,
-                keyspace=self.get_keyspace(),
+                keyspace=self._keyspace,
             )
         except Exception as e:
             msg = f"Error fetching database object: {e}"
@@ -457,25 +450,21 @@ def _initialize_database_options(self):
             raise ValueError(msg) from e
 
     def _initialize_collection_options(self, api_endpoint: str | None = None):
-        # Retrieve the database object
         database = self.get_database_object(api_endpoint=api_endpoint)
-        keyspace = self.get_keyspace()  # Cache the keyspace
-
-        # Get the list of collections
-        collection_list = list(database.list_collections(keyspace=keyspace))
+        keyspace = self._keyspace  # Use cached keyspace
 
-        def get_collection_metadata(col):
-            service = col.options.vector.service if col.options.vector else None
-            return {
+        # Get the list of collections and metadata using list comprehension
+        collection_list = database.list_collections(keyspace=keyspace)
+        return [
+            {
                 "name": col.name,
                 "records": self.collection_data(collection_name=col.name, database=database),
-                "provider": service.provider if service else None,
+                "provider": col.options.vector.service.provider if col.options.vector else None,
                 "icon": "",
-                "model": service.model_name if service else None,
+                "model": col.options.vector.service.model_name if col.options.vector else None,
             }
-
-        # Return the list of collections and their metadata
-        return [get_collection_metadata(col) for col in collection_list]
+            for col in collection_list
+        ]
 
     def reset_collection_list(self, build_config: dict):
         # Get the list of options we have based on the token provided
@@ -807,3 +796,8 @@ def get_retriever_kwargs(self):
             "search_type": self._map_search_type(),
             "search_kwargs": search_args,
         }
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._keyspace = self.keyspace.strip() if self.keyspace else None  # Cache keyspace
+        self._client = DataAPIClient(token=self.token, environment=self.environment)  # Initialize client once