jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

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

Browse source 
* 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
Download patch file Download diff file Expand all files Collapse all files

20
src/backend/base/langflow/components/vectorstores/AstraDB.py
View file

@ -137,7 +137,7 @@ class AstraVectorStoreComponent(LCVectorStoreComponent):
), ),
] ]
def build_vector_store(self): def _build_vector_store_no_ingest(self):
try: try:
from langchain_astradb import AstraDBVectorStore from langchain_astradb import AstraDBVectorStore
from langchain_astradb.utils.astradb import SetupMode from langchain_astradb.utils.astradb import SetupMode
@ -196,11 +196,13 @@ class AstraVectorStoreComponent(LCVectorStoreComponent):
except Exception as e: except Exception as e:
raise ValueError(f"Error initializing AstraDBVectorStore: {str(e)}") from 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: if hasattr(self, "ingest_data") and self.ingest_data:
logger.debug("Ingesting data into the Vector Store.") logger.debug("Ingesting data into the Vector Store.")
self._add_documents_to_vector_store(vector_store) self._add_documents_to_vector_store(vector_store)
self.status = self._astradb_collection_to_data(vector_store.collection)
return vector_store return vector_store
def _add_documents_to_vector_store(self, 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.") logger.debug("No documents to add to the Vector Store.")
def search_documents(self) -> list[Data]: 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 input: {self.search_input}")
logger.debug(f"Search type: {self.search_type}") logger.debug(f"Search type: {self.search_type}")
@ -253,13 +255,3 @@ class AstraVectorStoreComponent(LCVectorStoreComponent):
else: else:
logger.debug("No search input provided. Skipping search.") logger.debug("No search input provided. Skipping search.")
return [] 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

17
src/backend/base/langflow/components/vectorstores/Cassandra.py
View file

@ -102,9 +102,9 @@ class CassandraVectorStoreComponent(LCVectorStoreComponent):
] ]
def build_vector_store(self) -> Cassandra: 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: try:
import cassio import cassio
except ImportError: except ImportError:
@ -142,11 +142,12 @@ class CassandraVectorStoreComponent(LCVectorStoreComponent):
documents = [] documents = []
for _input in self.ingest_data or []: if ingest:
if isinstance(_input, Data): for _input in self.ingest_data or []:
documents.append(_input.to_lc_document()) if isinstance(_input, Data):
else: documents.append(_input.to_lc_document())
documents.append(_input) else:
documents.append(_input)
if documents: if documents:
table = Cassandra.from_documents( table = Cassandra.from_documents(
@ -172,7 +173,7 @@ class CassandraVectorStoreComponent(LCVectorStoreComponent):
return table return table
def search_documents(self) -> List[Data]: 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(): if self.search_query and isinstance(self.search_query, str) and self.search_query.strip():
try: try: