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

lfx/components/mongodb/mongodb_atlas.py

@@ -0,0 +1,213 @@
+import tempfile
+import time
+
+import certifi
+from langchain_community.vectorstores import MongoDBAtlasVectorSearch
+from pymongo.collection import Collection
+from pymongo.operations import SearchIndexModel
+
+from lfx.base.vectorstores.model import LCVectorStoreComponent, check_cached_vector_store
+from lfx.helpers.data import docs_to_data
+from lfx.io import BoolInput, DropdownInput, HandleInput, IntInput, SecretStrInput, StrInput
+from lfx.schema.data import Data
+
+
+class MongoVectorStoreComponent(LCVectorStoreComponent):
+    display_name = "MongoDB Atlas"
+    description = "MongoDB Atlas Vector Store with search capabilities"
+    name = "MongoDBAtlasVector"
+    icon = "MongoDB"
+    INSERT_MODES = ["append", "overwrite"]
+    SIMILARITY_OPTIONS = ["cosine", "euclidean", "dotProduct"]
+    QUANTIZATION_OPTIONS = ["scalar", "binary"]
+    inputs = [
+        SecretStrInput(name="mongodb_atlas_cluster_uri", display_name="MongoDB Atlas Cluster URI", required=True),
+        BoolInput(name="enable_mtls", display_name="Enable mTLS", value=False, advanced=True, required=True),
+        SecretStrInput(
+            name="mongodb_atlas_client_cert",
+            display_name="MongoDB Atlas Combined Client Certificate",
+            required=False,
+            info="Client Certificate combined with the private key in the following format:\n "
+            "-----BEGIN PRIVATE KEY-----\n...\n -----END PRIVATE KEY-----\n-----BEGIN CERTIFICATE-----\n"
+            "...\n-----END CERTIFICATE-----\n",
+        ),
+        StrInput(name="db_name", display_name="Database Name", required=True),
+        StrInput(name="collection_name", display_name="Collection Name", required=True),
+        StrInput(
+            name="index_name",
+            display_name="Index Name",
+            required=True,
+            info="The name of Atlas Search index, it should be a Vector Search.",
+        ),
+        *LCVectorStoreComponent.inputs,
+        DropdownInput(
+            name="insert_mode",
+            display_name="Insert Mode",
+            options=INSERT_MODES,
+            value=INSERT_MODES[0],
+            info="How to insert new documents into the collection.",
+            advanced=True,
+        ),
+        HandleInput(name="embedding", display_name="Embedding", input_types=["Embeddings"]),
+        IntInput(
+            name="number_of_results",
+            display_name="Number of Results",
+            info="Number of results to return.",
+            value=4,
+            advanced=True,
+        ),
+        StrInput(
+            name="index_field",
+            display_name="Index Field",
+            advanced=True,
+            required=True,
+            info="The field to index.",
+            value="embedding",
+        ),
+        StrInput(
+            name="filter_field", display_name="Filter Field", advanced=True, info="The field to filter the index."
+        ),
+        IntInput(
+            name="number_dimensions",
+            display_name="Number of Dimensions",
+            info="Embedding Context Length.",
+            value=1536,
+            advanced=True,
+            required=True,
+        ),
+        DropdownInput(
+            name="similarity",
+            display_name="Similarity",
+            options=SIMILARITY_OPTIONS,
+            value=SIMILARITY_OPTIONS[0],
+            info="The method used to measure the similarity between vectors.",
+            advanced=True,
+        ),
+        DropdownInput(
+            name="quantization",
+            display_name="Quantization",
+            options=QUANTIZATION_OPTIONS,
+            value=None,
+            info="Quantization reduces memory costs converting 32-bit floats to smaller data types",
+            advanced=True,
+        ),
+    ]
+
+    @check_cached_vector_store
+    def build_vector_store(self) -> MongoDBAtlasVectorSearch:
+        try:
+            from pymongo import MongoClient
+        except ImportError as e:
+            msg = "Please install pymongo to use MongoDB Atlas Vector Store"
+            raise ImportError(msg) from e
+
+        # Create temporary files for the client certificate
+        if self.enable_mtls:
+            client_cert_path = None
+            try:
+                client_cert = self.mongodb_atlas_client_cert.replace(" ", "\n")
+                client_cert = client_cert.replace("-----BEGIN\nPRIVATE\nKEY-----", "-----BEGIN PRIVATE KEY-----")
+                client_cert = client_cert.replace(
+                    "-----END\nPRIVATE\nKEY-----\n-----BEGIN\nCERTIFICATE-----",
+                    "-----END PRIVATE KEY-----\n-----BEGIN CERTIFICATE-----",
+                )
+                client_cert = client_cert.replace("-----END\nCERTIFICATE-----", "-----END CERTIFICATE-----")
+                with tempfile.NamedTemporaryFile(delete=False) as client_cert_file:
+                    client_cert_file.write(client_cert.encode("utf-8"))
+                    client_cert_path = client_cert_file.name
+
+            except Exception as e:
+                msg = f"Failed to write certificate to temporary file: {e}"
+                raise ValueError(msg) from e
+
+        try:
+            mongo_client: MongoClient = (
+                MongoClient(
+                    self.mongodb_atlas_cluster_uri,
+                    tls=True,
+                    tlsCertificateKeyFile=client_cert_path,
+                    tlsCAFile=certifi.where(),
+                )
+                if self.enable_mtls
+                else MongoClient(self.mongodb_atlas_cluster_uri)
+            )
+
+            collection = mongo_client[self.db_name][self.collection_name]
+
+        except Exception as e:
+            msg = f"Failed to connect to MongoDB Atlas: {e}"
+            raise ValueError(msg) from e
+
+        # Convert DataFrame to Data if needed using parent's method
+        self.ingest_data = self._prepare_ingest_data()
+
+        documents = []
+        for _input in self.ingest_data or []:
+            if isinstance(_input, Data):
+                documents.append(_input.to_lc_document())
+            else:
+                documents.append(_input)
+
+        if documents:
+            self.__insert_mode(collection)
+
+            return MongoDBAtlasVectorSearch.from_documents(
+                documents=documents, embedding=self.embedding, collection=collection, index_name=self.index_name
+            )
+        return MongoDBAtlasVectorSearch(embedding=self.embedding, collection=collection, index_name=self.index_name)
+
+    def search_documents(self) -> list[Data]:
+        from bson.objectid import ObjectId
+
+        vector_store = self.build_vector_store()
+
+        self.verify_search_index(vector_store._collection)
+
+        if self.search_query and isinstance(self.search_query, str):
+            docs = vector_store.similarity_search(
+                query=self.search_query,
+                k=self.number_of_results,
+            )
+            for doc in docs:
+                doc.metadata = {
+                    key: str(value) if isinstance(value, ObjectId) else value for key, value in doc.metadata.items()
+                }
+
+            data = docs_to_data(docs)
+            self.status = data
+            return data
+        return []
+
+    def __insert_mode(self, collection: Collection) -> None:
+        if self.insert_mode == "overwrite":
+            collection.delete_many({})  # Delete all documents while preserving collection structure
+
+    def verify_search_index(self, collection: Collection) -> None:
+        """Verify if the search index exists, if not, create it.
+
+        Args:
+            collection (Collection): The collection to verify the search index on.
+        """
+        indexes = collection.list_search_indexes()
+
+        index_names_types = {idx["name"]: idx["type"] for idx in indexes}
+        index_names = list(index_names_types.keys())
+        index_type = index_names_types.get(self.index_name)
+        if self.index_name not in index_names and index_type != "vectorSearch":
+            collection.create_search_index(self.__create_index_definition())
+
+            time.sleep(20)  # Give some time for index to be ready
+
+    def __create_index_definition(self) -> SearchIndexModel:
+        fields = [
+            {
+                "type": "vector",
+                "path": self.index_field,
+                "numDimensions": self.number_dimensions,
+                "similarity": self.similarity,
+                "quantization": self.quantization,
+            }
+        ]
+        if self.filter_field:
+            fields.append({"type": "filter", "path": self.filter_field})
+        return SearchIndexModel(definition={"fields": fields}, name=self.index_name, type="vectorSearch")

