Fix Cohere, SentenceTransformers

DouweM · DouweM · commit 4ec0b32d3e82 · 2025-11-28T22:54:03.000Z
diff --git a/pydantic_ai_slim/pydantic_ai/embeddings/cohere.py b/pydantic_ai_slim/pydantic_ai/embeddings/cohere.py
@@ -1,15 +1,20 @@
 from collections.abc import Sequence
 from dataclasses import dataclass, field
-from typing import Any, Literal, cast, overload
+from typing import Any, Literal, cast
 
-from pydantic_ai.embeddings.base import EmbeddingModel, EmbedInputType
-from pydantic_ai.embeddings.settings import EmbeddingSettings
 from pydantic_ai.exceptions import UnexpectedModelBehavior
-from pydantic_ai.providers import infer_provider
+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 cohere import AsyncClientV2
     from cohere.core.request_options import RequestOptions
-    from cohere.v2.client import EmbedInputType as CohereEmbedInputType
+    from cohere.types.embed_by_type_response import EmbedByTypeResponse
+    from cohere.types.embed_input_type import EmbedInputType as CohereEmbedInputType
     from cohere.v2.types.v2embed_request_truncate import V2EmbedRequestTruncate
 
     from pydantic_ai.providers.cohere import CohereProvider
@@ -73,7 +78,7 @@ def __init__(
         self,
         model_name: CohereEmbeddingModelName,
         *,
-        provider: Literal['cohere'] | CohereProvider = 'cohere',
+        provider: Literal['cohere'] | Provider[AsyncClientV2] | CohereProvider = 'cohere',
         settings: EmbeddingSettings | None = None,
     ):
         """Initialize an Cohere model.
@@ -92,7 +97,7 @@ def __init__(
             provider = infer_provider(provider)
         self._provider = provider
         self._client = provider.client
-        self._v1_client = provider.v1_client
+        self._v1_client = provider.v1_client if isinstance(provider, CohereProvider) else None
 
         super().__init__(settings=settings)
 
@@ -111,28 +116,15 @@ def system(self) -> str:
         """The embedding model provider."""
         return self._provider.name
 
-    @overload
-    async def embed(
-        self, documents: str, *, input_type: EmbedInputType, settings: EmbeddingSettings | None = None
-    ) -> list[float]:
-        pass
-
-    @overload
-    async def embed(
-        self, documents: Sequence[str], *, input_type: EmbedInputType, settings: EmbeddingSettings | None = None
-    ) -> list[list[float]]:
-        pass
-
     async def embed(
-        self, documents: Sequence[str], *, input_type: EmbedInputType, settings: EmbeddingSettings | None = None
-    ) -> list[float] | list[list[float]]:
-        documents, is_single_document, settings = self.prepare_embed(documents, settings)
-        embeddings = await self._embed(documents, input_type, cast(CohereEmbeddingSettings, settings))
-        return embeddings[0] if is_single_document else embeddings
+        self, documents: str | Sequence[str], *, input_type: EmbedInputType, settings: EmbeddingSettings | None = None
+    ) -> EmbeddingResult:
+        documents, settings = self.prepare_embed(documents, settings)
+        return await self._embed(documents, input_type, cast(CohereEmbeddingSettings, settings))
 
     async def _embed(
-        self, documents: Sequence[str], input_type: EmbedInputType, settings: CohereEmbeddingSettings
-    ) -> list[list[float]]:
+        self, documents: str | Sequence[str], input_type: EmbedInputType, settings: CohereEmbeddingSettings
+    ) -> EmbeddingResult:
         request_options = RequestOptions()
         if extra_headers := settings.get('extra_headers'):
             request_options['additional_headers'] = extra_headers
@@ -156,10 +148,18 @@ async def _embed(
         if embeddings is None:
             raise UnexpectedModelBehavior(
                 'The Cohere embeddings response did not have an `embeddings` field holding a list of floats',
-                str(response.data),
+                response,
             )
 
-        return embeddings
+        return EmbeddingResult(
+            embeddings=embeddings,
+            inputs=documents,
+            input_type=input_type,
+            usage=_map_usage(response),
+            model_name=self.model_name,
+            provider_name=self.system,
+            provider_response_id=response.id,
+        )
 
     async def max_input_tokens(self) -> int | None:
         return _MAX_INPUT_TOKENS.get(self.model_name)
@@ -173,3 +173,17 @@ async def count_tokens(self, text: str) -> int:
             offline=False,
         )
         return len(result.tokens)
+
+
+def _map_usage(response: EmbedByTypeResponse) -> RequestUsage:
+    u = response.meta
+    if u is None or u.billed_units is None:
+        return RequestUsage()
+    usage_data = u.billed_units.model_dump(exclude_none=True)
+    details = {k: int(v) for k, v in usage_data.items() if k != 'input_tokens' and isinstance(v, int | float) and v > 0}
+
+    # TODO (DouweM): Use RequestUsage.extract() once https://github.com/pydantic/genai-prices/blob/main/prices/providers/cohere.yml has been updated
+    return RequestUsage(
+        input_tokens=int(u.billed_units.input_tokens or 0),
+        details=details,
+    )
diff --git a/pydantic_ai_slim/pydantic_ai/embeddings/instrumented.py b/pydantic_ai_slim/pydantic_ai/embeddings/instrumented.py
@@ -1,17 +1,19 @@
 from __future__ import annotations
 
 import json
+import warnings
 from collections.abc import Callable, Iterator, Sequence
 from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, overload
+from typing import TYPE_CHECKING, Any
 from urllib.parse import urlparse
 
 from opentelemetry.util.types import AttributeValue
 
-from pydantic_ai.models.instrumented import ANY_ADAPTER, InstrumentationSettings
+from pydantic_ai.models.instrumented import ANY_ADAPTER, CostCalculationFailedWarning, InstrumentationSettings
 
 from .base import EmbeddingModel, EmbedInputType
+from .result import EmbeddingResult
 from .settings import EmbeddingSettings
 from .wrapper import WrapperEmbeddingModel
 
@@ -50,19 +52,9 @@ def __init__(
         super().__init__(wrapped)
         self.instrumentation_settings = options or InstrumentationSettings()
 
-    @overload
-    async def embed(
-        self, documents: str, *, input_type: EmbedInputType, settings: EmbeddingSettings | None = None
-    ) -> list[float]: ...
-
-    @overload
-    async def embed(
-        self, documents: Sequence[str], *, input_type: EmbedInputType, settings: EmbeddingSettings | None = None
-    ) -> list[list[float]]: ...
-
     async def embed(
         self, documents: str | Sequence[str], *, input_type: EmbedInputType, settings: EmbeddingSettings | None = None
-    ) -> list[float] | list[list[float]]:
+    ) -> EmbeddingResult:
         with self._instrument(documents, input_type, settings) as finish:
             result = await self.wrapped.embed(documents, input_type=input_type, settings=settings)
             finish(result)
@@ -74,7 +66,7 @@ def _instrument(
         documents: str | Sequence[str],
         input_type: EmbedInputType,
         settings: EmbeddingSettings | None,
-    ) -> Iterator[Callable[[list[float] | list[list[float]]], None]]:
+    ) -> Iterator[Callable[[EmbeddingResult], None]]:
         operation = 'embed'
         span_name = f'{operation} {self.model_name}'
 
@@ -111,31 +103,46 @@ def _instrument(
         try:
             with self.instrumentation_settings.tracer.start_as_current_span(span_name, attributes=attributes) as span:
 
-                def finish(result: list[float] | list[list[float]]):
+                def finish(result: EmbeddingResult):
                     if not span.is_recording():
                         return
 
-                    # Calculate output dimension
-                    if isinstance(result, list) and result:
-                        if isinstance(result[0], list):
-                            # Multiple embeddings
-                            output_dim = len(result[0]) if result[0] else 0
-                            num_outputs = len(result)
-                        else:
-                            # Single embedding
-                            output_dim = len(result)
-                            num_outputs = 1
+                    attributes_to_set: dict[str, AttributeValue] = {
+                        **result.usage.opentelemetry_attributes(),
+                        'gen_ai.response.model': result.model_name or self.model_name,
+                    }
+
+                    try:
+                        price_calculation = result.cost()
+                    except LookupError:
+                        # The cost of this provider/model is unknown, which is common.
+                        pass
+                    except Exception as e:
+                        warnings.warn(
+                            f'Failed to get cost from response: {type(e).__name__}: {e}', CostCalculationFailedWarning
+                        )
                     else:
-                        output_dim = 0
-                        num_outputs = 0
+                        attributes_to_set['operation.cost'] = float(price_calculation.total_price)
+
+                    # Calculate output dimension
+                    embeddings = result.embeddings
+                    if embeddings:
+                        output_dim = len(embeddings[0]) if embeddings[0] else 0
+                        num_outputs = len(embeddings)
+
+                        attributes_to_set.update(
+                            {
+                                'gen_ai.embedding.dimension': output_dim,
+                                'gen_ai.embedding.num_outputs': num_outputs,
+                            }
+                        )
+
+                    if result.provider_response_id is not None:
+                        attributes_to_set['gen_ai.response.id'] = result.provider_response_id
 
-                    attributes_to_set = {
-                        'gen_ai.embedding.dimension': output_dim,
-                        'gen_ai.embedding.num_outputs': num_outputs,
-                    }
                     span.set_attributes(attributes_to_set)
 
-                    # TODO (DouweM): Include cost as metric etc, just like on InstrumentedModel
+                    # TODO (DouweM): Record cost metric
 
                 yield finish
         finally:
diff --git a/pydantic_ai_slim/pydantic_ai/embeddings/result.py b/pydantic_ai_slim/pydantic_ai/embeddings/result.py
@@ -36,6 +36,8 @@ class EmbeddingResult:
 
     provider_response_id: str | None = None
 
+    # TODO (DouweM): Support `result[idx: int]` and `result[document: str]`
+
     def cost(self) -> genai_types.PriceCalculation:
         """Calculate the cost of the usage.
 
diff --git a/pydantic_ai_slim/pydantic_ai/embeddings/sentence_transformers.py b/pydantic_ai_slim/pydantic_ai/embeddings/sentence_transformers.py
@@ -2,13 +2,15 @@
 
 from collections.abc import Sequence
 from dataclasses import dataclass, field
-from typing import Any, cast, overload
+from typing import Any, cast
 
 import pydantic_ai._utils as _utils
-from pydantic_ai.embeddings.base import EmbeddingModel, EmbedInputType
-from pydantic_ai.embeddings.settings import EmbeddingSettings
 from pydantic_ai.exceptions import UnexpectedModelBehavior
 
+from .base import EmbeddingModel, EmbedInputType
+from .result import EmbeddingResult
+from .settings import EmbeddingSettings
+
 try:
     import numpy as np
     import torch
@@ -73,26 +75,18 @@ def system(self) -> str:
         """The embedding model provider/system identifier."""
         return 'sentence-transformers'
 
-    @overload
-    async def embed(
-        self, documents: str, *, input_type: EmbedInputType, settings: EmbeddingSettings | None = None
-    ) -> list[float]: ...
-
-    @overload
-    async def embed(
-        self, documents: Sequence[str], *, input_type: EmbedInputType, settings: EmbeddingSettings | None = None
-    ) -> list[list[float]]: ...
-
     async def embed(
         self, documents: str | Sequence[str], *, input_type: EmbedInputType, settings: EmbeddingSettings | None = None
     ) -> list[float] | list[list[float]]:
