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

feat: Add more icons and QOL improvements to Astra DB component (#6918)

Browse source 
* Test commit

* QoL updates for Astra DB component

* [autofix.ci] apply automated fixes

* Update astradb.py

* [autofix.ci] apply automated fixes

* [autofix.ci] apply automated fixes (attempt 2/3)

* Update Vector Store RAG.json

* [autofix.ci] apply automated fixes

* Keep ordering proper in Vector RAG template

* [autofix.ci] apply automated fixes

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
Eric Hare 2025-03-05 08:58:25 -08:00 • committed by GitHub
commit e1ee081d32
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 420 additions and 372 deletions
Download patch file Download diff file Expand all files Collapse all files

181
src/backend/base/langflow/components/vectorstores/astradb.py
View file

@ -37,25 +37,25 @@ class AstraDBVectorStoreComponent(LCVectorStoreComponent):
"data": { "data": {
"node": { "node": {
"name": "create_database", "name": "create_database",
"description": "", "description": "Please allow several minutes for creation to complete.",
"display_name": "Create new database", "display_name": "Create new database",
"field_order": ["new_database_name", "cloud_provider", "region"], "field_order": ["01_new_database_name", "02_cloud_provider", "03_region"],
"template": { "template": {
"new_database_name": StrInput( "01_new_database_name": StrInput(
name="new_database_name", name="new_database_name",
display_name="Name", display_name="Name",
info="Name of the new database to create in Astra DB.", info="Name of the new database to create in Astra DB.",
required=True, required=True,
), ),
"cloud_provider": DropdownInput( "02_cloud_provider": DropdownInput(
name="cloud_provider", name="cloud_provider",
display_name="Cloud provider", display_name="Cloud provider",
info="Cloud provider for the new database.", info="Cloud provider for the new database.",
options=["Amazon Web Services", "Google Cloud Platform", "Microsoft Azure"], options=[],
required=True, required=True,
real_time_refresh=True, real_time_refresh=True,
), ),
"region": DropdownInput( "03_region": DropdownInput(
name="region", name="region",
display_name="Region", display_name="Region",
info="Region for the new database.", info="Region for the new database.",
@ -76,37 +76,37 @@ class AstraDBVectorStoreComponent(LCVectorStoreComponent):
"data": { "data": {
"node": { "node": {
"name": "create_collection", "name": "create_collection",
"description": "", "description": "Please allow several seconds for creation to complete.",
"display_name": "Create new collection", "display_name": "Create new collection",
"field_order": [ "field_order": [
"new_collection_name", "01_new_collection_name",
"embedding_generation_provider", "02_embedding_generation_provider",
"embedding_generation_model", "03_embedding_generation_model",
"dimension", "04_dimension",
], ],
"template": { "template": {
"new_collection_name": StrInput( "01_new_collection_name": StrInput(
name="new_collection_name", name="new_collection_name",
display_name="Name", display_name="Name",
info="Name of the new collection to create in Astra DB.", info="Name of the new collection to create in Astra DB.",
required=True, required=True,
), ),
"embedding_generation_provider": DropdownInput( "02_embedding_generation_provider": DropdownInput(
name="embedding_generation_provider", name="embedding_generation_provider",
display_name="Embedding generation method", display_name="Embedding generation method",
info="Provider to use for generating embeddings.", info="Provider to use for generating embeddings.",
real_time_refresh=True, real_time_refresh=True,
required=True, required=True,
options=["Bring your own", "Nvidia"], options=[],
), ),
"embedding_generation_model": DropdownInput( "03_embedding_generation_model": DropdownInput(
name="embedding_generation_model", name="embedding_generation_model",
display_name="Embedding model", display_name="Embedding model",
info="Model to use for generating embeddings.", info="Model to use for generating embeddings.",
required=True, required=True,
options=[], options=[],
), ),
"dimension": IntInput( "04_dimension": IntInput(
name="dimension", name="dimension",
display_name="Dimensions (Required only for `Bring your own`)", display_name="Dimensions (Required only for `Bring your own`)",
info="Dimensions of the embeddings to generate.", info="Dimensions of the embeddings to generate.",
@ -254,6 +254,20 @@ class AstraDBVectorStoreComponent(LCVectorStoreComponent):
def map_cloud_providers(cls): def map_cloud_providers(cls):
# TODO: Programmatically fetch the regions for each cloud provider # TODO: Programmatically fetch the regions for each cloud provider
return { return {
"dev": {
"Google Cloud Platform": {
"id": "gcp",
"regions": ["us-central1"],
},
},
# TODO: Check test regions
"test": {
"Google Cloud Platform": {
"id": "gcp",
"regions": ["us-central1"],
},
},
"prod": {
"Amazon Web Services": { "Amazon Web Services": {
"id": "aws", "id": "aws",
"regions": ["us-east-2", "ap-south-1", "eu-west-1"], "regions": ["us-east-2", "ap-south-1", "eu-west-1"],
@ -266,6 +280,7 @@ class AstraDBVectorStoreComponent(LCVectorStoreComponent):
"id": "azure", "id": "azure",
"regions": ["westus3"], "regions": ["westus3"],
}, },
},
} }
@classmethod @classmethod
@ -309,10 +324,13 @@ class AstraDBVectorStoreComponent(LCVectorStoreComponent):
# Get the admin object # Get the admin object
admin_client = client.get_admin(token=token) admin_client = client.get_admin(token=token)
# Get the environment, set to prod if null like
my_env = environment or "prod"
# Call the create database function # Call the create database function
return await admin_client.async_create_database( return await admin_client.async_create_database(
name=new_database_name, name=new_database_name,
cloud_provider=cls.map_cloud_providers()[cloud_provider]["id"], cloud_provider=cls.map_cloud_providers()[my_env][cloud_provider]["id"],
region=region, region=region,
keyspace=keyspace, keyspace=keyspace,
wait_until_active=False, wait_until_active=False,
@ -506,18 +524,24 @@ class AstraDBVectorStoreComponent(LCVectorStoreComponent):
) )
# If there is no provider, use the vector store icon # If there is no provider, use the vector store icon
if not provider_name or provider_name == "bring your own": if not provider_name or provider_name == "Bring your own":
return "vectorstores" return "vectorstores"
# Special case for certain models # Map provider casings
# TODO: Add more icons case_map = {
if provider_name == "nvidia": "nvidia": "NVIDIA",
return "NVIDIA" "openai": "OpenAI",
if provider_name == "openai": "amazon bedrock": "AmazonBedrockEmbeddings",
return "OpenAI" "azure openai": "AzureOpenAiEmbeddings",
"cohere": "Cohere",
"jina ai": "JinaAI",
"mistral ai": "MistralAI",
"upstage": "Upstage",
"voyage ai": "VoyageAI",
}
# Title case on the provider for the icon if no special case # Adjust the casing on some like nvidia
return provider_name.title() return case_map[provider_name.lower()] if provider_name.lower() in case_map else provider_name.title()
def _initialize_collection_options(self, api_endpoint: str | None = None): def _initialize_collection_options(self, api_endpoint: str | None = None):
# Nothing to generate if we don't have an API endpoint yet # Nothing to generate if we don't have an API endpoint yet
@ -560,36 +584,41 @@ class AstraDBVectorStoreComponent(LCVectorStoreComponent):
# If the collection is set, allow user to see embedding options # If the collection is set, allow user to see embedding options
build_config["collection_name"]["dialog_inputs"]["fields"]["data"]["node"]["template"][ build_config["collection_name"]["dialog_inputs"]["fields"]["data"]["node"]["template"][
"embedding_generation_provider" "02_embedding_generation_provider"
]["options"] = ["Bring your own", "Nvidia", *[key for key in vectorize_providers if key != "Nvidia"]] ]["options"] = [
"Bring your own",
"Nvidia",
*[key for key in vectorize_providers if key not in ["Bring your own", "Nvidia"]],
]
# For all not Bring your own or Nvidia providers, add metadata saying configure in Astra DB Portal # For all not Bring your own or Nvidia providers, add metadata saying configure in Astra DB Portal
provider_options = build_config["collection_name"]["dialog_inputs"]["fields"]["data"]["node"]["template"][ provider_options = build_config["collection_name"]["dialog_inputs"]["fields"]["data"]["node"]["template"][
"embedding_generation_provider" "02_embedding_generation_provider"
]["options"] ]["options"]
# Go over each possible provider and add metadata to configure in Astra DB Portal # Go over each possible provider and add metadata to configure in Astra DB Portal
for provider in provider_options: for provider in provider_options:
# Skip Bring your own and Nvidia, automatically configured # Add the icon for the provider
if provider in {"Bring your own", "Nvidia"}: my_metadata = {"icon": self.get_provider_icon(provider_name=provider)}
build_config["collection_name"]["dialog_inputs"]["fields"]["data"]["node"]["template"][
"embedding_generation_provider"
]["options_metadata"].append({"icon": self.get_provider_icon(provider_name=provider.lower())})
continue
# Skip Bring your own and Nvidia, automatically configured
if provider not in {"Bring your own", "Nvidia"}:
# Add metadata to configure in Astra DB Portal # Add metadata to configure in Astra DB Portal
my_metadata[" "] = "Configure in Astra DB Portal"
# Add the metadata to the options metadata
build_config["collection_name"]["dialog_inputs"]["fields"]["data"]["node"]["template"][ build_config["collection_name"]["dialog_inputs"]["fields"]["data"]["node"]["template"][
"embedding_generation_provider" "02_embedding_generation_provider"
]["options_metadata"].append({" ": "Configure in Astra DB Portal"}) ]["options_metadata"].append(my_metadata)
# And allow the user to see the models based on a selected provider # And allow the user to see the models based on a selected provider
embedding_provider = build_config["collection_name"]["dialog_inputs"]["fields"]["data"]["node"]["template"][ embedding_provider = build_config["collection_name"]["dialog_inputs"]["fields"]["data"]["node"]["template"][
"embedding_generation_provider" "02_embedding_generation_provider"
]["value"] ]["value"]
# Set the options for the embedding model based on the provider # Set the options for the embedding model based on the provider
build_config["collection_name"]["dialog_inputs"]["fields"]["data"]["node"]["template"][ build_config["collection_name"]["dialog_inputs"]["fields"]["data"]["node"]["template"][
"embedding_generation_model" "03_embedding_generation_model"
]["options"] = vectorize_providers.get(embedding_provider, [[], []])[1] ]["options"] = vectorize_providers.get(embedding_provider, [[], []])[1]
return build_config return build_config
@ -617,6 +646,12 @@ class AstraDBVectorStoreComponent(LCVectorStoreComponent):
# Get the list of options we have based on the token provided # Get the list of options we have based on the token provided
database_options = self._initialize_database_options() database_options = self._initialize_database_options()
# Update the list of cloud providers
my_env = self.environment or "prod"
build_config["database_name"]["dialog_inputs"]["fields"]["data"]["node"]["template"]["02_cloud_provider"][
"options"
] = list(self.map_cloud_providers()[my_env].keys())
# If we retrieved options based on the token, show the dropdown # If we retrieved options based on the token, show the dropdown
build_config["database_name"]["options"] = [db["name"] for db in database_options] build_config["database_name"]["options"] = [db["name"] for db in database_options]
build_config["database_name"]["options_metadata"] = [ build_config["database_name"]["options_metadata"] = [
@ -652,68 +687,82 @@ class AstraDBVectorStoreComponent(LCVectorStoreComponent):
async def update_build_config(self, build_config: dict, field_value: str, field_name: str | None = None): async def update_build_config(self, build_config: dict, field_value: str, field_name: str | None = None):
# Callback for database creation # Callback for database creation
if field_name == "database_name" and isinstance(field_value, dict) and "new_database_name" in field_value: if field_name == "database_name" and isinstance(field_value, dict) and "01_new_database_name" in field_value:
try: try:
await self.create_database_api( await self.create_database_api(
new_database_name=field_value["new_database_name"], new_database_name=field_value["01_new_database_name"],
token=self.token, token=self.token,
keyspace=self.get_keyspace(), keyspace=self.get_keyspace(),
environment=self.environment, environment=self.environment,
cloud_provider=field_value["cloud_provider"], cloud_provider=field_value["02_cloud_provider"],
region=field_value["region"], region=field_value["03_region"],
) )
except Exception as e: except Exception as e:
msg = f"Error creating database: {e}" msg = f"Error creating database: {e}"
raise ValueError(msg) from e raise ValueError(msg) from e
# Add the new database to the list of options # Add the new database to the list of options
build_config["database_name"]["options"] += [field_value["new_database_name"]] build_config["database_name"]["options"] += [field_value["01_new_database_name"]]
build_config["database_name"]["options_metadata"] += [{"status": "PENDING"}] build_config["database_name"]["options_metadata"] += [{"status": "PENDING"}]
return self.reset_collection_list(build_config) return self.reset_collection_list(build_config)
# This is the callback required to update the list of regions for a cloud provider # This is the callback required to update the list of regions for a cloud provider
if field_name == "database_name" and isinstance(field_value, dict) and "new_database_name" not in field_value: if (
cloud_provider = field_value["cloud_provider"] field_name == "database_name"
build_config["database_name"]["dialog_inputs"]["fields"]["data"]["node"]["template"]["region"][ and isinstance(field_value, dict)
and "01_new_database_name" not in field_value
):
# Get the cloud provider and environment
cloud_provider = field_value["02_cloud_provider"]
my_env = self.environment or "prod"
# Update the list of regions for the cloud provider
build_config["database_name"]["dialog_inputs"]["fields"]["data"]["node"]["template"]["03_region"][
"options" "options"
] = self.map_cloud_providers()[cloud_provider]["regions"] ] = self.map_cloud_providers()[my_env][cloud_provider]["regions"]
return build_config return build_config
# Callback for the creation of collections # Callback for the creation of collections
if field_name == "collection_name" and isinstance(field_value, dict) and "new_collection_name" in field_value: if (
field_name == "collection_name"
and isinstance(field_value, dict)
and "01_new_collection_name" in field_value
):
try: try:
# Get the dimension if its a BYO provider # Get the dimension if its a BYO provider
dimension = ( dimension = (
field_value["dimension"] field_value["04_dimension"]
if field_value["embedding_generation_provider"] == "Bring your own" if field_value["02_embedding_generation_provider"] == "Bring your own"
else None else None
) )
# Create the collection # Create the collection
await self.create_collection_api( await self.create_collection_api(
new_collection_name=field_value["new_collection_name"], new_collection_name=field_value["01_new_collection_name"],
token=self.token, token=self.token,
api_endpoint=build_config["api_endpoint"]["value"], api_endpoint=build_config["api_endpoint"]["value"],
environment=self.environment, environment=self.environment,
keyspace=self.get_keyspace(), keyspace=self.get_keyspace(),
dimension=dimension, dimension=dimension,
embedding_generation_provider=field_value["embedding_generation_provider"], embedding_generation_provider=field_value["02_embedding_generation_provider"],
embedding_generation_model=field_value["embedding_generation_model"], embedding_generation_model=field_value["03_embedding_generation_model"],
) )
except Exception as e: except Exception as e:
msg = f"Error creating collection: {e}" msg = f"Error creating collection: {e}"
raise ValueError(msg) from e raise ValueError(msg) from e
# Add the new collection to the list of options # Add the new collection to the list of options
build_config["collection_name"]["value"] = field_value["new_collection_name"] build_config["collection_name"]["value"] = field_value["01_new_collection_name"]
build_config["collection_name"]["options"].append(field_value["new_collection_name"]) build_config["collection_name"]["options"].append(field_value["01_new_collection_name"])
# Get the provider and model for the new collection # Get the provider and model for the new collection
generation_provider = field_value["embedding_generation_provider"] generation_provider = field_value["02_embedding_generation_provider"]
provider = generation_provider if generation_provider != "Bring your own" else None provider = (
generation_model = field_value["embedding_generation_model"] generation_provider.lower() if generation_provider and generation_provider != "Bring your own" else None
)
generation_model = field_value["03_embedding_generation_model"]
model = generation_model if generation_model and generation_model != "Bring your own" else None model = generation_model if generation_model and generation_model != "Bring your own" else None
# Set the embedding choice # Set the embedding choice
@ -721,7 +770,7 @@ class AstraDBVectorStoreComponent(LCVectorStoreComponent):
build_config["embedding_model"]["advanced"] = bool(provider) build_config["embedding_model"]["advanced"] = bool(provider)
# Add the new collection to the list of options # Add the new collection to the list of options
icon = "NVIDIA" if provider == "Nvidia" else "vectorstores" icon = self.get_provider_icon(provider_name=generation_provider)
build_config["collection_name"]["options_metadata"] += [ build_config["collection_name"]["options_metadata"] += [
{"records": 0, "provider": provider, "icon": icon, "model": model} {"records": 0, "provider": provider, "icon": icon, "model": model}
] ]
@ -732,7 +781,7 @@ class AstraDBVectorStoreComponent(LCVectorStoreComponent):
if ( if (
field_name == "collection_name" field_name == "collection_name"
and isinstance(field_value, dict) and isinstance(field_value, dict)
and "new_collection_name" not in field_value and "01_new_collection_name" not in field_value
): ):
return self.reset_provider_options(build_config) return self.reset_provider_options(build_config)
@ -789,7 +838,7 @@ class AstraDBVectorStoreComponent(LCVectorStoreComponent):
{ {
"records": 0, "records": 0,
"provider": None, "provider": None,
"icon": "", "icon": "vectorstores",
"model": None, "model": None,
} }
) )
@ -806,12 +855,8 @@ class AstraDBVectorStoreComponent(LCVectorStoreComponent):
value_of_provider = build_config["collection_name"]["options_metadata"][index_of_name]["provider"] value_of_provider = build_config["collection_name"]["options_metadata"][index_of_name]["provider"]
# If we were able to determine the Vectorize provider, set it accordingly # If we were able to determine the Vectorize provider, set it accordingly
if value_of_provider: build_config["embedding_model"]["advanced"] = bool(value_of_provider)
build_config["embedding_model"]["advanced"] = True build_config["embedding_choice"]["value"] = "Astra Vectorize" if value_of_provider else "Embedding Model"
build_config["embedding_choice"]["value"] = "Astra Vectorize"
else:
build_config["embedding_model"]["advanced"] = False
build_config["embedding_choice"]["value"] = "Embedding Model"
return build_config return build_config

591
src/backend/base/langflow/initial_setup/starter_projects/Vector Store RAG.json
View file

File diff suppressed because one or more lines are too long