From 4d8e050f40969029354a10cac7c78d816d418806 Mon Sep 17 00:00:00 2001 From: Cristhian Zanforlin Lousa Date: Mon, 20 Jan 2025 09:22:35 -0300 Subject: [PATCH] feat: Add output dimensionality validation for Google AI embeddings (#5758) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 🐛 (google_generative_ai.py): add input validation to ensure output dimensionality is within valid range * [autofix.ci] apply automated fixes * 🐛 (google_generative_ai.py): refactor code to use constants for min and max output dimensionality values to improve readability and maintainability * [autofix.ci] apply automated fixes * 📝 (google_generative_ai.py): remove unnecessary error message detail about received dimensionality, as it is not relevant for users and does not affect functionality * [autofix.ci] apply automated fixes * 🐛 (google_generative_ai.py): enhance output dimensionality validation by ensuring checks only occur when output_dimensionality is not None --------- Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> Co-authored-by: Gabriel Luiz Freitas Almeida --- .../embeddings/google_generative_ai.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/src/backend/base/langflow/components/embeddings/google_generative_ai.py b/src/backend/base/langflow/components/embeddings/google_generative_ai.py index 94a3c5a43..2284d4055 100644 --- a/src/backend/base/langflow/components/embeddings/google_generative_ai.py +++ b/src/backend/base/langflow/components/embeddings/google_generative_ai.py @@ -9,6 +9,13 @@ from langchain_google_genai._common import GoogleGenerativeAIError from langflow.custom import Component from langflow.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" @@ -62,6 +69,12 @@ class GoogleGenerativeAIEmbeddingsComponent(Component): 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): @@ -111,6 +124,12 @@ class GoogleGenerativeAIEmbeddingsComponent(Component): 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],