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],