lfx-nightly 0.1.11.dev0__py3-none-any.whl

lfx/components/datastax/cassandra.py

@@ -0,0 +1,92 @@
+from lfx.base.memory.model import LCChatMemoryComponent
+from lfx.field_typing.constants import Memory
+from lfx.inputs.inputs import DictInput, MessageTextInput, SecretStrInput
+
+
+class CassandraChatMemory(LCChatMemoryComponent):
+    display_name = "Cassandra Chat Memory"
+    description = "Retrieves and store chat messages from Apache Cassandra."
+    name = "CassandraChatMemory"
+    icon = "Cassandra"
+
+    inputs = [
+        MessageTextInput(
+            name="database_ref",
+            display_name="Contact Points / Astra Database ID",
+            info="Contact points for the database (or Astra DB database ID)",
+            required=True,
+        ),
+        MessageTextInput(
+            name="username", display_name="Username", info="Username for the database (leave empty for Astra DB)."
+        ),
+        SecretStrInput(
+            name="token",
+            display_name="Password / Astra DB Token",
+            info="User password for the database (or Astra DB token).",
+            required=True,
+        ),
+        MessageTextInput(
+            name="keyspace",
+            display_name="Keyspace",
+            info="Table Keyspace (or Astra DB namespace).",
+            required=True,
+        ),
+        MessageTextInput(
+            name="table_name",
+            display_name="Table Name",
+            info="The name of the table (or Astra DB collection) where vectors will be stored.",
+            required=True,
+        ),
+        MessageTextInput(
+            name="session_id", display_name="Session ID", info="Session ID for the message.", advanced=True
+        ),
+        DictInput(
+            name="cluster_kwargs",
+            display_name="Cluster arguments",
+            info="Optional dictionary of additional keyword arguments for the Cassandra cluster.",
+            advanced=True,
+            is_list=True,
+        ),
+    ]
+
+    def build_message_history(self) -> Memory:
+        from langchain_community.chat_message_histories import CassandraChatMessageHistory
+
+        try:
+            import cassio
+        except ImportError as e:
+            msg = "Could not import cassio integration package. Please install it with `pip install cassio`."
+            raise ImportError(msg) from e
+
+        from uuid import UUID
+
+        database_ref = self.database_ref
+
+        try:
+            UUID(self.database_ref)
+            is_astra = True
+        except ValueError:
+            is_astra = False
+            if "," in self.database_ref:
+                # use a copy because we can't change the type of the parameter
+                database_ref = self.database_ref.split(",")
+
+        if is_astra:
+            cassio.init(
+                database_id=database_ref,
+                token=self.token,
+                cluster_kwargs=self.cluster_kwargs,
+            )
+        else:
+            cassio.init(
+                contact_points=database_ref,
+                username=self.username,
+                password=self.token,
+                cluster_kwargs=self.cluster_kwargs,
+            )
+
+        return CassandraChatMessageHistory(
+            session_id=self.session_id,
+            table_name=self.table_name,
+            keyspace=self.keyspace,
+        )

lfx/components/datastax/create_assistant.py

@@ -0,0 +1,58 @@
+from lfx.base.astra_assistants.util import get_patched_openai_client
+from lfx.custom.custom_component.component_with_cache import ComponentWithCache
+from lfx.inputs.inputs import MultilineInput, StrInput
+from lfx.log.logger import logger
+from lfx.schema.message import Message
+from lfx.template.field.base 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)

lfx/components/datastax/create_thread.py

@@ -0,0 +1,32 @@
+from lfx.base.astra_assistants.util import get_patched_openai_client
+from lfx.custom.custom_component.component_with_cache import ComponentWithCache
+from lfx.inputs.inputs import MultilineInput
+from lfx.schema.message import Message
+from lfx.template.field.base 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)

lfx/components/datastax/dotenv.py

@@ -0,0 +1,35 @@
+import io
+
+from dotenv import load_dotenv
+
+from lfx.custom.custom_component.component import Component
+from lfx.inputs.inputs import MultilineSecretInput
+from lfx.schema.message import Message
+from lfx.template.field.base 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

lfx/components/datastax/get_assistant.py

@@ -0,0 +1,37 @@
+from lfx.base.astra_assistants.util import get_patched_openai_client
+from lfx.custom.custom_component.component_with_cache import ComponentWithCache
+from lfx.inputs.inputs import MultilineInput, StrInput
+from lfx.schema.message import Message
+from lfx.template.field.base 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)

lfx/components/datastax/getenvvar.py

