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

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

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

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")
+	}
+}

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

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:**