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

lfx/components/google/google_bq_sql_executor.py

@@ -0,0 +1,157 @@
+import json
+import re
+from pathlib import Path
+
+from google.auth.exceptions import RefreshError
+from google.cloud import bigquery
+from google.oauth2.service_account import Credentials
+
+from lfx.custom import Component
+from lfx.io import BoolInput, FileInput, MessageTextInput, Output
+from lfx.schema.dataframe import DataFrame
+
+
+class BigQueryExecutorComponent(Component):
+    display_name = "BigQuery"
+    description = "Execute SQL queries on Google BigQuery."
+    name = "BigQueryExecutor"
+    icon = "Google"
+    beta: bool = True
+
+    inputs = [
+        FileInput(
+            name="service_account_json_file",
+            display_name="Upload Service Account JSON",
+            info="Upload the JSON file containing Google Cloud service account credentials.",
+            file_types=["json"],
+            required=True,
+        ),
+        MessageTextInput(
+            name="query",
+            display_name="SQL Query",
+            info="The SQL query to execute on BigQuery.",
+            required=True,
+            tool_mode=True,
+        ),
+        BoolInput(
+            name="clean_query",
+            display_name="Clean Query",
+            info="When enabled, this will automatically clean up your SQL query.",
+            value=False,
+            advanced=True,
+        ),
+    ]
+
+    outputs = [
+        Output(display_name="Query Results", name="query_results", method="execute_sql"),
+    ]
+
+    def _clean_sql_query(self, query: str) -> str:
+        """Clean SQL query by removing surrounding quotes and whitespace.
+
+        Also extracts SQL statements from text that might contain other content.
+
+        Args:
+            query: The SQL query to clean
+
+        Returns:
+            The cleaned SQL query
+        """
+        # First, try to extract SQL from code blocks
+        sql_pattern = r"```(?:sql)?\s*([\s\S]*?)\s*```"
+        sql_matches = re.findall(sql_pattern, query, re.IGNORECASE)
+
+        if sql_matches:
+            # If we found SQL in code blocks, use the first match
+            query = sql_matches[0]
+        else:
+            # If no code block, try to find SQL statements
+            # Look for common SQL keywords at the start of lines
+            sql_keywords = r"(?i)(SELECT|INSERT|UPDATE|DELETE|CREATE|ALTER|DROP|WITH|MERGE)"
+            lines = query.split("\n")
+            sql_lines = []
+            in_sql = False
+
+            for _line in lines:
+                line = _line.strip()
+                if re.match(sql_keywords, line):
+                    in_sql = True
+                if in_sql:
+                    sql_lines.append(line)
+                if line.endswith(";"):
+                    in_sql = False
+
+            if sql_lines:
+                query = "\n".join(sql_lines)
+
+        # Remove any backticks that might be at the start or end
+        query = query.strip("`")
+
+        # Then remove surrounding quotes (single or double) if they exist
+        query = query.strip()
+        if (query.startswith('"') and query.endswith('"')) or (query.startswith("'") and query.endswith("'")):
+            query = query[1:-1]
+
+        # Finally, clean up any remaining whitespace and ensure no backticks remain
+        query = query.strip()
+        # Remove any remaining backticks, but preserve them if they're part of a table/column name
+        # This regex will remove backticks that are not part of a valid identifier
+        return re.sub(r"`(?![a-zA-Z0-9_])|(?<![a-zA-Z0-9_])`", "", query)
+
+    def execute_sql(self) -> DataFrame:
+        try:
+            # First try to read the file
+            try:
+                service_account_path = Path(self.service_account_json_file)
+                with service_account_path.open() as f:
+                    credentials_json = json.load(f)
+                    project_id = credentials_json.get("project_id")
+                    if not project_id:
+                        msg = "No project_id found in service account credentials file."
+                        raise ValueError(msg)
+            except FileNotFoundError as e:
+                msg = f"Service account file not found: {e}"
+                raise ValueError(msg) from e
+            except json.JSONDecodeError as e:
+                msg = "Invalid JSON string for service account credentials"
+                raise ValueError(msg) from e
+
+            # Then try to load credentials
+            try:
+                credentials = Credentials.from_service_account_file(self.service_account_json_file)
+            except Exception as e:
+                msg = f"Error loading service account credentials: {e}"
+                raise ValueError(msg) from e
+
+        except ValueError:
+            raise
+        except Exception as e:
+            msg = f"Error executing BigQuery SQL query: {e}"
+            raise ValueError(msg) from e
+
+        try:
+            client = bigquery.Client(credentials=credentials, project=project_id)
+
+            # Check for empty or whitespace-only query before cleaning
+            if not str(self.query).strip():
+                msg = "No valid SQL query found in input text."
+                raise ValueError(msg)
+
+            # Always clean the query if it contains code block markers, quotes, or if clean_query is enabled
+            if "```" in str(self.query) or '"' in str(self.query) or "'" in str(self.query) or self.clean_query:
+                sql_query = self._clean_sql_query(str(self.query))
+            else:
+                sql_query = str(self.query).strip()  # At minimum, strip whitespace
+
+            query_job = client.query(sql_query)
+            results = query_job.result()
+            output_dict = [dict(row) for row in results]
+
+        except RefreshError as e:
+            msg = "Authentication error: Unable to refresh authentication token. Please try to reauthenticate."
+            raise ValueError(msg) from e
+        except Exception as e:
+            msg = f"Error executing BigQuery SQL query: {e}"
+            raise ValueError(msg) from e
+
+        return DataFrame(output_dict)

