Progress is made

Author: DouweM

pydantic_ai_slim/pydantic_ai/embeddings/__init__.py

@@ -1,17 +1,28 @@
-from collections.abc import Iterator, Sequence
+from collections.abc import Callable, Iterator, Sequence
 from contextlib import contextmanager
 from contextvars import ContextVar
 from dataclasses import dataclass
-from typing import Literal, overload
+from typing import Any, Literal, get_args, overload
 
 from typing_extensions import TypeAliasType
 
 from pydantic_ai import _utils
-from pydantic_ai.embeddings.embedding_model import EmbeddingModel
+from pydantic_ai.embeddings.base import EmbeddingModel
 from pydantic_ai.embeddings.settings import EmbeddingSettings, merge_embedding_settings
 from pydantic_ai.exceptions import UserError
+from pydantic_ai.models import OpenAIChatCompatibleProvider
 from pydantic_ai.models.instrumented import InstrumentationSettings
-from pydantic_ai.providers import infer_provider
+from pydantic_ai.providers import Provider, infer_provider
+
+__all__ = [
+    'Embedder',
+    'EmbeddingModel',
+    'EmbeddingSettings',
+    'merge_embedding_settings',
+    'KnownEmbeddingModelName',
+    'OpenAIEmbeddingsCompatibleProvider',
+    'infer_model',
+]
 
 KnownEmbeddingModelName = TypeAliasType(
     'KnownEmbeddingModelName',
@@ -21,13 +32,20 @@
         'openai:text-embedding-3-largecohere:embed-v4.0',
     ],
 )
-"""Known model names that can be used with the `model` parameter of [`Agent`][pydantic_ai.Agent].
+"""Known model names that can be used with the `model` parameter of [`Embedder`][pydantic_ai.embeddings.Embedder].
 
