jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

feat: Add output dimensionality validation for Google AI embeddings (#5758)

Browse source 
* 🐛 (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 <gabriel@langflow.org>
This commit is contained in:
Cristhian Zanforlin Lousa 2025-01-20 09:22:35 -03:00 • committed by GitHub
commit 4d8e050f40
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 19 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

19
src/backend/base/langflow/components/embeddings/google_generative_ai.py
View file

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