lfx-nightly 0.1.13.dev0__py3-none-any.whl → 0.2.0.dev26__py3-none-any.whl

lfx/components/datastax/astradb_graph.py

@@ -1,138 +1,36 @@
-import os
-
 import orjson
 
+from lfx.base.datastax.astradb_base import AstraDBBaseComponent
 from lfx.base.vectorstores.model import LCVectorStoreComponent, check_cached_vector_store
 from lfx.helpers.data import docs_to_data
 from lfx.inputs.inputs import (
-    BoolInput,
     DictInput,
     DropdownInput,
     FloatInput,
-    HandleInput,
     IntInput,
-    SecretStrInput,
     StrInput,
 )
 from lfx.schema.data import Data
 
 
-class AstraDBGraphVectorStoreComponent(LCVectorStoreComponent):
+class AstraDBGraphVectorStoreComponent(AstraDBBaseComponent, LCVectorStoreComponent):
     display_name: str = "Astra DB Graph"
     description: str = "Implementation of Graph Vector Store using Astra DB"
     name = "AstraDBGraph"
-    documentation: str = "https://docs.langflow.org/bundles-datastax#astra-db-graph"
+    documentation: str = "https://docs.langflow.org/bundles-datastax"
     icon: str = "AstraDB"
+    legacy: bool = True
+    replacement = ["datastax.GraphRAG"]
 
     inputs = [
-        SecretStrInput(
-            name="token",
-            display_name="Astra DB Application Token",
-            info="Authentication token for accessing Astra DB.",
-            value="ASTRA_DB_APPLICATION_TOKEN",
-            required=True,
-            advanced=os.getenv("ASTRA_ENHANCED", "false").lower() == "true",
-        ),
-        SecretStrInput(
-            name="api_endpoint",
-            display_name="Database" if os.getenv("ASTRA_ENHANCED", "false").lower() == "true" else "API Endpoint",
-            info="API endpoint URL for the Astra DB service.",
-            value="ASTRA_DB_API_ENDPOINT",
-            required=True,
-        ),
-        StrInput(
-            name="collection_name",
-            display_name="Collection Name",
-            info="The name of the collection within Astra DB where the vectors will be stored.",
-            required=True,
-        ),
+        *AstraDBBaseComponent.inputs,
+        *LCVectorStoreComponent.inputs,
         StrInput(
             name="metadata_incoming_links_key",
             display_name="Metadata incoming links key",
             info="Metadata key used for incoming links.",
             advanced=True,
         ),
-        *LCVectorStoreComponent.inputs,
-        StrInput(
-            name="keyspace",
-            display_name="Keyspace",
-            info="Optional keyspace within Astra DB to use for the collection.",
-            advanced=True,
-        ),
-        HandleInput(
-            name="embedding_model",
-            display_name="Embedding Model",
-            input_types=["Embeddings"],
-            info="Allows an embedding model configuration.",
-        ),
-        DropdownInput(
-            name="metric",
-            display_name="Metric",
-            info="Optional distance metric for vector comparisons in the vector store.",
-            options=["cosine", "dot_product", "euclidean"],
-            value="cosine",
-            advanced=True,
-        ),
-        IntInput(
-            name="batch_size",
-            display_name="Batch Size",
-            info="Optional number of data to process in a single batch.",
-            advanced=True,
-        ),
-        IntInput(
-            name="bulk_insert_batch_concurrency",
-            display_name="Bulk Insert Batch Concurrency",
-            info="Optional concurrency level for bulk insert operations.",
-            advanced=True,
-        ),
-        IntInput(
-            name="bulk_insert_overwrite_concurrency",
-            display_name="Bulk Insert Overwrite Concurrency",
-            info="Optional concurrency level for bulk insert operations that overwrite existing data.",
-            advanced=True,
-        ),
-        IntInput(
-            name="bulk_delete_concurrency",
-            display_name="Bulk Delete Concurrency",
-            info="Optional concurrency level for bulk delete operations.",
-            advanced=True,
-        ),
-        DropdownInput(
-            name="setup_mode",
-            display_name="Setup Mode",
-            info="Configuration mode for setting up the vector store, with options like 'Sync', or 'Off'.",
-            options=["Sync", "Off"],
-            advanced=True,
-            value="Sync",
-        ),
-        BoolInput(
-            name="pre_delete_collection",
-            display_name="Pre Delete Collection",
-            info="Boolean flag to determine whether to delete the collection before creating a new one.",
-            advanced=True,
-            value=False,
-        ),
-        StrInput(
-            name="metadata_indexing_include",
-            display_name="Metadata Indexing Include",
-            info="Optional list of metadata fields to include in the indexing.",
-            advanced=True,
-            list=True,
-        ),
-        StrInput(
-            name="metadata_indexing_exclude",
-            display_name="Metadata Indexing Exclude",
-            info="Optional list of metadata fields to exclude from the indexing.",
-            advanced=True,
-            list=True,
-        ),
-        StrInput(
-            name="collection_indexing_policy",
-            display_name="Collection Indexing Policy",
-            info='Optional JSON string for the "indexing" field of the collection. '
-            "See https://docs.datastax.com/en/astra-db-serverless/api-reference/collections.html#the-indexing-option",
-            advanced=True,
-        ),
         IntInput(
             name="number_of_results",
             display_name="Number of Results",
@@ -174,7 +72,6 @@ class AstraDBGraphVectorStoreComponent(LCVectorStoreComponent):
     @check_cached_vector_store
     def build_vector_store(self):
         try:
-            from astrapy.admin import parse_api_endpoint
             from langchain_astradb import AstraDBGraphVectorStore
             from langchain_astradb.utils.astradb import SetupMode
         except ImportError as e:
@@ -196,25 +93,14 @@ class AstraDBGraphVectorStoreComponent(LCVectorStoreComponent):
         try:
             self.log(f"Initializing Graph Vector Store {self.collection_name}")
 
-            # Handle environment parsing with try-except to avoid circular import
-            environment = None
-            if self.api_endpoint:
-                try:
-                    from astrapy.admin import parse_api_endpoint
-
-                    environment = parse_api_endpoint(self.api_endpoint).environment
-                except ImportError:
-                    self.log("Warning: Could not import parse_api_endpoint, using None for environment")
-                    environment = None
-
             vector_store = AstraDBGraphVectorStore(
                 embedding=self.embedding_model,
                 collection_name=self.collection_name,
                 metadata_incoming_links_key=self.metadata_incoming_links_key or "incoming_links",
                 token=self.token,
-                api_endpoint=self.api_endpoint,
-                namespace=self.keyspace or None,
-                environment=environment,
+                api_endpoint=self.get_api_endpoint(),
+                namespace=self.get_keyspace(),
+                environment=self.environment,
                 metric=self.metric or None,
                 batch_size=self.batch_size or None,
                 bulk_insert_batch_concurrency=self.bulk_insert_batch_concurrency or None,

lfx/components/datastax/astradb_tool.py

@@ -1,4 +1,3 @@
-import os
 from datetime import datetime, timezone
 from typing import Any
 
@@ -6,20 +5,25 @@ from astrapy import Collection, DataAPIClient, Database
 from langchain_core.tools import StructuredTool, Tool
 from pydantic import BaseModel, Field, create_model
 
+from lfx.base.datastax.astradb_base import AstraDBBaseComponent
 from lfx.base.langchain_utilities.model import LCToolComponent
-from lfx.io import BoolInput, DictInput, HandleInput, IntInput, SecretStrInput, StrInput, TableInput
+from lfx.io import BoolInput, DictInput, IntInput, StrInput, TableInput
 from lfx.log.logger import logger
 from lfx.schema.data import Data
 from lfx.schema.table import EditMode
 
 
-class AstraDBToolComponent(LCToolComponent):
+class AstraDBToolComponent(AstraDBBaseComponent, LCToolComponent):
     display_name: str = "Astra DB Tool"
     description: str = "Tool to run hybrid vector and metadata search on DataStax Astra DB Collection"
-    documentation: str = "https://docs.langflow.org/bundles-datastax#astra-db-tool"
+    documentation: str = "https://docs.langflow.org/bundles-datastax"
     icon: str = "AstraDB"
+    legacy: bool = True
+    name = "AstraDBTool"
+    replacement = ["datastax.AstraDB"]
 
     inputs = [
+        *AstraDBBaseComponent.inputs,
         StrInput(
             name="tool_name",
             display_name="Tool Name",
@@ -32,33 +36,6 @@ class AstraDBToolComponent(LCToolComponent):
             info="Describe the tool to LLM. Add any information that can help the LLM to use the tool.",
             required=True,
         ),
-        StrInput(
-            name="keyspace",
-            display_name="Keyspace Name",
-            info="The name of the keyspace within Astra where the collection is stored.",
-            value="default_keyspace",
-            advanced=True,
-        ),
-        StrInput(
-            name="collection_name",
-            display_name="Collection Name",
-            info="The name of the collection within Astra DB where the vectors will be stored.",
-            required=True,
-        ),
-        SecretStrInput(
-            name="token",
-            display_name="Astra DB Application Token",
-            info="Authentication token for accessing Astra DB.",
-            value="ASTRA_DB_APPLICATION_TOKEN",
-            required=True,
-        ),
-        SecretStrInput(
-            name="api_endpoint",
-            display_name="Database" if os.getenv("ASTRA_ENHANCED", "false").lower() == "true" else "API Endpoint",
-            info="API endpoint URL for the Astra DB service.",
-            value="ASTRA_DB_API_ENDPOINT",
-            required=True,
-        ),
         StrInput(
             name="projection_attributes",
             display_name="Projection Attributes",
@@ -175,7 +152,6 @@ class AstraDBToolComponent(LCToolComponent):
             advanced=False,
             value=False,
         ),
-        HandleInput(name="embedding", display_name="Embedding Model", input_types=["Embeddings"]),
         StrInput(
             name="semantic_search_instruction",
             display_name="Semantic Search Instruction",
@@ -190,26 +166,6 @@ class AstraDBToolComponent(LCToolComponent):
     _cached_db: Database | None = None
     _cached_collection: Collection | None = None
 
-    def _build_collection(self):
-        try:
-            from astrapy.admin import parse_api_endpoint
-        except ImportError as e:
-            msg = "Could not import Astra DB integration package. Please install it with `uv pip install astrapy`."
-            raise ImportError(msg) from e
-        if self._cached_collection:
-            return self._cached_collection
-
-        try:
-            environment = parse_api_endpoint(self.api_endpoint).environment
-            cached_client = DataAPIClient(self.token, environment=environment)
-            cached_db = cached_client.get_database(self.api_endpoint, keyspace=self.keyspace)
-            self._cached_collection = cached_db.get_collection(self.collection_name)
-        except Exception as e:
-            msg = f"Error building collection: {e}"
-            raise ValueError(msg) from e
-        else:
-            return self._cached_collection
-
     def create_args_schema(self) -> dict[str, BaseModel]:
         """DEPRECATED: This method is deprecated. Please use create_args_schema_v2 instead.
 
@@ -375,7 +331,6 @@ class AstraDBToolComponent(LCToolComponent):
 
     def run_model(self, **args) -> Data | list[Data]:
         """Run the query to get the data from the Astra DB collection."""
-        collection = self._build_collection()
         sort = {}
 
         # Build filters using the new method
@@ -405,6 +360,11 @@ class AstraDBToolComponent(LCToolComponent):
             find_options["projection"] = projection
 
         try:
+            database = self.get_database_object(api_endpoint=self.get_api_endpoint())
+            collection = database.get_collection(
+                name=self.collection_name,
+                keyspace=self.get_keyspace(),
+            )
             results = collection.find(**find_options)
         except Exception as e:
             msg = f"Error on Astra DB Tool {self.tool_name} request: {e}"