lfx/components/google/google_drive.py

@@ -0,0 +1,92 @@
+import json
+from json.decoder import JSONDecodeError
+
+from google.auth.exceptions import RefreshError
+from google.oauth2.credentials import Credentials
+from langchain_google_community import GoogleDriveLoader
+
+from lfx.custom.custom_component.component import Component
+from lfx.helpers.data import docs_to_data
+from lfx.inputs.inputs import MessageTextInput
+from lfx.io import SecretStrInput
+from lfx.schema.data import Data
+from lfx.template.field.base import Output
+
+
+class GoogleDriveComponent(Component):
+    display_name = "Google Drive Loader"
+    description = "Loads documents from Google Drive using provided credentials."
+    icon = "Google"
+    legacy: bool = True
+
+    inputs = [
+        SecretStrInput(
+            name="json_string",
+            display_name="JSON String of the Service Account Token",
+            info="JSON string containing OAuth 2.0 access token information for service account access",
+            required=True,
+        ),
+        MessageTextInput(
+            name="document_id", display_name="Document ID", info="Single Google Drive document ID", required=True
+        ),
+    ]
+
+    outputs = [
+        Output(display_name="Loaded Documents", name="docs", method="load_documents"),
+    ]
+
+    def load_documents(self) -> Data:
+        class CustomGoogleDriveLoader(GoogleDriveLoader):
+            creds: Credentials | None = None
+            """Credentials object to be passed directly."""
+
+            def _load_credentials(self):
+                """Load credentials from the provided creds attribute or fallback to the original method."""
+                if self.creds:
+                    return self.creds
+                msg = "No credentials provided."
+                raise ValueError(msg)
+
+            class Config:
+                arbitrary_types_allowed = True
+
+        json_string = self.json_string
+
+        document_ids = [self.document_id]
+        if len(document_ids) != 1:
+            msg = "Expected a single document ID"
+            raise ValueError(msg)
+
+        # TODO: Add validation to check if the document ID is valid
+
+        # Load the token information from the JSON string
+        try:
+            token_info = json.loads(json_string)
+        except JSONDecodeError as e:
+            msg = "Invalid JSON string"
+            raise ValueError(msg) from e
+
+        # Initialize the custom loader with the provided credentials and document IDs
+        loader = CustomGoogleDriveLoader(
+            creds=Credentials.from_authorized_user_info(token_info), document_ids=document_ids
+        )
+
+        # Load the documents
+        try:
+            docs = loader.load()
+        # catch google.auth.exceptions.RefreshError
+        except RefreshError as e:
+            msg = "Authentication error: Unable to refresh authentication token. Please try to reauthenticate."
+            raise ValueError(msg) from e
+        except Exception as e:
+            msg = f"Error loading documents: {e}"
+            raise ValueError(msg) from e
+
+        if len(docs) != 1:
+            msg = "Expected a single document to be loaded."
+            raise ValueError(msg)
+
+        data = docs_to_data(docs)
+        # Return the loaded documents
+        self.status = data
+        return Data(data={"text": data})

