feat: add `BedrockEmbeddingModel` for Nova, Cohere and Titan endpoints #4008

bitnahian · 2026-01-15T06:12:57Z

Adds to Support embeddings models #3252

Pre-Review Checklist

Any AI generated code has been reviewed line-by-line by the human PR author, who stands by it.
No breaking changes in accordance with the version policy.
Linting and type checking pass per make format and make typecheck.
PR title is fit for the release changelog.

Pre-Merge Checklist

New tests for any fix or new behavior, maintaining 100% coverage.
Updated documentation for new features and behaviors, including docstrings for API docs.

…coverage

…drock embeddings

… function

…ng_model test

…odel test

DouweM · 2026-01-16T16:58:15Z

docs/embeddings.md

+  export AWS_REGION='us-east-1'
+  ```
+- **AWS credentials file** (`~/.aws/credentials`)
+- **IAM roles** (when running on AWS infrastructure)


As we do in https://ai.pydantic.dev/embeddings/#vertex-ai, let's link to the existing Bedrock doc that is the canonical source on how to configure the provider, that also has more details like that we support the AWS_BEARER_TOKEN_BEDROCK env var. So I'd prefer for that to be the main place we explain all the options, and for this to have an example + link there for more details; if you want to specifically mention ~/.aws/credentials for example, let's add it there, not here.

DouweM · 2026-01-16T16:58:56Z

docs/embeddings.md

+
+#### Basic Usage
+
+```python {title="bedrock_embeddings.py" test="skip"}


Please don't skip testing examples unless we have no other option.

DouweM · 2026-01-16T16:59:36Z

docs/embeddings.md

+
+#### Supported Models
+
+Bedrock supports three families of embedding models:


This'll get outdated, let's link to Amazon's doc on this and have just 1 example here.

Edit: I see now that we specifically implement support these 3 families. So then it does make sense listing them, but I'd rather use subheadings so they show up in the ToC sidebar, and instead of one "Bedrock-Specific Settings", I think each model family section should have list its own settings, as they're model family specific right?

DouweM · 2026-01-16T17:00:28Z

docs/embeddings.md

