partners/voyageai: enable setting output dimension (#28740)

Voyage has introduced voyage-3-large and voyage-code-3, which feature
different output dimensions by leveraging a technique called "Matryoshka
Embeddings" (see blog -
https://blog.voyageai.com/2024/12/04/voyage-code-3/).
These two models are available in various sizes: [256, 512, 1024, 2048]
(https://docs.voyageai.com/docs/embeddings#model-choices).

This PR adds the option to set the required output dimension.

libs/partners/voyageai/langchain_voyageai/embeddings.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Any, Iterable, List, Optional
+from typing import Any, Iterable, List, Literal, Optional, cast
 
 import voyageai  # type: ignore
 from langchain_core.embeddings import Embeddings
@@ -37,8 +37,10 @@ class VoyageAIEmbeddings(BaseModel, Embeddings):
     _aclient: voyageai.client_async.AsyncClient = PrivateAttr()
     model: str
     batch_size: int
+
+    output_dimension: Optional[Literal[256, 512, 1024, 2048]] = None
     show_progress_bar: bool = False
-    truncation: Optional[bool] = None
+    truncation: bool = True
     voyage_api_key: SecretStr = Field(
         alias="api_key",
         default_factory=secret_from_env(
@@ -105,22 +107,26 @@ class VoyageAIEmbeddings(BaseModel, Embeddings):
 
         _iter = self._get_batch_iterator(texts)
         for i in _iter:
-            embeddings.extend(
-                self._client.embed(
-                    texts[i : i + self.batch_size],
-                    model=self.model,
-                    input_type="document",
-                    truncation=self.truncation,
-                ).embeddings
-            )
-
+            r = self._client.embed(
+                texts[i : i + self.batch_size],
+                model=self.model,
+                input_type="document",
+                truncation=self.truncation,
+                output_dimension=self.output_dimension,
+            ).embeddings
+            embeddings.extend(cast(Iterable[List[float]], r))
         return embeddings
 
     def embed_query(self, text: str) -> List[float]:
         """Embed query text."""
-        return self._client.embed(
-            [text], model=self.model, input_type="query", truncation=self.truncation
+        r = self._client.embed(
+            [text],
+            model=self.model,
+            input_type="query",
+            truncation=self.truncation,
+            output_dimension=self.output_dimension,
         ).embeddings[0]
+        return cast(List[float], r)
 
     async def aembed_documents(self, texts: List[str]) -> List[List[float]]:
         embeddings: List[List[float]] = []
@@ -132,8 +138,9 @@ class VoyageAIEmbeddings(BaseModel, Embeddings):
                 model=self.model,
                 input_type="document",
                 truncation=self.truncation,
+                output_dimension=self.output_dimension,
             )
-            embeddings.extend(r.embeddings)
+            embeddings.extend(cast(Iterable[List[float]], r.embeddings))
 
         return embeddings
 
@@ -143,5 +150,6 @@ class VoyageAIEmbeddings(BaseModel, Embeddings):
             model=self.model,
             input_type="query",
             truncation=self.truncation,
+            output_dimension=self.output_dimension,
         )
-        return r.embeddings[0]
+        return cast(List[float], r.embeddings[0])

libs/partners/voyageai/langchain_voyageai/rerank.py

@@ -16,8 +16,8 @@ from voyageai.object import RerankingObject  # type: ignore
 class VoyageAIRerank(BaseDocumentCompressor):
     """Document compressor that uses `VoyageAI Rerank API`."""
 
-    client: voyageai.Client = None
-    aclient: voyageai.AsyncClient = None
+    client: voyageai.Client = None  # type: ignore
+    aclient: voyageai.AsyncClient = None  # type: ignore
     """VoyageAI clients to use for compressing documents."""
     voyage_api_key: Optional[SecretStr] = None
     """VoyageAI API key. Must be specified directly or via environment variable 

libs/partners/voyageai/pyproject.toml

@@ -19,9 +19,9 @@ disallow_untyped_defs = "True"
 "Release Notes" = "https://github.com/langchain-ai/langchain/releases?q=tag%3A%22langchain-voyageai%3D%3D0%22&expanded=true"
 
 [tool.poetry.dependencies]
-python = ">=3.9,<4.0"
+python = ">=3.9,<3.13"
 langchain-core = "^0.3.15"
-voyageai = ">=0.2.1,<1"
+voyageai = ">=0.3.2,<1"
 pydantic = ">=2,<3"
 
 [tool.ruff.lint]

libs/partners/voyageai/tests/integration_tests/test_embeddings.py

@@ -51,3 +51,12 @@ async def test_langchain_voyageai_async_embedding_query() -> None:
     embedding = VoyageAIEmbeddings(model=MODEL)  # type: ignore[call-arg]
     output = await embedding.aembed_query(document)
     assert len(output) == 1024
+
+
+def test_langchain_voyageai_embedding_documents_with_output_dimension() -> None:
+    """Test voyage embeddings."""
+    documents = ["foo bar"]
+    embedding = VoyageAIEmbeddings(model="voyage-3-large", output_dimension=256)  # type: ignore[call-arg]
+    output = embedding.embed_documents(documents)
+    assert len(output) == 1
+    assert len(output[0]) == 256

libs/partners/voyageai/tests/unit_tests/test_embeddings.py

@@ -47,3 +47,15 @@ def test_initialization_voyage_1_batch_size() -> None:
     assert emb.batch_size == 15
     assert emb.model == "voyage-01"
     assert emb._client is not None
+
+
+def test_initialization_with_output_dimension() -> None:
+    emb = VoyageAIEmbeddings(
+        api_key="NOT_A_VALID_KEY",  # type: ignore
+        model="voyage-3-large",
+        output_dimension=256,
+        batch_size=10,
+    )
+    assert isinstance(emb, Embeddings)
+    assert emb.model == "voyage-3-large"
+    assert emb.output_dimension == 256