lfx-nightly 0.1.13.dev2__py3-none-any.whl → 0.1.13.dev3__py3-none-any.whl

lfx/base/mcp/util.py

@@ -35,13 +35,33 @@ HTTP_INTERNAL_SERVER_ERROR = 500
 HTTP_UNAUTHORIZED = 401
 HTTP_FORBIDDEN = 403
 
-# MCP Session Manager constants
-settings = get_settings_service().settings
-MAX_SESSIONS_PER_SERVER = (
-    settings.mcp_max_sessions_per_server
-)  # Maximum number of sessions per server to prevent resource exhaustion
-SESSION_IDLE_TIMEOUT = settings.mcp_session_idle_timeout  # 5 minutes idle timeout for sessions
-SESSION_CLEANUP_INTERVAL = settings.mcp_session_cleanup_interval  # Cleanup interval in seconds
+# MCP Session Manager constants - lazy loaded
+_mcp_settings_cache: dict[str, Any] = {}
+
+
+def _get_mcp_setting(key: str, default: Any = None) -> Any:
+    """Lazy load MCP settings from settings service."""
+    if key not in _mcp_settings_cache:
+        settings = get_settings_service().settings
+        _mcp_settings_cache[key] = getattr(settings, key, default)
+    return _mcp_settings_cache[key]
+
+
+def get_max_sessions_per_server() -> int:
+    """Get maximum number of sessions per server to prevent resource exhaustion."""
+    return _get_mcp_setting("mcp_max_sessions_per_server")
+
+
+def get_session_idle_timeout() -> int:
+    """Get 5 minutes idle timeout for sessions."""
+    return _get_mcp_setting("mcp_session_idle_timeout")
+
+
+def get_session_cleanup_interval() -> int:
+    """Get cleanup interval in seconds."""
+    return _get_mcp_setting("mcp_session_cleanup_interval")
+
+
 # RFC 7230 compliant header name pattern: token = 1*tchar
 # tchar = "!" / "#" / "$" / "%" / "&" / "'" / "*" / "+" / "-" / "." /
 #         "^" / "_" / "`" / "|" / "~" / DIGIT / ALPHA
@@ -432,7 +452,7 @@ class MCPSessionManager:
         """Periodically clean up idle sessions."""
         while True:
             try:
-                await asyncio.sleep(SESSION_CLEANUP_INTERVAL)
+                await asyncio.sleep(get_session_cleanup_interval())
                 await self._cleanup_idle_sessions()
             except asyncio.CancelledError:
                 break
@@ -450,7 +470,7 @@ class MCPSessionManager:
             sessions_to_remove = []
 
             for session_id, session_info in list(sessions.items()):
-                if current_time - session_info["last_used"] > SESSION_IDLE_TIMEOUT:
+                if current_time - session_info["last_used"] > get_session_idle_timeout():
                     sessions_to_remove.append(session_id)
 
             # Clean up idle sessions
@@ -573,7 +593,7 @@ class MCPSessionManager:
                 await self._cleanup_session_by_id(server_key, session_id)
 
         # Check if we've reached the maximum number of sessions for this server