@@ -0,0 +1,30 @@
+import os
+
+from lfx.custom.custom_component.component import Component
+from lfx.inputs.inputs import StrInput
+from lfx.schema.message import Message
+from lfx.template.field.base import Output
+
+
+class GetEnvVar(Component):
+    display_name = "Get Environment Variable"
+    description = "Gets the value of an environment variable from the system."
+    icon = "AstraDB"
+
+    inputs = [
+        StrInput(
+            name="env_var_name",
+            display_name="Environment Variable Name",
+            info="Name of the environment variable to get",
+        )
+    ]
+
+    outputs = [
+        Output(display_name="Environment Variable 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])

lfx/components/datastax/graph_rag.py

@@ -0,0 +1,141 @@
+import inspect
+from abc import ABC
+
+import graph_retriever.strategies as strategies_module
+from langchain_graph_retriever import GraphRetriever
+
+from lfx.base.vectorstores.model import LCVectorStoreComponent
+from lfx.helpers.data import docs_to_data
+from lfx.inputs.inputs import DropdownInput, HandleInput, MultilineInput, NestedDictInput, StrInput
+from lfx.schema.data import Data
+
+
+def traversal_strategies() -> list[str]:
+    """Retrieves a list of class names from the strategies_module.
+
+    This function uses the `inspect` module to get all the class members
+    from the `strategies_module` and returns their names as a list of strings.
+
+    Returns:
+        list[str]: A list of strategy class names.
+    """
+    classes = inspect.getmembers(strategies_module, inspect.isclass)
+    return [name for name, cls in classes if ABC not in cls.__bases__]
+
+
+class GraphRAGComponent(LCVectorStoreComponent):
+    """GraphRAGComponent is a component for performing Graph RAG traversal in a vector store.
+
+    Attributes:
+        display_name (str): The display name of the component.
+        description (str): A brief description of the component.
+        name (str): The name of the component.
+        icon (str): The icon representing the component.
+        inputs (list): A list of input configurations for the component.
+
+    Methods:
+        _build_search_args():
+            Builds the arguments required for the search operation.
+        search_documents() -> list[Data]:
+            Searches for documents using the specified strategy, edge definition, and query.
+        _edge_definition_from_input() -> tuple:
+            Processes the edge definition input and returns it as a tuple.
+    """
+
+    display_name: str = "Graph RAG"
+    description: str = "Graph RAG traversal for vector store."
+    name = "Graph RAG"
+    icon: str = "AstraDB"
+
+    inputs = [
+        HandleInput(
+            name="embedding_model",
+            display_name="Embedding Model",
+            input_types=["Embeddings"],
+            info="Specify the Embedding Model. Not required for Astra Vectorize collections.",
+            required=False,
+        ),
+        HandleInput(
+            name="vector_store",
+            display_name="Vector Store Connection",
+            input_types=["VectorStore"],
+            info="Connection to Vector Store.",
+        ),
+        StrInput(
+            name="edge_definition",
+            display_name="Edge Definition",
+            info="Edge definition for the graph traversal.",
+        ),
+        DropdownInput(
+            name="strategy",
+            display_name="Traversal Strategies",
+            options=traversal_strategies(),
+        ),
+        MultilineInput(
+            name="search_query",
+            display_name="Search Query",
+            tool_mode=True,
+        ),
+        NestedDictInput(
+            name="graphrag_strategy_kwargs",
+            display_name="Strategy Parameters",
+            info=(
+                "Optional dictionary of additional parameters for the retrieval strategy. "
+                "Please see https://datastax.github.io/graph-rag/reference/graph_retriever/strategies/ for details."
+            ),
+            advanced=True,
+        ),
+    ]
+
+    def search_documents(self) -> list[Data]:
+        """Searches for documents using the graph retriever based on the selected strategy, edge definition, and query.
+
+        Returns:
+            list[Data]: A list of retrieved documents.
+
+        Raises:
+            AttributeError: If there is an issue with attribute access.
+            TypeError: If there is a type mismatch.
+            ValueError: If there is a value error.
+        """
+        additional_params = self.graphrag_strategy_kwargs or {}
+
+        # Invoke the graph retriever based on the selected strategy, edge definition, and query
+        strategy_class = getattr(strategies_module, self.strategy)
+        retriever = GraphRetriever(
+            store=self.vector_store,
+            edges=[self._evaluate_edge_definition_input()],
+            strategy=strategy_class(**additional_params),
+        )
+
+        return docs_to_data(retriever.invoke(self.search_query))
+
+    def _edge_definition_from_input(self) -> tuple:
+        """Generates the edge definition from the input data.
+
+        Returns:
+            tuple: A tuple representing the edge definition.
+        """
+        values = self.edge_definition.split(",")
+        values = [value.strip() for value in values]
+
+        return tuple(values)
+
+    def _evaluate_edge_definition_input(self) -> tuple:
+        from graph_retriever.edges.metadata import Id
+
+        """Evaluates the edge definition, converting any function calls from strings.
+
+        Args:
+            edge_definition (tuple): The edge definition to evaluate.
+
+        Returns:
+            tuple: The evaluated edge definition.
+        """
+        evaluated_values = []
+        for value in self._edge_definition_from_input():
+            if value == "Id()":
+                evaluated_values.append(Id())  # Evaluate Id() as a function call
+            else:
+                evaluated_values.append(value)
+        return tuple(evaluated_values)