-        docs, is_single_document, settings = self.prepare_embed(documents, settings)
-        embeddings = await self._embed(docs, input_type, cast(SentenceTransformersEmbeddingSettings, settings))
-        return embeddings[0] if is_single_document else embeddings
+        docs, settings = self.prepare_embed(documents, settings)
+        return await self._embed(docs, input_type, cast(SentenceTransformersEmbeddingSettings, settings))
 
     async def _embed(
-        self, documents: Sequence[str], input_type: EmbedInputType, settings: SentenceTransformersEmbeddingSettings
-    ) -> list[list[float]]:
+        self,
+        documents: str | Sequence[str],
+        input_type: EmbedInputType,
+        settings: SentenceTransformersEmbeddingSettings,
+    ) -> EmbeddingResult:
         device = settings.get('sentence_transformers_device', None)
         normalize = settings.get('sentence_transformers_normalize_embeddings', False)
         batch_size = settings.get('sentence_transformers_batch_size', None)
@@ -111,7 +105,15 @@ async def _embed(
             normalize_embeddings=normalize,
             **{'batch_size': batch_size} if batch_size is not None else {},  # type: ignore[reportArgumentType]
         )
-        return np_embeddings.tolist()  # type: ignore[reportUnknownReturnType]
+        embeddings = np_embeddings.tolist()  # type: ignore[reportAttributeAccessIssue]
+
+        return EmbeddingResult(
+            embeddings=embeddings,  # type: ignore[reportUnknownArgumentType]
+            inputs=documents,
+            input_type=input_type,
+            model_name=self.model_name,
+            provider_name=self.system,
+        )
 
     async def max_input_tokens(self) -> int | None:
         model = await self._get_model()
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -34,6 +34,7 @@
     'IsBytes',
     'IsInt',
     'IsInstance',
+    'IsList',
     'TestEnv',
     'ClientWithHandler',
     'try_import',
@@ -62,8 +63,9 @@ def IsNow(*args: Any, **kwargs: Any) -> datetime: ...
     def IsStr(*args: Any, **kwargs: Any) -> str: ...
     def IsSameStr(*args: Any, **kwargs: Any) -> str: ...
     def IsBytes(*args: Any, **kwargs: Any) -> bytes: ...
+    def IsList(*args: T, **kwargs: Any) -> list[T]: ...
 else:
-    from dirty_equals import IsBytes, IsDatetime, IsFloat, IsInstance, IsInt, IsNow as _IsNow, IsStr
+    from dirty_equals import IsBytes, IsDatetime, IsFloat, IsInstance, IsInt, IsList, IsNow as _IsNow, IsStr
 
     def IsNow(*args: Any, **kwargs: Any):
         # Increase the default value of `delta` to 10 to reduce test flakiness on overburdened machines
diff --git a/tests/test_embeddings.py b/tests/test_embeddings.py
