pydantic/pydantic-ai
1
0
Fork 0
mirror of https://github.com/pydantic/pydantic-ai.git synced 2026-10-12 02:21:31 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
main
pydantic-ai/tests/test_embeddings.py
T

2654 lines
121 KiB
Python
Raw Permalink Blame History

from __future__ import annotations
import importlib
import io
import json
import os
import re
import threading
from collections.abc import Iterator, Sequence
from copy import deepcopy
from dataclasses import dataclass
from decimal import Decimal
from importlib.util import find_spec
from pathlib import Path
from typing import Any, Literal, get_args
from unittest.mock import AsyncMock, MagicMock, patch
from urllib.parse import urlparse
import anyio
import httpx
import httpx2
import pytest
from opentelemetry.trace import StatusCode
from pytest_mock import MockerFixture
import pydantic_ai.models
from pydantic_ai._utils import BaseExceptionGroup
from pydantic_ai.embeddings import (
Embedder,
EmbeddingResult,
EmbeddingSettings,
EmbedInputType,
InstrumentedEmbeddingModel,
KnownEmbeddingModelName,
TestEmbeddingModel,
WrapperEmbeddingModel,
infer_embedding_model,
)
from pydantic_ai.exceptions import ModelAPIError, ModelHTTPError, UserError
from pydantic_ai.models.instrumented import InstrumentationSettings
from pydantic_ai.usage import RequestUsage
from ._inline_snapshot import snapshot
from .conftest import IsDatetime, IsFloat, IsInt, IsList, IsStr, TestEnv, try_import
pytestmark = [
pytest.mark.usefixtures('allow_model_requests'),
]
with try_import() as logfire_imports_successful:
from logfire.testing import CaptureLogfire
with try_import() as openai_imports_successful:
from openai import AsyncOpenAI
from pydantic_ai.embeddings.openai import LatestOpenAIEmbeddingModelNames, OpenAIEmbeddingModel
from pydantic_ai.providers.openai import OpenAIProvider
with try_import() as cohere_imports_successful:
import cohere
from cohere import AsyncClientV2
from cohere.types.embed_by_type_response import EmbedByTypeResponse
from cohere.types.embed_by_type_response_embeddings import EmbedByTypeResponseEmbeddings
from pydantic_ai.embeddings import cohere as cohere_embeddings
from pydantic_ai.embeddings.cohere import (
CohereEmbeddingModel,
CohereEmbeddingSettings,
LatestCohereEmbeddingModelNames,
)
from pydantic_ai.providers.cohere import CohereProvider
with try_import() as bedrock_imports_successful:
from botocore.awsrequest import AWSPreparedRequest, AWSResponse, HTTPHeaders
from botocore.exceptions import ClientError, EndpointConnectionError, ReadTimeoutError
from urllib3 import HTTPResponse
from pydantic_ai.embeddings.bedrock import (
BedrockEmbeddingModel,
BedrockEmbeddingSettings,
LatestBedrockEmbeddingModelNames,
)
from pydantic_ai.providers.bedrock import BedrockProvider
with try_import() as google_imports_successful:
from pydantic_ai.embeddings.google import (
GoogleEmbeddingModel,
GoogleEmbeddingSettings,
LatestGoogleGLAEmbeddingModelNames,
LatestGoogleVertexEmbeddingModelNames,
)
from pydantic_ai.providers.google import GoogleProvider
from pydantic_ai.providers.google_cloud import GoogleCloudProvider
# `sentence_transformers` imports `torch`, and `voyageai` imports both when `sentence-transformers` is installed, which
# takes seconds: the tests that use them import them, so collecting this module stays cheap. See "Test cost" in
# `tests/AGENTS.md`.
voyageai_installed = find_spec('voyageai') is not None
sentence_transformers_installed = find_spec('sentence_transformers') is not None
@pytest.fixture(scope='module')
def voyageai_imported() -> None:
"""Pay for importing `voyageai` in shared setup, outside each test's time budget."""
importlib.import_module('pydantic_ai.embeddings.voyageai')
@pytest.fixture(scope='module')
def sentence_transformers_imported() -> None:
"""Pay for importing `sentence_transformers` and `torch` in shared setup, outside each test's time budget."""
importlib.import_module('pydantic_ai.embeddings.sentence_transformers')
@pytest.mark.skipif(not openai_imports_successful(), reason='openai not installed')
async def test_openai_embedding_model_blocks_requests_when_disabled():
model = OpenAIEmbeddingModel('text-embedding-3-small', provider=OpenAIProvider(api_key='test-key'))
with pydantic_ai.models.override_allow_model_requests(False):
with pytest.raises(RuntimeError, match='Model requests are not allowed'):
await model.embed('hello', input_type='query')
@pytest.mark.skipif(not cohere_imports_successful(), reason='cohere not installed')
async def test_cohere_embedding_model_blocks_requests_when_disabled():
model = CohereEmbeddingModel('embed-v4.0', provider=CohereProvider(api_key='test-key'))
with pydantic_ai.models.override_allow_model_requests(False):
with pytest.raises(RuntimeError, match='Model requests are not allowed'):
await model.embed('hello', input_type='query')
@pytest.mark.skipif(not google_imports_successful(), reason='google not installed')
async def test_google_embedding_model_blocks_requests_when_disabled():
model = GoogleEmbeddingModel('gemini-embedding-001', provider=GoogleProvider(api_key='test-key'))
with pydantic_ai.models.override_allow_model_requests(False):
with pytest.raises(RuntimeError, match='Model requests are not allowed'):
await model.embed('hello', input_type='query')
@pytest.mark.skipif(not bedrock_imports_successful(), reason='bedrock not installed')
async def test_bedrock_embedding_model_blocks_requests_when_disabled():
model = BedrockEmbeddingModel('amazon.titan-embed-text-v2:0', provider=BedrockProvider(bedrock_client=MagicMock()))
with pydantic_ai.models.override_allow_model_requests(False):
with pytest.raises(RuntimeError, match='Model requests are not allowed'):
await model.embed('hello', input_type='query')
@pytest.mark.skipif(not voyageai_installed, reason='voyageai not installed')
@pytest.mark.usefixtures('voyageai_imported')
async def test_voyageai_embedding_model_blocks_requests_when_disabled():
from pydantic_ai.embeddings.voyageai import VoyageAIEmbeddingModel
from pydantic_ai.providers.voyageai import VoyageAIProvider
model = VoyageAIEmbeddingModel('voyage-4', provider=VoyageAIProvider(api_key='test-key'))
with pydantic_ai.models.override_allow_model_requests(False):
with pytest.raises(RuntimeError, match='Model requests are not allowed'):
await model.embed('hello', input_type='query')
@pytest.mark.skipif(not google_imports_successful(), reason='google not installed')
async def test_google_embedding_model_blocks_count_tokens_when_disabled():
model = GoogleEmbeddingModel('gemini-embedding-001', provider=GoogleProvider(api_key='test-key'))
with pydantic_ai.models.override_allow_model_requests(False):
with pytest.raises(RuntimeError, match='Model requests are not allowed'):
await model.count_tokens('hello')
@pytest.mark.skipif(not cohere_imports_successful(), reason='cohere not installed')
async def test_cohere_embedding_model_blocks_count_tokens_when_disabled():
model = CohereEmbeddingModel('embed-v4.0', provider=CohereProvider(api_key='test-key'))
with pydantic_ai.models.override_allow_model_requests(False):
with pytest.raises(RuntimeError, match='Model requests are not allowed'):
await model.count_tokens('hello')
@pytest.mark.skipif(not openai_imports_successful(), reason='openai not installed')
async def test_embedder_blocks_requests_when_disabled():
"""Pins the guard on the public `Embedder` surface, which is what issue #6763 reports as leaking.
A guard that fires before the request is made can't be exercised through a VCR recording.
The per-model tests above prove each `embed()` guards; this proves the wrapper chain
(`Embedder` -> `InstrumentedEmbeddingModel` / `WrapperEmbeddingModel` -> concrete model)
surfaces the `RuntimeError` rather than swallowing or wrapping it.
"""
embedder = Embedder(OpenAIEmbeddingModel('text-embedding-3-small', provider=OpenAIProvider(api_key='test-key')))
with pydantic_ai.models.override_allow_model_requests(False):
with pytest.raises(RuntimeError, match='Model requests are not allowed'):
await embedder.embed_query('hello')
async def test_test_embedding_model_is_exempt_from_request_guard():
"""`ALLOW_MODEL_REQUESTS`'s docstring promises `TestEmbeddingModel` is unaffected; pin that promise.
Without this, adding the guard to `TestEmbeddingModel.embed` would break every user's test
suite while the whole file still passed, since the module-level `allow_model_requests`
fixture keeps the flag on for every other test here.
"""
embedder = Embedder(TestEmbeddingModel())
with pydantic_ai.models.override_allow_model_requests(False):
result = await embedder.embed_query('hello')
assert result.embeddings == snapshot([[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]])
async def test_test_embedding_model_counts_blank_input_as_zero_tokens():
"""Blank input reports no tokens on both methods.
The estimator behind them is shared with the other test models, and it counts blank text as one
token; `TestEmbeddingModel` guards that, so a blank input contributes nothing to the reported
usage instead of one phantom token per empty string in the batch.
"""
model = TestEmbeddingModel()
assert await model.count_tokens('') == snapshot(0)
result = await model.embed(['', 'hi there'], input_type='document')
assert result.usage.input_tokens == snapshot(2)
STSB_BERT_TINY_MODEL = 'sentence-transformers-testing/stsb-bert-tiny-safetensors'
# Pinned so a warm HF cache is served without revalidating files against the Hub.
# Keep in sync with the HF cache keys and warmup commands in .github/workflows/ci.yml;
# `test_stsb_model_pin_matches_ci` guards against drift.
STSB_BERT_TINY_REVISION = 'f3cb857cba53019a20df283396bcca179cf051a4'
def _hf_hub_unavailable(exc: BaseException) -> bool:
"""Whether an exception from a SentenceTransformer load means the HF Hub is
unavailable (worth skipping the real-model smoke tests) rather than a genuine
integration defect or a bad model pin (which must fail loudly). Shapes verified
against real outage and bad-pin probes.
"""
if isinstance(exc, httpx.TransportError):
# Transport-level httpx errors (e.g. connect timeouts) escape
# huggingface_hub 1.x unwrapped, and the only httpx traffic in a model
# load is Hub traffic. Status-level httpx errors are not outage proof.
return True
if isinstance(exc, RuntimeError):
# After a connection error, huggingface_hub 1.x retries on an httpx client
# it just closed, surfacing httpx's RuntimeError instead of its own error.
return 'client has been closed' in str(exc)
if type(exc).__module__.startswith('huggingface_hub'):
# huggingface_hub errors are OSError subclasses. Ones carrying an HTTP
# response are an outage only for transient statuses; 401/403/404 mean a
# bad pin or an auth problem. Ones without a response (e.g.
# LocalEntryNotFoundError) mean the Hub couldn't be reached at all.
status = getattr(getattr(exc, 'response', None), 'status_code', None)
return status is None or status == 429 or status >= 500
# transformers wraps connection failures in a plain OSError: "We couldn't
# connect to 'https://huggingface.co' to load the files [...]". Other OSError
# shapes (corrupt cache, invalid model identifier) must propagate.
return "couldn't connect" in str(exc).lower()
def test_hf_hub_unavailable_classifier():
"""Exercises the fixture's outage-classification guard on synthetic exceptions:
no HTTP is involved (hence no VCR), just shapes captured from real probes."""
class FakeResponse:
def __init__(self, status_code: int):
self.status_code = status_code
class FakeHubError(OSError):
def __init__(self, msg: str, status_code: int | None = None):
super().__init__(msg)
self.response = FakeResponse(status_code) if status_code is not None else None
FakeHubError.__module__ = 'huggingface_hub.errors'
# Hub unavailability: skip the real-model smoke tests.
assert _hf_hub_unavailable(httpx.ConnectTimeout('timed out'))
assert _hf_hub_unavailable(RuntimeError('Cannot send a request, as the client has been closed.'))
assert _hf_hub_unavailable(FakeHubError('504 Server Error', 504))
assert _hf_hub_unavailable(FakeHubError('429 Too Many Requests', 429))
assert _hf_hub_unavailable(FakeHubError('cannot find the requested files in the disk cache'))
assert _hf_hub_unavailable(OSError("We couldn't connect to 'https://huggingface.co' to load the files"))
# Anything else fails loudly: bad pins, auth problems, local defects.
assert not _hf_hub_unavailable(FakeHubError('401 Unauthorized', 401))
assert not _hf_hub_unavailable(FakeHubError('404 Repository Not Found', 404))
assert not _hf_hub_unavailable(
httpx.HTTPStatusError('404', request=httpx.Request('GET', 'https://x'), response=httpx.Response(404))
)
assert not _hf_hub_unavailable(RuntimeError('CUDA error: device-side assert triggered'))
assert not _hf_hub_unavailable(OSError('Unable to load weights from pytorch checkpoint file'))
assert not _hf_hub_unavailable(
OSError(
"xyz is not a local folder and is not a valid model identifier listed on 'https://huggingface.co/models'"
)
)
def test_stsb_model_pin_matches_ci():
"""Parses CI configuration only, no network or VCR involved: drift between this
pin and ci.yml silently reintroduces per-run Hub downloads, because the stale
cache key keeps exact-hitting, so CI never caches the new revision."""
ci_yml = Path(__file__).parent.parent / '.github' / 'workflows' / 'ci.yml'
if not ci_yml.is_file(): # pragma: lax no cover
pytest.skip('not running from a repo checkout')
stsb_lines = [line for line in ci_yml.read_text().splitlines() if 'stsb-bert-tiny' in line]
cache_keys = [line for line in stsb_lines if 'key:' in line]
warmups = [line for line in stsb_lines if 'snapshot_download' in line]
assert len(cache_keys) >= 2, 'expected a model cache key in both test jobs in ci.yml'
assert len(warmups) >= 2, 'expected a model warmup command in both test jobs in ci.yml'
assert all(STSB_BERT_TINY_MODEL in line for line in warmups)
for line in cache_keys + warmups:
assert STSB_BERT_TINY_REVISION in line, f'stale or missing revision pin in ci.yml: {line.strip()}'
stray = {sha for line in stsb_lines for sha in re.findall(r'[0-9a-f]{40}', line)} - {STSB_BERT_TINY_REVISION}
assert not stray, f'stale revision pins left in ci.yml: {stray}'
@pytest.mark.skipif(not openai_imports_successful(), reason='OpenAI not installed')
@pytest.mark.vcr
class TestOpenAI:
@pytest.fixture
def embedder(self, openai_api_key: str) -> Embedder:
return Embedder(OpenAIEmbeddingModel('text-embedding-3-small', provider=OpenAIProvider(api_key=openai_api_key)))
async def test_infer_model(self, openai_api_key: str):
with patch.dict(os.environ, {'OPENAI_API_KEY': openai_api_key}):
model = infer_embedding_model('openai:text-embedding-3-small')
assert isinstance(model, OpenAIEmbeddingModel)
assert model.model_name == 'text-embedding-3-small'
assert model.system == 'openai'
assert model.base_url == 'https://api.openai.com/v1/'
async def test_infer_model_azure(self):
with patch.dict(
os.environ,
{
'AZURE_OPENAI_API_KEY': 'azure-openai-api-key',
'AZURE_OPENAI_ENDPOINT': 'https://project-id.openai.azure.com/',
'OPENAI_API_VERSION': '2023-03-15-preview',
},
):
model = infer_embedding_model('azure:text-embedding-3-small')
assert isinstance(model, OpenAIEmbeddingModel)
assert model.model_name == 'text-embedding-3-small'
assert model.system == 'azure'
assert urlparse(model.base_url).hostname == 'project-id.openai.azure.com'
assert await model.max_input_tokens() is None
with pytest.raises(UserError, match='Counting tokens is not supported for non-OpenAI embedding models'):
await model.count_tokens('Hello, world!')
async def test_infer_model_gateway(self):
with patch.dict(
os.environ,
{
'PYDANTIC_AI_GATEWAY_API_KEY': 'test-api-key',
'PYDANTIC_AI_GATEWAY_BASE_URL': 'https://gateway.pydantic.dev/proxy',
},
):
model = infer_embedding_model('gateway/openai:text-embedding-3-small')
assert isinstance(model, OpenAIEmbeddingModel)
assert model.model_name == 'text-embedding-3-small'
assert model.system == 'openai'
assert urlparse(model.base_url).hostname == 'gateway.pydantic.dev'
async def test_infer_model_vllm(self):
with patch.dict(os.environ, {'VLLM_BASE_URL': 'http://localhost:8000/v1'}):
model = infer_embedding_model('vllm:intfloat/e5-mistral-7b-instruct')
assert isinstance(model, OpenAIEmbeddingModel)
assert model.model_name == 'intfloat/e5-mistral-7b-instruct'
assert model.system == 'vllm'
assert model.base_url == 'http://localhost:8000/v1/'
async def test_query(self, embedder: Embedder):
result = await embedder.embed_query('Hello, world!')
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=1536), length=1),
inputs=['Hello, world!'],
input_type='query',
usage=RequestUsage(input_tokens=4),
model_name='text-embedding-3-small',
timestamp=IsDatetime(),
provider_name='openai',
)
)
assert result.cost().total_price == snapshot(Decimal('8E-8'))
async def test_documents(self, embedder: Embedder):
result = await embedder.embed_documents(['hello', 'world'])
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=1536), length=2),
inputs=['hello', 'world'],
input_type='document',
usage=RequestUsage(input_tokens=2),
model_name='text-embedding-3-small',
timestamp=IsDatetime(),
provider_name='openai',
)
)
assert result.cost().total_price == snapshot(Decimal('4E-8'))
async def test_max_input_tokens(self, embedder: Embedder):
max_input_tokens = await embedder.max_input_tokens()
assert max_input_tokens == snapshot(8192)
async def test_count_tokens(self, embedder: Embedder):
count = await embedder.count_tokens('Hello, world!')
assert count == snapshot(4)
async def test_embed_error(self, openai_api_key: str):
model = OpenAIEmbeddingModel('nonexistent', provider=OpenAIProvider(api_key=openai_api_key))
embedder = Embedder(model)
with pytest.raises(ModelHTTPError, match='model_not_found'):
await embedder.embed_query('Hello, world!')
async def test_response_with_no_usage(self):
mock_client = AsyncMock()
mock_embedding_item = MagicMock()
mock_embedding_item.embedding = [0.1, 0.2, 0.3]
mock_response = MagicMock()
mock_response.data = [mock_embedding_item]
mock_response.usage = None
mock_response.model = 'test-model'
mock_client.embeddings.create.return_value = mock_response
provider = OpenAIProvider(openai_client=mock_client)
model = OpenAIEmbeddingModel('test-model', provider=provider)
result = await model.embed('test', input_type='query')
assert result == snapshot(
EmbeddingResult(
embeddings=[[0.1, 0.2, 0.3]],
inputs=['test'],
input_type='query',
model_name='test-model',
provider_name='openai',
timestamp=IsDatetime(),
)
)
@pytest.mark.skipif(not logfire_imports_successful(), reason='logfire not installed')
async def test_instrumentation(self, openai_api_key: str, capfire: CaptureLogfire):
model = OpenAIEmbeddingModel('text-embedding-3-small', provider=OpenAIProvider(api_key=openai_api_key))
embedder = Embedder(model, instrument=True)
await embedder.embed_query('Hello, world!', settings={'dimensions': 128})
spans = capfire.exporter.exported_spans_as_dict(parse_json_attributes=True)
span = next(span for span in spans if 'embeddings' in span['name'])
assert span == snapshot(
{
'name': 'embeddings text-embedding-3-small',
'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'parent': None,
'start_time': IsInt(),
'end_time': IsInt(),
'attributes': {
'gen_ai.operation.name': 'embeddings',
'gen_ai.provider.name': 'openai',
'gen_ai.request.model': 'text-embedding-3-small',
'input_type': 'query',
'server.address': 'api.openai.com',
'inputs_count': 1,
'embedding_settings': {'dimensions': 128},
'inputs': ['Hello, world!'],
'logfire.json_schema': {
'type': 'object',
'properties': {
'input_type': {'type': 'string'},
'inputs_count': {'type': 'integer'},
'embedding_settings': {'type': 'object'},
'inputs': {'type': ['array']},
'embeddings': {'type': 'array'},
},
},
'logfire.span_type': 'span',
'logfire.msg': 'embeddings text-embedding-3-small',
'gen_ai.usage.input_tokens': 4,
'operation.cost': 8e-08,
'gen_ai.response.model': 'text-embedding-3-small',
'gen_ai.embeddings.dimension.count': 128,
},
}
)
assert capfire.get_collected_metrics() == snapshot(
[
{
'name': 'gen_ai.client.token.usage',
'description': 'Measures number of input and output tokens used',
'unit': '{token}',
'data': {
'data_points': [
{
'attributes': {
'gen_ai.provider.name': 'openai',
'gen_ai.operation.name': 'embeddings',
'gen_ai.request.model': 'text-embedding-3-small',
'gen_ai.response.model': 'text-embedding-3-small',
'gen_ai.token.type': 'input',
},
'start_time_unix_nano': IsInt(),
'time_unix_nano': IsInt(),
'count': 1,
'sum': 4,
'scale': 20,
'zero_count': 0,
'positive': {'offset': 2097151, 'bucket_counts': [1]},
'negative': {'offset': 0, 'bucket_counts': [0]},
'flags': 0,
'min': 4,
'max': 4,
'exemplars': [],
}
],
'aggregation_temporality': 1,
},
},
{
'name': 'operation.cost',
'description': 'Monetary cost',
'unit': '{USD}',
'data': {
'data_points': [
{
'attributes': {
'gen_ai.provider.name': 'openai',
'gen_ai.operation.name': 'embeddings',
'gen_ai.request.model': 'text-embedding-3-small',
'gen_ai.response.model': 'text-embedding-3-small',
},
'start_time_unix_nano': IsInt(),
'time_unix_nano': IsInt(),
'count': 1,
'sum': 8e-08,
'scale': 20,
'zero_count': 0,
'positive': {'offset': -24720625, 'bucket_counts': [1]},
'negative': {'offset': 0, 'bucket_counts': [0]},
'flags': 0,
'min': 8e-08,
'max': 8e-08,
'exemplars': [],
}
],
'aggregation_temporality': 1,
},
},
]
)
@pytest.mark.skipif(not cohere_imports_successful(), reason='Cohere not installed')
def test_cohere_empty_billed_units():
"""An empty SDK billing object is a defensive case that a VCR recording cannot reliably produce."""
for billed_units in (cohere.ApiMetaBilledUnits(), cohere.ApiMetaBilledUnits(input_tokens=0, output_tokens=0)):
response = EmbedByTypeResponse(
id='test',
embeddings=EmbedByTypeResponseEmbeddings(float_=[[0.1]]),
meta=cohere.ApiMeta(billed_units=billed_units),
)
assert (
cohere_embeddings._map_usage( # pyright: ignore[reportPrivateUsage]
response, 'cohere', 'https://api.cohere.com', 'embed-v4.0'
)
== RequestUsage()
)
@pytest.mark.skipif(not cohere_imports_successful(), reason='Cohere not installed')
@pytest.mark.vcr
class TestCohere:
async def test_infer_model(self, co_api_key: str):
with patch.dict(os.environ, {'CO_API_KEY': co_api_key}):
model = infer_embedding_model('cohere:embed-v4.0')
assert isinstance(model, CohereEmbeddingModel)
assert model.model_name == 'embed-v4.0'
assert model.system == 'cohere'
assert model.base_url == 'https://api.cohere.com'
assert isinstance(model._provider, CohereProvider) # type: ignore[reportAttributeAccess]
async def test_query(self, co_api_key: str):
model = CohereEmbeddingModel('embed-v4.0', provider=CohereProvider(api_key=co_api_key))
embedder = Embedder(model)
result = await embedder.embed_query('Hello, world!')
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(
IsList(snapshot(-0.018445116), snapshot(0.008921167), snapshot(-0.0011377502), length=1536),
length=1,
),
inputs=['Hello, world!'],
input_type='query',
usage=RequestUsage(input_tokens=4),
model_name='embed-v4.0',
timestamp=IsDatetime(),
provider_name='cohere',
provider_response_id='0728b136-9b30-4fb5-bf9a-2c7cf36d51d3',
)
)
assert result.cost().total_price == snapshot(Decimal('4.8E-7'))
async def test_documents(self, co_api_key: str):
model = CohereEmbeddingModel('embed-v4.0', provider=CohereProvider(api_key=co_api_key))
embedder = Embedder(model)
result = await embedder.embed_documents(['hello', 'world'])
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=1536), length=2),
inputs=['hello', 'world'],
input_type='document',
usage=RequestUsage(input_tokens=2),
model_name='embed-v4.0',
timestamp=IsDatetime(),
provider_name='cohere',
provider_response_id='199299d7-f43d-45af-903c-347fff81bbe4',
)
)
assert result.cost().total_price == snapshot(Decimal('2.4E-7'))
async def test_max_input_tokens(self, co_api_key: str):
model = CohereEmbeddingModel('embed-v4.0', provider=CohereProvider(api_key=co_api_key))
embedder = Embedder(model)
max_input_tokens = await embedder.max_input_tokens()
assert max_input_tokens == snapshot(128000)
async def test_count_tokens(self, co_api_key: str):
model = CohereEmbeddingModel('embed-v4.0', provider=CohereProvider(api_key=co_api_key))
embedder = Embedder(model)
count = await embedder.count_tokens('Hello, world!')
assert count == snapshot(4)
async def test_embed_error(self, co_api_key: str):
model = CohereEmbeddingModel('nonexistent', provider=CohereProvider(api_key=co_api_key))
embedder = Embedder(model)
with pytest.raises(ModelHTTPError, match='not found,'):
await embedder.embed_query('Hello, world!')
async def test_embed_connect_error(self):
"""`cohere` doesn't wrap transport failures; a cassette can't replay one, so a mock transport raises it."""
def handler(request: httpx.Request) -> httpx.Response:
raise httpx.ConnectError('connection refused', request=request)
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http_client:
cohere_client = AsyncClientV2(api_key='test-key', base_url='http://localhost', httpx_client=http_client)
embedder = Embedder(
CohereEmbeddingModel('embed-v4.0', provider=CohereProvider(cohere_client=cohere_client))
)
with pytest.raises(ModelAPIError) as exc_info:
await embedder.embed_query('Hello, world!')
assert type(exc_info.value) is ModelAPIError
assert exc_info.value.message == 'connection refused'
assert isinstance(exc_info.value.__cause__, httpx.ConnectError)
async def test_query_with_cohere_truncate(self, co_api_key: str):
model = CohereEmbeddingModel('embed-v4.0', provider=CohereProvider(api_key=co_api_key))
embedder = Embedder(model)
settings: CohereEmbeddingSettings = {'cohere_truncate': 'END'}
result = await embedder.embed_query('Hello, world!', settings=settings)
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=1536), length=1),
inputs=['Hello, world!'],
input_type='query',
usage=RequestUsage(input_tokens=4),
model_name='embed-v4.0',
timestamp=IsDatetime(),
provider_name='cohere',
provider_response_id=IsStr(),
)
)
async def test_query_with_truncate(self, co_api_key: str):
model = CohereEmbeddingModel('embed-v4.0', provider=CohereProvider(api_key=co_api_key))
embedder = Embedder(model)
result = await embedder.embed_query('Hello, world!', settings={'truncate': True})
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=1536), length=1),
inputs=['Hello, world!'],
input_type='query',
usage=RequestUsage(input_tokens=4),
model_name='embed-v4.0',
timestamp=IsDatetime(),
provider_name='cohere',
provider_response_id=IsStr(),
)
)
@pytest.mark.skipif(not voyageai_installed, reason='VoyageAI not installed')
@pytest.mark.usefixtures('voyageai_imported')
@pytest.mark.vcr
class TestVoyageAI:
async def test_infer_model(self, voyage_api_key: str):
from pydantic_ai.embeddings.voyageai import VoyageAIEmbeddingModel
from pydantic_ai.providers.voyageai import VoyageAIProvider
with patch.dict(os.environ, {'VOYAGE_API_KEY': voyage_api_key}):
model = infer_embedding_model('voyageai:voyage-3.5')
assert isinstance(model, VoyageAIEmbeddingModel)
assert model.model_name == 'voyage-3.5'
assert model.system == 'voyageai'
assert model.base_url == 'https://api.voyageai.com/v1'
assert isinstance(model._provider, VoyageAIProvider) # type: ignore[reportAttributeAccess]
async def test_query(self, voyage_api_key: str):
from pydantic_ai.embeddings.voyageai import VoyageAIEmbeddingModel
from pydantic_ai.providers.voyageai import VoyageAIProvider
model = VoyageAIEmbeddingModel('voyage-3.5', provider=VoyageAIProvider(api_key=voyage_api_key))
embedder = Embedder(model)
result = await embedder.embed_query('Hello, world!')
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=1024), length=1),
inputs=['Hello, world!'],
input_type='query',
usage=RequestUsage(input_tokens=3),
model_name='voyage-3.5',
timestamp=IsDatetime(),
provider_name='voyageai',
)
)
async def test_query_voyage_4(self, voyage_api_key: str):
from pydantic_ai.embeddings.voyageai import VoyageAIEmbeddingModel
from pydantic_ai.providers.voyageai import VoyageAIProvider
model = VoyageAIEmbeddingModel('voyage-4', provider=VoyageAIProvider(api_key=voyage_api_key))
embedder = Embedder(model)
result = await embedder.embed_query('Hello, world!')
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=1024), length=1),
inputs=['Hello, world!'],
input_type='query',
usage=RequestUsage(input_tokens=3),
model_name='voyage-4',
timestamp=IsDatetime(),
provider_name='voyageai',
)
)
async def test_documents(self, voyage_api_key: str):
from pydantic_ai.embeddings.voyageai import VoyageAIEmbeddingModel
from pydantic_ai.providers.voyageai import VoyageAIProvider
model = VoyageAIEmbeddingModel('voyage-3.5', provider=VoyageAIProvider(api_key=voyage_api_key))
embedder = Embedder(model)
result = await embedder.embed_documents(['hello', 'world'])
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=1024), length=2),
inputs=['hello', 'world'],
input_type='document',
usage=RequestUsage(),
model_name='voyage-3.5',
timestamp=IsDatetime(),
provider_name='voyageai',
)
)
async def test_max_input_tokens(self, voyage_api_key: str):
from pydantic_ai.embeddings.voyageai import VoyageAIEmbeddingModel
from pydantic_ai.providers.voyageai import VoyageAIProvider
model = VoyageAIEmbeddingModel('voyage-3.5', provider=VoyageAIProvider(api_key=voyage_api_key))
embedder = Embedder(model)
max_input_tokens = await embedder.max_input_tokens()
assert max_input_tokens == snapshot(32000)
async def test_embed_error(self, voyage_api_key: str):
from pydantic_ai.embeddings.voyageai import VoyageAIEmbeddingModel
from pydantic_ai.providers.voyageai import VoyageAIProvider
model = VoyageAIEmbeddingModel('nonexistent', provider=VoyageAIProvider(api_key=voyage_api_key))
embedder = Embedder(model)
with pytest.raises(ModelAPIError, match='not supported'):
await embedder.embed_query('Hello, world!')
async def test_query_with_truncate(self, voyage_api_key: str):
from pydantic_ai.embeddings.voyageai import VoyageAIEmbeddingModel
from pydantic_ai.providers.voyageai import VoyageAIProvider
model = VoyageAIEmbeddingModel('voyage-3.5', provider=VoyageAIProvider(api_key=voyage_api_key))
embedder = Embedder(model)
result = await embedder.embed_query('Hello, world!', settings={'truncate': True})
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=1024), length=1),
inputs=['Hello, world!'],
input_type='query',
usage=RequestUsage(input_tokens=3),
model_name='voyage-3.5',
timestamp=IsDatetime(),
provider_name='voyageai',
)
)
async def test_query_with_voyageai_input_type(self, voyage_api_key: str):
from pydantic_ai.embeddings.voyageai import VoyageAIEmbeddingModel, VoyageAIEmbeddingSettings
from pydantic_ai.providers.voyageai import VoyageAIProvider
model = VoyageAIEmbeddingModel('voyage-3.5', provider=VoyageAIProvider(api_key=voyage_api_key))
embedder = Embedder(model)
settings: VoyageAIEmbeddingSettings = {'voyageai_input_type': 'none'}
result = await embedder.embed_query('Hello, world!', settings=settings)
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=1024), length=1),
inputs=['Hello, world!'],
input_type='query',
usage=RequestUsage(input_tokens=4),
model_name='voyage-3.5',
timestamp=IsDatetime(),
provider_name='voyageai',
)
)
@pytest.mark.skipif(not bedrock_imports_successful(), reason='Bedrock not installed')
@pytest.mark.vcr
class TestBedrock:
async def test_infer_model(self):
with patch.dict(
os.environ,
{
'AWS_ACCESS_KEY_ID': 'test-access-key',
'AWS_SECRET_ACCESS_KEY': 'test-secret-key',
'AWS_DEFAULT_REGION': 'us-east-1',
},
):
model = infer_embedding_model('bedrock:amazon.titan-embed-text-v2:0')
assert isinstance(model, BedrockEmbeddingModel)
assert model.model_name == 'amazon.titan-embed-text-v2:0'
assert model.system == 'bedrock'
assert model.base_url == 'https://bedrock-runtime.us-east-1.amazonaws.com'
async def test_titan_v1_minimal(self, bedrock_provider: BedrockProvider):
"""Test Titan V1 with default settings (fixed 1536 dimensions)."""
model = BedrockEmbeddingModel('amazon.titan-embed-text-v1', provider=bedrock_provider)
embedder = Embedder(model)
result = await embedder.embed_query('Hello, world!')
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=1536), length=1),
inputs=['Hello, world!'],
input_type='query',
model_name='amazon.titan-embed-text-v1',
provider_name='bedrock',
timestamp=IsDatetime(),
usage=RequestUsage(input_tokens=4),
)
)
async def test_titan_v2_minimal(self, bedrock_provider: BedrockProvider):
"""Test Titan V2 with default settings (1024 dimensions, normalize=True)."""
model = BedrockEmbeddingModel('amazon.titan-embed-text-v2:0', provider=bedrock_provider)
embedder = Embedder(model)
result = await embedder.embed_query('Hello, world!')
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=1024), length=1),
inputs=['Hello, world!'],
input_type='query',
model_name='amazon.titan-embed-text-v2:0',
provider_name='bedrock',
timestamp=IsDatetime(),
usage=RequestUsage(input_tokens=5),
)
)
async def test_titan_v2_with_dimensions(self, bedrock_provider: BedrockProvider):
"""Test Titan V2 with custom dimensions setting."""
model = BedrockEmbeddingModel('amazon.titan-embed-text-v2:0', provider=bedrock_provider)
embedder = Embedder(model, settings=BedrockEmbeddingSettings(dimensions=256))
result = await embedder.embed_query('Test embedding dimensions')
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=256), length=1),
inputs=['Test embedding dimensions'],
input_type='query',
model_name='amazon.titan-embed-text-v2:0',
provider_name='bedrock',
timestamp=IsDatetime(),
usage=RequestUsage(input_tokens=4),
)
)
async def test_titan_v2_with_normalize_false(self, bedrock_provider: BedrockProvider):
"""Test Titan V2 with normalize=False (override default)."""
model = BedrockEmbeddingModel('amazon.titan-embed-text-v2:0', provider=bedrock_provider)
embedder = Embedder(model, settings=BedrockEmbeddingSettings(bedrock_titan_normalize=False))
result = await embedder.embed_query('Test normalization disabled')
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=1024), length=1),
inputs=['Test normalization disabled'],
input_type='query',
model_name='amazon.titan-embed-text-v2:0',
provider_name='bedrock',
timestamp=IsDatetime(),
usage=RequestUsage(input_tokens=4),
)
)
async def test_titan_v2_with_normalize_true(self, bedrock_provider: BedrockProvider):
"""Test Titan V2 with explicit normalize=True setting."""
model = BedrockEmbeddingModel('amazon.titan-embed-text-v2:0', provider=bedrock_provider)
embedder = Embedder(model, settings=BedrockEmbeddingSettings(bedrock_titan_normalize=True))
result = await embedder.embed_query('Test explicit normalization')
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=1024), length=1),
inputs=['Test explicit normalization'],
input_type='query',
model_name='amazon.titan-embed-text-v2:0',
provider_name='bedrock',
timestamp=IsDatetime(),
usage=RequestUsage(input_tokens=4),
)
)
async def test_titan_v2_multiple_texts(self, bedrock_provider: BedrockProvider):
"""Test Titan V2 document embedding (multiple texts, sequential requests)."""
model = BedrockEmbeddingModel('amazon.titan-embed-text-v2:0', provider=bedrock_provider)
# Use max_concurrency=1 to ensure deterministic request order for VCRpy
embedder = Embedder(model, settings=BedrockEmbeddingSettings(bedrock_max_concurrency=1))
result = await embedder.embed_documents(['hello', 'world'])
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=1024), length=2),
inputs=['hello', 'world'],
input_type='document',
model_name='amazon.titan-embed-text-v2:0',
provider_name='bedrock',
timestamp=IsDatetime(),
usage=RequestUsage(input_tokens=4),
)
)
@pytest.mark.parametrize('max_concurrency', [0, -1])
async def test_titan_v2_rejects_invalid_max_concurrency(
self, bedrock_provider: BedrockProvider, max_concurrency: int
):
model = BedrockEmbeddingModel('amazon.titan-embed-text-v2:0', provider=bedrock_provider)
embedder = Embedder(model, settings=BedrockEmbeddingSettings(bedrock_max_concurrency=max_concurrency))
with (
anyio.fail_after(1),
pytest.raises(UserError, match=f'bedrock_max_concurrency must be >= 1, got {max_concurrency}'),
):
await embedder.embed_query('hello')
async def test_cohere_v3_minimal(self, bedrock_provider: BedrockProvider):
"""Test Cohere V3 with default settings (1024 dimensions, truncate=NONE)."""
model = BedrockEmbeddingModel('cohere.embed-english-v3', provider=bedrock_provider)
embedder = Embedder(model)
result = await embedder.embed_query('Hello, world!')
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=1024), length=1),
inputs=['Hello, world!'],
input_type='query',
model_name='cohere.embed-english-v3',
provider_name='bedrock',
timestamp=IsDatetime(),
usage=RequestUsage(input_tokens=4),
provider_response_id=IsStr(),
)
)
async def test_cohere_v3_with_input_type(self, bedrock_provider: BedrockProvider):
"""Test Cohere V3 with custom input_type setting."""
model = BedrockEmbeddingModel('cohere.embed-english-v3', provider=bedrock_provider)
embedder = Embedder(model, settings=BedrockEmbeddingSettings(bedrock_cohere_input_type='classification'))
result = await embedder.embed_query('Test input type setting')
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=1024), length=1),
inputs=['Test input type setting'],
input_type='query',
model_name='cohere.embed-english-v3',
provider_name='bedrock',
timestamp=IsDatetime(),
usage=RequestUsage(input_tokens=4),
provider_response_id=IsStr(),
)
)
async def test_cohere_v3_with_truncate(self, bedrock_provider: BedrockProvider):
"""Test Cohere V3 with custom truncate setting."""
model = BedrockEmbeddingModel('cohere.embed-english-v3', provider=bedrock_provider)
embedder = Embedder(model, settings=BedrockEmbeddingSettings(bedrock_cohere_truncate='END'))
result = await embedder.embed_query('Test truncation setting')
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=1024), length=1),
inputs=['Test truncation setting'],
input_type='query',
model_name='cohere.embed-english-v3',
provider_name='bedrock',
timestamp=IsDatetime(),
usage=RequestUsage(input_tokens=5),
provider_response_id=IsStr(),
)
)
async def test_cohere_v3_with_base_truncate(self, bedrock_provider: BedrockProvider):
"""Test Cohere V3 with base truncate=True setting (maps to END)."""
model = BedrockEmbeddingModel('cohere.embed-multilingual-v3', provider=bedrock_provider)
embedder = Embedder(model, settings=BedrockEmbeddingSettings(truncate=True))
result = await embedder.embed_query('Test base truncate setting')
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=1024), length=1),
inputs=['Test base truncate setting'],
input_type='query',
model_name='cohere.embed-multilingual-v3',
provider_name='bedrock',
timestamp=IsDatetime(),
usage=RequestUsage(input_tokens=6),
provider_response_id=IsStr(),
)
)
async def test_cohere_v4_minimal(self, bedrock_provider: BedrockProvider):
"""Test Cohere V4 with default settings (1536 dimensions, truncate=NONE)."""
model = BedrockEmbeddingModel('cohere.embed-v4:0', provider=bedrock_provider)
embedder = Embedder(model)
result = await embedder.embed_query('Hello, world!')
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=1536), length=1),
inputs=['Hello, world!'],
input_type='query',
model_name='cohere.embed-v4:0',
provider_name='bedrock',
timestamp=IsDatetime(),
usage=RequestUsage(input_tokens=4),
provider_response_id=IsStr(),
)
)
async def test_cohere_v4_with_dimensions(self, bedrock_provider: BedrockProvider):
"""Test Cohere V4 with custom dimensions setting."""
model = BedrockEmbeddingModel('cohere.embed-v4:0', provider=bedrock_provider)
embedder = Embedder(model, settings=BedrockEmbeddingSettings(dimensions=512))
result = await embedder.embed_query('Test dimensions setting')
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=512), length=1),
inputs=['Test dimensions setting'],
input_type='query',
model_name='cohere.embed-v4:0',
provider_name='bedrock',
timestamp=IsDatetime(),
usage=RequestUsage(input_tokens=3),
provider_response_id=IsStr(),
)
)
async def test_cohere_v4_with_max_tokens(self, bedrock_provider: BedrockProvider):
"""Test Cohere V4 with max_tokens setting."""
model = BedrockEmbeddingModel('cohere.embed-v4:0', provider=bedrock_provider)
embedder = Embedder(model, settings=BedrockEmbeddingSettings(bedrock_cohere_max_tokens=256))
result = await embedder.embed_query('Test max tokens setting')
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=1536), length=1),
inputs=['Test max tokens setting'],
input_type='query',
model_name='cohere.embed-v4:0',
provider_name='bedrock',
timestamp=IsDatetime(),
usage=RequestUsage(input_tokens=4),
provider_response_id=IsStr(),
)
)
async def test_cohere_v4_with_input_type(self, bedrock_provider: BedrockProvider):
"""Test Cohere V4 with custom input_type setting."""
model = BedrockEmbeddingModel('cohere.embed-v4:0', provider=bedrock_provider)
embedder = Embedder(model, settings=BedrockEmbeddingSettings(bedrock_cohere_input_type='clustering'))
result = await embedder.embed_query('Test input type setting')
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=1536), length=1),
inputs=['Test input type setting'],
input_type='query',
model_name='cohere.embed-v4:0',
provider_name='bedrock',
timestamp=IsDatetime(),
usage=RequestUsage(input_tokens=4),
provider_response_id=IsStr(),
)
)
async def test_cohere_v4_with_truncate(self, bedrock_provider: BedrockProvider):
"""Test Cohere V4 with custom truncate setting."""
model = BedrockEmbeddingModel('cohere.embed-v4:0', provider=bedrock_provider)
embedder = Embedder(model, settings=BedrockEmbeddingSettings(bedrock_cohere_truncate='END'))
result = await embedder.embed_query('Test truncation setting')
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=1536), length=1),
inputs=['Test truncation setting'],
input_type='query',
model_name='cohere.embed-v4:0',
provider_name='bedrock',
timestamp=IsDatetime(),
usage=RequestUsage(input_tokens=4),
provider_response_id=IsStr(),
)
)
async def test_cohere_v4_with_truncate_start(self, bedrock_provider: BedrockProvider):
"""Test Cohere V4 with truncate START setting."""
model = BedrockEmbeddingModel('cohere.embed-v4:0', provider=bedrock_provider)
embedder = Embedder(model, settings=BedrockEmbeddingSettings(bedrock_cohere_truncate='START'))
result = await embedder.embed_query('Test truncation start setting')
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=1536), length=1),
inputs=['Test truncation start setting'],
input_type='query',
model_name='cohere.embed-v4:0',
provider_name='bedrock',
timestamp=IsDatetime(),
usage=RequestUsage(input_tokens=5),
provider_response_id=IsStr(),
)
)
async def test_cohere_v4_truncate_priority(self, bedrock_provider: BedrockProvider):
"""Test that bedrock_cohere_truncate takes precedence over base truncate."""
model = BedrockEmbeddingModel('cohere.embed-v4:0', provider=bedrock_provider)
# Both settings provided - model-specific should win (START over END from truncate=True)
embedder = Embedder(model, settings=BedrockEmbeddingSettings(bedrock_cohere_truncate='START', truncate=True))
result = await embedder.embed_query('Test truncate priority')
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=1536), length=1),
inputs=['Test truncate priority'],
input_type='query',
model_name='cohere.embed-v4:0',
provider_name='bedrock',
timestamp=IsDatetime(),
usage=RequestUsage(input_tokens=3),
provider_response_id=IsStr(),
)
)
async def test_cohere_v4_batch_documents(self, bedrock_provider: BedrockProvider):
"""Test Cohere V4 batch embedding (multiple texts in single request)."""
model = BedrockEmbeddingModel('cohere.embed-v4:0', provider=bedrock_provider)
embedder = Embedder(model)
result = await embedder.embed_documents(['hello', 'world'])
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=1536), length=2),
inputs=['hello', 'world'],
input_type='document',
model_name='cohere.embed-v4:0',
provider_name='bedrock',
timestamp=IsDatetime(),
usage=RequestUsage(input_tokens=2),
provider_response_id=IsStr(),
)
)
async def test_nova_minimal(self, bedrock_provider: BedrockProvider):
"""Test Nova with default settings (3072 dimensions, truncate=NONE)."""
model = BedrockEmbeddingModel('amazon.nova-2-multimodal-embeddings-v1:0', provider=bedrock_provider)
embedder = Embedder(model)
result = await embedder.embed_query('Hello, world!')
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=3072), length=1),
inputs=['Hello, world!'],
input_type='query',
model_name='amazon.nova-2-multimodal-embeddings-v1:0',
provider_name='bedrock',
timestamp=IsDatetime(),
usage=RequestUsage(input_tokens=19),
)
)
async def test_nova_with_dimensions(self, bedrock_provider: BedrockProvider):
"""Test Nova with custom dimensions setting."""
model = BedrockEmbeddingModel('amazon.nova-2-multimodal-embeddings-v1:0', provider=bedrock_provider)
embedder = Embedder(model, settings=BedrockEmbeddingSettings(dimensions=256))
result = await embedder.embed_query('Test Nova dimensions')
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=256), length=1),
inputs=['Test Nova dimensions'],
input_type='query',
model_name='amazon.nova-2-multimodal-embeddings-v1:0',
provider_name='bedrock',
timestamp=IsDatetime(),
usage=RequestUsage(input_tokens=18),
)
)
async def test_nova_with_truncate(self, bedrock_provider: BedrockProvider):
"""Test Nova with custom truncate setting."""
model = BedrockEmbeddingModel('amazon.nova-2-multimodal-embeddings-v1:0', provider=bedrock_provider)
embedder = Embedder(
model,
settings=BedrockEmbeddingSettings(bedrock_nova_truncate='END'),
)
result = await embedder.embed_query('Test Nova truncate')
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=3072), length=1),
inputs=['Test Nova truncate'],
input_type='query',
model_name='amazon.nova-2-multimodal-embeddings-v1:0',
provider_name='bedrock',
timestamp=IsDatetime(),
usage=RequestUsage(input_tokens=19),
)
)
async def test_nova_with_truncate_start(self, bedrock_provider: BedrockProvider):
"""Test Nova with truncate START setting."""
model = BedrockEmbeddingModel('amazon.nova-2-multimodal-embeddings-v1:0', provider=bedrock_provider)
embedder = Embedder(
model,
settings=BedrockEmbeddingSettings(bedrock_nova_truncate='START'),
)
result = await embedder.embed_query('Test Nova truncate start')
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=3072), length=1),
inputs=['Test Nova truncate start'],
input_type='query',
model_name='amazon.nova-2-multimodal-embeddings-v1:0',
provider_name='bedrock',
timestamp=IsDatetime(),
usage=RequestUsage(input_tokens=20),
)
)
async def test_nova_with_base_truncate(self, bedrock_provider: BedrockProvider):
"""Test Nova with base truncate=True setting (maps to END)."""
model = BedrockEmbeddingModel('amazon.nova-2-multimodal-embeddings-v1:0', provider=bedrock_provider)
embedder = Embedder(model, settings=BedrockEmbeddingSettings(truncate=True))
result = await embedder.embed_query('Test base truncate')
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=3072), length=1),
inputs=['Test base truncate'],
input_type='query',
model_name='amazon.nova-2-multimodal-embeddings-v1:0',
provider_name='bedrock',
timestamp=IsDatetime(),
usage=RequestUsage(input_tokens=19),
)
)
async def test_nova_with_embedding_purpose(self, bedrock_provider: BedrockProvider):
"""Test Nova with custom embedding_purpose setting."""
model = BedrockEmbeddingModel('amazon.nova-2-multimodal-embeddings-v1:0', provider=bedrock_provider)
embedder = Embedder(
model,
settings=BedrockEmbeddingSettings(bedrock_nova_embedding_purpose='TEXT_RETRIEVAL'),
)
result = await embedder.embed_query('Test Nova settings')
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=3072), length=1),
inputs=['Test Nova settings'],
input_type='query',
model_name='amazon.nova-2-multimodal-embeddings-v1:0',
provider_name='bedrock',
timestamp=IsDatetime(),
usage=RequestUsage(input_tokens=22),
)
)
async def test_nova_multiple_texts(self, bedrock_provider: BedrockProvider):
"""Test Nova document embedding (multiple texts, sequential requests)."""
model = BedrockEmbeddingModel('amazon.nova-2-multimodal-embeddings-v1:0', provider=bedrock_provider)
# Use max_concurrency=1 to ensure deterministic request order for VCRpy
embedder = Embedder(model, settings=BedrockEmbeddingSettings(bedrock_max_concurrency=1))
result = await embedder.embed_documents(['hello', 'world'])
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=3072), length=2),
inputs=['hello', 'world'],
input_type='document',
model_name='amazon.nova-2-multimodal-embeddings-v1:0',
provider_name='bedrock',
timestamp=IsDatetime(),
usage=RequestUsage(input_tokens=6),
)
)
async def test_titan_v1_max_input_tokens(self, bedrock_provider: BedrockProvider):
model = BedrockEmbeddingModel('amazon.titan-embed-text-v1', provider=bedrock_provider)
embedder = Embedder(model)
max_input_tokens = await embedder.max_input_tokens()
assert max_input_tokens == snapshot(8192)
async def test_titan_v2_max_input_tokens(self, bedrock_provider: BedrockProvider):
model = BedrockEmbeddingModel('amazon.titan-embed-text-v2:0', provider=bedrock_provider)
embedder = Embedder(model)
max_input_tokens = await embedder.max_input_tokens()
assert max_input_tokens == snapshot(8192)
async def test_cohere_v3_max_input_tokens(self, bedrock_provider: BedrockProvider):
model = BedrockEmbeddingModel('cohere.embed-english-v3', provider=bedrock_provider)
embedder = Embedder(model)
max_input_tokens = await embedder.max_input_tokens()
assert max_input_tokens == snapshot(512)
async def test_cohere_v4_max_input_tokens(self, bedrock_provider: BedrockProvider):
model = BedrockEmbeddingModel('cohere.embed-v4:0', provider=bedrock_provider)
embedder = Embedder(model)
max_input_tokens = await embedder.max_input_tokens()
assert max_input_tokens == snapshot(128000)
async def test_nova_max_input_tokens(self, bedrock_provider: BedrockProvider):
model = BedrockEmbeddingModel('amazon.nova-2-multimodal-embeddings-v1:0', provider=bedrock_provider)
embedder = Embedder(model)
max_input_tokens = await embedder.max_input_tokens()
assert max_input_tokens == snapshot(8192)
async def test_base_url_property(self, bedrock_provider: BedrockProvider):
"""Test that base_url property returns the endpoint URL."""
model = BedrockEmbeddingModel('amazon.titan-embed-text-v2:0', provider=bedrock_provider)
assert model.base_url is not None
assert isinstance(model.base_url, str)
async def test_regional_prefix_model_name(self, bedrock_provider: BedrockProvider):
"""Test model with regional prefix (e.g., us.amazon.titan-embed-text-v2:0) is handled correctly."""
model = BedrockEmbeddingModel('us.amazon.titan-embed-text-v2:0', provider=bedrock_provider)
# Model name preserves the regional prefix
assert model.model_name == 'us.amazon.titan-embed-text-v2:0'
# But handler uses normalized name (without prefix)
assert model._handler.model_name == 'amazon.titan-embed-text-v2:0' # pyright: ignore[reportPrivateUsage]
# max_input_tokens() works correctly with regional prefix
max_tokens = await model.max_input_tokens()
assert max_tokens == snapshot(8192)
async def test_regional_prefix_embed(self, bedrock_provider: BedrockProvider):
"""Test embedding with a regional prefix model ID using Cohere v4.
Cross-region inference profiles are supported for Cohere models on Bedrock.
"""
model = BedrockEmbeddingModel('us.cohere.embed-v4:0', provider=bedrock_provider)
embedder = Embedder(model)
result = await embedder.embed_query('Hello from regional endpoint!')
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=1536), length=1),
inputs=['Hello from regional endpoint!'],
input_type='query',
model_name='us.cohere.embed-v4:0',
provider_name='bedrock',
timestamp=IsDatetime(),
usage=RequestUsage(input_tokens=5),
provider_response_id=IsStr(),
)
)
async def test_inference_profile_embed(self, bedrock_provider: BedrockProvider):
# When re-recording, set AWS_ACCOUNT_ID to your real account ID
account_id = os.getenv('AWS_ACCOUNT_ID', '123456789012')
inference_profile_arn = f'arn:aws:bedrock:us-east-1:{account_id}:application-inference-profile/otnfa2ysixqd'
settings: BedrockEmbeddingSettings = {'bedrock_inference_profile': inference_profile_arn}
model = BedrockEmbeddingModel('amazon.titan-embed-text-v2:0', provider=bedrock_provider, settings=settings)
result = await model.embed('Hello, world!', input_type='document')
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=1024), length=1),
inputs=['Hello, world!'],
input_type='document',
model_name='amazon.titan-embed-text-v2:0',
provider_name='bedrock',
timestamp=IsDatetime(),
usage=RequestUsage(input_tokens=5),
)
)
async def test_unsupported_model_error(self, bedrock_provider: BedrockProvider):
with pytest.raises(UserError, match='Unsupported Bedrock embedding model'):
BedrockEmbeddingModel('unsupported.model', provider=bedrock_provider)
async def test_unknown_model_max_tokens_returns_none(self, bedrock_provider: BedrockProvider):
"""Test that unknown models with valid prefixes return None for max_input_tokens."""
model = BedrockEmbeddingModel('amazon.titan-embed-text-v99:0', provider=bedrock_provider)
assert await model.max_input_tokens() is None
def test_model_with_string_provider(self, bedrock_provider: BedrockProvider):
"""Test BedrockEmbeddingModel can be created with string provider."""
with patch('pydantic_ai.embeddings.bedrock.infer_provider', return_value=bedrock_provider) as mock_infer:
model = BedrockEmbeddingModel('amazon.titan-embed-text-v2:0', provider='bedrock')
mock_infer.assert_called_once_with('bedrock')
assert model.model_name == 'amazon.titan-embed-text-v2:0'
async def test_client_error_with_status_code(self, bedrock_provider: BedrockProvider):
"""Test error handling when ClientError is raised with HTTP status code.
ResponseMetadata.HTTPHeaders is the nested dict that BedrockEmbeddingModel extracts
via metadata.get('HTTPHeaders') — verify it reaches ModelHTTPError.headers unchanged.
"""
model = BedrockEmbeddingModel('amazon.titan-embed-text-v2:0', provider=bedrock_provider)
error_response = {
'Error': {'Code': 'ValidationException', 'Message': 'Invalid input'},
'ResponseMetadata': {
'HTTPStatusCode': 400,
'HTTPHeaders': {'retry-after': '5', 'x-amzn-requestid': 'req-abc'},
},
}
with patch.object(
model.client,
'invoke_model',
side_effect=ClientError(error_response, 'InvokeModel'), # pyright: ignore[reportArgumentType]
):
with pytest.raises(ModelHTTPError) as exc_info:
await model.embed(['test'], input_type='query')
exc = exc_info.value
assert exc.status_code == 400
assert exc.headers is not None
assert exc.headers.get('retry-after') == '5'
assert exc.headers.get('x-amzn-requestid') == 'req-abc'
async def test_client_error_without_status_code(self, bedrock_provider: BedrockProvider):
"""Test error handling when ClientError is raised without HTTP status code."""
model = BedrockEmbeddingModel('amazon.titan-embed-text-v2:0', provider=bedrock_provider)
error_response = {
'Error': {'Code': 'UnknownError', 'Message': 'Something went wrong'},
'ResponseMetadata': {}, # No HTTPStatusCode
}
with patch.object(
model.client,
'invoke_model',
side_effect=ClientError(error_response, 'InvokeModel'), # pyright: ignore[reportArgumentType]
):
with pytest.raises(ModelAPIError) as exc_info:
await model.embed(['test'], input_type='query')
assert type(exc_info.value) is ModelAPIError
async def test_client_error_on_every_concurrent_request(self, bedrock_provider: BedrockProvider):
"""A throttle that fails every concurrent request raises one `ModelHTTPError`, not an `ExceptionGroup`."""
model = BedrockEmbeddingModel('amazon.titan-embed-text-v2:0', provider=bedrock_provider)
error_response = {
'Error': {'Code': 'ThrottlingException', 'Message': 'Too many requests'},
'ResponseMetadata': {'HTTPStatusCode': 429},
}
# Every request waits for the others before failing, so all three fail rather than being cancelled.
all_sent = threading.Barrier(3, timeout=5)
def invoke_model(**kwargs: Any) -> Any:
all_sent.wait()
raise ClientError(error_response, 'InvokeModel') # pyright: ignore[reportArgumentType]
with patch.object(model.client, 'invoke_model', side_effect=invoke_model):
with pytest.raises(ModelHTTPError) as exc_info:
await model.embed(['a', 'b', 'c'], input_type='document')
assert exc_info.value.status_code == 429
assert isinstance(exc_info.value.__cause__, ClientError)
group = exc_info.value.__context__
assert isinstance(group, BaseExceptionGroup)
assert exc_info.value in group.exceptions # pyright: ignore[reportUnknownMemberType]
@pytest.mark.parametrize('model_errors', [0, 2], ids=['no-model-error', 'model-error-among-others'])
async def test_mixed_errors_on_concurrent_requests(self, bedrock_provider: BedrockProvider, model_errors: int):
"""A model error is raised whenever one of the failed requests produced one, whatever order they failed in;
when none did, the group of the other errors propagates."""
model = BedrockEmbeddingModel('amazon.titan-embed-text-v2:0', provider=bedrock_provider)
error_response = {
'Error': {'Code': 'ThrottlingException', 'Message': 'Too many requests'},
'ResponseMetadata': {'HTTPStatusCode': 429},
}
all_sent = threading.Barrier(3, timeout=5)
calls = iter(range(3))
lock = threading.Lock()
def invoke_model(**kwargs: Any) -> Any:
with lock:
call = next(calls)
all_sent.wait()
if call < model_errors:
raise ClientError(error_response, 'InvokeModel') # pyright: ignore[reportArgumentType]
raise ValueError('unexpected response body')
with patch.object(model.client, 'invoke_model', side_effect=invoke_model):
if model_errors:
with pytest.raises(ModelHTTPError):
await model.embed(['a', 'b', 'c'], input_type='document')
else:
with pytest.raises(BaseExceptionGroup) as group_info:
await model.embed(['a', 'b', 'c'], input_type='document')
exceptions: tuple[BaseException, ...] = group_info.value.exceptions
assert all(isinstance(e, ValueError) for e in exceptions)
@pytest.mark.parametrize('error_type', ['read-timeout', 'endpoint-connection'])
async def test_transport_error(self, bedrock_provider: BedrockProvider, error_type: str):
"""botocore raises transport failures as `BotoCoreError`, which maps to `ModelAPIError` like `BedrockConverseModel`."""
model = BedrockEmbeddingModel('amazon.titan-embed-text-v2:0', provider=bedrock_provider)
error_cls = ReadTimeoutError if error_type == 'read-timeout' else EndpointConnectionError
error = error_cls(endpoint_url='https://bedrock.example')
with patch.object(model.client, 'invoke_model', side_effect=error):
with pytest.raises(ModelAPIError) as exc_info:
await model.embed(['test'], input_type='query')
assert exc_info.value.__cause__ is error
async def test_count_tokens_not_implemented(self, bedrock_provider: BedrockProvider):
"""Test that count_tokens raises NotImplementedError (Bedrock doesn't support it)."""
model = BedrockEmbeddingModel('amazon.titan-embed-text-v2:0', provider=bedrock_provider)
embedder = Embedder(model)
with pytest.raises(NotImplementedError):
await embedder.count_tokens('Hello, world!')
@pytest.mark.skipif(not logfire_imports_successful(), reason='logfire not installed')
async def test_instrumentation(self, bedrock_provider: BedrockProvider, capfire: CaptureLogfire):
model = BedrockEmbeddingModel('amazon.titan-embed-text-v2:0', provider=bedrock_provider)
embedder = Embedder(model, instrument=True)
await embedder.embed_query('Hello, world!', settings={'dimensions': 256})
spans = capfire.exporter.exported_spans_as_dict(parse_json_attributes=True)
span = next(span for span in spans if 'embeddings' in span['name'])
assert span == snapshot(
{
'name': 'embeddings amazon.titan-embed-text-v2:0',
'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'parent': None,
'start_time': IsInt(),
'end_time': IsInt(),
'attributes': {
'gen_ai.operation.name': 'embeddings',
'gen_ai.provider.name': 'bedrock',
'gen_ai.request.model': 'amazon.titan-embed-text-v2:0',
'input_type': 'query',
'server.address': 'bedrock-runtime.us-east-1.amazonaws.com',
'inputs_count': 1,
'embedding_settings': {'dimensions': 256},
'inputs': ['Hello, world!'],
'logfire.json_schema': {
'type': 'object',
'properties': {
'input_type': {'type': 'string'},
'inputs_count': {'type': 'integer'},
'embedding_settings': {'type': 'object'},
'inputs': {'type': ['array']},
'embeddings': {'type': 'array'},
},
},
'logfire.span_type': 'span',
'logfire.msg': 'embeddings amazon.titan-embed-text-v2:0',
'gen_ai.usage.input_tokens': 5,
'gen_ai.response.model': 'amazon.titan-embed-text-v2:0',
'operation.cost': 5e-07,
'gen_ai.embeddings.dimension.count': 256,
},
}
)
@dataclass
class _GoogleTaskPrefixCase:
id: str
model_name: str
input_type: Literal['query', 'document']
inputs: list[str]
settings: GoogleEmbeddingSettings
expected_texts: list[str]
expected_task_type: str | None
expected_warning: str | None = None
# The `GoogleEmbeddingSettings(...)` calls are only evaluated when the `google` extra is installed.
_GOOGLE_TASK_PREFIX_CASES: list[_GoogleTaskPrefixCase] = (
[
_GoogleTaskPrefixCase(
id='default-query',
model_name='gemini-embedding-2',
input_type='query',
inputs=['Hello, world!'],
settings=GoogleEmbeddingSettings(),
expected_texts=['task: search result | query: Hello, world!'],
expected_task_type=None,
),
_GoogleTaskPrefixCase(
id='asymmetric-query',
model_name='gemini-embedding-2',
input_type='query',
inputs=['Hello, world!'],
settings=GoogleEmbeddingSettings(google_task='question answering'),
expected_texts=['task: question answering | query: Hello, world!'],
expected_task_type=None,
),
_GoogleTaskPrefixCase(
id='asymmetric-document-with-title',
model_name='gemini-embedding-2',
input_type='document',
inputs=['hello'],
settings=GoogleEmbeddingSettings(google_task='search result', google_title='Greeting'),
expected_texts=['title: Greeting | text: hello'],
expected_task_type=None,
),
_GoogleTaskPrefixCase(
id='asymmetric-document-no-title',
model_name='gemini-embedding-2',
input_type='document',
inputs=['hello', 'world'],
settings=GoogleEmbeddingSettings(),
expected_texts=['title: none | text: hello', 'title: none | text: world'],
expected_task_type=None,
),
_GoogleTaskPrefixCase(
id='symmetric-query',
model_name='gemini-embedding-2',
input_type='query',
inputs=['hello'],
settings=GoogleEmbeddingSettings(google_task='classification'),
expected_texts=['task: classification | query: hello'],
expected_task_type=None,
),
_GoogleTaskPrefixCase(
id='symmetric-document-ignores-title',
model_name='gemini-embedding-2',
input_type='document',
inputs=['hello'],
settings=GoogleEmbeddingSettings(google_task='clustering', google_title='ignored'),
expected_texts=['task: clustering | query: hello'],
expected_task_type=None,
),
_GoogleTaskPrefixCase(
id='symmetric-sentence-similarity-ignores-title',
model_name='gemini-embedding-2',
input_type='document',
inputs=['hello'],
settings=GoogleEmbeddingSettings(google_task='sentence similarity', google_title='ignored'),
expected_texts=['task: sentence similarity | query: hello'],
expected_task_type=None,
),
_GoogleTaskPrefixCase(
id='asymmetric-document-empty-title',
model_name='gemini-embedding-2',
input_type='document',
inputs=['hello'],
settings=GoogleEmbeddingSettings(google_title=''),
expected_texts=['title: none | text: hello'],
expected_task_type=None,
),
_GoogleTaskPrefixCase(
id='raw-passthrough',
model_name='gemini-embedding-2',
input_type='document',
inputs=['title: custom | text: hello'],
settings=GoogleEmbeddingSettings(google_task='raw'),
expected_texts=['title: custom | text: hello'],
expected_task_type=None,
),
_GoogleTaskPrefixCase(
id='task-type-ignored-on-embedding-2',
model_name='gemini-embedding-2',
input_type='query',
inputs=['hello'],
settings=GoogleEmbeddingSettings(google_task='classification', google_task_type='RETRIEVAL_QUERY'),
expected_texts=['task: classification | query: hello'],
expected_task_type=None,
expected_warning='`google_task_type` is not supported by `gemini-embedding-2`',
),
_GoogleTaskPrefixCase(
id='task-ignored-on-other-model',
model_name='gemini-embedding-2-preview',
input_type='query',
inputs=['hello'],
settings=GoogleEmbeddingSettings(google_task='classification'),
expected_texts=['hello'],
expected_task_type='RETRIEVAL_QUERY',
expected_warning='`google_task` is only supported by `gemini-embedding-2`',
),
]
if google_imports_successful()
else []
)
@pytest.mark.skipif(not google_imports_successful(), reason='Google not installed')
@pytest.mark.vcr
@pytest.mark.parametrize('case', [pytest.param(c, id=c.id) for c in _GOOGLE_TASK_PREFIX_CASES])
async def test_google_task_prefix(case: _GoogleTaskPrefixCase, gemini_api_key: str, monkeypatch: pytest.MonkeyPatch):
"""`google_task` builds the right text prefix (and `task_type`) for `gemini-embedding-2`.
Spies on `embed_content` to assert the exact text sent to the API for each
(task, input_type, title) combination, plus the warn-and-ignore behavior when
`google_task`/`google_task_type` are used on the wrong model.
"""
provider = GoogleProvider(api_key=gemini_api_key)
model = GoogleEmbeddingModel(case.model_name, provider=provider)
embedder = Embedder(model)
captured: dict[str, Any] = {}
real_embed_content = provider.client.aio.models.embed_content
async def spy(**kwargs: Any) -> Any:
captured['contents'] = kwargs['contents']
captured['config'] = kwargs['config']
return await real_embed_content(**kwargs)
monkeypatch.setattr(provider.client.aio.models, 'embed_content', spy)
async def run() -> EmbeddingResult:
if case.input_type == 'query':
return await embedder.embed_query(case.inputs, settings=case.settings)
return await embedder.embed_documents(case.inputs, settings=case.settings)
if case.expected_warning is not None:
with pytest.warns(UserWarning, match=case.expected_warning):
result = await run()
else:
result = await run()
sent_texts = [part.text for content in captured['contents'] for part in content.parts]
assert sent_texts == case.expected_texts
assert captured['config'].task_type == case.expected_task_type
assert captured['config'].title is None
assert len(result.embeddings) == len(case.inputs)
# The prefix is applied internally; the user gets their original (non-prefixed) text back.
assert result.inputs == case.inputs
@pytest.mark.skipif(not google_imports_successful(), reason='Google not installed')
@pytest.mark.skipif(
not os.getenv('CI', False), reason='Requires properly configured local google vertex config to pass'
)
@pytest.mark.vcr
async def test_google_task_prefix_vertex(
allow_model_requests: None, vertex_provider: GoogleCloudProvider, monkeypatch: pytest.MonkeyPatch
): # pragma: lax no cover
"""`google_task` builds the same `gemini-embedding-2` prefix against Google Cloud (Vertex) as against the Gemini API."""
model = GoogleEmbeddingModel('gemini-embedding-2', provider=vertex_provider)
embedder = Embedder(model)
captured: dict[str, Any] = {}
real_embed_content = vertex_provider.client.aio.models.embed_content
async def spy(**kwargs: Any) -> Any:
captured['contents'] = kwargs['contents']
captured['config'] = kwargs['config']
return await real_embed_content(**kwargs)
monkeypatch.setattr(vertex_provider.client.aio.models, 'embed_content', spy)
result = await embedder.embed_query(
'Hello, world!', settings=GoogleEmbeddingSettings(google_task='question answering')
)
sent_texts = [part.text for content in captured['contents'] for part in content.parts]
assert sent_texts == ['task: question answering | query: Hello, world!']
assert captured['config'].task_type is None
assert captured['config'].title is None
assert len(result.embeddings) == 1
@pytest.mark.skipif(not google_imports_successful(), reason='Google not installed')
@pytest.mark.vcr
class TestGoogle:
@pytest.fixture
def embedder(self, gemini_api_key: str) -> Embedder:
return Embedder(
GoogleEmbeddingModel('gemini-embedding-2-preview', provider=GoogleProvider(api_key=gemini_api_key))
)
async def test_infer_model_google(self, gemini_api_key: str):
with patch.dict(os.environ, {'GOOGLE_API_KEY': gemini_api_key}):
model = infer_embedding_model('google:gemini-embedding-001')
assert isinstance(model, GoogleEmbeddingModel)
assert model.model_name == 'gemini-embedding-001'
assert model.system == 'google'
assert urlparse(model.base_url).hostname == 'generativelanguage.googleapis.com'
async def test_infer_model_google_cloud(self, env: TestEnv):
for name in {
'GOOGLE_APPLICATION_CREDENTIALS',
'GOOGLE_CLOUD_PROJECT',
'GOOGLE_CLOUD_LOCATION',
'GEMINI_API_KEY',
}:
env.remove(name)
env.set('GOOGLE_API_KEY', 'mock-api-key')
model = infer_embedding_model('google-cloud:gemini-embedding-001')
assert isinstance(model, GoogleEmbeddingModel)
assert model.model_name == 'gemini-embedding-001'
assert model.system == 'google-cloud'
async def test_model_with_string_provider(self, gemini_api_key: str):
with patch.dict(os.environ, {'GOOGLE_API_KEY': gemini_api_key}):
model = GoogleEmbeddingModel('gemini-embedding-001', provider='google')
assert isinstance(model, GoogleEmbeddingModel)
assert model.model_name == 'gemini-embedding-001'
assert model.system == 'google'
async def test_query(self, embedder: Embedder):
result = await embedder.embed_query('Hello, world!')
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=3072), length=1),
inputs=['Hello, world!'],
input_type='query',
usage=RequestUsage(),
model_name='gemini-embedding-2-preview',
timestamp=IsDatetime(),
provider_name='google',
)
)
async def test_documents(self, embedder: Embedder):
result = await embedder.embed_documents(['hello', 'world'])
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=3072), length=2),
inputs=['hello', 'world'],
input_type='document',
usage=RequestUsage(),
model_name='gemini-embedding-2-preview',
timestamp=IsDatetime(),
provider_name='google',
)
)
async def test_query_with_dimensions(self, embedder: Embedder):
result = await embedder.embed_query('Hello, world!', settings={'dimensions': 768})
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=768), length=1),
inputs=['Hello, world!'],
input_type='query',
usage=RequestUsage(),
model_name='gemini-embedding-2-preview',
timestamp=IsDatetime(),
provider_name='google',
)
)
async def test_max_input_tokens(self, embedder: Embedder):
max_input_tokens = await embedder.max_input_tokens()
assert max_input_tokens == snapshot(8192)
async def test_count_tokens(self, embedder: Embedder):
count = await embedder.count_tokens('Hello, world!')
assert count == snapshot(5)
async def test_embed_error(self, gemini_api_key: str):
model = GoogleEmbeddingModel('nonexistent-model', provider=GoogleProvider(api_key=gemini_api_key))
embedder = Embedder(model)
with pytest.raises(ModelHTTPError, match='not found'):
await embedder.embed_query('Hello, world!')
async def test_embed_connect_error(self):
"""`google.genai` doesn't wrap transport failures; a cassette can't replay one, so a mock transport raises it."""
def handler(request: httpx2.Request) -> httpx2.Response:
raise httpx2.ConnectError('connection refused', request=request)
async with httpx2.AsyncClient(transport=httpx2.MockTransport(handler)) as http_client:
provider = GoogleProvider(api_key='test-key', http_client=http_client, base_url='http://localhost')
embedder = Embedder(GoogleEmbeddingModel('gemini-embedding-001', provider=provider))
with pytest.raises(ModelAPIError) as exc_info:
await embedder.embed_query('Hello, world!')
assert type(exc_info.value) is ModelAPIError
assert exc_info.value.message == 'connection refused'
assert isinstance(exc_info.value.__cause__, httpx2.ConnectError)
async def test_count_tokens_error(self, gemini_api_key: str):
model = GoogleEmbeddingModel('nonexistent-model', provider=GoogleProvider(api_key=gemini_api_key))
embedder = Embedder(model)
with pytest.raises(ModelHTTPError, match='not found'):
await embedder.count_tokens('Hello, world!')
async def test_embed_error_no_http_response(self, gemini_api_key: str, mocker: MockerFixture):
"""An APIError with response=None (no HTTP response object) yields headers=None on ModelHTTPError.
This exercises the defensive `e.response is not None else None` branch in the embed
path of GoogleEmbeddingModel — the branch that handles non-HTTP errors where the SDK
raises an APIError without attaching an httpx.Response.
"""
from google.genai import errors
model = GoogleEmbeddingModel('gemini-embedding-2-preview', provider=GoogleProvider(api_key=gemini_api_key))
error_without_response = errors.APIError(503, {'error': {'code': 503, 'message': 'Unavailable'}})
mocker.patch.object(model._client.aio.models, 'embed_content', side_effect=error_without_response) # pyright: ignore[reportPrivateUsage]
with pytest.raises(ModelHTTPError) as exc_info:
await model.embed(['test'], input_type='query')
assert exc_info.value.status_code == 503
assert exc_info.value.headers is None
async def test_count_tokens_error_no_http_response(self, gemini_api_key: str, mocker: MockerFixture):
"""Same as above for the count_tokens path."""
from google.genai import errors
model = GoogleEmbeddingModel('gemini-embedding-2-preview', provider=GoogleProvider(api_key=gemini_api_key))
error_without_response = errors.APIError(503, {'error': {'code': 503, 'message': 'Unavailable'}})
mocker.patch.object(model._client.aio.models, 'count_tokens', side_effect=error_without_response) # pyright: ignore[reportPrivateUsage]
with pytest.raises(ModelHTTPError) as exc_info:
await model.count_tokens('test')
assert exc_info.value.status_code == 503
assert exc_info.value.headers is None
async def test_embed_error_with_http_response(self, gemini_api_key: str, mocker: MockerFixture):
"""An APIError with a real httpx.Response propagates its headers to ModelHTTPError.
The positive path `headers=dict(e.response.headers) if e.response is not None else None`
converts httpx.Headers to a plain lowercased dict. Verify the value reaches
ModelHTTPError.headers so a wrong-attribute regression (e.g. swapping response for None)
would be caught.
"""
import httpx
from google.genai import errors
model = GoogleEmbeddingModel('gemini-embedding-2-preview', provider=GoogleProvider(api_key=gemini_api_key))
req = httpx.Request('POST', 'https://generativelanguage.googleapis.com/v1beta/models')
resp = httpx.Response(429, headers={'retry-after': '10', 'x-goog-request-id': 'rid-1'}, request=req)
error_with_response = errors.APIError(429, {'error': {'code': 429, 'message': 'Rate limited'}})
error_with_response.response = resp
mocker.patch.object(model._client.aio.models, 'embed_content', side_effect=error_with_response) # pyright: ignore[reportPrivateUsage]
with pytest.raises(ModelHTTPError) as exc_info:
await model.embed(['test'], input_type='query')
exc = exc_info.value
assert exc.status_code == 429
assert exc.headers is not None
assert exc.headers.get('retry-after') == '10'
assert exc.headers.get('x-goog-request-id') == 'rid-1'
async def test_embed_error_low_status_code(self, gemini_api_key: str, mocker: MockerFixture):
"""An APIError with code < 400 is re-raised verbatim, not wrapped in ModelHTTPError.
GoogleEmbeddingModel only wraps errors with status_code >= 400. A code below
400 is a non-HTTP-error signal from the SDK; the original exception propagates.
This covers the `raise` (else) branch of `if (status_code := e.code) >= 400`.
"""
from google.genai import errors
model = GoogleEmbeddingModel('gemini-embedding-2-preview', provider=GoogleProvider(api_key=gemini_api_key))
low_code_error = errors.APIError(0, {'error': {'code': 0, 'message': 'Unknown'}})
mocker.patch.object(model._client.aio.models, 'embed_content', side_effect=low_code_error) # pyright: ignore[reportPrivateUsage]
with pytest.raises(errors.APIError) as exc_info:
await model.embed(['test'], input_type='query')
assert exc_info.value is low_code_error
async def test_count_tokens_error_low_status_code(self, gemini_api_key: str, mocker: MockerFixture):
"""Same as test_embed_error_low_status_code for the count_tokens path."""
from google.genai import errors
model = GoogleEmbeddingModel('gemini-embedding-2-preview', provider=GoogleProvider(api_key=gemini_api_key))
low_code_error = errors.APIError(0, {'error': {'code': 0, 'message': 'Unknown'}})
mocker.patch.object(model._client.aio.models, 'count_tokens', side_effect=low_code_error) # pyright: ignore[reportPrivateUsage]
with pytest.raises(errors.APIError) as exc_info:
await model.count_tokens('test')
assert exc_info.value is low_code_error
async def test_query_with_task_type(self, embedder: Embedder):
result = await embedder.embed_query(
'Hello, world!', settings=GoogleEmbeddingSettings(google_task_type='RETRIEVAL_QUERY')
)
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=3072), length=1),
inputs=['Hello, world!'],
input_type='query',
usage=RequestUsage(),
model_name='gemini-embedding-2-preview',
timestamp=IsDatetime(),
provider_name='google',
)
)
@pytest.mark.skipif(
not os.getenv('CI', False), reason='Requires properly configured local google vertex config to pass'
)
@pytest.mark.vcr()
async def test_vertex_query(
self, allow_model_requests: None, vertex_provider: GoogleProvider
): # pragma: lax no cover
model = GoogleEmbeddingModel('gemini-embedding-001', provider=vertex_provider)
embedder = Embedder(model)
result = await embedder.embed_query('Hello, world!')
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=3072), length=1),
inputs=['Hello, world!'],
input_type='query',
usage=RequestUsage(input_tokens=4),
model_name='gemini-embedding-001',
timestamp=IsDatetime(),
provider_name='google-cloud',
)
)
@pytest.mark.skipif(not logfire_imports_successful(), reason='logfire not installed')
async def test_instrumentation(self, gemini_api_key: str, capfire: CaptureLogfire):
model = GoogleEmbeddingModel('gemini-embedding-2-preview', provider=GoogleProvider(api_key=gemini_api_key))
embedder = Embedder(model, instrument=True)
await embedder.embed_query('Hello, world!', settings={'dimensions': 768})
spans = capfire.exporter.exported_spans_as_dict(parse_json_attributes=True)
span = next(span for span in spans if 'embeddings' in span['name'])
assert span == snapshot(
{
'name': 'embeddings gemini-embedding-2-preview',
'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'parent': None,
'start_time': IsInt(),
'end_time': IsInt(),
'attributes': {
'gen_ai.operation.name': 'embeddings',
'gen_ai.provider.name': 'google',
'gen_ai.request.model': 'gemini-embedding-2-preview',
'input_type': 'query',
'server.address': 'generativelanguage.googleapis.com',
'inputs_count': 1,
'embedding_settings': {'dimensions': 768},
'inputs': ['Hello, world!'],
'logfire.json_schema': {
'type': 'object',
'properties': {
'input_type': {'type': 'string'},
'inputs_count': {'type': 'integer'},
'embedding_settings': {'type': 'object'},
'inputs': {'type': ['array']},
'embeddings': {'type': 'array'},
},
},
'logfire.span_type': 'span',
'logfire.msg': 'embeddings gemini-embedding-2-preview',
'gen_ai.response.model': 'gemini-embedding-2-preview',
'gen_ai.embeddings.dimension.count': 768,
},
}
)
@pytest.mark.skipif(not sentence_transformers_installed, reason='SentenceTransformers not installed')
@pytest.mark.usefixtures('sentence_transformers_imported')
class TestSentenceTransformers:
def _load_stsb_bert_tiny_model(self):
from sentence_transformers import SentenceTransformer
# The pinned commit revision lets huggingface_hub serve every model file
# straight from a warm cache without revalidating it against the Hub.
# Construction still fires a few metadata requests (model card, repo tree,
# adapter probe); those tolerate Hub HTTP errors, and connection-level
# failures fall through to the skip below. CI warms the cache out-of-band
# (see ci.yml); a cold cache downloads the model here on first use.
try:
model = SentenceTransformer(STSB_BERT_TINY_MODEL, revision=STSB_BERT_TINY_REVISION)
except (OSError, RuntimeError, httpx.HTTPError) as e:
# Skip only when the Hub is unavailable (see `_hf_hub_unavailable`) so a
# HF outage never reds the whole suite; anything else (bad pin, auth
# problem, corrupt cache, dependency mismatch) fails loudly.
if not _hf_hub_unavailable(e):
raise
pytest.skip(f'sentence-transformers test model unavailable (HF Hub): {e}')
# If the model card didn't load, `model_id` is unset and the model reports
# its name as 'unknown'. Set it explicitly so the reported name is
# deterministic (a no-op when the card did load).
model.model_card_data.model_id = STSB_BERT_TINY_MODEL
model.model_card_data.generate_widget_examples = False # Disable widget examples generation for testing
return model
@pytest.fixture(scope='session')
def stsb_bert_tiny_model(self):
return self._load_stsb_bert_tiny_model()
def test_model_unavailable(self, monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(
'sentence_transformers.SentenceTransformer', MagicMock(side_effect=httpx.ConnectTimeout('offline'))
)
skip = MagicMock(side_effect=RuntimeError('skipped'))
monkeypatch.setattr(pytest, 'skip', skip)
with pytest.raises(RuntimeError, match='skipped'):
self._load_stsb_bert_tiny_model()
skip.assert_called_once_with('sentence-transformers test model unavailable (HF Hub): offline')
@pytest.fixture
def embedder(self, stsb_bert_tiny_model: Any) -> Embedder:
from pydantic_ai.embeddings.sentence_transformers import SentenceTransformerEmbeddingModel
return Embedder(SentenceTransformerEmbeddingModel(stsb_bert_tiny_model))
async def test_embed_is_exempt_from_request_guard(self, embedder: Embedder):
"""`ALLOW_MODEL_REQUESTS`'s docstring promises this model is unaffected; pin that promise.
Inference is local, so there's no provider call to block and no recording to make.
"""
with pydantic_ai.models.override_allow_model_requests(False):
result = await embedder.embed_query('hello')
assert result.embeddings
async def test_infer_model(self):
from pydantic_ai.embeddings.sentence_transformers import SentenceTransformerEmbeddingModel
model = infer_embedding_model('sentence-transformers:all-MiniLM-L6-v2')
assert isinstance(model, SentenceTransformerEmbeddingModel)
assert model.model_name == 'all-MiniLM-L6-v2'
assert model.system == 'sentence-transformers'
assert model.base_url is None
async def test_adapter_without_downloaded_model(self, monkeypatch: pytest.MonkeyPatch):
"""VCR cannot exercise a local adapter, so cover it without a downloaded model."""
import torch
from sentence_transformers import SentenceTransformer
import pydantic_ai.embeddings.sentence_transformers as sentence_transformers_module
from pydantic_ai.embeddings.sentence_transformers import (
SentenceTransformerEmbeddingModel,
SentenceTransformersEmbeddingSettings,
)
st_model = MagicMock(spec=SentenceTransformer)
st_model.model_card_data = MagicMock()
st_model.model_card_data.model_id = 'test-model'
st_model.encode_query.return_value.tolist.return_value = [[0.1, 0.2]]
st_model.encode_document.return_value.tolist.return_value = [[0.3, 0.4], [0.5, 0.6]]
st_model.get_max_seq_length.return_value = 512
st_model.tokenize.return_value = {'input_ids': torch.tensor([[1, 2, 3]])}
def preserve_model(model: SentenceTransformer) -> SentenceTransformer:
return model
monkeypatch.setattr(sentence_transformers_module, 'deepcopy', preserve_model)
model = SentenceTransformerEmbeddingModel(
st_model,
settings=SentenceTransformersEmbeddingSettings(
dimensions=2,
sentence_transformers_batch_size=2,
sentence_transformers_device='cpu',
sentence_transformers_normalize_embeddings=True,
),
)
embedder = Embedder(model)
query_result = await embedder.embed_query('hello')
assert query_result == snapshot(
EmbeddingResult(
embeddings=[[0.1, 0.2]],
inputs=['hello'],
input_type='query',
model_name='test-model',
provider_name='sentence-transformers',
timestamp=IsDatetime(),
)
)
st_model.encode_query.assert_called_once_with(
['hello'],
show_progress_bar=False,
convert_to_numpy=True,
convert_to_tensor=False,
device='cpu',
normalize_embeddings=True,
truncate_dim=2,
batch_size=2,
)
documents_result = await embedder.embed_documents(['hello', 'world'])
assert documents_result == snapshot(
EmbeddingResult(
embeddings=[[0.3, 0.4], [0.5, 0.6]],
inputs=['hello', 'world'],
input_type='document',
model_name='test-model',
provider_name='sentence-transformers',
timestamp=IsDatetime(),
)
)
st_model.encode_document.assert_called_once_with(
['hello', 'world'],
show_progress_bar=False,
convert_to_numpy=True,
convert_to_tensor=False,
device='cpu',
normalize_embeddings=True,
truncate_dim=2,
batch_size=2,
)
assert await embedder.max_input_tokens() == 512
assert await embedder.count_tokens('hello') == 3
loaded_model = MagicMock(spec=SentenceTransformer)
loaded_model.get_max_seq_length.return_value = 256
constructor = MagicMock(return_value=loaded_model)
monkeypatch.setattr(sentence_transformers_module, 'SentenceTransformer', constructor)
lazy_embedder = Embedder(SentenceTransformerEmbeddingModel('lazy-model'))
assert await lazy_embedder.max_input_tokens() == 256
constructor.assert_called_once_with('lazy-model')
async def test_query(self, embedder: Embedder):
result = await embedder.embed_query('Hello, world!')
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=128), length=1),
inputs=['Hello, world!'],
input_type='query',
model_name='sentence-transformers-testing/stsb-bert-tiny-safetensors',
timestamp=IsDatetime(),
provider_name='sentence-transformers',
)
)
async def test_documents(self, embedder: Embedder):
result = await embedder.embed_documents(['hello', 'world'])
assert result == snapshot(
EmbeddingResult(
embeddings=IsList(IsList(IsFloat(), length=128), length=2),
inputs=['hello', 'world'],
input_type='document',
model_name='sentence-transformers-testing/stsb-bert-tiny-safetensors',
timestamp=IsDatetime(),
provider_name='sentence-transformers',
)
)
async def test_max_input_tokens(self, embedder: Embedder):
max_input_tokens = await embedder.max_input_tokens()
assert max_input_tokens == snapshot(512)
async def test_count_tokens(self, embedder: Embedder):
count = await embedder.count_tokens('Hello, world!')
assert count == snapshot(6)
@pytest.mark.skipif(
not openai_imports_successful()
or not cohere_imports_successful()
or not google_imports_successful()
or not voyageai_installed
or not bedrock_imports_successful(),
reason='some embedding package was not installed',
)
@pytest.mark.usefixtures('voyageai_imported')
def test_known_embedding_model_names(): # pragma: lax no cover
from pydantic_ai.embeddings.voyageai import LatestVoyageAIEmbeddingModelNames
# Coverage seems to be misbehaving..?
def get_model_names(model_name_type: Any) -> Iterator[str]:
for arg in get_args(model_name_type):
if isinstance(arg, str):
yield arg
else:
yield from get_model_names(arg)
openai_names = [f'openai:{n}' for n in get_model_names(LatestOpenAIEmbeddingModelNames)]
cohere_names = [f'cohere:{n}' for n in get_model_names(LatestCohereEmbeddingModelNames)]
google_names = [f'google:{n}' for n in get_model_names(LatestGoogleGLAEmbeddingModelNames)]
google_cloud_names = [f'google-cloud:{n}' for n in get_model_names(LatestGoogleVertexEmbeddingModelNames)]
voyageai_names = [f'voyageai:{n}' for n in get_model_names(LatestVoyageAIEmbeddingModelNames)]
bedrock_names = [f'bedrock:{n}' for n in get_model_names(LatestBedrockEmbeddingModelNames)]
generated_names = sorted(
openai_names + cohere_names + google_names + google_cloud_names + voyageai_names + bedrock_names
)
known_model_names = sorted(get_args(KnownEmbeddingModelName.__value__))
if generated_names != known_model_names:
errors: list[str] = []
missing_names = set(generated_names) - set(known_model_names)
if missing_names:
errors.append(f'Missing names: {missing_names}')
extra_names = set(known_model_names) - set(generated_names)
if extra_names:
errors.append(f'Extra names: {extra_names}')
raise AssertionError('\n'.join(errors))
def test_infer_model_error():
with pytest.raises(ValueError, match='You must provide a provider prefix when specifying an embedding model name'):
infer_embedding_model('nonexistent')
async def test_instrument_all():
model = TestEmbeddingModel()
embedder = Embedder(model)
def get_model():
return embedder._get_model() # pyright: ignore[reportPrivateUsage]
Embedder.instrument_all(False)
assert get_model() is model
Embedder.instrument_all()
m = get_model()
assert isinstance(m, InstrumentedEmbeddingModel)
assert m.wrapped is model
assert m.instrumentation_settings.version == InstrumentationSettings().version
assert m.model_name == model.model_name
assert m.system == model.system
assert m.base_url == model.base_url
assert m.settings == model.settings
assert (await m.embed('Hello, world!', input_type='query')).embeddings == (
await model.embed('Hello, world!', input_type='query')
).embeddings
assert await m.max_input_tokens() == await model.max_input_tokens()
assert await m.count_tokens('Hello, world!') == await model.count_tokens('Hello, world!')
options = InstrumentationSettings(version=5)
Embedder.instrument_all(options)
m = get_model()
assert isinstance(m, InstrumentedEmbeddingModel)
assert m.wrapped is model
assert m.instrumentation_settings is options
Embedder.instrument_all(False)
assert get_model() is model
class ExplicitPortEmbeddingModel(TestEmbeddingModel):
@property
def base_url(self) -> str:
return 'https://example.com:8000/v1'
class MalformedPortEmbeddingModel(TestEmbeddingModel):
@property
def base_url(self) -> str:
return 'https://example.com:notaport/v1'
@pytest.mark.skipif(not logfire_imports_successful(), reason='logfire not installed')
@pytest.mark.parametrize(
'model_type,expected_server_attributes',
[
pytest.param(
ExplicitPortEmbeddingModel,
snapshot({'server.address': 'example.com', 'server.port': 8000}),
id='explicit-port',
),
pytest.param(MalformedPortEmbeddingModel, snapshot({}), id='malformed-port'),
],
)
async def test_instrumented_embedding_model_server_attributes(
model_type: type[TestEmbeddingModel], expected_server_attributes: dict[str, str | int], capfire: CaptureLogfire
):
"""A `base_url` whose port isn't an integer omits the server attributes instead of failing the request.
`urlparse` accepts the URL and only raises when `hostname`/`port` are read, so this is a unit test:
no real provider produces a `base_url` that survives client construction and fails at attribute-building.
"""
model = InstrumentedEmbeddingModel(model_type(), InstrumentationSettings())
await model.embed('Hello, world!', input_type='query')
[span] = capfire.exporter.exported_spans_as_dict()
assert {k: v for k, v in span['attributes'].items() if k.startswith('server.')} == expected_server_attributes
def test_override():
model = TestEmbeddingModel()
embedder = Embedder(model)
model2 = TestEmbeddingModel()
with embedder.override(model=model2):
assert embedder._get_model() is model2 # pyright: ignore[reportPrivateUsage]
with embedder.override():
assert embedder._get_model() is model # pyright: ignore[reportPrivateUsage]
assert embedder._get_model() is model # pyright: ignore[reportPrivateUsage]
def test_sync():
model = TestEmbeddingModel()
embedder = Embedder(model)
result = embedder.embed_query_sync('Hello, world!')
assert isinstance(result, EmbeddingResult)
result = embedder.embed_documents_sync(['hello', 'world'])
assert isinstance(result, EmbeddingResult)
result = embedder.embed_sync('Hello, world!', input_type='query')
assert isinstance(result, EmbeddingResult)
result = embedder.max_input_tokens_sync()
assert isinstance(result, int)
result = embedder.count_tokens_sync('Hello, world!')
assert isinstance(result, int)
async def test_settings():
model_settings: EmbeddingSettings = {'dimensions': 128, 'from_model': True} # pyright: ignore[reportAssignmentType]
model = TestEmbeddingModel(settings=model_settings)
assert model.settings == model_settings
await Embedder(model).embed_query('Hello, world!')
assert model.last_settings == snapshot({'dimensions': 128, 'from_model': True})
embedder_settings: EmbeddingSettings = {'dimensions': 256, 'from_embedder': True} # pyright: ignore[reportAssignmentType]
embedder = Embedder(model, settings=embedder_settings)
await embedder.embed_query('Hello, world!')
assert model.last_settings == snapshot({'dimensions': 256, 'from_model': True, 'from_embedder': True})
embed_settings: EmbeddingSettings = {'dimensions': 512, 'from_embed': True} # pyright: ignore[reportAssignmentType]
await embedder.embed_query('Hello, world!', settings=embed_settings)
assert model.last_settings == snapshot(
{'dimensions': 512, 'from_model': True, 'from_embedder': True, 'from_embed': True}
)
def test_embedder_compares_by_identity():
"""An `Embedder` is a stateful client, so two are equal only when they are the same object.
Field equality compared only `instrument`, so it made two `Embedder`s over different models equal
whenever their `instrument` matched, and left the class unhashable. No request is involved, so
this is not a VCR test.
"""
model = TestEmbeddingModel('small')
embedder = Embedder(model)
assert embedder == embedder
assert embedder != Embedder(model)
assert {embedder: 'small'}[embedder] == 'small'
def test_wrapper_embedding_model_deepcopy():
"""`deepcopy` builds the copy without `__init__`, so `wrapped` is unset when `__getattr__` runs.
`copy` probes the new instance for `__setstate__`; forwarding that to an unset `wrapped` recursed
until `RecursionError`. No request is involved, so this is not a VCR test.
"""
model = WrapperEmbeddingModel(TestEmbeddingModel('wrapped'))
copied = deepcopy(model)
assert copied is not model
assert copied.wrapped is not model.wrapped
assert copied.model_name == 'wrapped'
# Only the wrapped `TestEmbeddingModel` defines it, so reading it proves the copy still forwards.
assert copied.last_settings is None
def test_result():
result = EmbeddingResult(
embeddings=[[-1.0], [-0.5], [0.0], [0.5], [1.0]],
inputs=['a', 'b', 'c', 'd', 'e'],
input_type='document',
model_name='test',
timestamp=IsDatetime(),
provider_name='test',
)
assert result[0] == result['a'] == snapshot([-1.0])
assert result[1] == result['b'] == snapshot([-0.5])
assert result[2] == result['c'] == snapshot([0.0])
assert result[3] == result['d'] == snapshot([0.5])
assert result[4] == result['e'] == snapshot([1.0])
@pytest.mark.skipif(not logfire_imports_successful(), reason='logfire not installed')
async def test_limited_instrumentation(capfire: CaptureLogfire):
model = TestEmbeddingModel()
embedder = Embedder(model, instrument=InstrumentationSettings(include_content=False))
await embedder.embed_query('Hello, world!')
assert capfire.exporter.exported_spans_as_dict(parse_json_attributes=True) == snapshot(
[
{
'name': 'embeddings test',
'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
'parent': None,
'start_time': IsInt(),
'end_time': IsInt(),
'attributes': {
'gen_ai.operation.name': 'embeddings',
'gen_ai.provider.name': 'test',
'gen_ai.request.model': 'test',
'input_type': 'query',
'inputs_count': 1,
'logfire.json_schema': {
'type': 'object',
'properties': {
'input_type': {'type': 'string'},
'inputs_count': {'type': 'integer'},
'embedding_settings': {'type': 'object'},
},
},
'logfire.span_type': 'span',
'logfire.msg': 'embeddings test',
'gen_ai.usage.input_tokens': 2,
'gen_ai.response.model': 'test',
'gen_ai.embeddings.dimension.count': 8,
'gen_ai.response.id': IsStr(),
},
}
]
)
@pytest.mark.skipif(not logfire_imports_successful(), reason='logfire not installed')
@pytest.mark.parametrize('include_content', [True, False])
async def test_instrumentation_exception_honors_include_content(capfire: CaptureLogfire, include_content: bool):
"""A failing embedding request follows `include_content` like the agent's spans do.
A provider error carries its response body in the exception message, so the exception event and
the ERROR status description on the embedding span are withheld when content capture is off.
"""
class FailingEmbeddingModel(TestEmbeddingModel):
async def embed(
self, inputs: str | Sequence[str], *, input_type: EmbedInputType, settings: EmbeddingSettings | None = None
) -> EmbeddingResult:
raise ModelHTTPError(status_code=400, model_name='failing', body='invalid input: embed-secret')
embedder = Embedder(FailingEmbeddingModel(), instrument=InstrumentationSettings(include_content=include_content))
with pytest.raises(ModelHTTPError):
await embedder.embed('hello', input_type='document')
[span] = [span for span in capfire.exporter.exported_spans if span.status.status_code is StatusCode.ERROR]
[event] = [event for event in span.events if event.name == 'exception']
attributes = dict(event.attributes or {})
assert attributes['exception.type'] == 'pydantic_ai.exceptions.ModelHTTPError'
assert attributes['exception.escaped'] == 'False'
if include_content:
assert {'exception.message', 'exception.stacktrace'} <= set(attributes)
assert 'embed-secret' in str(attributes['exception.message'])
assert span.status.description is not None
else:
assert set(attributes) == {'exception.type', 'exception.escaped'}
assert span.status.description is None
assert 'embed-secret' not in str(capfire.exporter.exported_spans)
@pytest.mark.skipif(not openai_imports_successful(), reason='openai not installed')
async def test_openai_embedding_non_json_response_body_raises_model_api_error():
"""A 200 response body that can't be decoded as JSON surfaces as `ModelAPIError`, not the raw decode error.
A mock transport stands in for a cassette because no real provider returns such a body on demand.
https://github.com/pydantic/pydantic-ai/issues/9340
"""
def handler(request: httpx2.Request) -> httpx2.Response:
return httpx2.Response(200, content=b' ', headers={'content-type': 'application/json'})
async with AsyncOpenAI(
api_key='test',
base_url='http://localhost/v1',
max_retries=0,
http_client=httpx2.AsyncClient(transport=httpx2.MockTransport(handler)),
) as client:
model = OpenAIEmbeddingModel('text-embedding-3-small', provider=OpenAIProvider(openai_client=client))
with pytest.raises(ModelAPIError) as exc_info:
await model.embed('hello', input_type='query')
assert isinstance(exc_info.value.__cause__, json.JSONDecodeError)
assert exc_info.value.message.startswith('Failed to decode response as JSON')
@pytest.mark.skipif(not google_imports_successful(), reason='google not installed')
@pytest.mark.parametrize('call', ['embed', 'count_tokens'])
async def test_google_embedding_non_json_response_body_raises_model_api_error(call: str):
"""A 200 response body that can't be decoded as JSON surfaces as `ModelAPIError`, not the raw decode error.
A mock transport stands in for a cassette because no real provider returns such a body on demand.
https://github.com/pydantic/pydantic-ai/issues/9340
"""
def handler(request: httpx2.Request) -> httpx2.Response:
return httpx2.Response(200, content=b' ', headers={'content-type': 'application/json'})
async with httpx2.AsyncClient(transport=httpx2.MockTransport(handler)) as http_client:
provider = GoogleProvider(api_key='test-key', http_client=http_client, base_url='http://localhost')
model = GoogleEmbeddingModel('gemini-embedding-001', provider=provider)
with pytest.raises(ModelAPIError) as exc_info:
if call == 'count_tokens':
await model.count_tokens('hello')
else:
await model.embed('hello', input_type='query')
assert isinstance(exc_info.value.__cause__, json.JSONDecodeError)
assert exc_info.value.message.startswith('Failed to decode response as JSON')
@pytest.mark.skipif(not bedrock_imports_successful(), reason='Bedrock not installed')
@pytest.mark.parametrize('body', [b'', b'not json'], ids=['empty', 'not-json'])
async def test_bedrock_embedding_non_json_response_body_raises_model_api_error(body: bytes):
"""A 200 response body that can't be decoded as JSON surfaces as `ModelAPIError`, not the raw decode error.
The response is injected at `before-send` because no real endpoint returns such a body on demand.
https://github.com/pydantic/pydantic-ai/issues/9340
"""
provider = BedrockProvider(
region_name='us-east-1',
aws_access_key_id='AKIA6666666666666666',
aws_secret_access_key='6666666666666666666666666666666666666666',
)
model = BedrockEmbeddingModel('cohere.embed-english-v3', provider=provider)
def respond(request: AWSPreparedRequest, **_: object) -> AWSResponse:
return AWSResponse(request.url, 200, HTTPHeaders(), HTTPResponse(body=io.BytesIO(body), preload_content=False))
# botocore sends the response a `before-send` handler returns instead of the request, though the stubs type every
# handler as returning `None`.
model.client.meta.events.register_last('before-send.bedrock-runtime', respond) # pyright: ignore[reportArgumentType]
with pytest.raises(ModelAPIError) as exc_info:
await model.embed('hello', input_type='query')
assert isinstance(exc_info.value.__cause__, json.JSONDecodeError)
assert exc_info.value.message.startswith('Failed to decode response as JSON')
Reference in New Issue View Git Blame Copy Permalink