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

Cassandra: support for non-AstraDB databases (#2380)

Browse source 
* Added support for Contact points
* Reworded some parameters depending on astra / local cassandra usage
* Added username

Successor of https://github.com/langflow-ai/langflow/pull/2206 but with
the new framework
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-06-26 08:11:22 -07:00 • committed by GitHub
commit 070df4d302
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 55 additions and 16 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -4,6 +4,7 @@ from langchain_community.vectorstores import Cassandra
from langflow.base.vectorstores.model import LCVectorStoreComponent
from langflow.helpers.data import docs_to_data
from langflow.inputs import DictInput
from langflow.io import (
DataInput,
DropdownInput,
@ -23,24 +24,30 @@ class CassandraVectorStoreComponent(LCVectorStoreComponent):
icon = "Cassandra"
inputs = [
MessageTextInput(name="database_ref",
display_name="Contact Points / Astra Database ID",
info="Contact points for the database (or AstraDB database ID)",
required=True),
MessageTextInput(name="username",
display_name="Username",
info="Username for the database (leave empty for AstraDB)."),
SecretStrInput(
name="token",
display_name="Token",
info="Authentication token for accessing Cassandra on Astra DB.",
required=True,
),
MessageTextInput(name="database_id", display_name="Database ID", info="The Astra database ID.", required=True),
MessageTextInput(
name="table_name",
display_name="Table Name",
info="The name of the table where vectors will be stored.",
required=True,
display_name="Password / AstraDB Token",
info="User password for the database (or AstraDB token).",
required=True
),
MessageTextInput(
name="keyspace",
display_name="Keyspace",
info="Optional key space within Astra DB. The keyspace should already be created.",
advanced=False,
info="Table Keyspace (or AstraDB namespace).",
required=True,
),
MessageTextInput(
name="table_name",
display_name="Table Name",
info="The name of the table (or AstraDB collection) where vectors will be stored.",
required=True,
),
IntInput(
name="ttl_seconds",
@ -69,6 +76,13 @@ class CassandraVectorStoreComponent(LCVectorStoreComponent):
value="Sync",
advanced=True,
),
DictInput(
name="cluster_kwargs",
display_name="Cluster arguments",
info="Optional dictionary of additional keyword arguments for the Cassandra cluster.",
advanced=True,
is_list=True
),
MultilineInput(name="search_query", display_name="Search Query"),
DataInput(
name="ingest_data",
@ -96,10 +110,35 @@ class CassandraVectorStoreComponent(LCVectorStoreComponent):
"Could not import cassio integration package. " "Please install it with `pip install cassio`."
)
cassio.init(
database_id=self.database_id,
token=self.token,
)
from uuid import UUID
database_ref = self.database_ref
try:
UUID(self.database_ref)
is_astra = True
except ValueError:
is_astra = False
if "," in self.database_ref:
# use a copy because we can't change the type of the parameter
database_ref = self.database_ref.split(",")
if is_astra:
cassio.init(
database_id=database_ref,
token=self.token,
cluster_kwargs=self.cluster_kwargs,
)
else:
cassio.init(
contact_points=database_ref,
username=self.username,
password=self.token,
cluster_kwargs=self.cluster_kwargs,
)
if not self.ttl_seconds:
self.ttl_seconds = None
documents = []