feat: add Bedrock (Titan) embeddings to ai/embed

.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.

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
 				}
 			}

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
 }
@@ -44,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)
 	}
@@ -76,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"
@@ -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(ctx, 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)
 	}
 }
 

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,14 +148,40 @@ func TestEmbeddingConnector_IncompleteResponseFailsFast(t *testing.T) {
 	}
 }
 
-func TestEmbeddingConnector_BedrockUnsupported(t *testing.T) {
-	c := &EmbeddingConnector{}
+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{
+		"model":    "some-model",
+		"input":    "x",
+		"provider": "not-a-provider",
 	})
 	if err == nil {
-		t.Error("expected error for unsupported bedrock provider")
+		t.Error("expected error for unknown provider")
 	}
 }

packages/engine/internal/connector/provider_bedrock.go

@@ -223,3 +223,90 @@ 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
+
+	// 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")
+	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}
+		if req.Dimensions > 0 && supportsDimensions {
+			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
+}

packages/engine/internal/connector/provider_bedrock_embed_test.go

@@ -0,0 +1,124 @@
+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_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_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{
+		Model:  "cohere.embed-english-v3",
+		Inputs: []string{"x"},
+	})
+	if err == nil {
+		t.Error("expected error for unsupported (non-Titan) Bedrock embedding model")
+	}
+}

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