fix (astra/cassandra): avoid duplicated ingestion (#2455)

* astra/cassandra: avoid duplicated ingestion

* [autofix.ci] apply automated fixes

* [autofix.ci] apply automated fixes

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
This commit is contained in:
Nicolò Boschi 2024-07-02 15:11:17 +02:00 • committed by GitHub
commit 653be0b358
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 15 additions and 22 deletions

View file

@ -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

View file

@ -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: