chore: Updating Cassandra Vector Store parameters format

This commit is contained in:
joaoguilhermeS 2024-06-22 16:20:09 -03:00 • committed by Gabriel Luiz Freitas Almeida
commit 49ca6c5bf1

View file

@ -41,7 +41,7 @@ class CassandraVectorStoreComponent(LCVectorStoreComponent):
name="keyspace", name="keyspace",
display_name="Keyspace", display_name="Keyspace",
info="Optional key space within Astra DB. The keyspace should already be created.", info="Optional key space within Astra DB. The keyspace should already be created.",
advanced=True, advanced=False,
), ),
IntInput( IntInput(
name="ttl_seconds", name="ttl_seconds",
@ -70,18 +70,13 @@ class CassandraVectorStoreComponent(LCVectorStoreComponent):
value="Sync", value="Sync",
advanced=True, advanced=True,
), ),
HandleInput(name="embedding", display_name="Embedding", input_types=["Embeddings"]), MultilineInput(name="search_query", display_name="Search Query"),
DataInput( DataInput(
name="vector_store_inputs", name="ingest_data",
display_name="Vector Store Inputs", display_name="Ingest Data",
is_list=True, is_list=True,
), ),
BoolInput( HandleInput(name="embedding", display_name="Embedding", input_types=["Embeddings"]),
name="add_to_vector_store",
display_name="Add to Vector Store",
info="If true, the Vector Store Inputs will be added to the Vector Store.",
),
MultilineInput(name="search_input", display_name="Search Input"),
IntInput( IntInput(
name="number_of_results", name="number_of_results",
display_name="Number of Results", display_name="Number of Results",
@ -107,33 +102,25 @@ class CassandraVectorStoreComponent(LCVectorStoreComponent):
token=self.token, token=self.token,
) )
if self.add_to_vector_store: documents = []
documents = []
for _input in self.vector_store_inputs or []:
if isinstance(_input, Data):
documents.append(_input.to_lc_document())
else:
documents.append(_input)
if documents: for _input in self.ingest_data or []:
table = Cassandra.from_documents( if isinstance(_input, Data):
documents=documents, documents.append(_input.to_lc_document())
embedding=self.embedding,
table_name=self.table_name,
keyspace=self.keyspace,
ttl_seconds=self.ttl_seconds,
batch_size=self.batch_size,
body_index_options=self.body_index_options,
)
else: else:
table = Cassandra( documents.append(_input)
embedding=self.embedding,
table_name=self.table_name, if documents:
keyspace=self.keyspace, table = Cassandra.from_documents(
ttl_seconds=self.ttl_seconds, documents=documents,
body_index_options=self.body_index_options, embedding=self.embedding,
setup_mode=self.setup_mode, table_name=self.table_name,
) keyspace=self.keyspace,
ttl_seconds=self.ttl_seconds,
batch_size=self.batch_size,
body_index_options=self.body_index_options,
)
else: else:
table = Cassandra( table = Cassandra(
embedding=self.embedding, embedding=self.embedding,
@ -149,10 +136,10 @@ class CassandraVectorStoreComponent(LCVectorStoreComponent):
def search_documents(self) -> List[Data]: def search_documents(self) -> List[Data]:
vector_store = self._build_cassandra() vector_store = self._build_cassandra()
if self.search_input and isinstance(self.search_input, str) and self.search_input.strip(): if self.search_query and isinstance(self.search_query, str) and self.search_query.strip():
try: try:
docs = vector_store.similarity_search( docs = vector_store.similarity_search(
query=self.search_input, query=self.search_query,
k=self.number_of_results, k=self.number_of_results,
) )
except KeyError as e: except KeyError as e: