feat: RAG follow-ups — kb/query metadata filtering, recursive chunking, Cohere Bedrock embeddings (#153)

.changeset/rag-followups.md

@@ -0,0 +1,11 @@
+---
+"@mantle/engine": minor
+---
+
+RAG follow-ups on #153:
+
+- **`kb/query` metadata filtering** — an optional `filter` object restricts the search to rows whose metadata contains it, via the JSONB containment operator (`metadata @> $2::jsonb`). The whole filter binds as a single parameter (no injection surface); a custom `metadata_column` is supported. Scope retrieval to a source, team, or tag alongside the vector search.
+- **Separator-aware `recursive` chunking** — `text/chunk` gains `unit: recursive`, which walks a separator hierarchy (paragraph → line → sentence → word → character) so chunks break on natural boundaries instead of mid-sentence, then merges pieces up to `chunk_size` with `chunk_overlap` characters carried between chunks. The `rag-ingest` example now uses it.
+- **Cohere Bedrock embeddings** — `ai/embed` (provider `bedrock`) now supports `cohere.embed-english-v3` and `cohere.embed-multilingual-v3` in addition to Amazon Titan. Cohere batches natively (up to 96 texts per call) and takes an `input_type` (`search_document` default / `search_query` / `classification` / `clustering`).
+
+Token-accurate chunking remains a follow-up on #153.

packages/engine/internal/connector/chunk.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"strings"
+	"unicode/utf8"
 )
 
 // TextChunkConnector implements text/chunk: split a long document into
@@ -16,10 +17,20 @@ const (
 	chunkDefaultOverlap = 0
 )
 