-`KnownModelName` is provided as a concise way to specify a model.
+`KnownEmbeddingModelName` is provided as a concise way to specify an embedding model.
 """
 
+# For now, we assume that every chat completions-compatible provider also supports the embeddings endpoint.
+OpenAIEmbeddingsCompatibleProvider = OpenAIChatCompatibleProvider
+
 
-def infer_model(model: EmbeddingModel | KnownEmbeddingModelName | str) -> EmbeddingModel:
+def infer_model(
+    model: EmbeddingModel | KnownEmbeddingModelName | str,
+    *,
+    provider_factory: Callable[[str], Provider[Any]] = infer_provider,
+) -> EmbeddingModel:
     """Infer the model from the name."""
     if isinstance(model, EmbeddingModel):
         return model
@@ -37,14 +55,15 @@ def infer_model(model: EmbeddingModel | KnownEmbeddingModelName | str) -> Embedd
     except ValueError as e:
         raise ValueError('You must provide a provider prefix when specifying an embedding model name') from e
 
-    provider = infer_provider(provider_name)
+    provider = provider_factory(provider_name)
 
     model_kind = provider_name
     if model_kind.startswith('gateway/'):
-        model_kind = provider_name.removeprefix('gateway/')
+        from ..providers.gateway import normalize_gateway_provider
 
-    # TODO: extend the following list for other providers as appropriate
-    if model_kind in ('openai',):
+        model_kind = normalize_gateway_provider(model_kind)
+
+    if model_kind in get_args(OpenAIEmbeddingsCompatibleProvider.__value__):
         model_kind = 'openai'
 
     if model_kind == 'openai':
@@ -59,8 +78,10 @@ def infer_model(model: EmbeddingModel | KnownEmbeddingModelName | str) -> Embedd
         raise UserError(f'Unknown embeddings model: {model}')  # pragma: no cover
 
 
-@dataclass
+@dataclass(init=False)
 class Embedder:
+    """TODO: Docstring."""
+
     instrument: InstrumentationSettings | bool | None
     """Options to automatically instrument with OpenTelemetry."""
 

pydantic_ai_slim/pydantic_ai/embeddings/embedding_model.py renamed to pydantic_ai_slim/pydantic_ai/embeddings/base.py

@@ -2,7 +2,7 @@
 from collections.abc import Sequence
 from typing import overload
 
-from pydantic_ai.embeddings.settings import EmbeddingSettings
+from pydantic_ai.embeddings.settings import EmbeddingSettings, merge_embedding_settings
 
 
 class EmbeddingModel(ABC):
@@ -28,18 +28,22 @@ def settings(self) -> EmbeddingSettings | None:
         """Get the model settings."""
         return self._settings
 
+    @property
+    def base_url(self) -> str | None:
+        """The base URL for the provider API, if available."""
+        return None
+
     @property
     @abstractmethod
     def model_name(self) -> str:
         """The model name."""
         raise NotImplementedError()
 
-    # TODO: Add system?
-
     @property
-    def base_url(self) -> str | None:
-        """The base URL for the provider API, if available."""
-        return None
+    @abstractmethod
+    def system(self) -> str:
+        """The embedding model provider."""
+        raise NotImplementedError()
 
     @overload
     async def embed(self, documents: str, *, settings: EmbeddingSettings | None = None) -> list[float]:
@@ -49,7 +53,20 @@ async def embed(self, documents: str, *, settings: EmbeddingSettings | None = No
     async def embed(self, documents: Sequence[str], *, settings: EmbeddingSettings | None = None) -> list[list[float]]:
         pass
 
+    @abstractmethod
     async def embed(
         self, documents: str | Sequence[str], *, settings: EmbeddingSettings | None = None
     ) -> list[float] | list[list[float]]:
         raise NotImplementedError
+
+    def prepare_embed(
+        self, documents: str | Sequence[str], settings: EmbeddingSettings | None = None
+    ) -> tuple[Sequence[str], bool, EmbeddingSettings]:
+        """Prepare the documents and settings for the embedding."""
+        is_single_document = isinstance(documents, str)
+        if is_single_document:
+            documents = [documents]
+
+        settings = merge_embedding_settings(self._settings, settings) or {}
+
+        return documents, is_single_document, settings

pydantic_ai_slim/pydantic_ai/embeddings/cohere.py

@@ -1,15 +1,15 @@
 from collections.abc import Sequence
 from dataclasses import dataclass, field
-from typing import Literal, cast, overload
+from typing import Any, Literal, cast, overload
 
-from pydantic_ai.embeddings.embedding_model import EmbeddingModel
+from pydantic_ai.embeddings.base import EmbeddingModel
 from pydantic_ai.embeddings.settings import EmbeddingSettings
 from pydantic_ai.providers import Provider, infer_provider
 
-from .settings import merge_embedding_settings
-
 try:
     from cohere import AsyncClientV2
+    from cohere.core.request_options import RequestOptions
+    from cohere.types.embed_input_type import EmbedInputType
 except ImportError as _import_error:
     raise ImportError(
         'Please install `cohere` to use the Cohere embeddings model, '
@@ -18,16 +18,36 @@
 
 LatestCohereEmbeddingModelNames = Literal[
     'cohere:embed-v4.0',
-    # TODO: Add the others
+    'embed-english-v3.0embed-english-light-v3.0',
+    'embed-multilingual-v3.0',
+    'embed-multilingual-light-v3.0',
 ]
 """Latest Cohere embeddings models."""
 
 CohereEmbeddingModelName = str | LatestCohereEmbeddingModelNames
 """Possible Cohere embeddings model names."""
 
 
+class CohereEmbeddingSettings(EmbeddingSettings, total=False):
+    """Settings used for a Cohere embedding model request."""
+
+    # ALL FIELDS MUST BE `cohere_` PREFIXED SO YOU CAN MERGE THEM WITH OTHER MODELS.
+
+    # TODO: Possibly move to base EmbeddingSettings if supported by more providers
+    cohere_max_tokens: int
+    """The maximum number of tokens to generate before stopping."""
+
+    # We don't support `embedding_types for now because it doesn't affect the user-facing API today..
+    # cohere_embedding_types: Literal["float", "int8", "uint8", "binary", "ubinary", "base64"]
+
+    cohere_input_type: EmbedInputType
+    """The input type of the embedding."""
+
+
 @dataclass(init=False)
 class CohereEmbeddingModel(EmbeddingModel):
+    """Cohere embedding model."""
+
     _model_name: CohereEmbeddingModelName = field(repr=False)
     _provider: Provider[AsyncClientV2] = field(repr=False)
 
@@ -42,11 +62,10 @@ def __init__(
 
         Args:
             model_name: The name of the Cohere model to use. List of model names
-                available [here](https://docs.cohere.com/docs/models#command).
+                available [here](https://docs.cohere.com/docs/cohere-embed).
             provider: The provider to use for authentication and API access. Can be either the string
                 'cohere' or an instance of `Provider[AsyncClientV2]`. If not provided, a new provider will be
                 created using the other parameters.
-            profile: The model profile to use. Defaults to a profile picked by the provider based on the model name.
             settings: Model-specific settings that will be used as defaults for this model.
         """
         self._model_name = model_name