-        if len(sessions) >= MAX_SESSIONS_PER_SERVER:
+        if len(sessions) >= get_max_sessions_per_server():
             # Remove the oldest session
             oldest_session_id = min(sessions.keys(), key=lambda x: sessions[x]["last_used"])
             await logger.ainfo(

lfx/components/datastax/__init__.py

@@ -5,17 +5,19 @@ from typing import TYPE_CHECKING, Any
 from lfx.components._importing import import_mod
 
 if TYPE_CHECKING:
-    from .astra_assistant_manager import AstraAssistantManager
-    from .astra_db import AstraDBChatMemory
-    from .astra_vectorize import AstraVectorizeComponent
+    from .astradb_assistant_manager import AstraAssistantManager
+    from .astradb_chatmemory import AstraDBChatMemory
     from .astradb_cql import AstraDBCQLToolComponent
+    from .astradb_graph import AstraDBGraphVectorStoreComponent
     from .astradb_tool import AstraDBToolComponent
+    from .astradb_vectorize import AstraVectorizeComponent
     from .astradb_vectorstore import AstraDBVectorStoreComponent
     from .create_assistant import AssistantsCreateAssistant
     from .create_thread import AssistantsCreateThread
     from .dotenv import Dotenv
     from .get_assistant import AssistantsGetAssistantName
     from .getenvvar import GetEnvVar
+    from .graph_rag import GraphRAGComponent
     from .list_assistants import AssistantsListAssistants
     from .run import AssistantsRun
 
@@ -25,14 +27,16 @@ _dynamic_imports = {
     "AssistantsGetAssistantName": "get_assistant",
     "AssistantsListAssistants": "list_assistants",
     "AssistantsRun": "run",
-    "AstraAssistantManager": "astra_assistant_manager",
+    "AstraAssistantManager": "astradb_assistant_manager",
     "AstraDBCQLToolComponent": "astradb_cql",
-    "AstraDBChatMemory": "astra_db",
+    "AstraDBChatMemory": "astradb_chatmemory",
+    "AstraDBGraphVectorStoreComponent": "astradb_graph",
     "AstraDBToolComponent": "astradb_tool",
     "AstraDBVectorStoreComponent": "astradb_vectorstore",
-    "AstraVectorizeComponent": "astra_vectorize",
+    "AstraVectorizeComponent": "astradb_vectorize",
     "Dotenv": "dotenv",
     "GetEnvVar": "getenvvar",
+    "GraphRAGComponent": "graph_rag",
 }
 
 __all__ = [
@@ -44,11 +48,13 @@ __all__ = [
     "AstraAssistantManager",
     "AstraDBCQLToolComponent",
     "AstraDBChatMemory",
+    "AstraDBGraphVectorStoreComponent",
     "AstraDBToolComponent",
     "AstraDBVectorStoreComponent",
     "AstraVectorizeComponent",
     "Dotenv",
     "GetEnvVar",
+    "GraphRAGComponent",
 ]
 
 

lfx/components/datastax/{astra_assistant_manager.py → astradb_assistant_manager.py}

@@ -30,6 +30,7 @@ class AstraAssistantManager(ComponentWithCache):
     name = "Astra Assistant Agent"
     description = "Manages Assistant Interactions"
     icon = "AstraDB"
+    legacy = True
 
     inputs = [
         DropdownInput(

lfx/components/datastax/astradb_chatmemory.py

@@ -0,0 +1,40 @@
+from lfx.base.datastax.astradb_base import AstraDBBaseComponent
+from lfx.base.memory.model import LCChatMemoryComponent
+from lfx.field_typing.constants import Memory
+from lfx.inputs.inputs import MessageTextInput
+
+
+class AstraDBChatMemory(AstraDBBaseComponent, LCChatMemoryComponent):
+    display_name = "Astra DB Chat Memory"
+    description = "Retrieves and stores chat messages from Astra DB."
+    name = "AstraDBChatMemory"
+    icon: str = "AstraDB"
+
+    inputs = [
+        *AstraDBBaseComponent.inputs,
+        MessageTextInput(
+            name="session_id",
+            display_name="Session ID",
+            info="The session ID of the chat. If empty, the current session ID parameter will be used.",
+            advanced=True,
+        ),
+    ]
+
+    def build_message_history(self) -> Memory:
+        try:
+            from langchain_astradb.chat_message_histories import AstraDBChatMessageHistory
+        except ImportError as e:
+            msg = (
+                "Could not import langchain Astra DB integration package. "
+                "Please install it with `uv pip install langchain-astradb`."
+            )
+            raise ImportError(msg) from e
+
+        return AstraDBChatMessageHistory(
+            session_id=self.session_id,
+            collection_name=self.collection_name,
+            token=self.token,
+            api_endpoint=self.get_api_endpoint(),
+            namespace=self.get_keyspace(),
+            environment=self.environment,
+        )

lfx/components/datastax/astradb_cql.py

@@ -8,20 +8,22 @@ import requests
 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 DictInput, IntInput, SecretStrInput, StrInput, TableInput
+from lfx.io import DictInput, IntInput, StrInput, TableInput
 from lfx.log.logger import logger
 from lfx.schema.data import Data
 from lfx.schema.table import EditMode
 
 
-class AstraDBCQLToolComponent(LCToolComponent):
+class AstraDBCQLToolComponent(AstraDBBaseComponent, LCToolComponent):
     display_name: str = "Astra DB CQL"
     description: str = "Create a tool to get transactional data from DataStax Astra DB CQL Table"
     documentation: str = "https://docs.langflow.org/bundles-datastax#astra-db-cql"
     icon: str = "AstraDB"
 
     inputs = [
+        *AstraDBBaseComponent.inputs,
         StrInput(name="tool_name", display_name="Tool Name", info="The name of the tool.", required=True),
         StrInput(
             name="tool_description",
@@ -29,34 +31,6 @@ class AstraDBCQLToolComponent(LCToolComponent):
             info="The tool description to be passed to the model.",
             required=True,
         ),
-        StrInput(
-            name="keyspace",
-            display_name="Keyspace",
-            value="default_keyspace",
-            info="The keyspace name within Astra DB where the data is stored.",
-            required=True,
-            advanced=True,
-        ),
-        StrInput(
-            name="table_name",
-            display_name="Table Name",
-            info="The name of the table within Astra DB where the data is 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,
-        ),
-        StrInput(
-            name="api_endpoint",
-            display_name="API Endpoint",
-            info="API endpoint URL for the Astra DB service.",
-            value="ASTRA_DB_API_ENDPOINT",
-            required=True,
-        ),
         StrInput(
             name="projection_fields",
             display_name="Projection fields",
@@ -204,7 +178,7 @@ class AstraDBCQLToolComponent(LCToolComponent):
 
     def astra_rest(self, args):
         headers = {"Accept": "application/json", "X-Cassandra-Token": f"{self.token}"}
-        astra_url = f"{self.api_endpoint}/api/rest/v2/keyspaces/{self.keyspace}/{self.table_name}/"
+        astra_url = f"{self.get_api_endpoint()}/api/rest/v2/keyspaces/{self.get_keyspace()}/{self.collection_name}/"
         where = {}
 
         for param in self.tools_params:

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"
     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"
     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}"