lfx/components/google/google_drive_search.py

@@ -0,0 +1,152 @@
+import json
+
+from google.oauth2.credentials import Credentials
+from googleapiclient.discovery import build
+
+from lfx.custom.custom_component.component import Component
+from lfx.inputs.inputs import DropdownInput, MessageTextInput
+from lfx.io import SecretStrInput
+from lfx.schema.data import Data
+from lfx.template.field.base import Output
+
+
+class GoogleDriveSearchComponent(Component):
+    display_name = "Google Drive Search"
+    description = "Searches Google Drive files using provided credentials and query parameters."
+    icon = "Google"
+    legacy: bool = True
+
+    inputs = [
+        SecretStrInput(
+            name="token_string",
+            display_name="Token String",
+            info="JSON string containing OAuth 2.0 access token information for service account access",
+            required=True,
+        ),
+        DropdownInput(
+            name="query_item",
+            display_name="Query Item",
+            options=[
+                "name",
+                "fullText",
+                "mimeType",
+                "modifiedTime",
+                "viewedByMeTime",
+                "trashed",
+                "starred",
+                "parents",
+                "owners",
+                "writers",
+                "readers",
+                "sharedWithMe",
+                "createdTime",
+                "properties",
+                "appProperties",
+                "visibility",
+                "shortcutDetails.targetId",
+            ],
+            info="The field to query.",
+            required=True,
+        ),
+        DropdownInput(
+            name="valid_operator",
+            display_name="Valid Operator",
+            options=["contains", "=", "!=", "<=", "<", ">", ">=", "in", "has"],
+            info="Operator to use in the query.",
+            required=True,
+        ),
+        MessageTextInput(
+            name="search_term",
+            display_name="Search Term",
+            info="The value to search for in the specified query item.",
+            required=True,
+        ),
+        MessageTextInput(
+            name="query_string",
+            display_name="Query String",
+            info="The query string used for searching. You can edit this manually.",
+            value="",  # This will be updated with the generated query string
+        ),
+    ]
+
+    outputs = [
+        Output(display_name="Document URLs", name="doc_urls", method="search_doc_urls"),
+        Output(display_name="Document IDs", name="doc_ids", method="search_doc_ids"),
+        Output(display_name="Document Titles", name="doc_titles", method="search_doc_titles"),
+        Output(display_name="Data", name="Data", method="search_data"),
+    ]
+
+    def generate_query_string(self) -> str:
+        query_item = self.query_item
+        valid_operator = self.valid_operator
+        search_term = self.search_term
+
+        # Construct the query string
+        query = f"{query_item} {valid_operator} '{search_term}'"
+
+        # Update the editable query string input with the generated query
+        self.query_string = query
+
+        return query
+
+    def on_inputs_changed(self) -> None:
+        # Automatically regenerate the query string when inputs change
+        self.generate_query_string()
+
+    def generate_file_url(self, file_id: str, mime_type: str) -> str:
+        """Generates the appropriate Google Drive URL for a file based on its MIME type."""
+        return {
+            "application/vnd.google-apps.document": f"https://docs.google.com/document/d/{file_id}/edit",
+            "application/vnd.google-apps.spreadsheet": f"https://docs.google.com/spreadsheets/d/{file_id}/edit",
+            "application/vnd.google-apps.presentation": f"https://docs.google.com/presentation/d/{file_id}/edit",
+            "application/vnd.google-apps.drawing": f"https://docs.google.com/drawings/d/{file_id}/edit",
+            "application/pdf": f"https://drive.google.com/file/d/{file_id}/view?usp=drivesdk",
+        }.get(mime_type, f"https://drive.google.com/file/d/{file_id}/view?usp=drivesdk")
+
+    def search_files(self) -> dict:
+        # Load the token information from the JSON string
+        token_info = json.loads(self.token_string)
+        creds = Credentials.from_authorized_user_info(token_info)
+
+        # Use the query string from the input (which might have been edited by the user)
+        query = self.query_string or self.generate_query_string()
+
+        # Initialize the Google Drive API service
+        service = build("drive", "v3", credentials=creds)
+
+        # Perform the search
+        results = service.files().list(q=query, pageSize=5, fields="nextPageToken, files(id, name, mimeType)").execute()
+        items = results.get("files", [])
+
+        doc_urls = []
+        doc_ids = []
+        doc_titles_urls = []
+        doc_titles = []
+
+        if items:
+            for item in items:
+                # Directly use the file ID, title, and MIME type to generate the URL
+                file_id = item["id"]
+                file_title = item["name"]
+                mime_type = item["mimeType"]
+                file_url = self.generate_file_url(file_id, mime_type)
+
+                # Store the URL, ID, and title+URL in their respective lists
+                doc_urls.append(file_url)
+                doc_ids.append(file_id)
+                doc_titles.append(file_title)
+                doc_titles_urls.append({"title": file_title, "url": file_url})
+
+        return {"doc_urls": doc_urls, "doc_ids": doc_ids, "doc_titles_urls": doc_titles_urls, "doc_titles": doc_titles}
+
+    def search_doc_ids(self) -> list[str]:
+        return self.search_files()["doc_ids"]
+
+    def search_doc_urls(self) -> list[str]:
+        return self.search_files()["doc_urls"]
+
+    def search_doc_titles(self) -> list[str]:
+        return self.search_files()["doc_titles"]
+
+    def search_data(self) -> Data:
+        return Data(data={"text": self.search_files()["doc_titles_urls"]})