+embedder = Embedder('bedrock:eu.cohere.embed-english-v3')
+```
+
+The model automatically normalizes these prefixes when looking up `max_input_tokens()`.


I don't think this needs to be specified

DouweM · 2026-01-16T17:02:18Z

docs/embeddings.md

+
+#### Using a Custom Provider
+
+For advanced configuration, you can create a [`BedrockProvider`][pydantic_ai.providers.bedrock.BedrockProvider] directly:


I think we can have just one example here of building the provider yourself, and link to the existing bedrock docs for more details.

DouweM · 2026-01-16T17:24:23Z

pydantic_ai_slim/pydantic_ai/embeddings/bedrock.py

+                str(response_body),
+            )
+
+        return EmbeddingResult(


Lets simplify this method by returning only the embeddings + tokens + in some cases the response ID, and then build the EmbeddingResult where we call parse_response. Right now when we combine single-document runs, we deconstruct the individual EmbeddingResults and turn them into one combined one anyway, so let's skip the intermediate EmbeddingResult. That also means we no longer have to pass inputs, model_name etc into this method anymore.

DouweM · 2026-01-16T17:26:46Z

tests/conftest.py

+def mock_vcr_botocore_content_length(mocker: MockerFixture):
+    # VCR doesn't properly handle botocore's content-length verification when replaying responses.
+    # This causes IncompleteReadError when the recorded response body length doesn't match the Content-Length header.
+    # This happens because VCR decodes compressed responses but doesn't update the Content-Length header.


Should we instead fix this by updating the Content-Length in tests/json_body_serializer.py's serialize method?

DouweM · 2026-01-16T17:27:21Z

tests/test_embeddings.py

+                embeddings=IsList(IsList(IsFloat(), length=1024), length=1),
+                inputs=['Hello, world!'],
+                input_type='query',
+                usage=RequestUsage(input_tokens=IsInt()),


For usage, can we see the actual value rather than IsInt? I want to make sure it's not 0

DouweM · 2026-01-16T17:28:41Z

tests/test_embeddings.py



+@pytest.mark.skipif(not bedrock_imports_successful(), reason='Bedrock not installed')
+class TestBedrockHandlers:


I don't think we need to unit test this, assuming all the paths are covered by the integration test above?

DouweM · 2026-01-16T17:28:59Z

tests/test_embeddings.py

+        """Test error handling when ClientError is raised with HTTP status code."""
+        from botocore.exceptions import ClientError
+
+        from pydantic_ai.exceptions import ModelHTTPError


All imports at the top please

bitnahian · 2026-01-17T14:50:48Z

@DouweM Are we keeping the clients opinionated to work as text embedding models only? For instance, Nova has a tonne of multi-modal specific parameters and if we're restricting to text only, we shouldn't surface them.

bitnahian added 6 commits January 15, 2026 17:01

feat: Add first cut bedrock embedding model implementation

29d7e7a

test: Adding vcrpy tests for reasonable bedrock models

8c95b1b

fix: failing CI tests

de195de

docs: Add Bedrock embedding model section to documentation

d577dc3

refactor: Make API simpler and better override implementation

4b1ffe2

docs: Update BedrockProvider example fix E402

cb45092

dsfaccini added new models Support for new model(s) bedrock embeddings feature New feature request, or PR implementing a feature (enhancement) labels Jan 15, 2026

dsfaccini changed the title ~~feat: add BedrockEmbeddingModel for Nova, Cohere and Titan endpoints~~ feat: add BedrockEmbeddingModel for Nova, Cohere and Titan endpoints Jan 15, 2026

bitnahian added 8 commits January 16, 2026 10:16

Merge branch 'main' into bitnahian-bedrock-embeddings

4c580c9

refactor: Simplify Bedrock embedding model handling and improve test …

90679f2

…coverage

docs: Add regional prefixes section and note on token counting for Be…

a5acd31

…drock embeddings

refactor: Update Bedrock model inference to use infer_embedding_model…

fac0841

… function

refactor: Update TestBedrock to use bedrock_provider in infer_embeddi…

4661cb9

…ng_model test

refactor: Try to fix infer_model coverage

a45cfb1

refactor: Update TestBedrock to set environment variables for infer_m…

b15b0aa

…odel test

docs: Clarify note on token counting for Bedrock embedding models

3772a1e

dmontagu approved these changes Jan 16, 2026

View reviewed changes

DouweM requested changes Jan 16, 2026

View reviewed changes

DouweM self-assigned this Jan 16, 2026

DouweM added the awaiting author revision label Jan 16, 2026

refactor: Improve imports and update assertions for text input

8a30649

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

feat: add `BedrockEmbeddingModel` for Nova, Cohere and Titan endpoints #4008

feat: add `BedrockEmbeddingModel` for Nova, Cohere and Titan endpoints #4008

bitnahian commented Jan 15, 2026 •

edited

Loading

Uh oh!

DouweM Jan 16, 2026

Uh oh!

DouweM Jan 16, 2026

Uh oh!

DouweM Jan 16, 2026

Uh oh!

DouweM Jan 16, 2026

Uh oh!

DouweM Jan 16, 2026

Uh oh!

DouweM Jan 16, 2026

Uh oh!

DouweM Jan 16, 2026

Uh oh!

DouweM Jan 16, 2026

Uh oh!

DouweM Jan 16, 2026

Uh oh!

DouweM Jan 16, 2026

Uh oh!

DouweM Jan 16, 2026

Uh oh!

bitnahian commented Jan 17, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants


		#### Basic Usage

		```python {title="bedrock_embeddings.py" test="skip"}


		#### Supported Models

		Bedrock supports three families of embedding models:


		#### Using a Custom Provider

		For advanced configuration, you can create a [`BedrockProvider`][pydantic_ai.providers.bedrock.BedrockProvider] directly:



		@pytest.mark.skipif(not bedrock_imports_successful(), reason='Bedrock not installed')
		class TestBedrockHandlers:

feat: add BedrockEmbeddingModel for Nova, Cohere and Titan endpoints #4008

Are you sure you want to change the base?

feat: add BedrockEmbeddingModel for Nova, Cohere and Titan endpoints #4008

Conversation

bitnahian commented Jan 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Pre-Review Checklist

Pre-Merge Checklist

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bitnahian commented Jan 17, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

feat: add `BedrockEmbeddingModel` for Nova, Cohere and Titan endpoints #4008

feat: add `BedrockEmbeddingModel` for Nova, Cohere and Titan endpoints #4008

bitnahian commented Jan 15, 2026 •

edited

Loading