diff --git a/.changeset/ai-embed-connector.md b/.changeset/ai-embed-connector.md new file mode 100644 index 0000000..601a214 --- /dev/null +++ b/.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). diff --git a/packages/engine/internal/cli/serve.go b/packages/engine/internal/cli/serve.go index ff1c939..c22c839 100644 --- a/packages/engine/internal/cli/serve.go +++ b/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, diff --git a/packages/engine/internal/connector/connector.go b/packages/engine/internal/connector/connector.go index 75a69c9..b725ce1 100644 --- a/packages/engine/internal/connector/connector.go +++ b/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{}) diff --git a/packages/engine/internal/connector/embed.go b/packages/engine/internal/connector/embed.go new file mode 100644 index 0000000..f2786b1 --- /dev/null +++ b/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 +} diff --git a/packages/engine/internal/connector/embed_test.go b/packages/engine/internal/connector/embed_test.go new file mode 100644 index 0000000..cb9b2b7 --- /dev/null +++ b/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") + } +} diff --git a/packages/engine/internal/connector/provider.go b/packages/engine/internal/connector/provider.go index 4bd366e..9317c7b 100644 --- a/packages/engine/internal/connector/provider.go +++ b/packages/engine/internal/connector/provider.go @@ -10,6 +10,27 @@ type LLMProvider interface { ChatCompletion(ctx context.Context, req *ChatRequest) (*ChatResponse, error) } +// EmbeddingProvider implements a specific provider's text-embeddings API. +type EmbeddingProvider interface { + Embeddings(ctx context.Context, req *EmbeddingRequest) (*EmbeddingResponse, error) +} + +// EmbeddingRequest is the provider-agnostic embeddings request format. +type EmbeddingRequest struct { + Model string + Inputs []string + Dimensions int // 0 = provider default + Credential map[string]string +} + +// EmbeddingResponse is the provider-agnostic embeddings response format. +// Embeddings are returned in the same order as EmbeddingRequest.Inputs. +type EmbeddingResponse struct { + Embeddings [][]float64 + Model string + Usage ChatUsage // PromptTokens / TotalTokens are meaningful for embeddings +} + // ChatRequest is the provider-agnostic request format. type ChatRequest struct { Model string diff --git a/packages/engine/internal/connector/provider_openai.go b/packages/engine/internal/connector/provider_openai.go index 6059eb0..ae3a69e 100644 --- a/packages/engine/internal/connector/provider_openai.go +++ b/packages/engine/internal/connector/provider_openai.go @@ -85,19 +85,7 @@ func (p *OpenAIProvider) ChatCompletion(ctx context.Context, req *ChatRequest) ( httpReq.Header.Set("Content-Type", "application/json") // Apply credential-based auth. - if cred := req.Credential; cred != nil { - switch { - case cred["api_key"] != "": - httpReq.Header.Set("Authorization", "Bearer "+cred["api_key"]) - case cred["token"] != "": - httpReq.Header.Set("Authorization", "Bearer "+cred["token"]) - case cred["key"] != "": - httpReq.Header.Set("Authorization", "Bearer "+cred["key"]) - } - if orgID := cred["org_id"]; orgID != "" { - httpReq.Header.Set("OpenAI-Organization", orgID) - } - } + applyOpenAICredential(httpReq, req.Credential) client := p.Client if client == nil { @@ -153,3 +141,117 @@ func (p *OpenAIProvider) ChatCompletion(ctx context.Context, req *ChatRequest) ( return chatResp, nil } + +// applyOpenAICredential sets the Authorization (and optional organization) +// headers from a resolved credential map. Shared by chat and embeddings. +func applyOpenAICredential(httpReq *http.Request, cred map[string]string) { + if cred == nil { + return + } + switch { + case cred["api_key"] != "": + httpReq.Header.Set("Authorization", "Bearer "+cred["api_key"]) + case cred["token"] != "": + httpReq.Header.Set("Authorization", "Bearer "+cred["token"]) + case cred["key"] != "": + httpReq.Header.Set("Authorization", "Bearer "+cred["key"]) + } + if orgID := cred["org_id"]; orgID != "" { + httpReq.Header.Set("OpenAI-Organization", orgID) + } +} + +// embeddingsAPIResponse is the OpenAI /embeddings response envelope. +type embeddingsAPIResponse struct { + Model string `json:"model"` + Data []struct { + Index int `json:"index"` + Embedding []float64 `json:"embedding"` + } `json:"data"` + Usage struct { + PromptTokens int `json:"prompt_tokens"` + TotalTokens int `json:"total_tokens"` + } `json:"usage"` +} + +// Embeddings calls the OpenAI-compatible /embeddings endpoint. Works with +// OpenAI, Azure OpenAI, and any OpenAI-compatible server via BaseURL. +func (p *OpenAIProvider) Embeddings(ctx context.Context, req *EmbeddingRequest) (*EmbeddingResponse, error) { + reqBody := map[string]any{ + "model": req.Model, + "input": req.Inputs, + } + if req.Dimensions > 0 { + reqBody["dimensions"] = req.Dimensions + } + + reqJSON, err := json.Marshal(reqBody) + if err != nil { + return nil, fmt.Errorf("openai: marshaling embeddings request: %w", err) + } + + httpReq, err := http.NewRequestWithContext(ctx, "POST", p.BaseURL+"/embeddings", bytes.NewReader(reqJSON)) + if err != nil { + return nil, fmt.Errorf("openai: creating embeddings request: %w", err) + } + httpReq.Header.Set("Content-Type", "application/json") + applyOpenAICredential(httpReq, req.Credential) + + client := p.Client + if client == nil { + client = &http.Client{Timeout: 120 * time.Second} + } + + resp, err := client.Do(httpReq) + if err != nil { + return nil, fmt.Errorf("openai: %w", err) + } + defer resp.Body.Close() + + body, err := io.ReadAll(io.LimitReader(resp.Body, DefaultMaxResponseBytes)) + if err != nil { + return nil, fmt.Errorf("openai: reading embeddings response: %w", err) + } + + if resp.StatusCode != 200 { + slog.Warn("OpenAI embeddings API error", "status", resp.StatusCode, "body", truncate(string(body), 500)) + if resp.StatusCode == 429 { + return nil, &RetryableError{Err: fmt.Errorf("openai: rate limited (429)")} + } + return nil, fmt.Errorf("openai: embeddings API returned status %d", resp.StatusCode) + } + + var apiResp embeddingsAPIResponse + if err := json.Unmarshal(body, &apiResp); err != nil { + return nil, fmt.Errorf("openai: parsing embeddings response: %w", err) + } + // Reassemble in request order — the API tags each item with its index. + // Size by the request so a short or duplicate-indexed response fails fast + // rather than silently misaligning embeddings with their inputs. + out := make([][]float64, len(req.Inputs)) + seen := make([]bool, len(req.Inputs)) + for _, d := range apiResp.Data { + if d.Index < 0 || d.Index >= len(req.Inputs) { + return nil, fmt.Errorf("openai: embedding index %d out of range for %d inputs", d.Index, len(req.Inputs)) + } + if seen[d.Index] { + return nil, fmt.Errorf("openai: duplicate embedding index %d", d.Index) + } + seen[d.Index] = true + out[d.Index] = d.Embedding + } + for i, ok := range seen { + if !ok { + return nil, fmt.Errorf("openai: missing embedding for input %d (got %d of %d)", i, len(apiResp.Data), len(req.Inputs)) + } + } + + return &EmbeddingResponse{ + Embeddings: out, + Model: apiResp.Model, + Usage: ChatUsage{ + PromptTokens: apiResp.Usage.PromptTokens, + TotalTokens: apiResp.Usage.TotalTokens, + }, + }, nil +} diff --git a/packages/site/src/content/docs/rag-guide.md b/packages/site/src/content/docs/rag-guide.md new file mode 100644 index 0000000..71650c8 --- /dev/null +++ b/packages/site/src/content/docs/rag-guide.md @@ -0,0 +1,97 @@ +--- +title: Retrieval-augmented generation (RAG) +--- + +Build a semantic knowledge base and answer questions grounded in it, using the +`ai/embed` connector, a pgvector table, and `ai/completion`. Two example +workflows tie it together: + +- **`rag-ingest`** — embed a document and store it in the vector store. +- **`rag-ask`** — embed a question, find the nearest documents by cosine + distance, and have an LLM answer grounded in them (with citations). + +This is semantic retrieval — matches are by meaning, not keywords. For a simpler +keyword-only variant with no extensions, see the +[office-assistant MVP](/docs/office-assistant-mvp), which uses Postgres +full-text search. + +## Prerequisites + +- Credentials: an `openai` credential (referenced as `openai`) for `ai/embed` + and `ai/completion`, and a `postgres` credential (referenced as `kb-db`) for + the vector store. +- A Postgres database with the [pgvector](https://github.com/pgvector/pgvector) + extension. Apply the schema from + [`rag-kb-schema.sql`](https://github.com/dvflw/mantle/blob/main/packages/site/src/content/examples/rag-kb-schema.sql): + +```bash +psql "$KB_DATABASE_URL" -f packages/site/src/content/examples/rag-kb-schema.sql +``` + +Apply both workflows: + +```bash +mantle apply packages/site/src/content/examples/rag-ingest.yaml +mantle apply packages/site/src/content/examples/rag-ask.yaml +``` + +## How it works + +`ai/embed` returns each embedding both as a float array (`output.embedding`) and +as a **pgvector text literal** (`output.vector`, e.g. `"[0.1,0.2,...]"`). The +literal binds straight into a `::vector` column, so storing and querying vectors +needs no special encoding: + +```yaml +# ingest +args: + - "{{ inputs.content }}" + - "{{ steps['embed'].output.vector }}" # → $2::vector +``` + +```sql +-- ask: nearest neighbours by cosine distance (<=> is pgvector's operator) +SELECT content, 1 - (embedding <=> $1::vector) AS similarity +FROM kb_documents +ORDER BY embedding <=> $1::vector +LIMIT 5 +``` + +The retrieved rows are passed to `ai/completion` as JSON context, and the model +answers using only those passages. + +## Ingesting documents + +```bash +mantle run rag-ingest \ + --input title="Client C integration" \ + --input source="confluence/12345" \ + --input content="Team A connects to Client C via the X integration; Team B uses the direct API." +``` + +For a long document, split it into chunks and run `rag-ingest` once per chunk. +Re-ingesting identical content is a no-op (the schema dedupes on `md5(content)`). + +## Asking questions + +```bash +mantle run rag-ask \ + --input question="Who works with Client C and how?" \ + --output json +``` + +The `answer` step's `output.text` is the grounded response (the CLI prints step +outputs with `--output json` or `-v`). + +## Notes and limitations + +- **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. +- **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 + [#153](https://github.com/dvflw/mantle/issues/153). diff --git a/packages/site/src/content/docs/workflow-reference/connectors.md b/packages/site/src/content/docs/workflow-reference/connectors.md index 84f3a3c..a95ed9f 100644 --- a/packages/site/src/content/docs/workflow-reference/connectors.md +++ b/packages/site/src/content/docs/workflow-reference/connectors.md @@ -141,6 +141,56 @@ When running on AWS infrastructure with an IAM role attached (IRSA, instance pro **Authentication:** The AI connector reads the credential's `api_key` field (or `token` or `key` as fallbacks) and sends it as a Bearer token. If the credential includes an `org_id` field, it is sent as the `OpenAI-Organization` header. See the [Secrets Guide](/docs/secrets-guide) for how to create an `openai`-type credential. +## ai/embed + +Turns text into vector embeddings via an OpenAI-compatible embeddings API. Use it to build a semantic knowledge base (RAG) — see the [RAG guide](/docs/rag-guide). Like `ai/completion`, it participates in token-budget accounting. + +**Params:** + +| 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`). | +| `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. | + +**Output:** + +| Field | Type | Description | +|---|---|---| +| `embedding` | list | The first input's embedding as a float array. Present when there is at least one input. | +| `vector` | string | The first input's embedding as a pgvector text literal (`"[0.1,0.2,...]"`), ready to bind into a `::vector` column. | +| `embeddings` | list | All embeddings, one float array per input, in input order. | +| `vectors` | list | All embeddings as pgvector text literals, in input order. | +| `count` | number | Number of embeddings returned. | +| `dimensions` | number | Dimension of the returned vectors. | +| `usage.prompt_tokens` | number | Tokens counted for the input. | +| `usage.total_tokens` | number | Total tokens used. | + +**Example — embed and store in pgvector:** + +```yaml +- name: embed + action: ai/embed + credential: my-openai + params: + model: text-embedding-3-small + input: "{{ inputs.content }}" + +- name: store + action: postgres/query + credential: kb-db + depends_on: [embed] + params: + query: "INSERT INTO kb_documents (content, embedding) VALUES ($1, $2::vector)" + args: + - "{{ inputs.content }}" + - "{{ steps['embed'].output.vector }}" +``` + +**Authentication:** same as `ai/completion` — the credential's `api_key` (or `token`/`key`) is sent as a Bearer token. + ## slack/send Sends a message to a Slack channel via the [chat.postMessage](https://api.slack.com/methods/chat.postMessage) API. Requires a credential with a Slack Bot User OAuth Token. diff --git a/packages/site/src/content/examples/rag-ask.yaml b/packages/site/src/content/examples/rag-ask.yaml new file mode 100644 index 0000000..f759703 --- /dev/null +++ b/packages/site/src/content/examples/rag-ask.yaml @@ -0,0 +1,60 @@ +name: rag-ask +description: > + Retrieval-augmented generation over the pgvector knowledge base: embed the + question, find the nearest documents by cosine distance, and have an LLM answer + grounded in them. Pair with rag-ingest.yaml. Read the answer from the CLI + output (mantle run rag-ask --input question="..." --output json). + +inputs: + question: + type: string + description: The question to answer from the knowledge base + +steps: + # 1. Embed the question with the same model used at ingest time. + - name: embed-question + action: ai/embed + credential: openai + timeout: "30s" + params: + model: text-embedding-3-small + input: "{{ inputs.question }}" + + # 2. Nearest-neighbour search by cosine distance (<=> is pgvector's distance + # operator; similarity = 1 - distance). + - name: search + action: postgres/query + credential: kb-db + depends_on: + - embed-question + timeout: "15s" + params: + query: > + SELECT title, + source, + content, + 1 - (embedding <=> $1::vector) AS similarity + FROM kb_documents + ORDER BY embedding <=> $1::vector + LIMIT 5 + args: + - "{{ steps['embed-question'].output.vector }}" + + # 3. Answer grounded strictly in the retrieved passages. + - name: answer + action: ai/completion + credential: openai + depends_on: + - search + timeout: "60s" + params: + model: gpt-4o + system_prompt: > + Answer using ONLY the provided context passages. Cite the titles you + draw from. If the context does not contain the answer, say so plainly + rather than guessing. + prompt: > + Question: {{ inputs.question }} + + Context passages (JSON, most similar first): + {{ jsonEncode(steps['search'].output.rows) }} diff --git a/packages/site/src/content/examples/rag-ingest.yaml b/packages/site/src/content/examples/rag-ingest.yaml new file mode 100644 index 0000000..78cf256 --- /dev/null +++ b/packages/site/src/content/examples/rag-ingest.yaml @@ -0,0 +1,49 @@ +name: rag-ingest +description: > + Embed a document with ai/embed and store it in a pgvector knowledge base for + semantic retrieval. Pair with rag-ask.yaml. Requires the schema in + rag-kb-schema.sql. For long documents, split into chunks and run once per + chunk (native chunking is a planned follow-up). + +inputs: + title: + type: string + description: "Document title (optional)" + default: "" + source: + type: string + description: "Where the document came from, e.g. a URL or file path (optional)" + default: "" + content: + type: string + description: The document text to embed and store + +steps: + # 1. Produce an embedding for the document. output.vector is a pgvector text + # literal ("[0.1,0.2,...]") ready to bind into a ::vector column. + - name: embed + action: ai/embed + credential: openai + timeout: "30s" + params: + model: text-embedding-3-small + input: "{{ inputs.content }}" + + # 2. Store the document + embedding. ON CONFLICT makes re-ingesting identical + # content a no-op (dedupe_key is md5(content)). + - name: store + action: postgres/query + credential: kb-db + depends_on: + - embed + timeout: "15s" + params: + query: > + INSERT INTO kb_documents (source, title, content, embedding) + VALUES ($1, $2, $3, $4::vector) + ON CONFLICT (dedupe_key) DO NOTHING + args: + - "{{ inputs.source }}" + - "{{ inputs.title }}" + - "{{ inputs.content }}" + - "{{ steps['embed'].output.vector }}" diff --git a/packages/site/src/content/examples/rag-kb-schema.sql b/packages/site/src/content/examples/rag-kb-schema.sql new file mode 100644 index 0000000..f6851e6 --- /dev/null +++ b/packages/site/src/content/examples/rag-kb-schema.sql @@ -0,0 +1,31 @@ +-- Vector knowledge base for the RAG example (rag-ingest.yaml / rag-ask.yaml). +-- +-- Semantic retrieval with pgvector: documents are stored alongside an embedding +-- produced by the ai/embed connector, and queried by cosine distance. Requires +-- the pgvector extension (https://github.com/pgvector/pgvector). Point a Mantle +-- credential at this database and reference it as `credential: kb-db`. +-- +-- The embedding dimension must match the model: text-embedding-3-small = 1536, +-- text-embedding-3-large = 3072. Adjust vector(1536) if you change models. + +CREATE EXTENSION IF NOT EXISTS vector; + +CREATE TABLE IF NOT EXISTS kb_documents ( + id BIGINT GENERATED ALWAYS AS IDENTITY PRIMARY KEY, + source TEXT, + title TEXT, + content TEXT NOT NULL, + embedding vector(1536) NOT NULL, + -- Idempotency: re-ingesting identical content is a no-op (see the + -- ON CONFLICT clause in rag-ingest.yaml). + dedupe_key TEXT GENERATED ALWAYS AS (md5(content)) STORED, + created_at TIMESTAMPTZ NOT NULL DEFAULT now() +); + +CREATE UNIQUE INDEX IF NOT EXISTS idx_kb_documents_dedupe + ON kb_documents (dedupe_key); + +-- Approximate-nearest-neighbour index for cosine distance (pgvector >= 0.5). +-- For small corpora you can omit this and rely on an exact scan. +CREATE INDEX IF NOT EXISTS idx_kb_documents_embedding + ON kb_documents USING hnsw (embedding vector_cosine_ops);