-// chunkText splits text into fixed-size, optionally overlapping chunks. `unit`
-// is "chars" (default, Unicode-aware) or "words" (whitespace-separated). size
-// and overlap are counted in that unit. Pure — no I/O — so it is fully
-// unit-testable.
+// recursiveSeparators is the default hierarchy the "recursive" unit walks, from
+// coarsest (paragraph) to finest. The empty string is the terminal fallback: a
+// stretch with no separator left is hard-split by character window.
+var recursiveSeparators = []string{"\n\n", "\n", ". ", " ", ""}
+
+// chunkText splits text into overlapping chunks. `unit` selects the strategy:
+//   - "chars" (default, Unicode-aware) / "words": fixed-size sliding window,
+//     size and overlap counted in that unit.
+//   - "recursive": separator-aware split (paragraph → line → sentence → word →
+//     character) that prefers to break on natural boundaries, then merges
+//     adjacent pieces toward `size` with `overlap` characters carried between
+//     chunks. A chunk targets `size` but may reach up to `size + overlap`.
+//
+// Pure — no I/O — so it is fully unit-testable.
 func chunkText(text string, size, overlap int, unit string) ([]string, error) {
 	if strings.TrimSpace(text) == "" {
 		return nil, fmt.Errorf("text must not be empty")
@@ -46,9 +57,103 @@ func chunkText(text string, size, overlap int, unit string) ([]string, error) {
 		return sliceChunks(len(words), size, step, func(a, b int) string {
 			return strings.Join(words[a:b], " ")
 		}), nil
+	case "recursive":
+		return recursiveChunks(text, size, overlap, recursiveSeparators), nil
 	default:
-		return nil, fmt.Errorf("unknown unit %q (use chars or words)", unit)
+		return nil, fmt.Errorf("unknown unit %q (use chars, words, or recursive)", unit)
+	}
+}
+
+// recursiveChunks splits text on a hierarchy of separators so chunks break on
+// natural boundaries, then greedily merges the pieces (with overlap) into
+// windows of at most `size` characters. size/overlap are counted in runes.
+func recursiveChunks(text string, size, overlap int, seps []string) []string {
+	pieces := splitRecursive(text, size, seps)
+	return mergeSplits(pieces, size, overlap)
+}
+
+// splitRecursive breaks text into atomic pieces, each at most `size` runes where
+// possible, preferring the coarsest separator that appears. A piece still larger
+// than `size` after exhausting separators is hard-split by character window.
+// Separators are re-attached to the pieces so the text reconstructs faithfully.
+func splitRecursive(text string, size int, seps []string) []string {
+	if utf8.RuneCountInString(text) <= size {
+		if text == "" {
+			return nil
+		}
+		return []string{text}
+	}
+	if len(seps) == 0 {
+		return hardSplit(text, size)
+	}
+	sep, rest := seps[0], seps[1:]
+	if sep == "" {
+		return hardSplit(text, size)
+	}
+	parts := strings.Split(text, sep)
+	var out []string
+	for i, p := range parts {
+		piece := p
+		if i < len(parts)-1 { // keep the separator that we split on
+			piece += sep
+		}
+		if piece == "" {
+			continue
+		}
+		if utf8.RuneCountInString(piece) <= size {
+			out = append(out, piece)
+		} else {
+			out = append(out, splitRecursive(piece, size, rest)...)
+		}
 	}
+	return out
+}
+
+// hardSplit chops text into consecutive windows of `size` runes (no overlap),
+// used as the terminal fallback for a stretch with no usable separator.
+func hardSplit(text string, size int) []string {
+	runes := []rune(text)
+	return sliceChunks(len(runes), size, size, func(a, b int) string {
+		return string(runes[a:b])
+	})
+}
+
+// mergeSplits greedily concatenates pieces into chunks that target `size` runes.
+// When a chunk fills, it starts the next one with a tail of the previous pieces
+// totalling up to `overlap` runes, so consecutive chunks share context. Because
+// that carried tail is prepended before the next piece, a chunk may reach up to
+// `size + overlap` runes (this mirrors how splitters like LangChain's treat
+// chunk_size as a target rather than a hard cap).
+func mergeSplits(pieces []string, size, overlap int) []string {
+	var chunks []string
+	var cur []string
+	curLen := 0
+
+	flush := func() {
+		if len(cur) == 0 {
+			return
+		}
+		if c := strings.TrimSpace(strings.Join(cur, "")); c != "" {
+			chunks = append(chunks, c)
+		}
+	}
+
+	for _, p := range pieces {
+		pl := utf8.RuneCountInString(p)
+		if curLen+pl > size && len(cur) > 0 {
+			flush()
+			// Retain a trailing window of pieces up to `overlap` runes as the
+			// start of the next chunk.
+			for curLen > overlap && len(cur) > 0 {
+				curLen -= utf8.RuneCountInString(cur[0])
+				cur = cur[1:]
+			}
+		}
+		cur = append(cur, p)
+		curLen += pl
+	}
+	flush()
+	return chunks
 }
 
 // sliceChunks walks [0,n) in windows of `size` advancing by `step`, emitting

packages/engine/internal/connector/chunk_test.go

@@ -63,6 +63,71 @@ func TestChunkText_Unicode(t *testing.T) {
 	}
 }
 
+func TestChunkText_RecursiveSentences(t *testing.T) {
+	// With ". " available and a size that holds one sentence but not two, the
+	// recursive splitter breaks on sentence boundaries.
+	got, err := chunkText("Alpha beta. Gamma delta. Epsilon zeta.", 15, 0, "recursive")
+	if err != nil {
+		t.Fatalf("error: %v", err)
+	}
+	want := []string{"Alpha beta.", "Gamma delta.", "Epsilon zeta."}
+	if strings.Join(got, "|") != strings.Join(want, "|") {
+		t.Errorf("chunks = %v, want %v", got, want)
+	}
+}
+
+func TestChunkText_RecursiveFallbackHardSplit(t *testing.T) {
+	// No usable separator (one long token): falls through to a character window.
+	got, err := chunkText("abcdefgh", 3, 0, "recursive")
+	if err != nil {
+		t.Fatalf("error: %v", err)
+	}
+	want := []string{"abc", "def", "gh"}
+	if strings.Join(got, "|") != strings.Join(want, "|") {
+		t.Errorf("chunks = %v, want %v", got, want)
+	}
+}
+
+func TestChunkText_RecursiveOverlap(t *testing.T) {
+	// Overlap carries a trailing word into the next chunk, so adjacent chunks
+	// share the boundary token.
+	got, err := chunkText("one two three four five six", 9, 4, "recursive")
+	if err != nil {
+		t.Fatalf("error: %v", err)
+	}
+	if len(got) < 2 || got[0] != "one two" || got[1] != "two three" {
+		t.Errorf("chunks = %v, want first two [one two, two three]", got)
+	}
+}
+
+func TestChunkText_RecursiveShorterThanSize(t *testing.T) {
+	got, err := chunkText("short text", 100, 10, "recursive")
+	if err != nil {
+		t.Fatalf("error: %v", err)
+	}
+	if len(got) != 1 || got[0] != "short text" {
+		t.Errorf("chunks = %v, want [short text]", got)
+	}
+}
+
+func TestChunkText_RecursiveChunkSizeBound(t *testing.T) {
+	// Every chunk stays within size + overlap runes (the documented bound).
+	text := strings.Repeat("lorem ipsum dolor sit amet. ", 50)
+	size, overlap := 60, 15
+	got, err := chunkText(text, size, overlap, "recursive")
+	if err != nil {
+		t.Fatalf("error: %v", err)
+	}
+	if len(got) < 2 {
+		t.Fatalf("expected multiple chunks, got %d", len(got))
+	}
+	for i, c := range got {
+		if n := len([]rune(c)); n > size+overlap {
+			t.Errorf("chunk %d has %d runes, exceeds size+overlap (%d)", i, n, size+overlap)
+		}
+	}
+}
+
 func TestChunkText_Errors(t *testing.T) {
 	cases := []struct {
 		name          string

packages/engine/internal/connector/embed.go

@@ -60,6 +60,9 @@ func (c *EmbeddingConnector) Execute(ctx context.Context, params map[string]any)
 	if d, ok := extractInt(params["dimensions"]); ok && d > 0 {
 		req.Dimensions = d
 	}
+	if it, ok := params["input_type"].(string); ok {
+		req.InputType = it
+	}
 
 	workflow, _ := params["_workflow"].(string)
 	step, _ := params["_step"].(string)

packages/engine/internal/connector/kb.go

@@ -271,11 +271,61 @@ func prepareQuery(params map[string]any) (string, []any, error) {
 		cols = []string{contentCol}
 	}
 
+	// Optional metadata filter: a JSON object matched with the JSONB containment
+	// operator (@>). The whole object is bound as one parameter, so there is no
+	// injection surface — only the metadata column identifier is interpolated.
+	args := []any{vector}
+	where, err := kbFilterClause(params, len(args))
+	if err != nil {
+		return "", nil, err
+	}
+	if where.clause != "" {
+		args = append(args, where.arg)
+	}
+
 	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
+		"SELECT %s, (%s %s $1::vector) AS distance FROM %s%s ORDER BY %s %s $1::vector LIMIT %s",
+		selectList, embCol, op, table, where.clause, embCol, op, strconv.Itoa(topK))
+	return sql, args, nil
+}
+
+// kbFilter is the built WHERE fragment plus its bound JSONB argument.
+type kbFilter struct {
+	clause string // e.g. " WHERE metadata @> $2::jsonb", or "" when absent
+	arg    string // the JSON-encoded filter object (only set when clause != "")
+}
+
+// kbFilterClause builds an optional JSONB-containment WHERE clause from the
+// `filter` param. The filter must be a JSON object; it is matched against the
+// metadata column with @> (does the row's metadata contain these key/values).
+// An absent or empty filter yields no clause. boundParamCount is the number of
+// SQL parameters already bound (currently always 1, for the query vector $1), so
+// the filter binds at the next placeholder, $(boundParamCount+1).
+func kbFilterClause(params map[string]any, boundParamCount int) (kbFilter, error) {
+	raw, ok := params["filter"]
+	if !ok || raw == nil {
+		return kbFilter{}, nil
+	}
+	obj, ok := raw.(map[string]any)
+	if !ok {
+		return kbFilter{}, fmt.Errorf("filter must be a JSON object, got %T", raw)
+	}
+	if len(obj) == 0 {
+		return kbFilter{}, nil // empty object matches everything — treat as no filter
+	}
+	metaCol, err := kbIdentParam(params, "metadata_column", kbDefaultMetadataColumn)
+	if err != nil {
+		return kbFilter{}, err
+	}
+	b, err := json.Marshal(obj)
+	if err != nil {
+		return kbFilter{}, fmt.Errorf("marshaling filter: %w", err)
+	}
+	return kbFilter{
+		clause: fmt.Sprintf(" WHERE %s @> $%d::jsonb", metaCol, boundParamCount+1),
+		arg:    string(b),
+	}, nil
 }
 
 func kbStringList(raw any) ([]string, error) {

packages/engine/internal/connector/kb_test.go

@@ -185,6 +185,65 @@ func TestPrepareQuery_MetricsAndColumns(t *testing.T) {
 	}
 }
 
+func TestPrepareQuery_MetadataFilter(t *testing.T) {
+	sql, args, err := prepareQuery(map[string]any{
+		"table":   "kb_documents",
+		"vector":  "[0.1,0.2]",
+		"columns": []any{"content", "metadata"},
+		"filter":  map[string]any{"source": "confluence", "team": "a"},
+		"top_k":   3,
+	})
+	if err != nil {
+		t.Fatalf("prepareQuery error: %v", err)
+	}
+	// The filter binds as $2 and the WHERE sits between FROM and ORDER BY.
+	if !strings.Contains(sql, "FROM kb_documents WHERE metadata @> $2::jsonb ORDER BY") {
+		t.Errorf("unexpected where clause: %s", sql)
+	}
+	if len(args) != 2 {
+		t.Fatalf("args len = %d, want 2 (vector + filter)", len(args))
+	}
+	filterJSON := args[1].(string)
+	if !strings.Contains(filterJSON, `"source":"confluence"`) || !strings.Contains(filterJSON, `"team":"a"`) {
+		t.Errorf("filter arg = %q", filterJSON)
+	}
+}
+
+func TestPrepareQuery_FilterCustomColumnAndPlaceholder(t *testing.T) {
+	// A custom metadata_column is interpolated; the filter placeholder follows
+	// the vector ($1), i.e. $2.
+	sql, _, err := prepareQuery(map[string]any{
+		"table":           "kb_documents",
+		"vector":          "[1]",
+		"metadata_column": "meta",
+		"filter":          map[string]any{"k": "v"},
+	})
+	if err != nil {
+		t.Fatalf("prepareQuery error: %v", err)
+	}
+	if !strings.Contains(sql, "WHERE meta @> $2::jsonb") {
+		t.Errorf("unexpected filter clause: %s", sql)
+	}
+}
+
+func TestPrepareQuery_EmptyFilterOmitted(t *testing.T) {
+	// An empty object matches everything; treat it as no filter (no WHERE, no arg).
+	sql, args, err := prepareQuery(map[string]any{
+		"table":  "t",
+		"vector": "[1]",
+		"filter": map[string]any{},
+	})
+	if err != nil {
+		t.Fatalf("prepareQuery error: %v", err)
+	}
+	if strings.Contains(sql, "WHERE") {
+		t.Errorf("empty filter should not emit WHERE: %s", sql)
+	}
+	if len(args) != 1 {
+		t.Errorf("args len = %d, want 1 (vector only)", len(args))
+	}
+}
+
 func TestPrepareQuery_TopKCap(t *testing.T) {
 	sql, _, err := prepareQuery(map[string]any{"table": "t", "vector": "[1]", "top_k": 999999})
 	if err != nil {
@@ -207,6 +266,8 @@ func TestPrepareQuery_Errors(t *testing.T) {
 		{"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"}}},
+		{"filter not an object", map[string]any{"table": "t", "vector": "[1]", "filter": "source=x"}},
+		{"sql injection in metadata_column", map[string]any{"table": "t", "vector": "[1]", "metadata_column": "m @> '{}'; DROP", "filter": map[string]any{"k": "v"}}},
 	}
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {

packages/engine/internal/connector/provider.go

@@ -19,7 +19,8 @@ type EmbeddingProvider interface {
 type EmbeddingRequest struct {
 	Model      string
 	Inputs     []string
-	Dimensions int // 0 = provider default
+	Dimensions int    // 0 = provider default
+	InputType  string // provider-specific hint (Cohere: search_document/search_query/...); "" = provider default
 	Credential map[string]string
 }