feat: add ai/embed connector for RAG (#167)

Author: michaelmcnees

.changeset/ai-embed-connector.md

@@ -0,0 +1,5 @@
+---
+"@mantle/engine": minor
+---
+
+Add an `ai/embed` connector for text embeddings, enabling retrieval-augmented generation (RAG). It mirrors `ai/completion`'s provider selection, base-URL allowlist, metrics, and token-budget accounting, and supports OpenAI and OpenAI-compatible endpoints (Azure OpenAI, local servers) via `base_url`. Output includes a pgvector text literal (`output.vector`) so embeddings drop straight into a `postgres/query` arg cast with `::vector`. Ships with a pgvector RAG example (`rag-ingest.yaml`, `rag-ask.yaml`, `rag-kb-schema.sql`) and a guide. Bedrock embeddings and native `kb/*` / chunking are planned follow-ups (#153).

packages/engine/internal/cli/serve.go

@@ -71,6 +71,16 @@ 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
+			// applying it here would block embedding models.
+			if embConn, err := eng.Registry.Get("ai/embed"); err == nil {
+				if emb, ok := embConn.(*connector.EmbeddingConnector); ok {
+					emb.AllowedBaseURLs = cfg.Engine.AllowedBaseURLs
+				}
+			}
+
 			// Wire up Postgres-backed audit emitter with auth context enrichment.
 			auditor := &audit.PostgresEmitter{
 				DB:                  database,

packages/engine/internal/connector/connector.go

@@ -27,6 +27,7 @@ func NewRegistry() *Registry {
 	}
 	r.Register("http/request", &HTTPConnector{})
 	r.Register("ai/completion", &AIConnector{})
+	r.Register("ai/embed", &EmbeddingConnector{})
 	r.Register("slack/send", &SlackSendConnector{})
 	r.Register("slack/history", &SlackHistoryConnector{})
 	r.Register("postgres/query", &PostgresQueryConnector{})

packages/engine/internal/connector/embed.go

@@ -0,0 +1,186 @@
+package connector
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"strconv"
+	"strings"
+	"time"
+
+	"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).
+type EmbeddingConnector struct {
+	Client          *http.Client
+	AllowedBaseURLs []string
+	AllowedModels   []string // empty = all models allowed
+}
+
+// Execute embeds one or more input strings and returns the vectors, including
+// pgvector text literals for direct use in a postgres/query arg (e.g.
+// INSERT ... VALUES ($1::vector)).
+func (c *EmbeddingConnector) Execute(ctx context.Context, params map[string]any) (map[string]any, error) {
+	providerName, _ := params["provider"].(string)
+	if providerName == "" {
+		providerName = "openai"
+	}
+
+	model, _ := params["model"].(string)
+	if model == "" {
+		return nil, fmt.Errorf("ai/embed: model is required")
+	}
+	if len(c.AllowedModels) > 0 && !stringInSlice(model, c.AllowedModels) {
+		return nil, fmt.Errorf("ai/embed: model %q not in allowed list", model)
+	}
+
+	inputs, err := extractEmbedInputs(params)
+	if err != nil {
+		return nil, fmt.Errorf("ai/embed: %w", err)
+	}
+
+	provider, err := c.getProvider(providerName, params)
+	if err != nil {
+		return nil, fmt.Errorf("ai/embed: %w", err)
+	}
+
+	req := &EmbeddingRequest{Model: model, Inputs: inputs}
+	if cred, ok := params["_credential"].(map[string]string); ok {
+		req.Credential = cred
+	}
+	if d, ok := extractInt(params["dimensions"]); ok && d > 0 {
+		req.Dimensions = d
+	}
+
+	workflow, _ := params["_workflow"].(string)
+	step, _ := params["_step"].(string)
+
+	start := time.Now()
+	resp, err := provider.Embeddings(ctx, req)
+	duration := time.Since(start).Seconds()
+
+	metrics.AIRequestDuration.WithLabelValues(workflow, step, model, providerName).Observe(duration)
+	if err != nil {
+		metrics.AIRequestsTotal.WithLabelValues(workflow, step, model, providerName, "error").Inc()
+		return nil, fmt.Errorf("ai/embed: %w", err)
+	}
+	metrics.AITokensTotal.WithLabelValues(workflow, step, model, providerName, "prompt").Add(float64(resp.Usage.PromptTokens))
+	metrics.AIRequestsTotal.WithLabelValues(workflow, step, model, providerName, "success").Inc()
+
+	return embeddingResponseToOutput(resp), nil
+}
+
+// getProvider returns the EmbeddingProvider for the given provider name.
+func (c *EmbeddingConnector) getProvider(name string, params map[string]any) (EmbeddingProvider, error) {
+	switch name {
+	case "openai":
+		baseURL := "https://api.openai.com/v1"
+		if u, ok := params["base_url"].(string); ok && u != "" {
+			baseURL = u
+		}
+		if len(c.AllowedBaseURLs) > 0 && !stringInSlice(baseURL, c.AllowedBaseURLs) {
+			return nil, fmt.Errorf("base_url %q not in allowed list", baseURL)
+		}
+		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)")
+	default:
+		return nil, fmt.Errorf("unknown provider %q (available: openai)", name)
+	}
+}
+
+// extractEmbedInputs reads the `input` param, accepting a single string or an
+// array of strings.
+func extractEmbedInputs(params map[string]any) ([]string, error) {
+	raw, ok := params["input"]
+	if !ok {
+		return nil, fmt.Errorf("input is required")
+	}
+	switch v := raw.(type) {
+	case string:
+		if v == "" {
+			return nil, fmt.Errorf("input must not be empty")
+		}
+		return []string{v}, nil
+	case []string:
+		if len(v) == 0 {
+			return nil, fmt.Errorf("input must not be empty")
+		}
+		return v, nil
+	case []any:
+		out := make([]string, 0, len(v))
+		for i, item := range v {
+			s, ok := item.(string)
+			if !ok {
+				return nil, fmt.Errorf("input[%d] must be a string, got %T", i, item)
+			}
+			out = append(out, s)
+		}
+		if len(out) == 0 {
+			return nil, fmt.Errorf("input must not be empty")
+		}
+		return out, nil
+	default:
+		return nil, fmt.Errorf("input must be a string or array of strings, got %T", raw)
+	}
+}
+
+// embeddingResponseToOutput converts an EmbeddingResponse to the connector's
+// output map. It always includes the full `embeddings`/`vectors` arrays, and
+// for the common single-input case surfaces `embedding`/`vector`/`dimensions`.
+func embeddingResponseToOutput(resp *EmbeddingResponse) map[string]any {
+	vectors := make([]string, len(resp.Embeddings))
+	for i, e := range resp.Embeddings {
+		vectors[i] = formatPGVector(e)
+	}
+
+	out := map[string]any{
+		"model":      resp.Model,
+		"embeddings": resp.Embeddings,
+		"vectors":    vectors,
+		"count":      len(resp.Embeddings),
+		"usage": map[string]any{
+			"prompt_tokens": resp.Usage.PromptTokens,
+			"total_tokens":  resp.Usage.TotalTokens,
+		},
+	}
+	if len(resp.Embeddings) > 0 {
+		out["embedding"] = resp.Embeddings[0]
+		out["vector"] = vectors[0]
+		out["dimensions"] = len(resp.Embeddings[0])
+	}
+	return out
+}
+
+// formatPGVector renders a float slice as a pgvector text literal: "[1,2,3]".
+// This can be passed straight into a postgres/query arg and cast with ::vector.
+func formatPGVector(v []float64) string {
+	var b strings.Builder
+	b.WriteByte('[')
+	for i, f := range v {
+		if i > 0 {
+			b.WriteByte(',')
+		}
+		b.WriteString(strconv.FormatFloat(f, 'g', -1, 64))
+	}
+	b.WriteByte(']')
+	return b.String()
+}
+
+func stringInSlice(s string, list []string) bool {
+	for _, v := range list {
+		if v == s {
+			return true
+		}
+	}
+	return false
+}

packages/engine/internal/connector/embed_test.go

@@ -0,0 +1,157 @@
+package connector
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+)
+
+// embedTestServer returns an httptest server that emits `n` deterministic
+// embeddings echoing the request order via the index field.
+func embedTestServer(t *testing.T, dims int) *httptest.Server {
+	t.Helper()
+	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path != "/embeddings" {
+			t.Errorf("path = %s, want /embeddings", r.URL.Path)
+		}
+		var body struct {
+			Model string   `json:"model"`
+			Input []string `json:"input"`
+		}
+		_ = json.NewDecoder(r.Body).Decode(&body)
+
+		data := make([]map[string]any, len(body.Input))
+		for i := range body.Input {
+			vec := make([]float64, dims)
+			for j := range vec {
+				vec[j] = float64(i) + float64(j)/10.0
+			}
+			// Return out of order to exercise index reassembly.
+			data[len(body.Input)-1-i] = map[string]any{"index": i, "embedding": vec}
+		}
+		w.Header().Set("Content-Type", "application/json")
+		_ = json.NewEncoder(w).Encode(map[string]any{
+			"model": body.Model,
+			"data":  data,
+			"usage": map[string]any{"prompt_tokens": 7, "total_tokens": 7},
+		})
+	}))
+}
+
+func TestEmbeddingConnector_SingleInput(t *testing.T) {
+	server := embedTestServer(t, 3)
+	defer server.Close()
+
+	c := &EmbeddingConnector{Client: server.Client()}
+	out, err := c.Execute(context.Background(), map[string]any{
+		"model":       "text-embedding-3-small",
+		"input":       "hello world",
+		"base_url":    server.URL,
+		"_credential": map[string]string{"api_key": "sk-test"},
+	})
+	if err != nil {
+		t.Fatalf("Execute() error: %v", err)
+	}
+	if got := out["count"].(int); got != 1 {
+		t.Errorf("count = %d, want 1", got)
+	}
+	if got := out["dimensions"].(int); got != 3 {
+		t.Errorf("dimensions = %d, want 3", got)
+	}
+	// First input's vector is [0, 0.1, 0.2] → pgvector literal.
+	if got := out["vector"].(string); got != "[0,0.1,0.2]" {
+		t.Errorf("vector = %q, want %q", got, "[0,0.1,0.2]")
+	}
+	usage := out["usage"].(map[string]any)
+	if usage["total_tokens"].(int) != 7 {
+		t.Errorf("total_tokens = %v, want 7", usage["total_tokens"])
+	}
+}
+
+func TestEmbeddingConnector_MultiInputPreservesOrder(t *testing.T) {
+	server := embedTestServer(t, 2)
+	defer server.Close()
+
+	c := &EmbeddingConnector{Client: server.Client()}
+	out, err := c.Execute(context.Background(), map[string]any{
+		"model":    "text-embedding-3-small",
+		"input":    []any{"a", "b", "c"},
+		"base_url": server.URL,
+	})
+	if err != nil {
+		t.Fatalf("Execute() error: %v", err)
+	}
+	if out["count"].(int) != 3 {
+		t.Fatalf("count = %d, want 3", out["count"])
+	}
+	vectors := out["vectors"].([]string)
+	// Index reassembly: input i → vector [i, i+0.1].
+	want := []string{"[0,0.1]", "[1,1.1]", "[2,2.1]"}
+	for i, w := range want {
+		if vectors[i] != w {
+			t.Errorf("vectors[%d] = %q, want %q", i, vectors[i], w)
+		}
+	}
+}
+
+func TestEmbeddingConnector_Validation(t *testing.T) {
+	c := &EmbeddingConnector{}
+	if _, err := c.Execute(context.Background(), map[string]any{"input": "x"}); err == nil {
+		t.Error("expected error when model is missing")
+	}
+	if _, err := c.Execute(context.Background(), map[string]any{"model": "m"}); err == nil {
+		t.Error("expected error when input is missing")
+	}
+	if _, err := c.Execute(context.Background(), map[string]any{"model": "m", "input": ""}); err == nil {
+		t.Error("expected error when input is empty")
+	}
+}
+
+func TestEmbeddingConnector_ModelAllowlist(t *testing.T) {
+	c := &EmbeddingConnector{AllowedModels: []string{"text-embedding-3-large"}}
+	_, err := c.Execute(context.Background(), map[string]any{
+		"model": "text-embedding-3-small",
+		"input": "x",
+	})
+	if err == nil {
+		t.Error("expected error for disallowed model")
+	}
+}
+
+func TestEmbeddingConnector_IncompleteResponseFailsFast(t *testing.T) {
+	// Server returns only one embedding for two inputs — must error rather
+	// than silently returning a misaligned/short vectors list.
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		_ = json.NewEncoder(w).Encode(map[string]any{
+			"model": "text-embedding-3-small",
+			"data":  []map[string]any{{"index": 0, "embedding": []float64{0.1, 0.2}}},
+			"usage": map[string]any{"prompt_tokens": 3, "total_tokens": 3},
+		})
+	}))
+	defer server.Close()
+
+	c := &EmbeddingConnector{Client: server.Client()}
+	_, err := c.Execute(context.Background(), map[string]any{
+		"model":    "text-embedding-3-small",
+		"input":    []any{"a", "b"},
+		"base_url": server.URL,
+	})
+	if err == nil {
+		t.Error("expected error when the provider returns fewer embeddings than inputs")
+	}
+}
+
+func TestEmbeddingConnector_BedrockUnsupported(t *testing.T) {
+	c := &EmbeddingConnector{}
+	_, err := c.Execute(context.Background(), map[string]any{
+		"model":    "amazon.titan-embed-text-v2:0",
+		"input":    "x",
+		"provider": "bedrock",
+	})
+	if err == nil {
+		t.Error("expected error for unsupported bedrock provider")
+	}
+}