lfx/components/needle/__init__.py

@@ -0,0 +1,3 @@
+from .needle import NeedleComponent
+
+__all__ = ["NeedleComponent"]

lfx/components/needle/needle.py

@@ -0,0 +1,104 @@
+from langchain_community.retrievers.needle import NeedleRetriever
+
+from lfx.custom.custom_component.component import Component
+from lfx.io import IntInput, MessageTextInput, Output, SecretStrInput
+from lfx.schema.message import Message
+from lfx.utils.constants import MESSAGE_SENDER_AI
+
+
+class NeedleComponent(Component):
+    display_name = "Needle Retriever"
+    description = "A retriever that uses the Needle API to search collections."
+    documentation = "https://docs.needle-ai.com"
+    icon = "Needle"
+    name = "needle"
+
+    inputs = [
+        SecretStrInput(
+            name="needle_api_key",
+            display_name="Needle API Key",
+            info="Your Needle API key.",
+            required=True,
+        ),
+        MessageTextInput(
+            name="collection_id",
+            display_name="Collection ID",
+            info="The ID of the Needle collection.",
+            required=True,
+        ),
+        MessageTextInput(
+            name="query",
+            display_name="User Query",
+            info="Enter your question here. In tool mode, you can also specify top_k parameter (min: 20).",
+            required=True,
+            tool_mode=True,
+        ),
+        IntInput(
+            name="top_k",
+            display_name="Top K Results",
+            info="Number of search results to return (min: 20).",
+            value=20,
+            required=True,
+        ),
+    ]
+
+    outputs = [Output(display_name="Result", name="result", type_="Message", method="run")]
+
+    def run(self) -> Message:
+        # Extract query and top_k
+        query_input = self.query
+        actual_query = query_input.get("query", "") if isinstance(query_input, dict) else query_input
+
+        # Parse top_k from tool input or use default, always enforcing minimum of 20
+        try:
+            if isinstance(query_input, dict) and "top_k" in query_input:
+                agent_top_k = query_input.get("top_k")
+                # Check if agent_top_k is not None before converting to int
+                top_k = max(20, int(agent_top_k)) if agent_top_k is not None else max(20, self.top_k)
+            else:
+                top_k = max(20, self.top_k)
+        except (ValueError, TypeError):
+            top_k = max(20, self.top_k)
+
+        # Validate required inputs
+        if not self.needle_api_key or not self.needle_api_key.strip():
+            error_msg = "The Needle API key cannot be empty."
+            raise ValueError(error_msg)
+        if not self.collection_id or not self.collection_id.strip():
+            error_msg = "The Collection ID cannot be empty."
+            raise ValueError(error_msg)
+        if not actual_query or not actual_query.strip():
+            error_msg = "The query cannot be empty."
+            raise ValueError(error_msg)
+
+        try:
+            # Initialize the retriever and get documents
+            retriever = NeedleRetriever(
+                needle_api_key=self.needle_api_key,
+                collection_id=self.collection_id,
+                top_k=top_k,
+            )
+
+            docs = retriever.get_relevant_documents(actual_query)
+
+            # Format the response
+            if not docs:
+                text_content = "No relevant documents found for the query."
+            else:
+                context = "\n\n".join([f"Document {i + 1}:\n{doc.page_content}" for i, doc in enumerate(docs)])
+                text_content = f"Question: {actual_query}\n\nContext:\n{context}"
+
+            # Return formatted message
+            return Message(
+                text=text_content,
+                type="assistant",
+                sender=MESSAGE_SENDER_AI,
+                additional_kwargs={
+                    "source_documents": [{"page_content": doc.page_content, "metadata": doc.metadata} for doc in docs],
+                    "top_k_used": top_k,
+                },
+            )
+
+        except Exception as e:
+            error_msg = f"Error processing query: {e!s}"
+            raise ValueError(error_msg) from e