@@ -84,21 +103,26 @@ async def embed(self, documents: Sequence[str], *, settings: EmbeddingSettings |
     async def embed(
         self, documents: Sequence[str], *, settings: EmbeddingSettings | None = None
     ) -> list[float] | list[list[float]]:
-        input_is_string = isinstance(documents, str)
-        if input_is_string:
-            documents = [documents]
+        documents, is_single_document, settings = self.prepare_embed(documents, settings)
+        embeddings = await self._embed(documents, cast(CohereEmbeddingSettings, settings))
+        return embeddings[0] if is_single_document else embeddings
+
+    async def _embed(self, documents: Sequence[str], settings: CohereEmbeddingSettings) -> list[list[float]]:
+        request_options = RequestOptions()
+        if extra_headers := settings.get('extra_headers'):
+            request_options['additional_headers'] = extra_headers
+        if extra_body := settings.get('extra_body'):
+            request_options['additional_body_parameters'] = cast(dict[str, Any], extra_body)
 
-        settings = merge_embedding_settings(self._settings, settings) or {}
         response = await self._client.embed(
             model=self.model_name,
-            input_type=settings.get('input_type', 'search_document'),
-            texts=cast(Sequence[str], documents),
-            output_dimension=settings.get('output_dimension'),
+            texts=documents,
+            output_dimension=settings.get('dimensions'),
+            input_type=settings.get('cohere_input_type', 'search_document'),
+            max_tokens=settings.get('cohere_max_tokens'),
+            request_options=request_options,
         )
         embeddings = response.embeddings.float_
         assert embeddings is not None, 'This is a bug in cohere?'
 
-        if input_is_string:
-            return embeddings[0]
-
         return embeddings

pydantic_ai_slim/pydantic_ai/embeddings/openai.py

@@ -2,11 +2,11 @@
 from dataclasses import dataclass, field
 from typing import Literal, overload
 
-from pydantic_ai.embeddings.embedding_model import EmbeddingModel
+from pydantic_ai.embeddings.base import EmbeddingModel
 from pydantic_ai.embeddings.settings import EmbeddingSettings
 from pydantic_ai.providers import Provider, infer_provider
 
-from .settings import merge_embedding_settings
+from . import OpenAIEmbeddingsCompatibleProvider
 
 try:
     from openai import NOT_GIVEN, AsyncOpenAI
@@ -21,27 +21,34 @@
 """Possible OpenAI embeddings model names."""
 
 
+class OpenAIEmbeddingSettings(EmbeddingSettings, total=False):
+    """Settings used for an OpenAI embedding model request."""
+
+    # ALL FIELDS MUST BE `openai_` PREFIXED SO YOU CAN MERGE THEM WITH OTHER MODELS.
+
+
 @dataclass(init=False)
 class OpenAIEmbeddingModel(EmbeddingModel):
+    """OpenAI embedding model."""
+
     _model_name: OpenAIEmbeddingModelName = field(repr=False)
     _provider: Provider[AsyncOpenAI] = field(repr=False)
 
     def __init__(
         self,
         model_name: OpenAIEmbeddingModelName,
         *,
-        provider: Literal['openai'] | Provider[AsyncOpenAI] = 'openai',
+        provider: OpenAIEmbeddingsCompatibleProvider | Literal['openai'] | Provider[AsyncOpenAI] = 'openai',
         settings: EmbeddingSettings | None = None,
     ):
         """Initialize an OpenAI model.
 
         Args:
             model_name: The name of the OpenAI model to use. List of model names
-                available [here](https://docs.OpenAI.com/docs/models#command).
+                available [here](https://platform.openai.com/docs/guides/embeddings#embedding-models).
             provider: The provider to use for authentication and API access. Can be either the string
                 'OpenAI' or an instance of `Provider[AsyncClientV2]`. If not provided, a new provider will be
                 created using the other parameters.
-            profile: The model profile to use. Defaults to a profile picked by the provider based on the model name.
             settings: Model-specific settings that will be used as defaults for this model.
         """
         self._model_name = model_name
@@ -78,18 +85,16 @@ async def embed(self, documents: Sequence[str], *, settings: EmbeddingSettings |
     async def embed(
         self, documents: str | Sequence[str], *, settings: EmbeddingSettings | None = None
     ) -> list[float] | list[list[float]]:
-        input_is_string = isinstance(documents, str)
-        if input_is_string:
-            documents = [documents]
+        documents, is_single_document, settings = self.prepare_embed(documents, settings)
+        embeddings = await self._embed(documents, settings)
+        return embeddings[0] if is_single_document else embeddings
 
-        settings = merge_embedding_settings(self._settings, settings) or {}
+    async def _embed(self, documents: Sequence[str], settings: OpenAIEmbeddingSettings) -> list[list[float]]:
         response = await self._client.embeddings.create(
             input=documents,  # pyright: ignore[reportArgumentType]  # Sequence[str] not compatible with SequenceNotStr[str] :/
             model=self.model_name,
-            dimensions=settings.get('output_dimension') or NOT_GIVEN,
+            dimensions=settings.get('dimensions') or NOT_GIVEN,
+            extra_headers=settings.get('extra_headers'),
+            extra_body=settings.get('extra_body'),
         )
-        result = [item.embedding for item in response.data]
-
-        if input_is_string:
-            return result[0]
-        return result
+        return [item.embedding for item in response.data]

pydantic_ai_slim/pydantic_ai/embeddings/settings.py

@@ -1,46 +1,47 @@
-from typing import Literal, TypedDict
+from typing_extensions import TypedDict
 
 
 class EmbeddingSettings(TypedDict, total=False):
-    # TODO: May want to add extra_headres, extra_query, extra_body, timeout, etc.
+    """Settings to configure an embedding model.
 
-    output_dimension: int
-    """The maximum number of tokens to generate before stopping.
+    Here we include only settings which apply to multiple models / model providers,
+    though not all of these settings are supported by all models.
+    """
+
+    dimensions: int
+    """The number of dimensions the resulting output embeddings should have.
 
     Supported by:
 
-    * Cohere
     * OpenAI
+    * Cohere
     """
 
-    max_tokens: int
-    """The maximum number of tokens to generate before stopping.
+    extra_headers: dict[str, str]
+    """Extra headers to send to the model.
 
     Supported by:
 
+    * OpenAI
     * Cohere
     """
 
-    # We don't support embedding_types for now because it doesn't affect the user-facing API today..
-    # embedding_types: Literal["float", "int8", "uint8", "binary", "ubinary", "base64"]
-
-    input_type: Literal['search_document', 'search_query', 'classification', 'clustering', 'image']
-    """The input type of the embedding.
+    extra_body: object
+    """Extra body to send to the model.
 
     Supported by:
 
-    * Cohere (See `cohere.EmbedInputType`)
+    * OpenAI
+    * Cohere
     """
 
-    # TODO: Add more?
-
 
 def merge_embedding_settings(
     base: EmbeddingSettings | None, overrides: EmbeddingSettings | None
 ) -> EmbeddingSettings | None:
     """Merge two sets of embedding settings, preferring the overrides.
 
-    A common use case is: merge_embedding_settings(<agent settings>, <run settings>)
+    A common use case is: merge_embedding_settings(<embedder settings>, <run settings>)
     """
     # Note: we may want merge recursively if/when we add non-primitive values
     if base and overrides: