From 232866b7dc9313edcddfdcfb769eafdc586e7d4f Mon Sep 17 00:00:00 2001 From: Saeed Seyfi Date: Sat, 11 Oct 2025 20:20:13 +0200 Subject: [PATCH] feat(vertex): add default dimensions parameter to VertexAIEmbeddings Add optional 'dimensions' constructor argument to VertexAIEmbeddings class. This allows users to set a default output dimensionality that will be used for all embedding requests unless explicitly overridden. - Add 'dimensions' field with default value None - Update embed() method to use constructor default when dimensions not specified - Maintain backward compatibility - existing code continues to work unchanged - Add comprehensive unit tests for default dimensions behavior --- .../langchain_google_vertexai/embeddings.py | 9 +- .../tests/unit_tests/test_embeddings.py | 100 ++++++++++++++++++ 2 files changed, 107 insertions(+), 2 deletions(-) diff --git a/libs/vertexai/langchain_google_vertexai/embeddings.py b/libs/vertexai/langchain_google_vertexai/embeddings.py index 7f808be8a..2fa381770 100644 --- a/libs/vertexai/langchain_google_vertexai/embeddings.py +++ b/libs/vertexai/langchain_google_vertexai/embeddings.py @@ -44,6 +44,9 @@ class VertexAIEmbeddings(BaseModel, Embeddings): "the environment." max_retries: int = 6 """The maximum number of retries to make when generating.""" + dimensions: Optional[int] = None + """Default output dimensionality for embeddings. If not specified, uses the + model's default. Can be overridden per request in embed() method.""" @model_validator(mode="before") @classmethod @@ -122,7 +125,8 @@ def embed( QUESTION_ANSWERING FACT_VERIFICATION dimensions: [int] optional. Output embeddings dimensions. - Only supported on preview models. + Only supported on preview models. If not provided, uses the + default dimensions specified in the constructor. title: [str] optional. Title for the text. Only applicable when TaskType is RETRIEVAL_DOCUMENT @@ -131,10 +135,11 @@ def embed( """ if len(texts) == 0: return [] + effective_dimensions = dimensions if dimensions is not None else self.dimensions embeddings = self._get_embeddings_with_retry( texts=texts, embeddings_type=embeddings_task_type, - dimensions=dimensions, + dimensions=effective_dimensions, title=title, ) return embeddings diff --git a/libs/vertexai/tests/unit_tests/test_embeddings.py b/libs/vertexai/tests/unit_tests/test_embeddings.py index 62d463e4c..33dee9b39 100644 --- a/libs/vertexai/tests/unit_tests/test_embeddings.py +++ b/libs/vertexai/tests/unit_tests/test_embeddings.py @@ -94,3 +94,103 @@ def test_embed_parameters(mock_get_embeddings, mock_client): dimensions=128, title="test-title", ) + + +@patch("langchain_google_vertexai.embeddings.genai.Client") +@patch.object(VertexAIEmbeddings, "_get_embeddings_with_retry") +def test_default_dimensions_used_when_not_specified(mock_get_embeddings, mock_client): + """Test that constructor dimensions are used when not specified in embed().""" + mock_client.return_value = MagicMock() + embeddings = VertexAIEmbeddings(model="text-embedding-004", dimensions=256) + texts = ["hello", "world"] + + mock_get_embeddings.return_value = [[0.001] * 256 for _ in texts] + + embeddings.embed(texts) + + mock_get_embeddings.assert_called_once_with( + texts=texts, + embeddings_type=None, + dimensions=256, + title=None, + ) + + +@patch("langchain_google_vertexai.embeddings.genai.Client") +@patch.object(VertexAIEmbeddings, "_get_embeddings_with_retry") +def test_explicit_dimensions_override_default(mock_get_embeddings, mock_client): + """Test that explicit dimensions in embed() override constructor default.""" + mock_client.return_value = MagicMock() + embeddings = VertexAIEmbeddings(model="text-embedding-004", dimensions=256) + texts = ["hello", "world"] + + mock_get_embeddings.return_value = [[0.001] * 512 for _ in texts] + + embeddings.embed(texts, dimensions=512) + + mock_get_embeddings.assert_called_once_with( + texts=texts, + embeddings_type=None, + dimensions=512, + title=None, + ) + + +@patch("langchain_google_vertexai.embeddings.genai.Client") +@patch.object(VertexAIEmbeddings, "_get_embeddings_with_retry") +def test_no_default_dimensions_works_as_before(mock_get_embeddings, mock_client): + """Test backward compatibility when no default dimensions specified.""" + mock_client.return_value = MagicMock() + embeddings = VertexAIEmbeddings(model="text-embedding-004") + texts = ["hello", "world"] + + mock_get_embeddings.return_value = [[0.001] * 768 for _ in texts] + + embeddings.embed(texts) + + mock_get_embeddings.assert_called_once_with( + texts=texts, + embeddings_type=None, + dimensions=None, + title=None, + ) + + +@patch("langchain_google_vertexai.embeddings.genai.Client") +@patch.object(VertexAIEmbeddings, "_get_embeddings_with_retry") +def test_default_dimensions_used_in_embed_documents(mock_get_embeddings, mock_client): + """Test that constructor dimensions are used in embed_documents().""" + mock_client.return_value = MagicMock() + embeddings = VertexAIEmbeddings(model="text-embedding-004", dimensions=128) + texts = ["hello", "world"] + + mock_get_embeddings.return_value = [[0.001] * 128 for _ in texts] + + embeddings.embed_documents(texts) + + mock_get_embeddings.assert_called_once_with( + texts=texts, + embeddings_type="RETRIEVAL_DOCUMENT", + dimensions=128, + title=None, + ) + + +@patch("langchain_google_vertexai.embeddings.genai.Client") +@patch.object(VertexAIEmbeddings, "_get_embeddings_with_retry") +def test_default_dimensions_used_in_embed_query(mock_get_embeddings, mock_client): + """Test that constructor dimensions are used in embed_query().""" + mock_client.return_value = MagicMock() + embeddings = VertexAIEmbeddings(model="text-embedding-004", dimensions=128) + text = "hello" + + mock_get_embeddings.return_value = [[0.001] * 128] + + embeddings.embed_query(text) + + mock_get_embeddings.assert_called_once_with( + texts=[text], + embeddings_type="RETRIEVAL_QUERY", + dimensions=128, + title=None, + )