feat: add datastax components bundle (#8184)

* feat: add datastax components bundle

* Update __init__.py

* Remove old astra assistants folder

* Remove old tools

* Update __init__.py

* Update test_assistants_components.py

---------

Co-authored-by: Edwin Jose <edwin.jose@datastax.com>

src/backend/base/langflow/components/datastax/__init__.py

@@ -0,0 +1,23 @@
+from .astra_assistant_manager import AstraAssistantManager
+from .astradb_cql import AstraDBCQLToolComponent
+from .astradb_tool import AstraDBToolComponent
+from .create_assistant import AssistantsCreateAssistant
+from .create_thread import AssistantsCreateThread
+from .dotenv import Dotenv
+from .get_assistant import AssistantsGetAssistantName
+from .getenvvar import GetEnvVar
+from .list_assistants import AssistantsListAssistants
+from .run import AssistantsRun
+
+__all__ = [
+    "AssistantsCreateAssistant",
+    "AssistantsCreateThread",
+    "AssistantsGetAssistantName",
+    "AssistantsListAssistants",
+    "AssistantsRun",
+    "AstraAssistantManager",
+    "AstraDBCQLToolComponent",
+    "AstraDBToolComponent",
+    "Dotenv",
+    "GetEnvVar",
+]

src/backend/base/langflow/components/datastax/astra_assistant_manager.py

@@ -0,0 +1,306 @@
+import asyncio
+from asyncio import to_thread
+from typing import TYPE_CHECKING, Any, cast
+
+from astra_assistants.astra_assistants_manager import AssistantManager
+from langchain_core.agents import AgentFinish
+from loguru import logger
+
+from langflow.base.agents.events import ExceptionWithMessageError, process_agent_events
+from langflow.base.astra_assistants.util import (
+    get_patched_openai_client,
+    litellm_model_names,
+    sync_upload,
+    wrap_base_tool_as_tool_interface,
+)
+from langflow.custom.custom_component.component_with_cache import ComponentWithCache
+from langflow.inputs import DropdownInput, FileInput, HandleInput, MultilineInput
+from langflow.memory import delete_message
+from langflow.schema.content_block import ContentBlock
+from langflow.schema.message import Message
+from langflow.template import Output
+from langflow.utils.constants import MESSAGE_SENDER_AI
+
+if TYPE_CHECKING:
+    from langflow.schema.log import SendMessageFunctionType
+
+
+class AstraAssistantManager(ComponentWithCache):
+    display_name = "Astra Assistant Agent"
+    name = "Astra Assistant Agent"
+    description = "Manages Assistant Interactions"
+    icon = "AstraDB"
+
+    inputs = [
+        DropdownInput(
+            name="model_name",
+            display_name="Model",
+            advanced=False,
+            options=litellm_model_names,
+            value="gpt-4o-mini",
+        ),
+        MultilineInput(
+            name="instructions",
+            display_name="Agent Instructions",
+            info="Instructions for the assistant, think of these as the system prompt.",
+        ),
+        HandleInput(
+            name="input_tools",
+            display_name="Tools",
+            input_types=["Tool"],
+            is_list=True,
+            required=False,
+            info="These are the tools that the agent can use to help with tasks.",
+        ),
+        # DropdownInput(
+        #    display_name="Tools",
+        #    name="tool",
+        #    options=tool_names,
+        # ),
+        MultilineInput(
+            name="user_message", display_name="User Message", info="User message to pass to the run.", tool_mode=True
+        ),
+        FileInput(
+            name="file",
+            display_name="File(s) for retrieval",
+            list=True,
+            info="Files to be sent with the message.",
+            required=False,
+            show=True,
+            file_types=[
+                "txt",
+                "md",
+                "mdx",
+                "csv",
+                "json",
+                "yaml",
+                "yml",
+                "xml",
+                "html",
+                "htm",
+                "pdf",
+                "docx",
+                "py",
+                "sh",
+                "sql",
+                "js",
+                "ts",
+                "tsx",
+                "jpg",
+                "jpeg",
+                "png",
+                "bmp",
+                "image",
+                "zip",
+                "tar",
+                "tgz",
+                "bz2",
+                "gz",
+                "c",
+                "cpp",
+                "cs",
+                "css",
+                "go",
+                "java",
+                "php",
+                "rb",
+                "tex",
+                "doc",
+                "docx",
+                "ppt",
+                "pptx",
+                "xls",
+                "xlsx",
+                "jsonl",
+            ],
+        ),
+        MultilineInput(
+            name="input_thread_id",
+            display_name="Thread ID (optional)",
+            info="ID of the thread",
+            advanced=True,
+        ),
+        MultilineInput(
+            name="input_assistant_id",
+            display_name="Assistant ID (optional)",
+            info="ID of the assistant",
+            advanced=True,
+        ),
+        MultilineInput(
+            name="env_set",
+            display_name="Environment Set",
+            info="Dummy input to allow chaining with Dotenv Component.",
+            advanced=True,
+        ),
+    ]
+
+    outputs = [
+        Output(display_name="Assistant Response", name="assistant_response", method="get_assistant_response"),
+        Output(display_name="Tool output", name="tool_output", method="get_tool_output", hidden=True),
+        Output(display_name="Thread Id", name="output_thread_id", method="get_thread_id", hidden=True),
+        Output(display_name="Assistant Id", name="output_assistant_id", method="get_assistant_id", hidden=True),
+        Output(display_name="Vector Store Id", name="output_vs_id", method="get_vs_id", hidden=True),
+    ]
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self.lock = asyncio.Lock()
+        self.initialized: bool = False
+        self._assistant_response: Message = None  # type: ignore[assignment]
+        self._tool_output: Message = None  # type: ignore[assignment]
+        self._thread_id: Message = None  # type: ignore[assignment]
+        self._assistant_id: Message = None  # type: ignore[assignment]
+        self._vs_id: Message = None  # type: ignore[assignment]
+        self.client = get_patched_openai_client(self._shared_component_cache)
+        self.input_tools: list[Any]
+
+    async def get_assistant_response(self) -> Message:
+        await self.initialize()
+        self.status = self._assistant_response
+        return self._assistant_response
+
+    async def get_vs_id(self) -> Message:
+        await self.initialize()
+        self.status = self._vs_id
+        return self._vs_id
+
+    async def get_tool_output(self) -> Message:
+        await self.initialize()
+        self.status = self._tool_output
+        return self._tool_output
+
+    async def get_thread_id(self) -> Message:
+        await self.initialize()
+        self.status = self._thread_id
+        return self._thread_id
+
+    async def get_assistant_id(self) -> Message:
+        await self.initialize()
+        self.status = self._assistant_id
+        return self._assistant_id
+
+    async def initialize(self) -> None:
+        async with self.lock:
+            if not self.initialized:
+                await self.process_inputs()
+                self.initialized = True
+
+    async def process_inputs(self) -> None:
+        logger.info(f"env_set is {self.env_set}")
+        logger.info(self.input_tools)
+        tools = []
+        tool_obj = None
+        if self.input_tools is None:
+            self.input_tools = []
+        for tool in self.input_tools:
+            tool_obj = wrap_base_tool_as_tool_interface(tool)
+            tools.append(tool_obj)
+
+        assistant_id = None
+        thread_id = None
+        if self.input_assistant_id:
+            assistant_id = self.input_assistant_id
+        if self.input_thread_id:
+            thread_id = self.input_thread_id
+
+        if hasattr(self, "graph"):
+            session_id = self.graph.session_id
+        elif hasattr(self, "_session_id"):
+            session_id = self._session_id
+        else:
+            session_id = None
+
+        agent_message = Message(
+            sender=MESSAGE_SENDER_AI,
+            sender_name=self.display_name or "Astra Assistant",
+            properties={"icon": "Bot", "state": "partial"},
+            content_blocks=[ContentBlock(title="Assistant Steps", contents=[])],
+            session_id=session_id,
+        )
+
+        assistant_manager = AssistantManager(
+            instructions=self.instructions,
+            model=self.model_name,
+            name="managed_assistant",
+            tools=tools,
+            client=self.client,
+            thread_id=thread_id,
+            assistant_id=assistant_id,
+        )
+
+        if self.file:
+            file = await to_thread(sync_upload, self.file, assistant_manager.client)
+            vector_store = assistant_manager.client.beta.vector_stores.create(name="my_vs", file_ids=[file.id])
+            assistant_tools = assistant_manager.assistant.tools
+            assistant_tools += [{"type": "file_search"}]
+            assistant = assistant_manager.client.beta.assistants.update(
+                assistant_manager.assistant.id,
+                tools=assistant_tools,
+                tool_resources={"file_search": {"vector_store_ids": [vector_store.id]}},
+            )
+            assistant_manager.assistant = assistant
+
+        async def step_iterator():
+            # Initial event
+            yield {"event": "on_chain_start", "name": "AstraAssistant", "data": {"input": {"text": self.user_message}}}
+
+            content = self.user_message
+            result = await assistant_manager.run_thread(content=content, tool=tool_obj)
+
+            # Tool usage if present
+            if "output" in result and "arguments" in result:
+                yield {"event": "on_tool_start", "name": "tool", "data": {"input": {"text": str(result["arguments"])}}}
+                yield {"event": "on_tool_end", "name": "tool", "data": {"output": result["output"]}}
+
+            if "file_search" in result and result["file_search"] is not None:
+                yield {"event": "on_tool_start", "name": "tool", "data": {"input": {"text": self.user_message}}}
+                file_search_str = ""
+                for chunk in result["file_search"].to_dict().get("chunks", []):
+                    file_search_str += f"## Chunk ID: `{chunk['chunk_id']}`\n"
+                    file_search_str += f"**Content:**\n\n```\n{chunk['content']}\n```\n\n"
+                    if "score" in chunk:
+                        file_search_str += f"**Score:** {chunk['score']}\n\n"
+                    if "file_id" in chunk:
+                        file_search_str += f"**File ID:** `{chunk['file_id']}`\n\n"
+                    if "file_name" in chunk:
+                        file_search_str += f"**File Name:** `{chunk['file_name']}`\n\n"
+                    if "bytes" in chunk:
+                        file_search_str += f"**Bytes:** {chunk['bytes']}\n\n"
+                    if "search_string" in chunk:
+                        file_search_str += f"**Search String:** {chunk['search_string']}\n\n"
+                yield {"event": "on_tool_end", "name": "tool", "data": {"output": file_search_str}}
+
+            if "text" not in result:
+                msg = f"No text in result, {result}"
+                raise ValueError(msg)
+
+            self._assistant_response = Message(text=result["text"])
+            if "decision" in result:
+                self._tool_output = Message(text=str(result["decision"].is_complete))
+            else:
+                self._tool_output = Message(text=result["text"])
+            self._thread_id = Message(text=assistant_manager.thread.id)
+            self._assistant_id = Message(text=assistant_manager.assistant.id)
+
+            # Final event - format it like AgentFinish to match the expected format
+            yield {
+                "event": "on_chain_end",
+                "name": "AstraAssistant",
+                "data": {"output": AgentFinish(return_values={"output": result["text"]}, log="")},
+            }
+
+        try:
+            if hasattr(self, "send_message"):
+                processed_result = await process_agent_events(
+                    step_iterator(),
+                    agent_message,
+                    cast("SendMessageFunctionType", self.send_message),
+                )
+                self.status = processed_result
+        except ExceptionWithMessageError as e:
+            msg_id = e.agent_message.id
+            await delete_message(id_=msg_id)
+            await self._send_message_event(e.agent_message, category="remove_message")
+            raise
+        except Exception:
+            raise

src/backend/base/langflow/components/datastax/astradb_cql.py

@@ -0,0 +1,314 @@
+import json
+import urllib
+from datetime import datetime, timezone
+from http import HTTPStatus
+from typing import Any
+
+import requests
+from langchain.pydantic_v1 import BaseModel, Field, create_model
+from langchain_core.tools import StructuredTool, Tool
+
+from langflow.base.langchain_utilities.model import LCToolComponent
+from langflow.io import DictInput, IntInput, SecretStrInput, StrInput, TableInput
+from langflow.logging import logger
+from langflow.schema import Data
+from langflow.schema.table import EditMode
+
+
+class AstraDBCQLToolComponent(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/Components/components-tools#astra-db-cql-tool"
+    icon: str = "AstraDB"
+
+    inputs = [
+        StrInput(name="tool_name", display_name="Tool Name", info="The name of the tool.", required=True),
+        StrInput(
+            name="tool_description",
+            display_name="Tool Description",
+            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",
+            info="Attributes to return separated by comma.",
+            required=True,
+            value="*",
+            advanced=True,
+        ),
+        TableInput(
+            name="tools_params",
+            display_name="Tools Parameters",
+            info="Define the structure for the tool parameters. Describe the parameters "
+            "in a way the LLM can understand how to use them. Add the parameters "
+            "respecting the table schema (Partition Keys, Clustering Keys and Indexed Fields).",
+            required=False,
+            table_schema=[
+                {
+                    "name": "name",
+                    "display_name": "Name",
+                    "type": "str",
+                    "description": "Name of the field/parameter to be used by the model.",
+                    "default": "field",
+                    "edit_mode": EditMode.INLINE,
+                },
+                {
+                    "name": "field_name",
+                    "display_name": "Field Name",
+                    "type": "str",
+                    "description": "Specify the column name to be filtered on the table. "
+                    "Leave empty if the attribute name is the same as the name of the field.",
+                    "default": "",
+                    "edit_mode": EditMode.INLINE,
+                },
+                {
+                    "name": "description",
+                    "display_name": "Description",
+                    "type": "str",
+                    "description": "Describe the purpose of the parameter.",
+                    "default": "description of tool parameter",
+                    "edit_mode": EditMode.POPOVER,
+                },
+                {
+                    "name": "mandatory",
+                    "display_name": "Is Mandatory",
+                    "type": "boolean",
+                    "edit_mode": EditMode.INLINE,
+                    "description": ("Indicate if the field is mandatory."),
+                    "options": ["True", "False"],
+                    "default": "False",
+                },
+                {
+                    "name": "is_timestamp",
+                    "display_name": "Is Timestamp",
+                    "type": "boolean",
+                    "edit_mode": EditMode.INLINE,
+                    "description": ("Indicate if the field is a timestamp."),
+                    "options": ["True", "False"],
+                    "default": "False",
+                },
+                {
+                    "name": "operator",
+                    "display_name": "Operator",
+                    "type": "str",
+                    "description": "Set the operator for the field. "
+                    "https://docs.datastax.com/en/astra-db-serverless/api-reference/documents.html#operators",
+                    "default": "$eq",
+                    "options": ["$gt", "$gte", "$lt", "$lte", "$eq", "$ne", "$in", "$nin", "$exists", "$all", "$size"],
+                    "edit_mode": EditMode.INLINE,
+                },
+            ],
+            value=[],
+        ),
+        DictInput(
+            name="partition_keys",
+            display_name="DEPRECATED: Partition Keys",
+            is_list=True,
+            info="Field name and description to the model",
+            required=False,
+            advanced=True,
+        ),
+        DictInput(
+            name="clustering_keys",
+            display_name="DEPRECATED: Clustering Keys",
+            is_list=True,
+            info="Field name and description to the model",
+            required=False,
+            advanced=True,
+        ),
+        DictInput(
+            name="static_filters",
+            display_name="Static Filters",
+            is_list=True,
+            advanced=True,
+            info="Field name and value. When filled, it will not be generated by the LLM.",
+        ),
+        IntInput(
+            name="number_of_results",
+            display_name="Number of Results",
+            info="Number of results to return.",
+            advanced=True,
+            value=5,
+        ),
+    ]
+
+    def parse_timestamp(self, timestamp_str: str) -> str:
+        """Parse a timestamp string into Astra DB REST API format.
+
+        Args:
+            timestamp_str (str): Input timestamp string
+
+        Returns:
+            str: Formatted timestamp string in YYYY-MM-DDTHH:MI:SS.000Z format
+
+        Raises:
+            ValueError: If the timestamp cannot be parsed
+        """
+        # Common datetime formats to try
+        formats = [
+            "%Y-%m-%d",  # 2024-03-21
+            "%Y-%m-%dT%H:%M:%S",  # 2024-03-21T15:30:00
+            "%Y-%m-%dT%H:%M:%S%z",  # 2024-03-21T15:30:00+0000
+            "%Y-%m-%d %H:%M:%S",  # 2024-03-21 15:30:00
+            "%d/%m/%Y",  # 21/03/2024
+            "%Y/%m/%d",  # 2024/03/21
+        ]
+
+        for fmt in formats:
+            try:
+                # Parse the date string
+                date_obj = datetime.strptime(timestamp_str, fmt).astimezone()
+
+                # If the parsed date has no timezone info, assume UTC
+                if date_obj.tzinfo is None:
+                    date_obj = date_obj.replace(tzinfo=timezone.utc)
+
+                # Convert to UTC and format
+                utc_date = date_obj.astimezone(timezone.utc)
+                return utc_date.strftime("%Y-%m-%dT%H:%M:%S.000Z")
+            except ValueError:
+                continue
+
+        msg = f"Could not parse date: {timestamp_str}"
+        logger.error(msg)
+        raise ValueError(msg)
+
+    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}/"
+        where = {}
+
+        for param in self.tools_params:
+            field_name = param["field_name"] if param["field_name"] else param["name"]
+            field_value = None
+
+            if field_name in self.static_filters:
+                field_value = self.static_filters[field_name]
+            elif param["name"] in args:
+                field_value = args[param["name"]]
+
+            if field_value is None:
+                continue
+
+            if param["is_timestamp"] == True:  # noqa: E712
+                try:
+                    field_value = self.parse_timestamp(field_value)
+                except ValueError as e:
+                    msg = f"Error parsing timestamp: {e} - Use the prompt to specify the date in the correct format"
+                    logger.error(msg)
+                    raise ValueError(msg) from e
+
+            if param["operator"] == "$exists":
+                where[field_name] = {**where.get(field_name, {}), param["operator"]: True}
+            elif param["operator"] in ["$in", "$nin", "$all"]:
+                where[field_name] = {
+                    **where.get(field_name, {}),
+                    param["operator"]: field_value.split(",") if isinstance(field_value, str) else field_value,
+                }
+            else:
+                where[field_name] = {**where.get(field_name, {}), param["operator"]: field_value}
+
+        url = f"{astra_url}?page-size={self.number_of_results}"
+        url += f"&where={json.dumps(where)}"
+
+        if self.projection_fields != "*":
+            url += f"&fields={urllib.parse.quote(self.projection_fields.replace(' ', ''))}"
+
+        res = requests.request("GET", url=url, headers=headers, timeout=10)
+
+        if int(res.status_code) >= HTTPStatus.BAD_REQUEST:
+            msg = f"Error on Astra DB CQL Tool {self.tool_name} request: {res.text}"
+            logger.error(msg)
+            raise ValueError(msg)
+
+        try:
+            res_data = res.json()
+            return res_data["data"]
+        except ValueError:
+            return res.status_code
+
+    def create_args_schema(self) -> dict[str, BaseModel]:
+        args: dict[str, tuple[Any, Field]] = {}
+
+        for param in self.tools_params:
+            field_name = param["field_name"] if param["field_name"] else param["name"]
+            if field_name not in self.static_filters:
+                if param["mandatory"]:
+                    args[param["name"]] = (str, Field(description=param["description"]))
+                else:
+                    args[param["name"]] = (str | None, Field(description=param["description"], default=None))
+
+        model = create_model("ToolInput", **args, __base__=BaseModel)
+        return {"ToolInput": model}
+
+    def build_tool(self) -> Tool:
+        """Builds a Astra DB CQL Table tool.
+
+        Args:
+            name (str, optional): The name of the tool.
+
+        Returns:
+            Tool: The built AstraDB tool.
+        """
+        schema_dict = self.create_args_schema()
+        return StructuredTool.from_function(
+            name=self.tool_name,
+            args_schema=schema_dict["ToolInput"],
+            description=self.tool_description,
+            func=self.run_model,
+            return_direct=False,
+        )
+
+    def projection_args(self, input_str: str) -> dict:
+        elements = input_str.split(",")
+        result = {}
+
+        for element in elements:
+            if element.startswith("!"):
+                result[element[1:]] = False
+            else:
+                result[element] = True
+
+        return result
+
+    def run_model(self, **args) -> Data | list[Data]:
+        results = self.astra_rest(args)
+        data: list[Data] = []
+
+        if isinstance(results, list):
+            data = [Data(data=doc) for doc in results]
+        else:
+            self.status = results
+            return []
+
+        self.status = data
+        return data

src/backend/base/langflow/components/datastax/astradb_tool.py

@@ -0,0 +1,414 @@
+import os
+from datetime import datetime, timezone
+from typing import Any
+
+from astrapy import Collection, DataAPIClient, Database
+from astrapy.admin import parse_api_endpoint
+from langchain.pydantic_v1 import BaseModel, Field, create_model
+from langchain_core.tools import StructuredTool, Tool
+
+from langflow.base.langchain_utilities.model import LCToolComponent
+from langflow.io import BoolInput, DictInput, HandleInput, IntInput, SecretStrInput, StrInput, TableInput
+from langflow.logging import logger
+from langflow.schema import Data
+from langflow.schema.table import EditMode
+
+
+class AstraDBToolComponent(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/Components/components-tools#astra-db-tool"
+    icon: str = "AstraDB"
+
+    inputs = [
+        StrInput(
+            name="tool_name",
+            display_name="Tool Name",
+            info="The name of the tool to be passed to the LLM.",
+            required=True,
+        ),
+        StrInput(
+            name="tool_description",
+            display_name="Tool Description",
+            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",
+            info="Attributes to be returned by the tool separated by comma.",
+            required=True,
+            value="*",
+            advanced=True,
+        ),
+        TableInput(
+            name="tools_params_v2",
+            display_name="Tools Parameters",
+            info="Define the structure for the tool parameters. Describe the parameters "
+            "in a way the LLM can understand how to use them.",
+            required=False,
+            table_schema=[
+                {
+                    "name": "name",
+                    "display_name": "Name",
+                    "type": "str",
+                    "description": "Specify the name of the output field/parameter for the model.",
+                    "default": "field",
+                    "edit_mode": EditMode.INLINE,
+                },
+                {
+                    "name": "attribute_name",
+                    "display_name": "Attribute Name",
+                    "type": "str",
+                    "description": "Specify the attribute name to be filtered on the collection. "
+                    "Leave empty if the attribute name is the same as the name of the field.",
+                    "default": "",
+                    "edit_mode": EditMode.INLINE,
+                },
+                {
+                    "name": "description",
+                    "display_name": "Description",
+                    "type": "str",
+                    "description": "Describe the purpose of the output field.",
+                    "default": "description of field",
+                    "edit_mode": EditMode.POPOVER,
+                },
+                {
+                    "name": "metadata",
+                    "display_name": "Is Metadata",
+                    "type": "boolean",
+                    "edit_mode": EditMode.INLINE,
+                    "description": ("Indicate if the field is included in the metadata field."),
+                    "options": ["True", "False"],
+                    "default": "False",
+                },
+                {
+                    "name": "mandatory",
+                    "display_name": "Is Mandatory",
+                    "type": "boolean",
+                    "edit_mode": EditMode.INLINE,
+                    "description": ("Indicate if the field is mandatory."),
+                    "options": ["True", "False"],
+                    "default": "False",
+                },
+                {
+                    "name": "is_timestamp",
+                    "display_name": "Is Timestamp",
+                    "type": "boolean",
+                    "edit_mode": EditMode.INLINE,
+                    "description": ("Indicate if the field is a timestamp."),
+                    "options": ["True", "False"],
+                    "default": "False",
+                },
+                {
+                    "name": "operator",
+                    "display_name": "Operator",
+                    "type": "str",
+                    "description": "Set the operator for the field. "
+                    "https://docs.datastax.com/en/astra-db-serverless/api-reference/documents.html#operators",
+                    "default": "$eq",
+                    "options": ["$gt", "$gte", "$lt", "$lte", "$eq", "$ne", "$in", "$nin", "$exists", "$all", "$size"],
+                    "edit_mode": EditMode.INLINE,
+                },
+            ],
+            value=[],
+        ),
+        DictInput(
+            name="tool_params",
+            info="DEPRECATED: Attributes to filter and description to the model. "
+            "Add ! for mandatory (e.g: !customerId)",
+            display_name="Tool params",
+            is_list=True,
+            advanced=True,
+        ),
+        DictInput(
+            name="static_filters",
+            info="Attributes to filter and correspoding value",
+            display_name="Static filters",
+            advanced=True,
+            is_list=True,
+        ),
+        IntInput(
+            name="number_of_results",
+            display_name="Number of Results",
+            info="Number of results to return.",
+            advanced=True,
+            value=5,
+        ),
+        BoolInput(
+            name="use_search_query",
+            display_name="Semantic Search",
+            info="When this parameter is activated, the search query parameter will be used to search the collection.",
+            advanced=False,
+            value=False,
+        ),
+        BoolInput(
+            name="use_vectorize",
+            display_name="Use Astra DB Vectorize",
+            info="When this parameter is activated, Astra DB Vectorize method will be used to generate the embeddings.",
+            advanced=False,
+            value=False,
+        ),
+        HandleInput(name="embedding", display_name="Embedding Model", input_types=["Embeddings"]),
+        StrInput(
+            name="semantic_search_instruction",
+            display_name="Semantic Search Instruction",
+            info="The instruction to use for the semantic search.",
+            required=True,
+            value="Search query to find relevant documents.",
+            advanced=True,
+        ),
+    ]
+
+    _cached_client: DataAPIClient | None = None
+    _cached_db: Database | None = None
+    _cached_collection: Collection | None = None
+
+    def _build_collection(self):
+        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.
+
+        It is keep only for backward compatibility.
+        """
+        logger.warning("This is the old way to define the tool parameters. Please use the new way.")
+        args: dict[str, tuple[Any, Field] | list[str]] = {}
+
+        for key in self.tool_params:
+            if key.startswith("!"):  # Mandatory
+                args[key[1:]] = (str, Field(description=self.tool_params[key]))
+            else:  # Optional
+                args[key] = (str | None, Field(description=self.tool_params[key], default=None))
+
+        if self.use_search_query:
+            args["search_query"] = (
+                str | None,
+                Field(description="Search query to find relevant documents.", default=None),
+            )
+
+        model = create_model("ToolInput", **args, __base__=BaseModel)
+        return {"ToolInput": model}
+
+    def create_args_schema_v2(self) -> dict[str, BaseModel]:
+        """Create the tool input schema using the new tool parameters configuration."""
+        args: dict[str, tuple[Any, Field] | list[str]] = {}
+
+        for tool_param in self.tools_params_v2:
+            if tool_param["mandatory"]:
+                args[tool_param["name"]] = (str, Field(description=tool_param["description"]))
+            else:
+                args[tool_param["name"]] = (str | None, Field(description=tool_param["description"], default=None))
+
+        if self.use_search_query:
+            args["search_query"] = (
+                str,
+                Field(description=self.semantic_search_instruction),
+            )
+
+        model = create_model("ToolInput", **args, __base__=BaseModel)
+        return {"ToolInput": model}
+
+    def build_tool(self) -> Tool:
+        """Builds an Astra DB Collection tool.
+
+        Returns:
+            Tool: The built Astra DB tool.
+        """
+        schema_dict = self.create_args_schema() if len(self.tool_params.keys()) > 0 else self.create_args_schema_v2()
+
+        tool = StructuredTool.from_function(
+            name=self.tool_name,
+            args_schema=schema_dict["ToolInput"],
+            description=self.tool_description,
+            func=self.run_model,
+            return_direct=False,
+        )
+        self.status = "Astra DB Tool created"
+
+        return tool
+
+    def projection_args(self, input_str: str) -> dict | None:
+        """Build the projection arguments for the AstraDB query."""
+        elements = input_str.split(",")
+        result = {}
+
+        if elements == ["*"]:
+            return None
+
+        # Force the projection to exclude the $vector field as it is not required by the tool
+        result["$vector"] = False
+
+        # Fields with ! as prefix should be removed from the projection
+        for element in elements:
+            if element.startswith("!"):
+                result[element[1:]] = False
+            else:
+                result[element] = True
+
+        return result
+
+    def parse_timestamp(self, timestamp_str: str) -> datetime:
+        """Parse a timestamp string into Astra DB REST API format.
+
+        Args:
+            timestamp_str (str): Input timestamp string
+
+        Returns:
+            datetime: Datetime object
+
+        Raises:
+            ValueError: If the timestamp cannot be parsed
+        """
+        # Common datetime formats to try
+        formats = [
+            "%Y-%m-%d",  # 2024-03-21
+            "%Y-%m-%dT%H:%M:%S",  # 2024-03-21T15:30:00
+            "%Y-%m-%dT%H:%M:%S%z",  # 2024-03-21T15:30:00+0000
+            "%Y-%m-%d %H:%M:%S",  # 2024-03-21 15:30:00
+            "%d/%m/%Y",  # 21/03/2024
+            "%Y/%m/%d",  # 2024/03/21
+        ]
+
+        for fmt in formats:
+            try:
+                # Parse the date string
+                date_obj = datetime.strptime(timestamp_str, fmt).astimezone()
+
+                # If the parsed date has no timezone info, assume UTC
+                if date_obj.tzinfo is None:
+                    date_obj = date_obj.replace(tzinfo=timezone.utc)
+
+                # Convert to UTC and format
+                return date_obj.astimezone(timezone.utc)
+
+            except ValueError:
+                continue
+
+        msg = f"Could not parse date: {timestamp_str}"
+        logger.error(msg)
+        raise ValueError(msg)
+
+    def build_filter(self, args: dict, filter_settings: list) -> dict:
+        """Build filter dictionary for AstraDB query.
+
+        Args:
+            args: Dictionary of arguments from the tool
+            filter_settings: List of filter settings from tools_params_v2
+        Returns:
+            Dictionary containing the filter conditions
+        """
+        filters = {**self.static_filters}
+
+        for key, value in args.items():
+            # Skip search_query as it's handled separately
+            if key == "search_query":
+                continue
+
+            filter_setting = next((x for x in filter_settings if x["name"] == key), None)
+            if filter_setting and value is not None:
+                field_name = filter_setting["attribute_name"] if filter_setting["attribute_name"] else key
+                filter_key = field_name if not filter_setting["metadata"] else f"metadata.{field_name}"
+                if filter_setting["operator"] == "$exists":
+                    filters[filter_key] = {**filters.get(filter_key, {}), filter_setting["operator"]: True}
+                elif filter_setting["operator"] in ["$in", "$nin", "$all"]:
+                    filters[filter_key] = {
+                        **filters.get(filter_key, {}),
+                        filter_setting["operator"]: value.split(",") if isinstance(value, str) else value,
+                    }
+                elif filter_setting["is_timestamp"] == True:  # noqa: E712
+                    try:
+                        filters[filter_key] = {
+                            **filters.get(filter_key, {}),
+                            filter_setting["operator"]: self.parse_timestamp(value),
+                        }
+                    except ValueError as e:
+                        msg = f"Error parsing timestamp: {e} - Use the prompt to specify the date in the correct format"
+                        logger.error(msg)
+                        raise ValueError(msg) from e
+                else:
+                    filters[filter_key] = {**filters.get(filter_key, {}), filter_setting["operator"]: value}
+        return filters
+
+    def run_model(self, **args) -> Data | list[Data]:
+        """Run the query to get the data from the AstraDB collection."""
+        collection = self._build_collection()
+        sort = {}
+
+        # Build filters using the new method
+        filters = self.build_filter(args, self.tools_params_v2)
+
+        # Build the vector search on
+        if self.use_search_query and args["search_query"] is not None and args["search_query"] != "":
+            if self.use_vectorize:
+                sort["$vectorize"] = args["search_query"]
+            else:
+                if self.embedding is None:
+                    msg = "Embedding model is not set. Please set the embedding model or use Astra DB Vectorize."
+                    logger.error(msg)
+                    raise ValueError(msg)
+                embedding_query = self.embedding.embed_query(args["search_query"])
+                sort["$vector"] = embedding_query
+            del args["search_query"]
+
+        find_options = {
+            "filter": filters,
+            "limit": self.number_of_results,
+            "sort": sort,
+        }
+
+        projection = self.projection_args(self.projection_attributes)
+        if projection and len(projection) > 0:
+            find_options["projection"] = projection
+
+        try:
+            results = collection.find(**find_options)
+        except Exception as e:
+            msg = f"Error on Astra DB Tool {self.tool_name} request: {e}"
+            logger.error(msg)
+            raise ValueError(msg) from e
+
+        logger.info(f"Tool {self.tool_name} executed`")
+
+        data: list[Data] = [Data(data=doc) for doc in results]
+        self.status = data
+        return data

src/backend/base/langflow/components/datastax/create_assistant.py

@@ -0,0 +1,59 @@
+from loguru import logger
+
+from langflow.base.astra_assistants.util import get_patched_openai_client
+from langflow.custom.custom_component.component_with_cache import ComponentWithCache
+from langflow.inputs import MultilineInput, StrInput
+from langflow.schema.message import Message
+from langflow.template import Output
+
+
+class AssistantsCreateAssistant(ComponentWithCache):
+    icon = "AstraDB"
+    display_name = "Create Assistant"
+    description = "Creates an Assistant and returns it's id"
+
+    inputs = [
+        StrInput(
+            name="assistant_name",
+            display_name="Assistant Name",
+            info="Name for the assistant being created",
+        ),
+        StrInput(
+            name="instructions",
+            display_name="Instructions",
+            info="Instructions for the assistant, think of these as the system prompt.",
+        ),
+        StrInput(
+            name="model",
+            display_name="Model name",
+            info=(
+                "Model for the assistant.\n\n"
+                "Environment variables for provider credentials can be set with the Dotenv Component.\n\n"
+                "Models are supported via LiteLLM, "
+                "see (https://docs.litellm.ai/docs/providers) for supported model names and env vars."
+            ),
+            # refresh_model=True
+        ),
+        MultilineInput(
+            name="env_set",
+            display_name="Environment Set",
+            info="Dummy input to allow chaining with Dotenv Component.",
+        ),
+    ]
+
+    outputs = [
+        Output(display_name="Assistant ID", name="assistant_id", method="process_inputs"),
+    ]
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self.client = get_patched_openai_client(self._shared_component_cache)
+
+    def process_inputs(self) -> Message:
+        logger.info(f"env_set is {self.env_set}")
+        assistant = self.client.beta.assistants.create(
+            name=self.assistant_name,
+            instructions=self.instructions,
+            model=self.model,
+        )
+        return Message(text=assistant.id)

src/backend/base/langflow/components/datastax/create_thread.py

@@ -0,0 +1,32 @@
+from langflow.base.astra_assistants.util import get_patched_openai_client
+from langflow.custom.custom_component.component_with_cache import ComponentWithCache
+from langflow.inputs import MultilineInput
+from langflow.schema.message import Message
+from langflow.template import Output
+
+
+class AssistantsCreateThread(ComponentWithCache):
+    display_name = "Create Assistant Thread"
+    description = "Creates a thread and returns the thread id"
+    icon = "AstraDB"
+    inputs = [
+        MultilineInput(
+            name="env_set",
+            display_name="Environment Set",
+            info="Dummy input to allow chaining with Dotenv Component.",
+        ),
+    ]
+
+    outputs = [
+        Output(display_name="Thread ID", name="thread_id", method="process_inputs"),
+    ]
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self.client = get_patched_openai_client(self._shared_component_cache)
+
+    def process_inputs(self) -> Message:
+        thread = self.client.beta.threads.create()
+        thread_id = thread.id
+
+        return Message(text=thread_id)

src/backend/base/langflow/components/datastax/dotenv.py

@@ -0,0 +1,35 @@
+import io
+
+from dotenv import load_dotenv
+
+from langflow.custom import Component
+from langflow.inputs import MultilineSecretInput
+from langflow.schema.message import Message
+from langflow.template import Output
+
+
+class Dotenv(Component):
+    display_name = "Dotenv"
+    description = "Load .env file into env vars"
+    icon = "AstraDB"
+    inputs = [
+        MultilineSecretInput(
+            name="dotenv_file_content",
+            display_name="Dotenv file content",
+            info="Paste the content of your .env file directly, since contents are sensitive, "
+            "using a Global variable set as 'password' is recommended",
+        )
+    ]
+
+    outputs = [
+        Output(display_name="env_set", name="env_set", method="process_inputs"),
+    ]
+
+    def process_inputs(self) -> Message:
+        fake_file = io.StringIO(self.dotenv_file_content)
+        result = load_dotenv(stream=fake_file, override=True)
+
+        message = Message(text="No variables found in .env")
+        if result:
+            message = Message(text="Loaded .env")
+        return message

src/backend/base/langflow/components/datastax/get_assistant.py

@@ -0,0 +1,37 @@
+from langflow.base.astra_assistants.util import get_patched_openai_client
+from langflow.custom.custom_component.component_with_cache import ComponentWithCache
+from langflow.inputs import MultilineInput, StrInput
+from langflow.schema.message import Message
+from langflow.template import Output
+
+
+class AssistantsGetAssistantName(ComponentWithCache):
+    display_name = "Get Assistant name"
+    description = "Assistant by id"
+    icon = "AstraDB"
+    inputs = [
+        StrInput(
+            name="assistant_id",
+            display_name="Assistant ID",
+            info="ID of the assistant",
+        ),
+        MultilineInput(
+            name="env_set",
+            display_name="Environment Set",
+            info="Dummy input to allow chaining with Dotenv Component.",
+        ),
+    ]
+
+    outputs = [
+        Output(display_name="Assistant Name", name="assistant_name", method="process_inputs"),
+    ]
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self.client = get_patched_openai_client(self._shared_component_cache)
+
+    def process_inputs(self) -> Message:
+        assistant = self.client.beta.assistants.retrieve(
+            assistant_id=self.assistant_id,
+        )
+        return Message(text=assistant.name)

src/backend/base/langflow/components/datastax/getenvvar.py

@@ -0,0 +1,30 @@
+import os
+
+from langflow.custom import Component
+from langflow.inputs import StrInput
+from langflow.schema.message import Message
+from langflow.template import Output
+
+
+class GetEnvVar(Component):
+    display_name = "Get env var"
+    description = "Get env var"
+    icon = "AstraDB"
+
+    inputs = [
+        StrInput(
+            name="env_var_name",
+            display_name="Env var name",
+            info="Name of the environment variable to get",
+        )
+    ]
+
+    outputs = [
+        Output(display_name="Env var value", name="env_var_value", method="process_inputs"),
+    ]
+
+    def process_inputs(self) -> Message:
+        if self.env_var_name not in os.environ:
+            msg = f"Environment variable {self.env_var_name} not set"
+            raise ValueError(msg)
+        return Message(text=os.environ[self.env_var_name])

src/backend/base/langflow/components/datastax/list_assistants.py

@@ -0,0 +1,25 @@
+from langflow.base.astra_assistants.util import get_patched_openai_client
+from langflow.custom.custom_component.component_with_cache import ComponentWithCache
+from langflow.schema.message import Message
+from langflow.template.field.base import Output
+
+
+class AssistantsListAssistants(ComponentWithCache):
+    display_name = "List Assistants"
+    description = "Returns a list of assistant id's"
+    icon = "AstraDB"
+    outputs = [
+        Output(display_name="Assistants", name="assistants", method="process_inputs"),
+    ]
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self.client = get_patched_openai_client(self._shared_component_cache)
+
+    def process_inputs(self) -> Message:
+        assistants = self.client.beta.assistants.list().data
+        id_list = [assistant.id for assistant in assistants]
+        return Message(
+            # get text from list
+            text="\n".join(id_list)
+        )

src/backend/base/langflow/components/datastax/run.py

@@ -0,0 +1,89 @@
+from typing import Any
+
+from openai.lib.streaming import AssistantEventHandler
+
+from langflow.base.astra_assistants.util import get_patched_openai_client
+from langflow.custom.custom_component.component_with_cache import ComponentWithCache
+from langflow.inputs import MultilineInput
+from langflow.schema import dotdict
+from langflow.schema.message import Message
+from langflow.template import Output
+
+
+class AssistantsRun(ComponentWithCache):
+    display_name = "Run Assistant"
+    description = "Executes an Assistant Run against a thread"
+    icon = "AstraDB"
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self.client = get_patched_openai_client(self._shared_component_cache)
+        self.thread_id = None
+
+    def update_build_config(
+        self,
+        build_config: dotdict,
+        field_value: Any,
+        field_name: str | None = None,
+    ) -> None:
+        if field_name == "thread_id":
+            if field_value is None:
+                thread = self.client.beta.threads.create()
+                self.thread_id = thread.id
+            build_config["thread_id"] = field_value
+
+    inputs = [
+        MultilineInput(
+            name="assistant_id",
+            display_name="Assistant ID",
+            info=(
+                "The ID of the assistant to run. \n\n"
+                "Can be retrieved using the List Assistants component or created with the Create Assistant component."
+            ),
+        ),
+        MultilineInput(
+            name="user_message",
+            display_name="User Message",
+            info="User message to pass to the run.",
+        ),
+        MultilineInput(
+            name="thread_id",
+            display_name="Thread ID",
+            required=False,
+            info="Thread ID to use with the run. If not provided, a new thread will be created.",
+        ),
+        MultilineInput(
+            name="env_set",
+            display_name="Environment Set",
+            info="Dummy input to allow chaining with Dotenv Component.",
+        ),
+    ]
+
+    outputs = [Output(display_name="Assistant Response", name="assistant_response", method="process_inputs")]
+
+    def process_inputs(self) -> Message:
+        text = ""
+
+        if self.thread_id is None:
+            thread = self.client.beta.threads.create()
+            self.thread_id = thread.id
+
+        # add the user message
+        self.client.beta.threads.messages.create(thread_id=self.thread_id, role="user", content=self.user_message)
+
+        class EventHandler(AssistantEventHandler):
+            def __init__(self) -> None:
+                super().__init__()
+
+            def on_exception(self, exception: Exception) -> None:
+                raise exception
+
+        event_handler = EventHandler()
+        with self.client.beta.threads.runs.create_and_stream(
+            thread_id=self.thread_id,
+            assistant_id=self.assistant_id,
+            event_handler=event_handler,
+        ) as stream:
+            for part in stream.text_deltas:
+                text += part
+        return Message(text=text)

src/backend/base/langflow/components/tools/__init__.py

@@ -35,13 +35,9 @@ from .yahoo_finance import YfinanceToolComponent
 
 with warnings.catch_warnings():
     warnings.simplefilter("ignore", LangChainDeprecationWarning)
-    from .astradb import AstraDBToolComponent
-    from .astradb_cql import AstraDBCQLToolComponent
 
 __all__ = [
     "ArXivComponent",
-    "AstraDBCQLToolComponent",
-    "AstraDBToolComponent",
     "BingSearchAPIComponent",
     "CalculatorComponent",
     "CalculatorToolComponent",

src/backend/tests/integration/components/assistants/test_assistants_components.py

@@ -1,5 +1,5 @@
 import pytest
-from langflow.components.astra_assistants import (
+from langflow.components.datastax import (
     AssistantsCreateAssistant,
     AssistantsCreateThread,
     AssistantsGetAssistantName,

src/frontend/src/utils/styleUtils.ts

@@ -243,7 +243,7 @@ export const SIDEBAR_BUNDLES = [
   { display_name: "AssemblyAI", name: "assemblyai", icon: "AssemblyAI" },
   {
     display_name: "DataStax",
-    name: "astra_assistants",
+    name: "datastax",
     icon: "AstraDB",
   },
   { display_name: "Olivya", name: "olivya", icon: "Olivya" },