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..4950169 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 } @@ -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) } } diff --git a/packages/engine/internal/connector/embed_test.go b/packages/engine/internal/connector/embed_test.go index cb9b2b7..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,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") } } diff --git a/packages/engine/internal/connector/provider_bedrock.go b/packages/engine/internal/connector/provider_bedrock.go index 5383255..039674a 100644 --- a/packages/engine/internal/connector/provider_bedrock.go +++ b/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 +} 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..da48c74 --- /dev/null +++ b/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") + } +} 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:**