feat: add Bedrock (Titan) embeddings to ai/embed#168
Conversation
Extends the ai/embed connector with a Bedrock provider using the Amazon Titan text-embedding models (amazon.titan-embed-text-v2:0 and amazon.titan-embed-text-v1). - BedrockEmbeddingProvider uses the Bedrock InvokeModel API (distinct from the Converse API used for chat). Titan embeds one input per call, so multi-input requests are issued sequentially and reassembled in order; token counts are summed for budget accounting. Retryable Bedrock errors reuse classifyBedrockError. - ai/embed getProvider builds the Bedrock client from AWS config (region from params/credential/DefaultRegion), mirroring ai/completion; serve wires cfg.AWS.Region + NewAWSConfig into the embeddings connector. - honors `dimensions` on Titan v2. - Tests: provider-level Titan tests (single/multi input ordering, token sum, request body, unsupported-model error) with a mock InvokeModel client; connector unknown-provider test. - Docs: ai/embed connector reference + RAG guide updated for Bedrock. Cohere Bedrock embedding models remain a follow-up on #153. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WnKgCRfMvFcZAkoVbr8k79
|
Warning Review limit reached
Next review available in: 48 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (4)
📝 WalkthroughWalkthroughThis PR adds Bedrock Titan embedding support to the ChangesBedrock embeddings support
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant ServeCLI
participant EmbeddingConnector
participant BedrockEmbeddingProvider
participant BedrockRuntime
ServeCLI->>EmbeddingConnector: configure AWS region/config defaults
EmbeddingConnector->>EmbeddingConnector: getProvider("bedrock")
EmbeddingConnector->>BedrockEmbeddingProvider: build with AWS config
EmbeddingConnector->>BedrockEmbeddingProvider: Embeddings(inputs)
loop per input string
BedrockEmbeddingProvider->>BedrockRuntime: InvokeModel(Titan request)
BedrockRuntime-->>BedrockEmbeddingProvider: embedding + tokens
end
BedrockEmbeddingProvider-->>EmbeddingConnector: EmbeddingResponse
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: fbe454eba7
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if req.Dimensions > 0 { | ||
| body.Dimensions = req.Dimensions | ||
| } |
There was a problem hiding this comment.
When callers use the documented amazon.titan-embed-text-v1 model and pass dimensions (for example by reusing a shared embedding config), this branch serializes dimensions into the InvokeModel body even though Titan G1/Text v1 only accepts inputText. That makes the new Bedrock path fail instead of returning the default v1 embedding; either reject dimensions for v1 up front or only set this field for amazon.titan-embed-text-v2:0.
Useful? React with 👍 / 👎.
Addresses PR review (Codex P2). Titan v1 (G1) rejects any InvokeModel body field other than inputText, so serializing `dimensions` for it — e.g. from an embedding config shared with v2 — would fail the request. Only set dimensions for amazon.titan-embed-text-v2 and ignore it for v1 (which returns its fixed-size embedding). Adds a regression test. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WnKgCRfMvFcZAkoVbr8k79
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
packages/engine/internal/connector/embed.go (1)
83-109: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winThread
ctxinto Bedrock provider setup.
EmbeddingConnector.Executealready has a request context, butgetProviderreplaces it withcontext.Background()before callingNewAWSConfig. That drops cancellation and deadlines during AWS config/credential resolution. Passctxthrough togetProviderandconfigFunc.🔧 Proposed fix
-func (c *EmbeddingConnector) getProvider(name string, params map[string]any) (EmbeddingProvider, error) { +func (c *EmbeddingConnector) getProvider(ctx context.Context, name string, params map[string]any) (EmbeddingProvider, error) { @@ - awsCfg, err := configFunc(context.Background(), cred, defaultRegion) + awsCfg, err := configFunc(ctx, cred, defaultRegion)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/engine/internal/connector/embed.go` around lines 83 - 109, The Bedrock path in EmbeddingConnector.getProvider is discarding the request context by calling NewAWSConfig with context.Background(), so cancellation and deadlines from EmbeddingConnector.Execute are lost. Thread the existing ctx from Execute into getProvider, and pass that ctx through to configFunc in the bedrock case instead of creating a new background context.
🧹 Nitpick comments (2)
packages/engine/internal/connector/provider_bedrock_embed_test.go (1)
21-81: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider adding coverage for the v1 model and the dimensions edge case.
Current tests only cover
amazon.titan-embed-text-v2:0with dimensions set. Adding a case foramazon.titan-embed-text-v1(dimensions omitted) and one foramazon.titan-embed-text-v1+ non-zeroDimensionswould pin the behavior discussed inprovider_bedrock.goand catch regressions if the model-prefix check is tightened.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/engine/internal/connector/provider_bedrock_embed_test.go` around lines 21 - 81, Add test coverage in BedrockEmbeddingProvider_Embeddings for the Titan v1 path in provider_bedrock.go: add a case using amazon.titan-embed-text-v1 with Dimensions omitted to verify the request body behavior, and another case using the same model with a non-zero Dimensions value to pin the edge-case handling. Reuse BedrockEmbeddingProvider, Embeddings, and the titanEmbedRequest/request-body assertions so the tests cover both the model-prefix check and dimension-specific behavior.packages/engine/internal/connector/embed_test.go (1)
147-157: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winNo test exercises the new
getProvider("bedrock", ...)path inembed.go.The old Bedrock-specific test was replaced with a generic unknown-provider test, but nothing now covers region resolution (
c.DefaultRegionvsparams["region"]), theAWSConfigFuncoverride, or the[bedrock]: %werror wrapping added inembed.go. SinceAWSConfigFuncis injectable, a test with a stub function avoiding real AWS calls would be straightforward.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/engine/internal/connector/embed_test.go` around lines 147 - 157, Add a dedicated test for the EmbeddingConnector bedrock provider path in embed_test.go, since TestEmbeddingConnector_UnknownProvider no longer covers it. Exercise EmbeddingConnector.Execute/getProvider with provider set to "bedrock" and verify region selection from c.DefaultRegion versus params["region"], using a stubbed AWSConfigFunc to avoid real AWS calls. Also assert that failures from the bedrock branch are wrapped with the "[bedrock]: %w" format so the error path is covered.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@packages/engine/internal/connector/provider_bedrock.go`:
- Around line 264-268: Validate Titan v2 embedding dimensions before building
the request in provider_bedrock.go’s Titan request path. The current logic in
the embedding request loop forwards any positive req.Dimensions value into
titanEmbedRequest, but Titan v2 only supports 256, 512, or 1024. Add validation
in the Bedrock provider flow before the request is sent, and reject or normalize
unsupported values using the existing embedding request handling around
req.Inputs, req.Dimensions, and titanEmbedRequest.
---
Outside diff comments:
In `@packages/engine/internal/connector/embed.go`:
- Around line 83-109: The Bedrock path in EmbeddingConnector.getProvider is
discarding the request context by calling NewAWSConfig with
context.Background(), so cancellation and deadlines from
EmbeddingConnector.Execute are lost. Thread the existing ctx from Execute into
getProvider, and pass that ctx through to configFunc in the bedrock case instead
of creating a new background context.
---
Nitpick comments:
In `@packages/engine/internal/connector/embed_test.go`:
- Around line 147-157: Add a dedicated test for the EmbeddingConnector bedrock
provider path in embed_test.go, since TestEmbeddingConnector_UnknownProvider no
longer covers it. Exercise EmbeddingConnector.Execute/getProvider with provider
set to "bedrock" and verify region selection from c.DefaultRegion versus
params["region"], using a stubbed AWSConfigFunc to avoid real AWS calls. Also
assert that failures from the bedrock branch are wrapped with the "[bedrock]:
%w" format so the error path is covered.
In `@packages/engine/internal/connector/provider_bedrock_embed_test.go`:
- Around line 21-81: Add test coverage in BedrockEmbeddingProvider_Embeddings
for the Titan v1 path in provider_bedrock.go: add a case using
amazon.titan-embed-text-v1 with Dimensions omitted to verify the request body
behavior, and another case using the same model with a non-zero Dimensions value
to pin the edge-case handling. Reuse BedrockEmbeddingProvider, Embeddings, and
the titanEmbedRequest/request-body assertions so the tests cover both the
model-prefix check and dimension-specific behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 8d9f22ee-6b7c-4e09-b43b-2a4ed896a5bd
📒 Files selected for processing (8)
.changeset/ai-embed-bedrock.mdpackages/engine/internal/cli/serve.gopackages/engine/internal/connector/embed.gopackages/engine/internal/connector/embed_test.gopackages/engine/internal/connector/provider_bedrock.gopackages/engine/internal/connector/provider_bedrock_embed_test.gopackages/site/src/content/docs/rag-guide.mdpackages/site/src/content/docs/workflow-reference/connectors.md
Addresses CodeRabbit review on the Bedrock embeddings PR: - Thread the request context through EmbeddingConnector.getProvider into NewAWSConfig instead of context.Background(), so cancellation/deadlines aren't dropped during AWS config resolution (Major). - Validate Titan v2 dimensions up front — only 256, 512, or 1024 are accepted, so reject other values with a clear error instead of a cryptic Bedrock runtime failure (Minor). - Tests: cover the ai/embed bedrock provider path via a stubbed AWSConfigFunc (region resolution + [bedrock] error wrapping) and reject invalid v2 dimensions. (v1 dimension-ignoring is already covered by TestBedrockEmbeddingProvider_V1IgnoresDimensions.) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WnKgCRfMvFcZAkoVbr8k79
|
Addressed the review in
Build, vet, gofmt, and golangci-lint are clean. Generated by Claude Code |
2 participants
Summary
Continues RAG (#153): adds a Bedrock provider to the
ai/embedconnector via the Amazon Titan text-embedding models, so teams already using Claude-via-Bedrock can generate embeddings in the same AWS account without an OpenAI key.Changes
BedrockEmbeddingProvider(provider_bedrock.go) — uses the BedrockInvokeModelAPI (distinct from the Converse API used for chat). Supportsamazon.titan-embed-text-v2:0andamazon.titan-embed-text-v1. Titan embeds one input per call, so multi-input requests are issued sequentially and reassembled in order; input token counts are summed for budget accounting; retryable errors reuseclassifyBedrockError. Honorsdimensionson Titan v2.ai/embedconnector —getProvidernow builds the Bedrock client from AWS config (region fromparams.region/ credential /DefaultRegion), mirroringai/completion. AddedAWSConfigFunc+DefaultRegionfields.servewirescfg.AWS.Region+NewAWSConfiginto the embeddings connector alongside the chat connector.ai/embedconnector reference and the RAG guide updated for the Bedrock provider.Scope
Titan only for now. Cohere Bedrock embedding models (which use a batch
textsbody plus aninput_typedistinction for documents vs queries) remain a follow-up on #153, as do nativekb/*connectors + chunking.Testing
InvokeModelclient; connector unknown-provider test. Existing OpenAI embed + Bedrock chat tests still pass.go build,go vet,gofmt,golangci-lint runclean.Related Issues
Part of #153 (RAG subsystem).
Generated by Claude Code
Summary by CodeRabbit
New Features
ai/embed, including Titan text-embedding models.Documentation
Bug Fixes