lfx/components/google/google_generative_ai.py

@@ -0,0 +1,147 @@
+from typing import Any
+
+import requests
+from pydantic.v1 import SecretStr
+
+from lfx.base.models.google_generative_ai_constants import GOOGLE_GENERATIVE_AI_MODELS
+from lfx.base.models.model import LCModelComponent
+from lfx.field_typing import LanguageModel
+from lfx.field_typing.range_spec import RangeSpec
+from lfx.inputs.inputs import BoolInput, DropdownInput, FloatInput, IntInput, SecretStrInput, SliderInput
+from lfx.log.logger import logger
+from lfx.schema.dotdict import dotdict
+
+
+class GoogleGenerativeAIComponent(LCModelComponent):
+    display_name = "Google Generative AI"
+    description = "Generate text using Google Generative AI."
+    icon = "GoogleGenerativeAI"
+    name = "GoogleGenerativeAIModel"
+
+    inputs = [
+        *LCModelComponent.get_base_inputs(),
+        IntInput(
+            name="max_output_tokens", display_name="Max Output Tokens", info="The maximum number of tokens to generate."
+        ),
+        DropdownInput(
+            name="model_name",
+            display_name="Model",
+            info="The name of the model to use.",
+            options=GOOGLE_GENERATIVE_AI_MODELS,
+            value="gemini-1.5-pro",
+            refresh_button=True,
+            combobox=True,
+        ),
+        SecretStrInput(
+            name="api_key",
+            display_name="Google API Key",
+            info="The Google API Key to use for the Google Generative AI.",
+            required=True,
+            real_time_refresh=True,
+        ),
+        FloatInput(
+            name="top_p",
+            display_name="Top P",
+            info="The maximum cumulative probability of tokens to consider when sampling.",
+            advanced=True,
+        ),
+        SliderInput(
+            name="temperature",
+            display_name="Temperature",
+            value=0.1,
+            range_spec=RangeSpec(min=0, max=1, step=0.01),
+            info="Controls randomness. Lower values are more deterministic, higher values are more creative.",
+        ),
+        IntInput(
+            name="n",
+            display_name="N",
+            info="Number of chat completions to generate for each prompt. "
+            "Note that the API may not return the full n completions if duplicates are generated.",
+            advanced=True,
+        ),
+        IntInput(
+            name="top_k",
+            display_name="Top K",
+            info="Decode using top-k sampling: consider the set of top_k most probable tokens. Must be positive.",
+            advanced=True,
+        ),
+        BoolInput(
+            name="tool_model_enabled",
+            display_name="Tool Model Enabled",
+            info="Whether to use the tool model.",
+            value=False,
+        ),
+    ]
+
+    def build_model(self) -> LanguageModel:  # type: ignore[type-var]
+        try:
+            from langchain_google_genai import ChatGoogleGenerativeAI
+        except ImportError as e:
+            msg = "The 'langchain_google_genai' package is required to use the Google Generative AI model."
+            raise ImportError(msg) from e
+
+        google_api_key = self.api_key
+        model = self.model_name
+        max_output_tokens = self.max_output_tokens
+        temperature = self.temperature
+        top_k = self.top_k
+        top_p = self.top_p
+        n = self.n
+
+        return ChatGoogleGenerativeAI(
+            model=model,
+            max_output_tokens=max_output_tokens or None,
+            temperature=temperature,
+            top_k=top_k or None,
+            top_p=top_p or None,
+            n=n or 1,
+            google_api_key=SecretStr(google_api_key).get_secret_value(),
+        )
+
+    def get_models(self, *, tool_model_enabled: bool | None = None) -> list[str]:
+        try:
+            import google.generativeai as genai
+
+            genai.configure(api_key=self.api_key)
+            model_ids = [
+                model.name.replace("models/", "")
+                for model in genai.list_models()
+                if "generateContent" in model.supported_generation_methods
+            ]
+            model_ids.sort(reverse=True)
+        except (ImportError, ValueError) as e:
+            logger.exception(f"Error getting model names: {e}")
+            model_ids = GOOGLE_GENERATIVE_AI_MODELS
+        if tool_model_enabled:
+            try:
+                from langchain_google_genai.chat_models import ChatGoogleGenerativeAI
+            except ImportError as e:
+                msg = "langchain_google_genai is not installed."
+                raise ImportError(msg) from e
+            for model in model_ids:
+                model_with_tool = ChatGoogleGenerativeAI(
+                    model=self.model_name,
+                    google_api_key=self.api_key,
+                )
+                if not self.supports_tool_calling(model_with_tool):
+                    model_ids.remove(model)
+        return model_ids
+
+    def update_build_config(self, build_config: dotdict, field_value: Any, field_name: str | None = None):
+        if field_name in {"base_url", "model_name", "tool_model_enabled", "api_key"} and field_value:
+            try:
+                if len(self.api_key) == 0:
+                    ids = GOOGLE_GENERATIVE_AI_MODELS
+                else:
+                    try:
+                        ids = self.get_models(tool_model_enabled=self.tool_model_enabled)
+                    except (ImportError, ValueError, requests.exceptions.RequestException) as e:
+                        logger.exception(f"Error getting model names: {e}")
+                        ids = GOOGLE_GENERATIVE_AI_MODELS
+                build_config.setdefault("model_name", {})
+                build_config["model_name"]["options"] = ids
+                build_config["model_name"].setdefault("value", ids[0])
+            except Exception as e:
+                msg = f"Error getting model names: {e}"
+                raise ValueError(msg) from e
+        return build_config