lfx/components/notdiamond/__init__.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from lfx.components._importing import import_mod
+
+if TYPE_CHECKING:
+    from lfx.components.notdiamond.notdiamond import NotDiamondComponent
+
+_dynamic_imports = {
+    "NotDiamondComponent": "notdiamond",
+}
+
+__all__ = [
+    "NotDiamondComponent",
+]
+
+
+def __getattr__(attr_name: str) -> Any:
+    """Lazily import notdiamond components on attribute access."""
+    if attr_name not in _dynamic_imports:
+        msg = f"module '{__name__}' has no attribute '{attr_name}'"
+        raise AttributeError(msg)
+    try:
+        result = import_mod(attr_name, _dynamic_imports[attr_name], __spec__.parent)
+    except (ModuleNotFoundError, ImportError, AttributeError) as e:
+        msg = f"Could not import '{attr_name}' from '{__name__}': {e}"
+        raise AttributeError(msg) from e
+    globals()[attr_name] = result
+    return result
+
+
+def __dir__() -> list[str]:
+    return list(__all__)

lfx/components/notdiamond/notdiamond.py

@@ -0,0 +1,228 @@
+import warnings
+
+import requests
+from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage
+from pydantic.v1 import SecretStr
+
+from lfx.base.models.chat_result import get_chat_result
+from lfx.base.models.model_utils import get_model_name
+from lfx.custom.custom_component.component import Component
+from lfx.io import (
+    BoolInput,
+    DropdownInput,
+    HandleInput,
+    MessageInput,
+    MessageTextInput,
+    Output,
+    SecretStrInput,
+    StrInput,
+)
+from lfx.schema.message import Message
+
+ND_MODEL_MAPPING = {
+    "gpt-4o": {"provider": "openai", "model": "gpt-4o"},
+    "gpt-4o-mini": {"provider": "openai", "model": "gpt-4o-mini"},
+    "gpt-4-turbo": {"provider": "openai", "model": "gpt-4-turbo-2024-04-09"},
+    "claude-3-5-haiku-20241022": {"provider": "anthropic", "model": "claude-3-5-haiku-20241022"},
+    "claude-3-5-sonnet-20241022": {"provider": "anthropic", "model": "claude-3-5-sonnet-20241022"},
+    "anthropic.claude-3-5-sonnet-20241022-v2:0": {"provider": "anthropic", "model": "claude-3-5-sonnet-20241022"},
+    "anthropic.claude-3-5-haiku-20241022-v1:0": {"provider": "anthropic", "model": "claude-3-5-haiku-20241022"},
+    "gemini-1.5-pro": {"provider": "google", "model": "gemini-1.5-pro-latest"},
+    "gemini-1.5-flash": {"provider": "google", "model": "gemini-1.5-flash-latest"},
+    "llama-3.1-sonar-large-128k-online": {"provider": "perplexity", "model": "llama-3.1-sonar-large-128k-online"},
+    "meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo": {
+        "provider": "togetherai",
+        "model": "Meta-Llama-3.1-70B-Instruct-Turbo",
+    },
+    "meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo": {
+        "provider": "togetherai",
+        "model": "Meta-Llama-3.1-405B-Instruct-Turbo",
+    },
+    "mistral-large-latest": {"provider": "mistral", "model": "mistral-large-2407"},
+}
+
+
+class NotDiamondComponent(Component):
+    display_name = "Not Diamond Router"
+    description = "Call the right model at the right time with the world's most powerful AI model router."
+    documentation: str = "https://docs.notdiamond.ai/"
+    icon = "NotDiamond"
+    name = "NotDiamond"
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._selected_model_name = None
+
+    inputs = [
+        MessageInput(name="input_value", display_name="Input", required=True),
+        MessageTextInput(
+            name="system_message",
+            display_name="System Message",
+            info="System message to pass to the model.",
+            advanced=False,
+        ),
+        HandleInput(
+            name="models",
+            display_name="Language Models",
+            input_types=["LanguageModel"],
+            required=True,
+            is_list=True,
+            info="Link the models you want to route between.",
+        ),
+        SecretStrInput(
+            name="api_key",
+            display_name="Not Diamond API Key",
+            info="The Not Diamond API Key to use for routing.",
+            advanced=False,
+            value="NOTDIAMOND_API_KEY",
+            required=True,
+        ),
+        StrInput(
+            name="preference_id",
+            display_name="Preference ID",
+            info="The ID of the router preference that was configured via the Dashboard.",
+            advanced=False,
+        ),
+        DropdownInput(
+            name="tradeoff",
+            display_name="Tradeoff",
+            info="The tradeoff between cost and latency for the router to determine the best LLM for a given query.",
+            advanced=False,
+            options=["quality", "cost", "latency"],
+            value="quality",
+        ),
+        BoolInput(
+            name="hash_content",
+            display_name="Hash Content",
+            info="Whether to hash the content before being sent to the NotDiamond API.",
+            advanced=False,
+            value=False,
+        ),
+    ]
+
+    outputs = [
+        Output(display_name="Output", name="output", method="model_select"),
+        Output(
+            display_name="Selected Model",
+            name="selected_model",
+            method="get_selected_model",
+            required_inputs=["output"],
+        ),
+    ]
+
+    def get_selected_model(self) -> str:
+        return self._selected_model_name
+
+    def model_select(self) -> Message:
+        api_key = SecretStr(self.api_key).get_secret_value() if self.api_key else None
+        input_value = self.input_value
+        system_message = self.system_message
+        messages = self._format_input(input_value, system_message)
+
+        selected_models = []
+        mapped_selected_models = []
+        for model in self.models:
+            model_name = get_model_name(model)
+
+            if model_name in ND_MODEL_MAPPING:
+                selected_models.append(model)
+                mapped_selected_models.append(ND_MODEL_MAPPING[model_name])
+
+        payload = {
+            "messages": messages,
+            "llm_providers": mapped_selected_models,
+            "hash_content": self.hash_content,
+        }
+
+        if self.tradeoff != "quality":
+            payload["tradeoff"] = self.tradeoff
+
+        if self.preference_id and self.preference_id != "":
+            payload["preference_id"] = self.preference_id
+
+        header = {
+            "Authorization": f"Bearer {api_key}",
+            "accept": "application/json",
+            "content-type": "application/json",
+        }
+
+        response = requests.post(
+            "https://api.notdiamond.ai/v2/modelRouter/modelSelect",
+            json=payload,
+            headers=header,
+            timeout=10,
+        )
+
+        result = response.json()
+        chosen_model = self.models[0]  # By default there is a fallback model
+        self._selected_model_name = get_model_name(chosen_model)
+
+        if "providers" not in result:
+            # No provider returned by NotDiamond API, likely failed. Fallback to first model.
+            return self._call_get_chat_result(chosen_model, input_value, system_message)
+
+        providers = result["providers"]
+
+        if len(providers) == 0:
+            # No provider returned by NotDiamond API, likely failed. Fallback to first model.
+            return self._call_get_chat_result(chosen_model, input_value, system_message)
+
+        nd_result = providers[0]
+
+        for nd_model, selected_model in zip(mapped_selected_models, selected_models, strict=False):
+            if nd_model["provider"] == nd_result["provider"] and nd_model["model"] == nd_result["model"]:
+                chosen_model = selected_model
+                self._selected_model_name = get_model_name(chosen_model)
+                break
+
+        return self._call_get_chat_result(chosen_model, input_value, system_message)
+
+    def _call_get_chat_result(self, chosen_model, input_value, system_message):
+        return get_chat_result(
+            runnable=chosen_model,
+            input_value=input_value,
+            system_message=system_message,
+        )
+
+    def _format_input(
+        self,
+        input_value: str | Message,
+        system_message: str | None = None,
+    ):
+        messages: list[BaseMessage] = []
+        if not input_value and not system_message:
+            msg = "The message you want to send to the router is empty."
+            raise ValueError(msg)
+        system_message_added = False
+        if input_value:
+            if isinstance(input_value, Message):
+                with warnings.catch_warnings():
+                    warnings.simplefilter("ignore")
+                    if "prompt" in input_value:
+                        prompt = input_value.load_lc_prompt()
+                        if system_message:
+                            prompt.messages = [
+                                SystemMessage(content=system_message),
+                                *prompt.messages,  # type: ignore[has-type]
+                            ]
+                            system_message_added = True
+                        messages.extend(prompt.messages)
+                    else:
+                        messages.append(input_value.to_lc_message())
+            else:
+                messages.append(HumanMessage(content=input_value))
+
+        if system_message and not system_message_added:
+            messages.insert(0, SystemMessage(content=system_message))
+
+        # Convert Langchain messages to OpenAI format
+        openai_messages = []
+        for msg in messages:
+            if isinstance(msg, HumanMessage):
+                openai_messages.append({"role": "user", "content": msg.content})
+            elif isinstance(msg, AIMessage):
+                openai_messages.append({"role": "assistant", "content": msg.content})
+            elif isinstance(msg, SystemMessage):
+                openai_messages.append({"role": "system", "content": msg.content})
+
+        return openai_messages

lfx/components/novita/__init__.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from lfx.components._importing import import_mod
+
+if TYPE_CHECKING:
+    from lfx.components.novita.novita import NovitaModelComponent
+
+_dynamic_imports = {
+    "NovitaModelComponent": "novita",
+}
+
+__all__ = ["NovitaModelComponent"]
+
+
+def __getattr__(attr_name: str) -> Any:
+    """Lazily import novita components on attribute access."""
+    if attr_name not in _dynamic_imports:
+        msg = f"module '{__name__}' has no attribute '{attr_name}'"
+        raise AttributeError(msg)
+    try:
+        result = import_mod(attr_name, _dynamic_imports[attr_name], __spec__.parent)
+    except (ModuleNotFoundError, ImportError, AttributeError) as e:
+        msg = f"Could not import '{attr_name}' from '{__name__}': {e}"
+        raise AttributeError(msg) from e
+    globals()[attr_name] = result
+    return result
+
+
+def __dir__() -> list[str]:
+    return list(__all__)