diff --git a/src/backend/base/langflow/components/vectorstores/AstraDB.py b/src/backend/base/langflow/components/vectorstores/AstraDB.py index aabbb91a7..eb9011e5f 100644 --- a/src/backend/base/langflow/components/vectorstores/AstraDB.py +++ b/src/backend/base/langflow/components/vectorstores/AstraDB.py @@ -137,7 +137,7 @@ class AstraVectorStoreComponent(LCVectorStoreComponent): ), ] - def build_vector_store(self): + def _build_vector_store_no_ingest(self): try: from langchain_astradb import AstraDBVectorStore from langchain_astradb.utils.astradb import SetupMode @@ -196,11 +196,13 @@ class AstraVectorStoreComponent(LCVectorStoreComponent): except Exception as e: raise ValueError(f"Error initializing AstraDBVectorStore: {str(e)}") from e + return vector_store + + def build_vector_store(self): + vector_store = self._build_vector_store_no_ingest() if hasattr(self, "ingest_data") and self.ingest_data: logger.debug("Ingesting data into the Vector Store.") self._add_documents_to_vector_store(vector_store) - - self.status = self._astradb_collection_to_data(vector_store.collection) return vector_store def _add_documents_to_vector_store(self, vector_store): @@ -221,7 +223,7 @@ class AstraVectorStoreComponent(LCVectorStoreComponent): logger.debug("No documents to add to the Vector Store.") def search_documents(self) -> list[Data]: - vector_store = self.build_vector_store() + vector_store = self._build_vector_store_no_ingest() logger.debug(f"Search input: {self.search_input}") logger.debug(f"Search type: {self.search_type}") @@ -253,13 +255,3 @@ class AstraVectorStoreComponent(LCVectorStoreComponent): else: logger.debug("No search input provided. Skipping search.") return [] - - def _astradb_collection_to_data(self, collection): - data = [] - data_dict = collection.find() - if data_dict and "data" in data_dict: - data_dict = data_dict["data"].get("documents", []) - - for item in data_dict: - data.append(Data(content=item["content"])) - return data diff --git a/src/backend/base/langflow/components/vectorstores/Cassandra.py b/src/backend/base/langflow/components/vectorstores/Cassandra.py index 053e06ec0..90800e70b 100644 --- a/src/backend/base/langflow/components/vectorstores/Cassandra.py +++ b/src/backend/base/langflow/components/vectorstores/Cassandra.py @@ -102,9 +102,9 @@ class CassandraVectorStoreComponent(LCVectorStoreComponent): ] def build_vector_store(self) -> Cassandra: - return self._build_cassandra() + return self._build_cassandra(ingest=True) - def _build_cassandra(self) -> Cassandra: + def _build_cassandra(self, ingest: bool) -> Cassandra: try: import cassio except ImportError: @@ -142,11 +142,12 @@ class CassandraVectorStoreComponent(LCVectorStoreComponent): documents = [] - for _input in self.ingest_data or []: - if isinstance(_input, Data): - documents.append(_input.to_lc_document()) - else: - documents.append(_input) + if ingest: + for _input in self.ingest_data or []: + if isinstance(_input, Data): + documents.append(_input.to_lc_document()) + else: + documents.append(_input) if documents: table = Cassandra.from_documents( @@ -172,7 +173,7 @@ class CassandraVectorStoreComponent(LCVectorStoreComponent): return table def search_documents(self) -> List[Data]: - vector_store = self._build_cassandra() + vector_store = self._build_cassandra(ingest=False) if self.search_query and isinstance(self.search_query, str) and self.search_query.strip(): try: