diff --git a/.changeset/kb-connectors.md b/.changeset/kb-connectors.md new file mode 100644 index 0000000..ed177f4 --- /dev/null +++ b/.changeset/kb-connectors.md @@ -0,0 +1,5 @@ +--- +"@mantle/engine": minor +--- + +Add `kb/upsert` and `kb/query` connectors for retrieval-augmented generation. They are thin pgvector helpers over a Postgres database (the step's `credential`) that compose with `ai/embed`: `kb/upsert` stores document text + embedding (+ optional JSONB metadata), taking `ai/embed`'s pgvector literal directly and supporting single or batch (chunk) inserts with idempotent `ON CONFLICT`; `kb/query` runs distance-ordered nearest-neighbour search (cosine/l2/inner_product) and returns the closest rows. Table and column names are validated as SQL identifiers to prevent injection. The RAG example (`rag-ingest.yaml`, `rag-ask.yaml`, `rag-kb-schema.sql`) and guide now use these connectors. They do not manage schema (you create the pgvector table); native chunking and metadata filtering remain follow-ups on #153. diff --git a/packages/engine/internal/connector/connector.go b/packages/engine/internal/connector/connector.go index b725ce1..bd84cb0 100644 --- a/packages/engine/internal/connector/connector.go +++ b/packages/engine/internal/connector/connector.go @@ -31,6 +31,8 @@ func NewRegistry() *Registry { r.Register("slack/send", &SlackSendConnector{}) r.Register("slack/history", &SlackHistoryConnector{}) r.Register("postgres/query", &PostgresQueryConnector{}) + r.Register("kb/upsert", &KBUpsertConnector{}) + r.Register("kb/query", &KBQueryConnector{}) r.Register("email/send", &EmailSendConnector{}) r.Register("email/receive", &EmailReceiveConnector{}) r.Register("email/move", &EmailMoveConnector{}) diff --git a/packages/engine/internal/connector/kb.go b/packages/engine/internal/connector/kb.go new file mode 100644 index 0000000..4de94f0 --- /dev/null +++ b/packages/engine/internal/connector/kb.go @@ -0,0 +1,351 @@ +package connector + +import ( + "context" + "encoding/json" + "fmt" + "regexp" + "strconv" + "strings" + "time" + + "github.com/jackc/pgx/v5" +) + +// The kb/* connectors are thin pgvector helpers over a Postgres database (the +// step's `credential` is a postgres credential). They compose with ai/embed: +// take the pgvector literal from ai/embed's `output.vector` / `output.vectors` +// and store or query it, hiding the ::vector casts, distance operators, and +// multi-row unnest inserts. They do not manage schema — create the table +// yourself (see the RAG guide / rag-kb-schema.sql). + +// kbIdentRe matches a safe SQL identifier, optionally schema-qualified. Table +// and column names are interpolated into SQL (they cannot be bound as +// parameters), so every identifier is validated against this before use. +var kbIdentRe = regexp.MustCompile(`^[A-Za-z_][A-Za-z0-9_]*(\.[A-Za-z_][A-Za-z0-9_]*)?$`) + +const ( + kbDefaultContentColumn = "content" + kbDefaultEmbeddingColumn = "embedding" + kbDefaultMetadataColumn = "metadata" + kbDefaultTopK = 5 + kbMaxTopK = 1000 + // kbConnectTimeout bounds the Postgres connection handshake so a stalled + // server can't hang a step that has no timeout of its own. It only shortens + // an existing deadline, never extends one. + kbConnectTimeout = 15 * time.Second +) + +// kbIdentParam reads an identifier param, applies a default when absent, and +// validates it to prevent SQL injection through table/column names. +func kbIdentParam(params map[string]any, key, def string) (string, error) { + v, ok := params[key].(string) + if !ok || v == "" { + v = def + } + if v == "" { + return "", fmt.Errorf("%s is required", key) + } + if !kbIdentRe.MatchString(v) { + return "", fmt.Errorf("%s %q is not a valid SQL identifier", key, v) + } + return v, nil +} + +// kbStrings normalizes a "single or many" text param: `singleKey` (a string) or +// `manyKey` (an array of strings). Exactly one must be present. +func kbStrings(params map[string]any, singleKey, manyKey string) ([]string, error) { + single, hasSingle := params[singleKey] + many, hasMany := params[manyKey] + if hasSingle == hasMany { // both or neither + return nil, fmt.Errorf("exactly one of %s or %s is required", singleKey, manyKey) + } + if hasSingle { + s, ok := single.(string) + if !ok || s == "" { + return nil, fmt.Errorf("%s must be a non-empty string", singleKey) + } + return []string{s}, nil + } + switch v := many.(type) { + case []string: + if len(v) == 0 { + return nil, fmt.Errorf("%s must not be empty", manyKey) + } + 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("%s[%d] must be a string, got %T", manyKey, i, item) + } + out = append(out, s) + } + if len(out) == 0 { + return nil, fmt.Errorf("%s must not be empty", manyKey) + } + return out, nil + default: + return nil, fmt.Errorf("%s must be an array of strings, got %T", manyKey, many) + } +} + +// kbMetadata normalizes the optional metadata param into JSON strings aligned +// with the rows. Returns present=false when no metadata param is set. +func kbMetadata(params map[string]any, rows int) (jsons []string, present bool, err error) { + single, hasSingle := params["metadata"] + many, hasMany := params["metadatas"] + if !hasSingle && !hasMany { + return nil, false, nil + } + if hasSingle && hasMany { + return nil, false, fmt.Errorf("set only one of metadata or metadatas") + } + + var objs []any + if hasSingle { + objs = []any{single} + } else { + arr, ok := many.([]any) + if !ok { + return nil, false, fmt.Errorf("metadatas must be an array of objects, got %T", many) + } + objs = arr + } + if len(objs) != rows { + return nil, false, fmt.Errorf("metadata count %d does not match content count %d", len(objs), rows) + } + + out := make([]string, len(objs)) + for i, o := range objs { + b, mErr := json.Marshal(o) + if mErr != nil { + return nil, false, fmt.Errorf("marshaling metadata[%d]: %w", i, mErr) + } + out[i] = string(b) + } + return out, true, nil +} + +// prepareUpsert validates params and builds the INSERT ... SELECT unnest(...) +// statement plus its args. Pure (no DB) so it is fully unit-testable. +func prepareUpsert(params map[string]any) (string, []any, error) { + table, err := kbIdentParam(params, "table", "") + if err != nil { + return "", nil, err + } + contentCol, err := kbIdentParam(params, "content_column", kbDefaultContentColumn) + if err != nil { + return "", nil, err + } + embCol, err := kbIdentParam(params, "embedding_column", kbDefaultEmbeddingColumn) + if err != nil { + return "", nil, err + } + metaCol, err := kbIdentParam(params, "metadata_column", kbDefaultMetadataColumn) + if err != nil { + return "", nil, err + } + + contents, err := kbStrings(params, "content", "contents") + if err != nil { + return "", nil, err + } + vectors, err := kbStrings(params, "vector", "vectors") + if err != nil { + return "", nil, err + } + if len(vectors) != len(contents) { + return "", nil, fmt.Errorf("vector count %d does not match content count %d", len(vectors), len(contents)) + } + + metas, hasMeta, err := kbMetadata(params, len(contents)) + if err != nil { + return "", nil, err + } + + var conflict string + if ct, ok := params["conflict_target"].(string); ok && ct != "" { + // Accept a single column or a comma-separated list (composite key), + // validating each part as an identifier. + parts := strings.Split(ct, ",") + for i, p := range parts { + p = strings.TrimSpace(p) + if !kbIdentRe.MatchString(p) { + return "", nil, fmt.Errorf("conflict_target %q is not a valid SQL identifier", ct) + } + parts[i] = p + } + conflict = fmt.Sprintf(" ON CONFLICT (%s) DO NOTHING", strings.Join(parts, ", ")) + } + + var sb strings.Builder + args := []any{contents, vectors} + if hasMeta { + fmt.Fprintf(&sb, + "INSERT INTO %s (%s, %s, %s) SELECT c, e::vector, m::jsonb FROM unnest($1::text[], $2::text[], $3::text[]) AS u(c, e, m)%s", + table, contentCol, embCol, metaCol, conflict) + args = append(args, metas) + } else { + fmt.Fprintf(&sb, + "INSERT INTO %s (%s, %s) SELECT c, e::vector FROM unnest($1::text[], $2::text[]) AS u(c, e)%s", + table, contentCol, embCol, conflict) + } + return sb.String(), args, nil +} + +// kbDistanceOperator maps a metric name to its pgvector operator. +func kbDistanceOperator(metric string) (string, error) { + switch metric { + case "", "cosine": + return "<=>", nil + case "l2", "euclidean": + return "<->", nil + case "inner_product", "ip": + return "<#>", nil + default: + return "", fmt.Errorf("unknown metric %q (use cosine, l2/euclidean, or inner_product/ip)", metric) + } +} + +// prepareQuery validates params and builds the nearest-neighbour SELECT plus +// its args ($1 = query vector). Pure (no DB) so it is fully unit-testable. +func prepareQuery(params map[string]any) (string, []any, error) { + table, err := kbIdentParam(params, "table", "") + if err != nil { + return "", nil, err + } + embCol, err := kbIdentParam(params, "embedding_column", kbDefaultEmbeddingColumn) + if err != nil { + return "", nil, err + } + + vector, ok := params["vector"].(string) + if !ok || vector == "" { + return "", nil, fmt.Errorf("vector is required (a pgvector literal, e.g. from ai/embed output.vector)") + } + + metric, _ := params["metric"].(string) + op, err := kbDistanceOperator(metric) + if err != nil { + return "", nil, err + } + + topK := kbDefaultTopK + if v, ok := extractInt(params["top_k"]); ok { + if v <= 0 { + return "", nil, fmt.Errorf("top_k must be positive, got %d", v) + } + topK = v + } + if topK > kbMaxTopK { + topK = kbMaxTopK + } + + // Columns to return: default to the content column; callers may request more. + var cols []string + if raw, ok := params["columns"]; ok { + list, lErr := kbStringList(raw) + if lErr != nil { + return "", nil, fmt.Errorf("columns: %w", lErr) + } + for _, c := range list { + if !kbIdentRe.MatchString(c) { + return "", nil, fmt.Errorf("column %q is not a valid SQL identifier", c) + } + } + cols = list + } else { + contentCol, cErr := kbIdentParam(params, "content_column", kbDefaultContentColumn) + if cErr != nil { + return "", nil, cErr + } + cols = []string{contentCol} + } + + selectList := strings.Join(cols, ", ") + sql := fmt.Sprintf( + "SELECT %s, (%s %s $1::vector) AS distance FROM %s ORDER BY %s %s $1::vector LIMIT %s", + selectList, embCol, op, table, embCol, op, strconv.Itoa(topK)) + return sql, []any{vector}, nil +} + +func kbStringList(raw any) ([]string, error) { + switch v := raw.(type) { + case []string: + 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("[%d] must be a string, got %T", i, item) + } + out = append(out, s) + } + return out, nil + default: + return nil, fmt.Errorf("must be an array of strings, got %T", raw) + } +} + +// KBUpsertConnector implements kb/upsert: store document text + embedding +// (and optional JSONB metadata) into a pgvector table, one or many rows. +type KBUpsertConnector struct{} + +func (c *KBUpsertConnector) Execute(ctx context.Context, params map[string]any) (map[string]any, error) { + connURL, err := extractPostgresURL(params) + if err != nil { + return nil, fmt.Errorf("kb/upsert: %w", err) + } + + query, args, err := prepareUpsert(params) + if err != nil { + return nil, fmt.Errorf("kb/upsert: %w", err) + } + + connectCtx, cancel := context.WithTimeout(ctx, kbConnectTimeout) + conn, err := pgx.Connect(connectCtx, connURL) + cancel() + if err != nil { + return nil, fmt.Errorf("kb/upsert: connecting: %w", err) + } + defer conn.Close(ctx) + + out, err := executeExec(ctx, conn, query, args) + if err != nil { + return nil, fmt.Errorf("kb/upsert: %w", err) + } + return out, nil +} + +// KBQueryConnector implements kb/query: nearest-neighbour search over a +// pgvector table for a query embedding, returning the closest rows. +type KBQueryConnector struct{} + +func (c *KBQueryConnector) Execute(ctx context.Context, params map[string]any) (map[string]any, error) { + connURL, err := extractPostgresURL(params) + if err != nil { + return nil, fmt.Errorf("kb/query: %w", err) + } + + query, args, err := prepareQuery(params) + if err != nil { + return nil, fmt.Errorf("kb/query: %w", err) + } + + connectCtx, cancel := context.WithTimeout(ctx, kbConnectTimeout) + conn, err := pgx.Connect(connectCtx, connURL) + cancel() + if err != nil { + return nil, fmt.Errorf("kb/query: connecting: %w", err) + } + defer conn.Close(ctx) + + out, err := executeSelect(ctx, conn, query, args) + if err != nil { + return nil, fmt.Errorf("kb/query: %w", err) + } + return out, nil +} diff --git a/packages/engine/internal/connector/kb_test.go b/packages/engine/internal/connector/kb_test.go new file mode 100644 index 0000000..fb46419 --- /dev/null +++ b/packages/engine/internal/connector/kb_test.go @@ -0,0 +1,196 @@ +package connector + +import ( + "context" + "strings" + "testing" +) + +// A malformed connection string so pgx.Connect fails immediately at parse time +// (no network), letting us assert the connector didn't mutate its params. +const kbBadCredURL = "postgres://bad host" + +// The engine resolves a step's params once and reuses that map across retry +// attempts, so a connector must not mutate it — deleting _credential would make +// attempt 2 fail with a missing credential instead of retrying. +func TestKBUpsert_DoesNotMutateCredential(t *testing.T) { + params := map[string]any{ + "_credential": map[string]string{"url": kbBadCredURL}, + "table": "kb_documents", + "content": "x", + "vector": "[1]", + } + if _, err := (&KBUpsertConnector{}).Execute(context.Background(), params); err == nil { + t.Fatal("expected a connection error") + } + if _, ok := params["_credential"]; !ok { + t.Error("_credential was removed from params (would break retries)") + } +} + +func TestKBQuery_DoesNotMutateCredential(t *testing.T) { + params := map[string]any{ + "_credential": map[string]string{"url": kbBadCredURL}, + "table": "kb_documents", + "vector": "[1]", + } + if _, err := (&KBQueryConnector{}).Execute(context.Background(), params); err == nil { + t.Fatal("expected a connection error") + } + if _, ok := params["_credential"]; !ok { + t.Error("_credential was removed from params (would break retries)") + } +} + +func TestPrepareUpsert_SingleRow(t *testing.T) { + sql, args, err := prepareUpsert(map[string]any{ + "table": "kb_documents", + "content": "hello", + "vector": "[0.1,0.2]", + }) + if err != nil { + t.Fatalf("prepareUpsert error: %v", err) + } + if !strings.Contains(sql, "INSERT INTO kb_documents (content, embedding)") { + t.Errorf("unexpected sql: %s", sql) + } + if !strings.Contains(sql, "e::vector") || !strings.Contains(sql, "unnest($1::text[], $2::text[])") { + t.Errorf("missing unnest/cast: %s", sql) + } + if len(args) != 2 { + t.Fatalf("args len = %d, want 2", len(args)) + } + contents := args[0].([]string) + vectors := args[1].([]string) + if len(contents) != 1 || contents[0] != "hello" || vectors[0] != "[0.1,0.2]" { + t.Errorf("args = %v / %v", contents, vectors) + } +} + +func TestPrepareUpsert_BatchWithMetadataAndConflict(t *testing.T) { + sql, args, err := prepareUpsert(map[string]any{ + "table": "kb_documents", + "contents": []any{"a", "b"}, + "vectors": []any{"[1]", "[2]"}, + "metadatas": []any{map[string]any{"src": "x"}, map[string]any{"src": "y"}}, + "metadata_column": "meta", + "conflict_target": "dedupe_key", + "embedding_column": "emb", + }) + if err != nil { + t.Fatalf("prepareUpsert error: %v", err) + } + if !strings.Contains(sql, "INSERT INTO kb_documents (content, emb, meta)") { + t.Errorf("unexpected columns: %s", sql) + } + if !strings.Contains(sql, "m::jsonb") || !strings.Contains(sql, "unnest($1::text[], $2::text[], $3::text[])") { + t.Errorf("missing metadata handling: %s", sql) + } + if !strings.HasSuffix(sql, "ON CONFLICT (dedupe_key) DO NOTHING") { + t.Errorf("missing on conflict: %s", sql) + } + if len(args) != 3 { + t.Fatalf("args len = %d, want 3", len(args)) + } + metas := args[2].([]string) + if len(metas) != 2 || !strings.Contains(metas[0], `"src":"x"`) { + t.Errorf("metadata args = %v", metas) + } +} + +func TestPrepareUpsert_Errors(t *testing.T) { + cases := []struct { + name string + params map[string]any + }{ + {"missing table", map[string]any{"content": "x", "vector": "[1]"}}, + {"missing content", map[string]any{"table": "t", "vector": "[1]"}}, + {"both content forms", map[string]any{"table": "t", "content": "x", "contents": []any{"y"}, "vector": "[1]"}}, + {"count mismatch", map[string]any{"table": "t", "contents": []any{"a", "b"}, "vectors": []any{"[1]"}}}, + {"metadata mismatch", map[string]any{"table": "t", "content": "a", "vector": "[1]", "metadatas": []any{map[string]any{}, map[string]any{}}}}, + {"both metadata forms", map[string]any{"table": "t", "content": "a", "vector": "[1]", "metadata": map[string]any{}, "metadatas": []any{map[string]any{}}}}, + {"both vector forms", map[string]any{"table": "t", "content": "a", "vector": "[1]", "vectors": []any{"[1]"}}}, + {"sql injection in table", map[string]any{"table": "kb; DROP TABLE users;--", "content": "x", "vector": "[1]"}}, + {"sql injection in column", map[string]any{"table": "t", "content_column": "c)); DROP", "content": "x", "vector": "[1]"}}, + {"sql injection in conflict", map[string]any{"table": "t", "content": "x", "vector": "[1]", "conflict_target": "a) DO UPDATE"}}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if _, _, err := prepareUpsert(tc.params); err == nil { + t.Errorf("expected error for %s", tc.name) + } + }) + } +} + +func TestPrepareQuery_Default(t *testing.T) { + sql, args, err := prepareQuery(map[string]any{ + "table": "kb_documents", + "vector": "[0.1,0.2]", + }) + if err != nil { + t.Fatalf("prepareQuery error: %v", err) + } + // Default: cosine operator, content column, top_k 5. + if !strings.Contains(sql, "SELECT content, (embedding <=> $1::vector) AS distance") { + t.Errorf("unexpected select: %s", sql) + } + if !strings.Contains(sql, "ORDER BY embedding <=> $1::vector LIMIT 5") { + t.Errorf("unexpected order/limit: %s", sql) + } + if len(args) != 1 || args[0].(string) != "[0.1,0.2]" { + t.Errorf("args = %v", args) + } +} + +func TestPrepareQuery_MetricsAndColumns(t *testing.T) { + sql, _, err := prepareQuery(map[string]any{ + "table": "kb_documents", + "vector": "[1]", + "metric": "l2", + "columns": []any{"title", "content", "source"}, + "embedding_column": "emb", + "top_k": 3, + }) + if err != nil { + t.Fatalf("prepareQuery error: %v", err) + } + if !strings.Contains(sql, "SELECT title, content, source, (emb <-> $1::vector) AS distance") { + t.Errorf("unexpected select: %s", sql) + } + if !strings.HasSuffix(sql, "LIMIT 3") { + t.Errorf("unexpected limit: %s", sql) + } +} + +func TestPrepareQuery_TopKCap(t *testing.T) { + sql, _, err := prepareQuery(map[string]any{"table": "t", "vector": "[1]", "top_k": 999999}) + if err != nil { + t.Fatalf("error: %v", err) + } + if !strings.HasSuffix(sql, "LIMIT 1000") { + t.Errorf("top_k not capped: %s", sql) + } +} + +func TestPrepareQuery_Errors(t *testing.T) { + cases := []struct { + name string + params map[string]any + }{ + {"missing table", map[string]any{"vector": "[1]"}}, + {"missing vector", map[string]any{"table": "t"}}, + {"bad metric", map[string]any{"table": "t", "vector": "[1]", "metric": "manhattan"}}, + {"negative top_k", map[string]any{"table": "t", "vector": "[1]", "top_k": -1}}, + {"zero top_k", map[string]any{"table": "t", "vector": "[1]", "top_k": 0}}, + {"sql injection in table", map[string]any{"table": "t; DROP TABLE x", "vector": "[1]"}}, + {"sql injection in column", map[string]any{"table": "t", "vector": "[1]", "columns": []any{"content", "x); DROP"}}}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if _, _, err := prepareQuery(tc.params); err == nil { + t.Errorf("expected error for %s", tc.name) + } + }) + } +} diff --git a/packages/site/src/content/docs/rag-guide.md b/packages/site/src/content/docs/rag-guide.md index b3fec9a..9a46c8a 100644 --- a/packages/site/src/content/docs/rag-guide.md +++ b/packages/site/src/content/docs/rag-guide.md @@ -37,28 +37,41 @@ 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: +`ai/embed` returns each embedding as a **pgvector text literal** +(`output.vector`, e.g. `"[0.1,0.2,...]"`). The `kb/upsert` and `kb/query` +connectors take that literal and handle the pgvector SQL for you — the `::vector` +casts, the cosine-distance operator, `ORDER BY`/`LIMIT`, JSONB metadata, and +multi-row inserts: ```yaml -# ingest -args: - - "{{ inputs.content }}" - - "{{ steps['embed'].output.vector }}" # → $2::vector +# ingest — store content + embedding + metadata +- action: kb/upsert + credential: kb-db + params: + table: kb_documents + content: "{{ inputs.content }}" + vector: "{{ steps['embed'].output.vector }}" + metadata: { title: "{{ inputs.title }}", source: "{{ inputs.source }}" } + conflict_target: dedupe_key # idempotent re-ingest + +# ask — nearest-neighbour search +- action: kb/query + credential: kb-db + params: + table: kb_documents + vector: "{{ steps['embed-question'].output.vector }}" + columns: [content, metadata] + top_k: 5 ``` -```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 -``` +`kb/query` returns the requested columns plus a `distance` field (lower = closer; +for cosine, similarity = `1 - distance`). The retrieved rows go to +`ai/completion` as JSON context, and the model answers using only those passages. -The retrieved rows are passed to `ai/completion` as JSON context, and the model -answers using only those passages. +`kb/*` are thin sugar — they run against a pgvector database you provide (the +step `credential`) and don't manage schema, so you still create the table +yourself. If you'd rather write the SQL directly, `postgres/query` works too +(that's what these connectors generate). ## Ingesting documents @@ -69,8 +82,10 @@ mantle run rag-ingest \ --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)`). +For a long document, split it into chunks and pass arrays to `ai/embed` +(`input`) and `kb/upsert` (`contents` + `vectors`, plus `metadatas`) to store +them all in one step. Re-ingesting identical content is a no-op (the schema +dedupes on `md5(content)`). ## Asking questions @@ -94,7 +109,8 @@ outputs with `--output json` or `-v`). `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 - [#153](https://github.com/dvflw/mantle/issues/153). +- **You provide the vector store.** `kb/*` run against a pgvector database you + create and point a credential at; they don't provision the extension or table. +- **Chunking is manual.** Split long documents yourself and pass arrays to + `ai/embed` / `kb/upsert`. A native chunking helper and metadata filtering on + `kb/query` 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 bf9eb06..25c2cd1 100644 --- a/packages/site/src/content/docs/workflow-reference/connectors.md +++ b/packages/site/src/content/docs/workflow-reference/connectors.md @@ -320,6 +320,71 @@ Executes a parameterized SQL query against an external Postgres database. The co mantle secrets create --name my-database --type generic --field url=postgres://user:pass@host:5432/dbname?sslmode=require ``` +## kb/upsert + +Stores document text and a vector embedding (plus optional JSONB metadata) into a [pgvector](https://github.com/pgvector/pgvector) table. A thin helper over `postgres/query` for [RAG](/docs/rag-guide): it takes the pgvector literal from `ai/embed`'s `output.vector` and handles the `::vector` cast and multi-row inserts. Uses a `postgres` credential; you create the table (it does not manage schema). Table and column names are validated as SQL identifiers. + +**Params:** + +| Param | Type | Required | Description | +|---|---|---|---| +| `table` | string | Yes | Target table name (optionally `schema.table`). | +| `content` / `contents` | string / list | Yes | Document text — a single string, or an array for a batch (e.g. chunks). | +| `vector` / `vectors` | string / list | Yes | Matching pgvector literal(s) from `ai/embed`. Count must match `content`/`contents`. | +| `metadata` / `metadatas` | object / list | No | Optional JSONB metadata per row. If provided, the count must match. | +| `content_column` | string | No | Content column name. Default `content`. | +| `embedding_column` | string | No | Vector column name. Default `embedding`. | +| `metadata_column` | string | No | Metadata column name. Default `metadata`. | +| `conflict_target` | string | No | Column/constraint for `ON CONFLICT (...) DO NOTHING` (idempotent upsert). | + +**Output:** `rows_affected` (number). + +**Example:** + +```yaml +- name: store + action: kb/upsert + credential: kb-db + depends_on: [embed] + params: + table: kb_documents + content: "{{ inputs.content }}" + vector: "{{ steps['embed'].output.vector }}" + metadata: { title: "{{ inputs.title }}", source: "{{ inputs.source }}" } + conflict_target: dedupe_key +``` + +## kb/query + +Nearest-neighbour search over a pgvector table for a query embedding. A thin helper over `postgres/query`: it builds the distance-ordered `SELECT` so you don't write the pgvector SQL. Uses a `postgres` credential; identifiers are validated. + +**Params:** + +| Param | Type | Required | Description | +|---|---|---|---| +| `table` | string | Yes | Table to search (optionally `schema.table`). | +| `vector` | string | Yes | The query embedding as a pgvector literal (from `ai/embed`). | +| `top_k` | integer | No | Number of rows to return. Default `5`, capped at `1000`. | +| `columns` | list | No | Columns to return. Default `[content]`. | +| `metric` | string | No | Distance metric: `cosine` (default), `l2`, or `inner_product`. | +| `embedding_column` | string | No | Vector column name. Default `embedding`. | + +**Output:** `rows` (list of maps — the requested columns plus a `distance` field, nearest first) and `row_count`. For `cosine`, similarity is `1 - distance`. + +**Example:** + +```yaml +- name: search + action: kb/query + credential: kb-db + depends_on: [embed-question] + params: + table: kb_documents + vector: "{{ steps['embed-question'].output.vector }}" + columns: [content, metadata] + top_k: 5 +``` + ## email/send Sends an email via SMTP. Supports plaintext and HTML content. diff --git a/packages/site/src/content/examples/rag-ask.yaml b/packages/site/src/content/examples/rag-ask.yaml index f759703..082097f 100644 --- a/packages/site/src/content/examples/rag-ask.yaml +++ b/packages/site/src/content/examples/rag-ask.yaml @@ -1,9 +1,9 @@ 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). + question, find the nearest documents with kb/query (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: @@ -20,25 +20,21 @@ steps: model: text-embedding-3-small input: "{{ inputs.question }}" - # 2. Nearest-neighbour search by cosine distance (<=> is pgvector's distance - # operator; similarity = 1 - distance). + # 2. Nearest-neighbour search. Returns the requested columns plus a `distance` + # field (lower = closer; for cosine, similarity = 1 - distance). - name: search - action: postgres/query + action: kb/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 }}" + table: kb_documents + vector: "{{ steps['embed-question'].output.vector }}" + columns: + - content + - metadata + top_k: 5 # 3. Answer grounded strictly in the retrieved passages. - name: answer @@ -50,11 +46,11 @@ steps: 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. + Answer using ONLY the provided context passages. Cite the source or + title from each passage's metadata when you use it. If the context does + not contain the answer, say so plainly rather than guessing. prompt: > Question: {{ inputs.question }} - Context passages (JSON, most similar first): + Context passages (JSON, nearest 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 index 78cf256..5072732 100644 --- a/packages/site/src/content/examples/rag-ingest.yaml +++ b/packages/site/src/content/examples/rag-ingest.yaml @@ -1,9 +1,9 @@ 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). + Embed a document with ai/embed and store it in a pgvector knowledge base with + kb/upsert for semantic retrieval. Pair with rag-ask.yaml. Requires the schema + in rag-kb-schema.sql. For long documents, split into chunks and pass arrays to + ai/embed (input) and kb/upsert (contents/vectors). inputs: title: @@ -19,8 +19,8 @@ inputs: 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. + # 1. Produce an embedding. output.vector is a pgvector text literal that + # kb/upsert stores directly. - name: embed action: ai/embed credential: openai @@ -29,21 +29,19 @@ steps: 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)). + # 2. Store the document, embedding, and metadata. conflict_target makes + # re-ingesting identical content a no-op (dedupe_key is md5(content)). - name: store - action: postgres/query + action: kb/upsert 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 }}" + table: kb_documents + content: "{{ inputs.content }}" + vector: "{{ steps['embed'].output.vector }}" + metadata: + title: "{{ inputs.title }}" + source: "{{ inputs.source }}" + conflict_target: dedupe_key diff --git a/packages/site/src/content/examples/rag-kb-schema.sql b/packages/site/src/content/examples/rag-kb-schema.sql index f6851e6..ac2f324 100644 --- a/packages/site/src/content/examples/rag-kb-schema.sql +++ b/packages/site/src/content/examples/rag-kb-schema.sql @@ -1,23 +1,25 @@ -- 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`. +-- Semantic retrieval with pgvector, driven by the kb/upsert and kb/query +-- connectors. Requires the pgvector extension +-- (https://github.com/pgvector/pgvector). Point a Mantle postgres 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. +-- text-embedding-3-large = 3072, Bedrock titan-embed-text-v2 = 1024 (default). +-- 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). + -- Arbitrary per-document metadata (title, source, ...), written by + -- kb/upsert's `metadata` param and returned by kb/query's `columns`. + metadata JSONB NOT NULL DEFAULT '{}'::jsonb, + -- Idempotency: re-ingesting identical content is a no-op via kb/upsert's + -- conflict_target (dedupe_key). dedupe_key TEXT GENERATED ALWAYS AS (md5(content)) STORED, created_at TIMESTAMPTZ NOT NULL DEFAULT now() ); @@ -25,7 +27,9 @@ CREATE TABLE IF NOT EXISTS kb_documents ( 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. +-- Approximate-nearest-neighbour index for cosine distance (pgvector >= 0.5), +-- which is kb/query's default metric. For small corpora you can omit this and +-- rely on an exact scan. If you query with metric: l2 or inner_product, add a +-- matching index (vector_l2_ops / vector_ip_ops) or those queries seq-scan. CREATE INDEX IF NOT EXISTS idx_kb_documents_embedding ON kb_documents USING hnsw (embedding vector_cosine_ops);