diff --git a/src/backend/base/langflow/components/vectorstores/Cassandra.py b/src/backend/base/langflow/components/vectorstores/Cassandra.py index 0e4e9eb84..14ca572f4 100644 --- a/src/backend/base/langflow/components/vectorstores/Cassandra.py +++ b/src/backend/base/langflow/components/vectorstores/Cassandra.py @@ -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 = []