feat(vertexai): add default dimensions parameter to VertexAIEmbeddings (#1242)

## Summary

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.

## Changes

- Add `dimensions` field to `VertexAIEmbeddings` with default value
`None`
- Update `embed()` method to use constructor default when dimensions not
specified in method call
- Explicit dimensions in method calls override the constructor default
- Full backward compatibility maintained - existing code continues to
work unchanged

## Usage Example

```python
# Set default dimensions in constructor
embeddings = VertexAIEmbeddings(model="text-embedding-004", dimensions=128)

# Uses 128 dimensions (from constructor)
embeddings.embed(["hello", "world"])

# Overrides with 256 dimensions
embeddings.embed(["hello", "world"], dimensions=256)

# Backward compatible - no dimensions specified
embeddings = VertexAIEmbeddings(model="text-embedding-004")
embeddings.embed(["hello", "world"])  # Uses model's default
```

## Test Plan

- [x] Added 5 new unit tests covering all scenarios
- [x] Test that constructor dimensions are used when not specified in
embed()
- [x] Test that explicit dimensions override constructor default
- [x] Test backward compatibility when no default dimensions specified
- [x] Test default dimensions work with embed_documents() and
embed_query()
- [x] All 10 unit tests pass
- [x] Code quality checks pass (lint, format, mypy)

---------

Co-authored-by: Mason Daugherty <mason@langchain.dev>

Author: saeedseyfi

libs/vertexai/langchain_google_vertexai/embeddings.py

@@ -54,6 +54,9 @@ class VertexAIEmbeddings(BaseModel, Embeddings):
 
     max_retries: int = 6
     """The maximum number of retries to make when generating."""
+    dimensions: int | None = 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
@@ -130,10 +133,11 @@ def embed(
 
                 The following are only supported on preview models:
                     `QUESTION_ANSWERING`, `FACT_VERIFICATION`.
-            dimensions: Optional output embeddings dimensions.
+            dimensions: Output embeddings dimensions.
 
-                Only supported on preview models.
-            title: Optional title for the text.
+                Only supported on preview models. If not provided, uses the
+                default dimensions specified in the constructor.
+            title: Title for the text.
 
                 Only applicable when `TaskType` is `RETRIEVAL_DOCUMENT`.
 
@@ -142,10 +146,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

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,
+    )