From fbe454eba78ca5cc54a78a650a804f1b34f1d3a9 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 2 Jul 2026 04:41:04 +0000 Subject: [PATCH 1/3] feat: add Bedrock (Titan) embeddings to ai/embed (#153) 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 Claude-Session: https://claude.ai/code/session_01WnKgCRfMvFcZAkoVbr8k79 --- .changeset/ai-embed-bedrock.md | 5 ++ packages/engine/internal/cli/serve.go | 11 ++- packages/engine/internal/connector/embed.go | 32 +++++--- .../engine/internal/connector/embed_test.go | 8 +- .../internal/connector/provider_bedrock.go | 75 +++++++++++++++++ .../connector/provider_bedrock_embed_test.go | 81 +++++++++++++++++++ packages/site/src/content/docs/rag-guide.md | 9 ++- .../docs/workflow-reference/connectors.md | 9 ++- 8 files changed, 207 insertions(+), 23 deletions(-) create mode 100644 .changeset/ai-embed-bedrock.md create mode 100644 packages/engine/internal/connector/provider_bedrock_embed_test.go diff --git a/.changeset/ai-embed-bedrock.md b/.changeset/ai-embed-bedrock.md new file mode 100644 index 0000000..57fa17d --- /dev/null +++ b/.changeset/ai-embed-bedrock.md @@ -0,0 +1,5 @@ +--- +"@mantle/engine": minor +--- + +Add Bedrock support to the `ai/embed` connector via the Amazon Titan text-embedding models (`amazon.titan-embed-text-v2:0`, `amazon.titan-embed-text-v1`). Set `provider: bedrock` with a `region` and an `aws` credential; the connector uses Bedrock `InvokeModel` (embedding one input per call, reassembled in order) and honors `dimensions` on Titan v2. `serve` wires the AWS region/config into the embeddings connector alongside the chat connector. Cohere Bedrock embedding models are a planned follow-up. diff --git a/packages/engine/internal/cli/serve.go b/packages/engine/internal/cli/serve.go index c22c839..23f0eeb 100644 --- a/packages/engine/internal/cli/serve.go +++ b/packages/engine/internal/cli/serve.go @@ -71,12 +71,17 @@ func newServeCommand() *cobra.Command { } } - // Apply the base-URL allowlist to the embeddings connector too, so - // ai/embed can't be pointed at arbitrary hosts. The model allowlist - // is intentionally not shared: allowed_models lists chat models, and + // Apply the base-URL allowlist and AWS defaults to the embeddings + // connector too, so ai/embed can't be pointed at arbitrary hosts and + // its Bedrock provider gets a region. The model allowlist is + // intentionally not shared: allowed_models lists chat models, and // applying it here would block embedding models. if embConn, err := eng.Registry.Get("ai/embed"); err == nil { if emb, ok := embConn.(*connector.EmbeddingConnector); ok { + if cfg.AWS.Region != "" { + emb.DefaultRegion = cfg.AWS.Region + emb.AWSConfigFunc = connector.NewAWSConfig + } emb.AllowedBaseURLs = cfg.Engine.AllowedBaseURLs } } diff --git a/packages/engine/internal/connector/embed.go b/packages/engine/internal/connector/embed.go index f2786b1..298244d 100644 --- a/packages/engine/internal/connector/embed.go +++ b/packages/engine/internal/connector/embed.go @@ -8,16 +8,20 @@ import ( "strings" "time" + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/service/bedrockruntime" "github.com/dvflw/mantle/internal/metrics" ) // EmbeddingConnector implements the ai/embed action: it turns text into vector // embeddings via a provider's embeddings API. It mirrors AIConnector's provider -// selection, base-URL/model allowlists, and metrics. Because the action lives -// under the "ai/" prefix, the engine's token-budget accounting applies to it -// automatically (via the usage.total_tokens it returns). +// selection, base-URL/model allowlists, AWS config, and metrics. Because the +// action lives under the "ai/" prefix, the engine's token-budget accounting +// applies to it automatically (via the usage.total_tokens it returns). type EmbeddingConnector struct { Client *http.Client + AWSConfigFunc func(ctx context.Context, cred map[string]string, defaultRegion string) (aws.Config, error) + DefaultRegion string AllowedBaseURLs []string AllowedModels []string // empty = all models allowed } @@ -88,13 +92,23 @@ func (c *EmbeddingConnector) getProvider(name string, params map[string]any) (Em } return &OpenAIProvider{Client: c.Client, BaseURL: baseURL}, nil case "bedrock": - // Bedrock embeddings use InvokeModel with per-model request shapes - // (Titan, Cohere), distinct from the Converse API used for chat. - // Tracked as a follow-up; OpenAI-compatible endpoints (including Azure - // OpenAI and local servers) are reachable today via base_url. - return nil, fmt.Errorf("bedrock embeddings are not yet supported; use provider: openai (optionally with base_url)") + cred, _ := params["_credential"].(map[string]string) + region, _ := params["region"].(string) + defaultRegion := c.DefaultRegion + if region != "" { + defaultRegion = region + } + configFunc := c.AWSConfigFunc + if configFunc == nil { + configFunc = NewAWSConfig + } + awsCfg, err := configFunc(context.Background(), cred, defaultRegion) + if err != nil { + return nil, fmt.Errorf("[bedrock]: %w", err) + } + return &BedrockEmbeddingProvider{Client: bedrockruntime.NewFromConfig(awsCfg)}, nil default: - return nil, fmt.Errorf("unknown provider %q (available: openai)", name) + return nil, fmt.Errorf("unknown provider %q (available: openai, bedrock)", name) } } diff --git a/packages/engine/internal/connector/embed_test.go b/packages/engine/internal/connector/embed_test.go index cb9b2b7..fcaa689 100644 --- a/packages/engine/internal/connector/embed_test.go +++ b/packages/engine/internal/connector/embed_test.go @@ -144,14 +144,14 @@ func TestEmbeddingConnector_IncompleteResponseFailsFast(t *testing.T) { } } -func TestEmbeddingConnector_BedrockUnsupported(t *testing.T) { +func TestEmbeddingConnector_UnknownProvider(t *testing.T) { c := &EmbeddingConnector{} _, err := c.Execute(context.Background(), map[string]any{ - "model": "amazon.titan-embed-text-v2:0", + "model": "some-model", "input": "x", - "provider": "bedrock", + "provider": "not-a-provider", }) if err == nil { - t.Error("expected error for unsupported bedrock provider") + t.Error("expected error for unknown provider") } } diff --git a/packages/engine/internal/connector/provider_bedrock.go b/packages/engine/internal/connector/provider_bedrock.go index 5383255..13b1f84 100644 --- a/packages/engine/internal/connector/provider_bedrock.go +++ b/packages/engine/internal/connector/provider_bedrock.go @@ -223,3 +223,78 @@ func classifyBedrockError(err error) error { return fmt.Errorf("bedrock: API request failed") } + +// BedrockInvokeAPI abstracts the Bedrock InvokeModel call (used for embeddings) +// for testability. +type BedrockInvokeAPI interface { + InvokeModel(ctx context.Context, input *bedrockruntime.InvokeModelInput, opts ...func(*bedrockruntime.Options)) (*bedrockruntime.InvokeModelOutput, error) +} + +// BedrockEmbeddingProvider implements EmbeddingProvider using the AWS Bedrock +// InvokeModel API. It currently supports the Amazon Titan text-embedding models +// (amazon.titan-embed-text-v1 and amazon.titan-embed-text-v2:0). Titan embeds a +// single input per call, so multi-input requests are issued sequentially and +// reassembled in order. +type BedrockEmbeddingProvider struct { + Client BedrockInvokeAPI +} + +// titanEmbedRequest is the Amazon Titan text-embeddings InvokeModel body. +// Dimensions is only honoured by titan-embed-text-v2; omitempty keeps it out of +// v1 requests. +type titanEmbedRequest struct { + InputText string `json:"inputText"` + Dimensions int `json:"dimensions,omitempty"` +} + +type titanEmbedResponse struct { + Embedding []float64 `json:"embedding"` + InputTextTokenCount int `json:"inputTextTokenCount"` +} + +// Embeddings implements EmbeddingProvider for Bedrock Titan models. +func (p *BedrockEmbeddingProvider) Embeddings(ctx context.Context, req *EmbeddingRequest) (*EmbeddingResponse, error) { + if !strings.HasPrefix(req.Model, "amazon.titan-embed") { + return nil, fmt.Errorf("bedrock: unsupported embedding model %q (supported: amazon.titan-embed-text-v1, amazon.titan-embed-text-v2:0)", req.Model) + } + + out := make([][]float64, len(req.Inputs)) + totalTokens := 0 + + for i, text := range req.Inputs { + body := titanEmbedRequest{InputText: text} + if req.Dimensions > 0 { + body.Dimensions = req.Dimensions + } + bodyJSON, err := json.Marshal(body) + if err != nil { + return nil, fmt.Errorf("bedrock: marshaling embedding request: %w", err) + } + + resp, err := p.Client.InvokeModel(ctx, &bedrockruntime.InvokeModelInput{ + ModelId: aws.String(req.Model), + Body: bodyJSON, + ContentType: aws.String("application/json"), + Accept: aws.String("application/json"), + }) + if err != nil { + return nil, classifyBedrockError(err) + } + + var er titanEmbedResponse + if err := json.Unmarshal(resp.Body, &er); err != nil { + return nil, fmt.Errorf("bedrock: parsing embedding response: %w", err) + } + if len(er.Embedding) == 0 { + return nil, fmt.Errorf("bedrock: empty embedding returned for input %d", i) + } + out[i] = er.Embedding + totalTokens += er.InputTextTokenCount + } + + return &EmbeddingResponse{ + Embeddings: out, + Model: req.Model, + Usage: ChatUsage{PromptTokens: totalTokens, TotalTokens: totalTokens}, + }, nil +} diff --git a/packages/engine/internal/connector/provider_bedrock_embed_test.go b/packages/engine/internal/connector/provider_bedrock_embed_test.go new file mode 100644 index 0000000..d1bb03a --- /dev/null +++ b/packages/engine/internal/connector/provider_bedrock_embed_test.go @@ -0,0 +1,81 @@ +package connector + +import ( + "context" + "encoding/json" + "testing" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/service/bedrockruntime" +) + +// mockBedrockInvokeClient implements BedrockInvokeAPI for testing. +type mockBedrockInvokeClient struct { + InvokeFunc func(ctx context.Context, input *bedrockruntime.InvokeModelInput, opts ...func(*bedrockruntime.Options)) (*bedrockruntime.InvokeModelOutput, error) +} + +func (m *mockBedrockInvokeClient) InvokeModel(ctx context.Context, input *bedrockruntime.InvokeModelInput, opts ...func(*bedrockruntime.Options)) (*bedrockruntime.InvokeModelOutput, error) { + return m.InvokeFunc(ctx, input, opts...) +} + +func TestBedrockEmbeddingProvider_Titan(t *testing.T) { + var calls int + mock := &mockBedrockInvokeClient{ + InvokeFunc: func(ctx context.Context, input *bedrockruntime.InvokeModelInput, opts ...func(*bedrockruntime.Options)) (*bedrockruntime.InvokeModelOutput, error) { + calls++ + if aws.ToString(input.ModelId) != "amazon.titan-embed-text-v2:0" { + t.Errorf("ModelId = %q", aws.ToString(input.ModelId)) + } + // Verify the request body carries inputText and the requested dimensions. + var body titanEmbedRequest + if err := json.Unmarshal(input.Body, &body); err != nil { + t.Fatalf("unmarshaling request body: %v", err) + } + if body.InputText == "" { + t.Error("inputText is empty") + } + if body.Dimensions != 256 { + t.Errorf("dimensions = %d, want 256", body.Dimensions) + } + resp, _ := json.Marshal(titanEmbedResponse{ + Embedding: []float64{0.5, -0.25}, + InputTextTokenCount: 4, + }) + return &bedrockruntime.InvokeModelOutput{Body: resp}, nil + }, + } + + p := &BedrockEmbeddingProvider{Client: mock} + resp, err := p.Embeddings(context.Background(), &EmbeddingRequest{ + Model: "amazon.titan-embed-text-v2:0", + Inputs: []string{"one", "two", "three"}, + Dimensions: 256, + }) + if err != nil { + t.Fatalf("Embeddings() error: %v", err) + } + // Titan embeds one input per call. + if calls != 3 { + t.Errorf("InvokeModel calls = %d, want 3", calls) + } + if len(resp.Embeddings) != 3 { + t.Fatalf("embeddings = %d, want 3", len(resp.Embeddings)) + } + if resp.Embeddings[0][0] != 0.5 || resp.Embeddings[0][1] != -0.25 { + t.Errorf("embedding[0] = %v", resp.Embeddings[0]) + } + if resp.Usage.TotalTokens != 12 { // 4 per input * 3 + t.Errorf("TotalTokens = %d, want 12", resp.Usage.TotalTokens) + } +} + +func TestBedrockEmbeddingProvider_UnsupportedModel(t *testing.T) { + p := &BedrockEmbeddingProvider{Client: &mockBedrockInvokeClient{}} + _, err := p.Embeddings(context.Background(), &EmbeddingRequest{ + Model: "cohere.embed-english-v3", + Inputs: []string{"x"}, + }) + if err == nil { + t.Error("expected error for unsupported (non-Titan) Bedrock embedding model") + } +} diff --git a/packages/site/src/content/docs/rag-guide.md b/packages/site/src/content/docs/rag-guide.md index 71650c8..b3fec9a 100644 --- a/packages/site/src/content/docs/rag-guide.md +++ b/packages/site/src/content/docs/rag-guide.md @@ -88,9 +88,12 @@ outputs with `--output json` or `-v`). - **Embedding dimension must match the model.** The schema uses `vector(1536)` for `text-embedding-3-small`; use `vector(3072)` for `text-embedding-3-large`. Embed queries and documents with the same model. -- **Provider support:** `ai/embed` currently supports `openai` and - OpenAI-compatible endpoints (Azure OpenAI, local servers) via `base_url`. - Bedrock embeddings are a planned follow-up. +- **Provider support:** `ai/embed` supports `openai` (and OpenAI-compatible + endpoints like Azure OpenAI or local servers via `base_url`) and `bedrock` + with the Amazon Titan text-embedding models (`amazon.titan-embed-text-v2:0`, + `amazon.titan-embed-text-v1`). To use Bedrock, set `provider: bedrock`, a + `region`, and an `aws` credential; adjust the schema's `vector(...)` dimension + to match (Titan v2 defaults to 1024). Cohere Bedrock models are a follow-up. - **No native chunking or `kb/*` convenience steps yet.** You compose RAG from `ai/embed` + `postgres/query` + `ai/completion` today; native chunking and `kb/upsert` / `kb/query` connectors are tracked by diff --git a/packages/site/src/content/docs/workflow-reference/connectors.md b/packages/site/src/content/docs/workflow-reference/connectors.md index a95ed9f..bf9eb06 100644 --- a/packages/site/src/content/docs/workflow-reference/connectors.md +++ b/packages/site/src/content/docs/workflow-reference/connectors.md @@ -149,11 +149,12 @@ Turns text into vector embeddings via an OpenAI-compatible embeddings API. Use i | Param | Type | Required | Description | |---|---|---|---| -| `provider` | string | No | Embeddings provider: `openai` (default). Bedrock is not yet supported. | -| `model` | string | Yes | Embedding model (e.g., `text-embedding-3-small`, `text-embedding-3-large`). | +| `provider` | string | No | Embeddings provider: `openai` (default) or `bedrock`. | +| `model` | string | Yes | Embedding model. OpenAI: `text-embedding-3-small`, `text-embedding-3-large`. Bedrock: `amazon.titan-embed-text-v2:0`, `amazon.titan-embed-text-v1`. | | `input` | string \| list | Yes | A single string or an array of strings to embed. | -| `dimensions` | integer | No | Requested embedding dimension, when the model supports truncation. | -| `base_url` | string | No | Override the API base URL. Defaults to `https://api.openai.com/v1`. Use for Azure OpenAI or OpenAI-compatible servers. | +| `dimensions` | integer | No | Requested embedding dimension, when the model supports it (OpenAI `text-embedding-3-*`, Bedrock `titan-embed-text-v2`). | +| `region` | string | No | AWS region for the Bedrock provider. Only used when `provider` is `bedrock`. | +| `base_url` | string | No | Override the API base URL (OpenAI provider). Defaults to `https://api.openai.com/v1`. Use for Azure OpenAI or OpenAI-compatible servers. | **Output:** From cdff8d27ad38b1619ecc9baf0ae0348b50f3d00f Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 2 Jul 2026 04:46:41 +0000 Subject: [PATCH 2/3] fix: gate Titan embedding dimensions to v2 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 Claude-Session: https://claude.ai/code/session_01WnKgCRfMvFcZAkoVbr8k79 --- .../internal/connector/provider_bedrock.go | 7 +++++- .../connector/provider_bedrock_embed_test.go | 25 +++++++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/packages/engine/internal/connector/provider_bedrock.go b/packages/engine/internal/connector/provider_bedrock.go index 13b1f84..2fabe45 100644 --- a/packages/engine/internal/connector/provider_bedrock.go +++ b/packages/engine/internal/connector/provider_bedrock.go @@ -261,9 +261,14 @@ func (p *BedrockEmbeddingProvider) Embeddings(ctx context.Context, req *Embeddin out := make([][]float64, len(req.Inputs)) totalTokens := 0 + // dimensions is only accepted by Titan v2; Titan v1 (G1) rejects any field + // other than inputText, so ignore it there rather than failing a request + // that reuses a shared embedding config. + supportsDimensions := strings.Contains(req.Model, "titan-embed-text-v2") + for i, text := range req.Inputs { body := titanEmbedRequest{InputText: text} - if req.Dimensions > 0 { + if req.Dimensions > 0 && supportsDimensions { body.Dimensions = req.Dimensions } bodyJSON, err := json.Marshal(body) diff --git a/packages/engine/internal/connector/provider_bedrock_embed_test.go b/packages/engine/internal/connector/provider_bedrock_embed_test.go index d1bb03a..d1ad9ce 100644 --- a/packages/engine/internal/connector/provider_bedrock_embed_test.go +++ b/packages/engine/internal/connector/provider_bedrock_embed_test.go @@ -69,6 +69,31 @@ func TestBedrockEmbeddingProvider_Titan(t *testing.T) { } } +func TestBedrockEmbeddingProvider_V1IgnoresDimensions(t *testing.T) { + mock := &mockBedrockInvokeClient{ + InvokeFunc: func(ctx context.Context, input *bedrockruntime.InvokeModelInput, opts ...func(*bedrockruntime.Options)) (*bedrockruntime.InvokeModelOutput, error) { + var body titanEmbedRequest + if err := json.Unmarshal(input.Body, &body); err != nil { + t.Fatalf("unmarshaling request body: %v", err) + } + // Titan v1 must not receive dimensions even when requested. + if body.Dimensions != 0 { + t.Errorf("dimensions = %d, want 0 (omitted) for titan v1", body.Dimensions) + } + resp, _ := json.Marshal(titanEmbedResponse{Embedding: []float64{1}, InputTextTokenCount: 1}) + return &bedrockruntime.InvokeModelOutput{Body: resp}, nil + }, + } + p := &BedrockEmbeddingProvider{Client: mock} + if _, err := p.Embeddings(context.Background(), &EmbeddingRequest{ + Model: "amazon.titan-embed-text-v1", + Inputs: []string{"x"}, + Dimensions: 512, + }); err != nil { + t.Fatalf("Embeddings() error: %v", err) + } +} + func TestBedrockEmbeddingProvider_UnsupportedModel(t *testing.T) { p := &BedrockEmbeddingProvider{Client: &mockBedrockInvokeClient{}} _, err := p.Embeddings(context.Background(), &EmbeddingRequest{ From 33a948c619ab46e27a556ec6fb8276e922bc05f3 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 2 Jul 2026 04:53:41 +0000 Subject: [PATCH 3/3] fix: thread ctx + validate Titan v2 dimensions in ai/embed MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 Claude-Session: https://claude.ai/code/session_01WnKgCRfMvFcZAkoVbr8k79 --- packages/engine/internal/connector/embed.go | 6 ++-- .../engine/internal/connector/embed_test.go | 30 +++++++++++++++++++ .../internal/connector/provider_bedrock.go | 7 +++++ .../connector/provider_bedrock_embed_test.go | 18 +++++++++++ 4 files changed, 58 insertions(+), 3 deletions(-) diff --git a/packages/engine/internal/connector/embed.go b/packages/engine/internal/connector/embed.go index 298244d..4950169 100644 --- a/packages/engine/internal/connector/embed.go +++ b/packages/engine/internal/connector/embed.go @@ -48,7 +48,7 @@ func (c *EmbeddingConnector) Execute(ctx context.Context, params map[string]any) return nil, fmt.Errorf("ai/embed: %w", err) } - provider, err := c.getProvider(providerName, params) + provider, err := c.getProvider(ctx, providerName, params) if err != nil { return nil, fmt.Errorf("ai/embed: %w", err) } @@ -80,7 +80,7 @@ func (c *EmbeddingConnector) Execute(ctx context.Context, params map[string]any) } // getProvider returns the EmbeddingProvider for the given provider name. -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) { switch name { case "openai": baseURL := "https://api.openai.com/v1" @@ -102,7 +102,7 @@ func (c *EmbeddingConnector) getProvider(name string, params map[string]any) (Em if configFunc == nil { configFunc = NewAWSConfig } - awsCfg, err := configFunc(context.Background(), cred, defaultRegion) + awsCfg, err := configFunc(ctx, cred, defaultRegion) if err != nil { return nil, fmt.Errorf("[bedrock]: %w", err) } diff --git a/packages/engine/internal/connector/embed_test.go b/packages/engine/internal/connector/embed_test.go index fcaa689..2bbf49b 100644 --- a/packages/engine/internal/connector/embed_test.go +++ b/packages/engine/internal/connector/embed_test.go @@ -3,9 +3,13 @@ package connector import ( "context" "encoding/json" + "errors" "net/http" "net/http/httptest" + "strings" "testing" + + "github.com/aws/aws-sdk-go-v2/aws" ) // embedTestServer returns an httptest server that emits `n` deterministic @@ -144,6 +148,32 @@ func TestEmbeddingConnector_IncompleteResponseFailsFast(t *testing.T) { } } +func TestEmbeddingConnector_BedrockRegionAndErrorWrapping(t *testing.T) { + var gotRegion string + c := &EmbeddingConnector{ + DefaultRegion: "us-west-2", + AWSConfigFunc: func(ctx context.Context, cred map[string]string, defaultRegion string) (aws.Config, error) { + gotRegion = defaultRegion + return aws.Config{}, errors.New("boom") + }, + } + _, err := c.Execute(context.Background(), map[string]any{ + "model": "amazon.titan-embed-text-v2:0", + "input": "x", + "provider": "bedrock", + "region": "eu-central-1", // params region overrides DefaultRegion + }) + if err == nil { + t.Fatal("expected error from stubbed AWS config") + } + if !strings.Contains(err.Error(), "[bedrock]") { + t.Errorf("error not wrapped with [bedrock]: %v", err) + } + if gotRegion != "eu-central-1" { + t.Errorf("resolved region = %q, want eu-central-1 (params override)", gotRegion) + } +} + func TestEmbeddingConnector_UnknownProvider(t *testing.T) { c := &EmbeddingConnector{} _, err := c.Execute(context.Background(), map[string]any{ diff --git a/packages/engine/internal/connector/provider_bedrock.go b/packages/engine/internal/connector/provider_bedrock.go index 2fabe45..039674a 100644 --- a/packages/engine/internal/connector/provider_bedrock.go +++ b/packages/engine/internal/connector/provider_bedrock.go @@ -265,6 +265,13 @@ func (p *BedrockEmbeddingProvider) Embeddings(ctx context.Context, req *Embeddin // other than inputText, so ignore it there rather than failing a request // that reuses a shared embedding config. supportsDimensions := strings.Contains(req.Model, "titan-embed-text-v2") + if req.Dimensions > 0 && supportsDimensions { + switch req.Dimensions { + case 256, 512, 1024: + default: + return nil, fmt.Errorf("bedrock: %s supports dimensions 256, 512, or 1024, got %d", req.Model, req.Dimensions) + } + } for i, text := range req.Inputs { body := titanEmbedRequest{InputText: text} diff --git a/packages/engine/internal/connector/provider_bedrock_embed_test.go b/packages/engine/internal/connector/provider_bedrock_embed_test.go index d1ad9ce..da48c74 100644 --- a/packages/engine/internal/connector/provider_bedrock_embed_test.go +++ b/packages/engine/internal/connector/provider_bedrock_embed_test.go @@ -94,6 +94,24 @@ func TestBedrockEmbeddingProvider_V1IgnoresDimensions(t *testing.T) { } } +func TestBedrockEmbeddingProvider_V2InvalidDimensions(t *testing.T) { + mock := &mockBedrockInvokeClient{ + InvokeFunc: func(ctx context.Context, input *bedrockruntime.InvokeModelInput, opts ...func(*bedrockruntime.Options)) (*bedrockruntime.InvokeModelOutput, error) { + t.Fatal("InvokeModel should not be called for invalid dimensions") + return nil, nil + }, + } + p := &BedrockEmbeddingProvider{Client: mock} + _, err := p.Embeddings(context.Background(), &EmbeddingRequest{ + Model: "amazon.titan-embed-text-v2:0", + Inputs: []string{"x"}, + Dimensions: 999, // only 256/512/1024 are valid + }) + if err == nil { + t.Error("expected error for unsupported Titan v2 dimensions") + } +} + func TestBedrockEmbeddingProvider_UnsupportedModel(t *testing.T) { p := &BedrockEmbeddingProvider{Client: &mockBedrockInvokeClient{}} _, err := p.Embeddings(context.Background(), &EmbeddingRequest{