VoyageAI embeddings support

docs/api/embeddings.md

@@ -14,6 +14,8 @@
 
 ::: pydantic_ai.embeddings.google
 
+::: pydantic_ai.embeddings.voyageai
+
 ::: pydantic_ai.embeddings.sentence_transformers
 
 ::: pydantic_ai.embeddings.test

docs/api/providers.md

@@ -20,6 +20,8 @@
 
 ::: pydantic_ai.providers.cohere
 
+::: pydantic_ai.providers.voyageai.VoyageAIProvider
+
 ::: pydantic_ai.providers.cerebras.CerebrasProvider
 
 ::: pydantic_ai.providers.mistral.MistralProvider

docs/embeddings.md

@@ -342,6 +342,61 @@ embedder = Embedder(
 )
 ```
 
+### VoyageAI
+
+[`VoyageAIEmbeddingModel`][pydantic_ai.embeddings.voyageai.VoyageAIEmbeddingModel] provides access to VoyageAI's embedding models, which are optimized for retrieval with specialized models for code, finance, and legal domains.
+
+#### Install
+
+To use VoyageAI embedding models, you need to install `pydantic-ai-slim` with the `voyageai` optional group:
+
+```bash
+pip/uv-add "pydantic-ai-slim[voyageai]"
+```
+
+#### Configuration
+
+To use `VoyageAIEmbeddingModel`, go to [dash.voyageai.com](https://dash.voyageai.com/) to generate an API key. Once you have the API key, you can set it as an environment variable:
+
+```bash
+export VOYAGE_API_KEY='your-api-key'
+```
+
+You can then use the model:
+
+```python {title="voyageai_embeddings.py"}
+from pydantic_ai import Embedder
+
+embedder = Embedder('voyageai:voyage-3.5')
+
+
+async def main():
+    result = await embedder.embed_query('Hello world')
+    print(len(result.embeddings[0]))
+    #> 1024
+```
+
+_(This example is complete, it can be run "as is" — you'll need to add `asyncio.run(main())` to run `main`)_
+
+See the [VoyageAI Embeddings documentation](https://docs.voyageai.com/docs/embeddings) for available models.
+
+#### VoyageAI-Specific Settings
+
+VoyageAI models support additional settings via [`VoyageAIEmbeddingSettings`][pydantic_ai.embeddings.voyageai.VoyageAIEmbeddingSettings]:
+
+```python {title="voyageai_settings.py"}
+from pydantic_ai import Embedder
+from pydantic_ai.embeddings.voyageai import VoyageAIEmbeddingSettings
+
+embedder = Embedder(
+    'voyageai:voyage-3.5',
+    settings=VoyageAIEmbeddingSettings(
+        dimensions=512,  # Reduce output dimensions
+        voyageai_input_type='document',  # Override input type for all requests
+    ),
+)
+```
+
 ### Sentence Transformers (Local)
 
 [`SentenceTransformerEmbeddingModel`][pydantic_ai.embeddings.sentence_transformers.SentenceTransformerEmbeddingModel] runs embeddings locally using the [sentence-transformers](https://www.sbert.net/) library. This is ideal for:
@@ -418,7 +473,10 @@ embedder = Embedder(model)
 
 ## Settings
 
-[`EmbeddingSettings`][pydantic_ai.embeddings.EmbeddingSettings] provides common configuration options that work across providers.
+[`EmbeddingSettings`][pydantic_ai.embeddings.EmbeddingSettings] provides common configuration options that work across providers:
+
+- `dimensions`: Reduce the output embedding dimensions (supported by OpenAI, Google, Cohere, VoyageAI)
+- `truncate`: When `True`, truncate input text that exceeds the model's context length instead of raising an error (supported by Cohere, VoyageAI)
 
 Settings can be specified at the embedder level (applied to all calls) or per-call:
 

pydantic_ai_slim/pydantic_ai/embeddings/__init__.py

@@ -48,6 +48,16 @@
         'cohere:embed-english-light-v3.0',
         'cohere:embed-multilingual-v3.0',
         'cohere:embed-multilingual-light-v3.0',
+        'voyageai:voyage-4-large',
+        'voyageai:voyage-4',
+        'voyageai:voyage-4-lite',
+        'voyageai:voyage-3-large',
+        'voyageai:voyage-3.5',
+        'voyageai:voyage-3.5-lite',
+        'voyageai:voyage-code-3',
+        'voyageai:voyage-finance-2',
+        'voyageai:voyage-law-2',
+        'voyageai:voyage-code-2',
     ],
 )
 """Known model names that can be used with the `model` parameter of [`Embedder`][pydantic_ai.embeddings.Embedder].
@@ -104,6 +114,10 @@ def infer_embedding_model(
         from .sentence_transformers import SentenceTransformerEmbeddingModel
 
         return SentenceTransformerEmbeddingModel(model_name)
+    elif model_kind == 'voyageai':
+        from .voyageai import VoyageAIEmbeddingModel
+
+        return VoyageAIEmbeddingModel(model_name, provider=provider)
     else:
         raise UserError(f'Unknown embeddings model: {model}')  # pragma: no cover
 

pydantic_ai_slim/pydantic_ai/embeddings/cohere.py

@@ -76,6 +76,8 @@ class CohereEmbeddingSettings(EmbeddingSettings, total=False):
     - `'NONE'` (default): Raise an error if input exceeds max tokens.
     - `'END'`: Truncate the end of the input text.
     - `'START'`: Truncate the start of the input text.
+
+    Note: This setting overrides the standard `truncate` boolean setting when specified.
     """
 
 
@@ -159,14 +161,22 @@ async def embed(
         if extra_body := settings.get('extra_body'):  # pragma: no cover
             request_options['additional_body_parameters'] = cast(dict[str, Any], extra_body)
 
+        # Determine truncation strategy: cohere_truncate takes precedence over truncate
+        if 'cohere_truncate' in settings:
+            truncate = settings['cohere_truncate']
+        elif settings.get('truncate'):
+            truncate = 'END'
+        else:
+            truncate = 'NONE'
+
         try:
             response = await self._client.embed(
                 model=self.model_name,
                 texts=inputs,
                 output_dimension=settings.get('dimensions'),
                 input_type=cohere_input_type,
                 max_tokens=settings.get('cohere_max_tokens'),
-                truncate=settings.get('cohere_truncate', 'NONE'),
+                truncate=truncate,
                 request_options=request_options,
             )
         except ApiError as e:

pydantic_ai_slim/pydantic_ai/embeddings/settings.py

@@ -23,6 +23,27 @@ class EmbeddingSettings(TypedDict, total=False):
     * Cohere
     * Google
     * Sentence Transformers
+    * VoyageAI
+    """
+
+    truncate: bool
+    """Whether to truncate inputs that exceed the model's context length.
+
+    Defaults to `False`. If `True`, inputs that are too long will be truncated.
+    If `False`, an error will be raised for inputs that exceed the context length.
+
+    For more control over truncation, you can use
+    [`max_input_tokens()`][pydantic_ai.embeddings.Embedder.max_input_tokens] and
+    [`count_tokens()`][pydantic_ai.embeddings.Embedder.count_tokens] to implement
+    your own truncation logic.
+
+    Provider-specific truncation settings (e.g., `cohere_truncate`) take precedence
+    if specified.
+
+    Supported by:
+
+    * Cohere
+    * VoyageAI
     """
 
     extra_headers: dict[str, str]

pydantic_ai_slim/pydantic_ai/embeddings/voyageai.py

@@ -0,0 +1,188 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass, field
+from typing import Literal, cast
+
+from pydantic_ai.exceptions import ModelAPIError
+from pydantic_ai.providers import Provider, infer_provider
+from pydantic_ai.usage import RequestUsage
+
+from .base import EmbeddingModel, EmbedInputType
+from .result import EmbeddingResult
+from .settings import EmbeddingSettings
+
+try:
+    from voyageai.client_async import AsyncClient
+    from voyageai.error import VoyageError
+except ImportError as _import_error:
+    raise ImportError(
+        'Please install `voyageai` to use the VoyageAI embeddings model, '
+        'you can use the `voyageai` optional group — `pip install "pydantic-ai-slim[voyageai]"`'
+    ) from _import_error
+
+LatestVoyageAIEmbeddingModelNames = Literal[
+    'voyage-4-large',
+    'voyage-4',
+    'voyage-4-lite',
+    'voyage-3-large',
+    'voyage-3.5',
+    'voyage-3.5-lite',
+    'voyage-code-3',
+    'voyage-finance-2',
+    'voyage-law-2',
+    'voyage-code-2',
+]
+"""Latest VoyageAI embedding models.
+
+See [VoyageAI Embeddings](https://docs.voyageai.com/docs/embeddings)
+for available models and their capabilities.
+"""
+
+VoyageAIEmbeddingModelName = str | LatestVoyageAIEmbeddingModelNames
+"""Possible VoyageAI embedding model names."""
+
+VoyageAIEmbedInputType = Literal['query', 'document', 'none']
+"""VoyageAI embedding input types.
+
+- `'query'`: For search queries; prepends retrieval-optimized prefix.
+- `'document'`: For documents; prepends document retrieval prefix.
+- `'none'`: Direct embedding without any prefix.
+"""
+
+
+class VoyageAIEmbeddingSettings(EmbeddingSettings, total=False):
+    """Settings used for a VoyageAI embedding model request.
+
+    All fields from [`EmbeddingSettings`][pydantic_ai.embeddings.EmbeddingSettings] are supported,
+    plus VoyageAI-specific settings prefixed with `voyageai_`.
+    """
+
+    # ALL FIELDS MUST BE `voyageai_` PREFIXED SO YOU CAN MERGE THEM WITH OTHER MODELS.
+
+    voyageai_input_type: VoyageAIEmbedInputType
+    """The VoyageAI-specific input type for the embedding.
+
+    Overrides the standard `input_type` argument. Options include:
+    `'query'`, `'document'`, or `'none'` for direct embedding without prefix.
+    """
+
+
+_MAX_INPUT_TOKENS: dict[VoyageAIEmbeddingModelName, int] = {
+    'voyage-4-large': 32000,
+    'voyage-4': 32000,
+    'voyage-4-lite': 32000,
+    'voyage-3-large': 32000,
+    'voyage-3.5': 32000,
+    'voyage-3.5-lite': 32000,
+    'voyage-code-3': 32000,
+    'voyage-finance-2': 32000,
+    'voyage-law-2': 16000,
+    'voyage-code-2': 16000,
+}
+
+
+@dataclass(init=False)
+class VoyageAIEmbeddingModel(EmbeddingModel):
+    """VoyageAI embedding model implementation.
+
+    VoyageAI provides state-of-the-art embedding models optimized for
+    retrieval, with specialized models for code, finance, and legal domains.
+
+    Example:
+    ```python
+    from pydantic_ai.embeddings.voyageai import VoyageAIEmbeddingModel
+
+    model = VoyageAIEmbeddingModel('voyage-3.5')
+    ```
+    """
+
+    _model_name: VoyageAIEmbeddingModelName = field(repr=False)
+    _provider: Provider[AsyncClient] = field(repr=False)
+
+    def __init__(
+        self,
+        model_name: VoyageAIEmbeddingModelName,
+        *,
+        provider: Literal['voyageai'] | Provider[AsyncClient] = 'voyageai',
+        settings: EmbeddingSettings | None = None,
+    ):
+        """Initialize a VoyageAI embedding model.
+
+        Args:
+            model_name: The name of the VoyageAI model to use.
+                See [VoyageAI models](https://docs.voyageai.com/docs/embeddings)
+                for available options.
+            provider: The provider to use for authentication and API access. Can be:
+
+                - `'voyageai'` (default): Uses the standard VoyageAI API
+                - A [`VoyageAIProvider`][pydantic_ai.providers.voyageai.VoyageAIProvider] instance
+                  for custom configuration
+            settings: Model-specific [`EmbeddingSettings`][pydantic_ai.embeddings.EmbeddingSettings]
+                to use as defaults for this model.
+        """
+        self._model_name = model_name
+
+        if isinstance(provider, str):
+            provider = infer_provider(provider)
+        self._provider = provider
+
+        super().__init__(settings=settings)
+
+    @property
+    def base_url(self) -> str:
+        """The base URL for the provider API."""
+        return self._provider.base_url
+
+    @property
+    def model_name(self) -> VoyageAIEmbeddingModelName:
+        """The embedding model name."""
+        return self._model_name
+
+    @property
+    def system(self) -> str:
+        """The embedding model provider."""
+        return self._provider.name
+
+    async def embed(
+        self,
+        inputs: str | Sequence[str],
+        *,
+        input_type: EmbedInputType,
+        settings: EmbeddingSettings | None = None,
+    ) -> EmbeddingResult:
+        inputs, settings = self.prepare_embed(inputs, settings)
+        settings = cast(VoyageAIEmbeddingSettings, settings)
+
+        voyageai_input_type: VoyageAIEmbedInputType = settings.get(
+            'voyageai_input_type', 'document' if input_type == 'document' else 'query'
+        )
+        # Convert 'none' string to None for the API
+        api_input_type = None if voyageai_input_type == 'none' else voyageai_input_type
+
+        try:
+            response = await self._provider.client.embed(
+                texts=list(inputs),
+                model=self.model_name,
+                input_type=api_input_type,
+                truncation=settings.get('truncate', False),
+                output_dimension=settings.get('dimensions'),
+            )
+        except VoyageError as e:
+            raise ModelAPIError(model_name=self.model_name, message=str(e)) from e
+
+        return EmbeddingResult(
+            embeddings=response.embeddings,
+            inputs=inputs,
+            input_type=input_type,
+            usage=_map_usage(response.total_tokens),
+            model_name=self.model_name,
+            provider_name=self.system,
+        )
+
+    async def max_input_tokens(self) -> int | None:
+        return _MAX_INPUT_TOKENS.get(self.model_name)
+
+
+def _map_usage(total_tokens: int) -> RequestUsage:
+    return RequestUsage(input_tokens=total_tokens)

pydantic_ai_slim/pydantic_ai/providers/__init__.py

@@ -161,6 +161,10 @@ def infer_provider_class(provider: str) -> type[Provider[Any]]:  # noqa: C901
         from .sentence_transformers import SentenceTransformersProvider
 
         return SentenceTransformersProvider
+    elif provider == 'voyageai':
+        from .voyageai import VoyageAIProvider
+
+        return VoyageAIProvider
     else:  # pragma: no cover
         raise ValueError(f'Unknown provider: {provider}')