lfx/components/google/google_generative_ai_embeddings.py

@@ -0,0 +1,141 @@
+# from lfx.field_typing import Data
+
+# TODO: remove ignore once the google package is published with types
+from google.ai.generativelanguage_v1beta.types import BatchEmbedContentsRequest
+from langchain_core.embeddings import Embeddings
+from langchain_google_genai import GoogleGenerativeAIEmbeddings
+from langchain_google_genai._common import GoogleGenerativeAIError
+
+from lfx.custom.custom_component.component import Component
+from lfx.io import MessageTextInput, Output, SecretStrInput
+
+MIN_DIMENSION_ERROR = "Output dimensionality must be at least 1"
+MAX_DIMENSION_ERROR = (
+    "Output dimensionality cannot exceed 768. Google's embedding models only support dimensions up to 768."
+)
+MAX_DIMENSION = 768
+MIN_DIMENSION = 1
+
+
+class GoogleGenerativeAIEmbeddingsComponent(Component):
+    display_name = "Google Generative AI Embeddings"
+    description = (
+        "Connect to Google's generative AI embeddings service using the GoogleGenerativeAIEmbeddings class, "
+        "found in the langchain-google-genai package."
+    )
+    documentation: str = "https://python.langchain.com/v0.2/docs/integrations/text_embedding/google_generative_ai/"
+    icon = "GoogleGenerativeAI"
+    name = "Google Generative AI Embeddings"
+
+    inputs = [
+        SecretStrInput(name="api_key", display_name="API Key", required=True),
+        MessageTextInput(name="model_name", display_name="Model Name", value="models/text-embedding-004"),
+    ]
+
+    outputs = [
+        Output(display_name="Embeddings", name="embeddings", method="build_embeddings"),
+    ]
+
+    def build_embeddings(self) -> Embeddings:
+        if not self.api_key:
+            msg = "API Key is required"
+            raise ValueError(msg)
+
+        class HotaGoogleGenerativeAIEmbeddings(GoogleGenerativeAIEmbeddings):
+            def __init__(self, *args, **kwargs) -> None:
+                super(GoogleGenerativeAIEmbeddings, self).__init__(*args, **kwargs)
+
+            def embed_documents(
+                self,
+                texts: list[str],
+                *,
+                batch_size: int = 100,
+                task_type: str | None = None,
+                titles: list[str] | None = None,
+                output_dimensionality: int | None = 768,
+            ) -> list[list[float]]:
+                """Embed a list of strings.
+
+                Google Generative AI currently sets a max batch size of 100 strings.
+
+                Args:
+                    texts: List[str] The list of strings to embed.
+                    batch_size: [int] The batch size of embeddings to send to the model
+                    task_type: task_type (https://ai.google.dev/api/rest/v1/TaskType)
+                    titles: An optional list of titles for texts provided.
+                    Only applicable when TaskType is RETRIEVAL_DOCUMENT.
+                    output_dimensionality: Optional reduced dimension for the output embedding.
+                    https://ai.google.dev/api/rest/v1/models/batchEmbedContents#EmbedContentRequest
+                Returns:
+                    List of embeddings, one for each text.
+                """
+                if output_dimensionality is not None and output_dimensionality < MIN_DIMENSION:
+                    raise ValueError(MIN_DIMENSION_ERROR)
+                if output_dimensionality is not None and output_dimensionality > MAX_DIMENSION:
+                    error_msg = MAX_DIMENSION_ERROR.format(output_dimensionality)
+                    raise ValueError(error_msg)
+
+                embeddings: list[list[float]] = []
+                batch_start_index = 0
+                for batch in GoogleGenerativeAIEmbeddings._prepare_batches(texts, batch_size):
+                    if titles:
+                        titles_batch = titles[batch_start_index : batch_start_index + len(batch)]
+                        batch_start_index += len(batch)
+                    else:
+                        titles_batch = [None] * len(batch)  # type: ignore[list-item]
+
+                    requests = [
+                        self._prepare_request(
+                            text=text,
+                            task_type=task_type,
+                            title=title,
+                            output_dimensionality=output_dimensionality,
+                        )
+                        for text, title in zip(batch, titles_batch, strict=True)
+                    ]
+
+                    try:
+                        result = self.client.batch_embed_contents(
+                            BatchEmbedContentsRequest(requests=requests, model=self.model)
+                        )
+                    except Exception as e:
+                        msg = f"Error embedding content: {e}"
+                        raise GoogleGenerativeAIError(msg) from e
+                    embeddings.extend([list(e.values) for e in result.embeddings])
+                return embeddings
+
+            def embed_query(
+                self,
+                text: str,
+                task_type: str | None = None,
+                title: str | None = None,
+                output_dimensionality: int | None = 768,
+            ) -> list[float]:
+                """Embed a text.
+
+                Args:
+                    text: The text to embed.
+                    task_type: task_type (https://ai.google.dev/api/rest/v1/TaskType)
+                    title: An optional title for the text.
+                    Only applicable when TaskType is RETRIEVAL_DOCUMENT.
+                    output_dimensionality: Optional reduced dimension for the output embedding.
+                    https://ai.google.dev/api/rest/v1/models/batchEmbedContents#EmbedContentRequest
+
+                Returns:
+                    Embedding for the text.
+                """
+                if output_dimensionality is not None and output_dimensionality < MIN_DIMENSION:
+                    raise ValueError(MIN_DIMENSION_ERROR)
+                if output_dimensionality is not None and output_dimensionality > MAX_DIMENSION:
+                    error_msg = MAX_DIMENSION_ERROR.format(output_dimensionality)
+                    raise ValueError(error_msg)
+
+                task_type = task_type or "RETRIEVAL_QUERY"
+                return self.embed_documents(
+                    [text],
+                    task_type=task_type,
+                    titles=[title] if title else None,
+                    output_dimensionality=output_dimensionality,
+                )[0]
+
+        return HotaGoogleGenerativeAIEmbeddings(model=self.model_name, google_api_key=self.api_key)