From 1fc0beb7ba0f8e6ac615b31ad5328d31618fb756 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 14:06:48 +0000 Subject: [PATCH] feat: add text/chunk connector for RAG ingest (#153) Completes the RAG ingest pipeline: text/chunk -> ai/embed (batch) -> kb/upsert (batch). - text/chunk splits a document into overlapping chunks by characters (Unicode-aware) or words, with configurable chunk_size/chunk_overlap. output.chunks feeds straight into ai/embed's input and kb/upsert's contents. Pure chunkText core with thorough unit tests (char/word, overlap, exact-fit, unicode, shorter-than-size, error cases). - kb/upsert: a single `metadata` object is now broadcast to every row of a batch, so chunked ingest can share one title/source. (`metadatas` still matches per-row.) - rag-ingest example rewritten to chunk -> embed -> upsert; text/chunk documented in the connector reference; RAG guide updated. rag-ingest passes `mantle validate`. Smarter (separator-aware/token) chunking and kb/query metadata filtering remain follow-ups on #153. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01WnKgCRfMvFcZAkoVbr8k79 --- .changeset/text-chunk-connector.md | 5 + packages/engine/internal/connector/chunk.go | 93 ++++++++++++++++ .../engine/internal/connector/chunk_test.go | 102 ++++++++++++++++++ .../engine/internal/connector/connector.go | 1 + packages/engine/internal/connector/kb.go | 35 +++--- packages/engine/internal/connector/kb_test.go | 22 ++++ packages/site/src/content/docs/rag-guide.md | 15 +-- .../docs/workflow-reference/connectors.md | 44 ++++++++ .../site/src/content/examples/rag-ingest.yaml | 40 ++++--- 9 files changed, 322 insertions(+), 35 deletions(-) create mode 100644 .changeset/text-chunk-connector.md create mode 100644 packages/engine/internal/connector/chunk.go create mode 100644 packages/engine/internal/connector/chunk_test.go diff --git a/.changeset/text-chunk-connector.md b/.changeset/text-chunk-connector.md new file mode 100644 index 0000000..06b6bc1 --- /dev/null +++ b/.changeset/text-chunk-connector.md @@ -0,0 +1,5 @@ +--- +"@mantle/engine": minor +--- + +Add a `text/chunk` connector that splits a long document into overlapping chunks for embedding, completing the RAG ingest pipeline: `text/chunk` → `ai/embed` (batch) → `kb/upsert` (batch). Chunk by characters (Unicode-aware) or words with configurable size and overlap; `output.chunks` feeds straight into `ai/embed`'s `input` and `kb/upsert`'s `contents`. Also enhances `kb/upsert` so a single `metadata` object is broadcast to every row of a batch (so chunked ingest can share one title/source). The RAG example now chunks documents; smarter (separator-aware/token) chunking and `kb/query` metadata filtering remain follow-ups on #153. diff --git a/packages/engine/internal/connector/chunk.go b/packages/engine/internal/connector/chunk.go new file mode 100644 index 0000000..2ce719b --- /dev/null +++ b/packages/engine/internal/connector/chunk.go @@ -0,0 +1,93 @@ +package connector + +import ( + "context" + "fmt" + "strings" +) + +// TextChunkConnector implements text/chunk: split a long document into +// overlapping chunks for embedding. It composes with ai/embed (pass the chunks +// as `input`) and kb/upsert (pass them as `contents`): chunk → embed → upsert. +type TextChunkConnector struct{} + +const ( + chunkDefaultSize = 1000 + 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. +func chunkText(text string, size, overlap int, unit string) ([]string, error) { + if strings.TrimSpace(text) == "" { + return nil, fmt.Errorf("text must not be empty") + } + if size <= 0 { + return nil, fmt.Errorf("chunk_size must be positive, got %d", size) + } + if overlap < 0 { + return nil, fmt.Errorf("chunk_overlap must be >= 0, got %d", overlap) + } + if overlap >= size { + return nil, fmt.Errorf("chunk_overlap (%d) must be less than chunk_size (%d)", overlap, size) + } + step := size - overlap + + switch unit { + case "", "chars": + runes := []rune(text) + return sliceChunks(len(runes), size, step, func(a, b int) string { + return string(runes[a:b]) + }), nil + case "words": + words := strings.Fields(text) + return sliceChunks(len(words), size, step, func(a, b int) string { + return strings.Join(words[a:b], " ") + }), nil + default: + return nil, fmt.Errorf("unknown unit %q (use chars or words)", unit) + } +} + +// sliceChunks walks [0,n) in windows of `size` advancing by `step`, emitting +// each window via slice(start, end). It stops once a window reaches the end so +// overlap never produces a trailing duplicate. +func sliceChunks(n, size, step int, slice func(a, b int) string) []string { + var chunks []string + for start := 0; start < n; start += step { + end := start + size + if end >= n { + end = n + chunks = append(chunks, slice(start, end)) + break + } + chunks = append(chunks, slice(start, end)) + } + return chunks +} + +func (c *TextChunkConnector) Execute(_ context.Context, params map[string]any) (map[string]any, error) { + text, _ := params["text"].(string) + + size := chunkDefaultSize + if v, ok := extractInt(params["chunk_size"]); ok { + size = v + } + overlap := chunkDefaultOverlap + if v, ok := extractInt(params["chunk_overlap"]); ok { + overlap = v + } + unit, _ := params["unit"].(string) + + chunks, err := chunkText(text, size, overlap, unit) + if err != nil { + return nil, fmt.Errorf("text/chunk: %w", err) + } + + return map[string]any{ + "chunks": chunks, + "count": len(chunks), + }, nil +} diff --git a/packages/engine/internal/connector/chunk_test.go b/packages/engine/internal/connector/chunk_test.go new file mode 100644 index 0000000..44d14ca --- /dev/null +++ b/packages/engine/internal/connector/chunk_test.go @@ -0,0 +1,102 @@ +package connector + +import ( + "context" + "strings" + "testing" +) + +func TestChunkText_Chars(t *testing.T) { + // "abcdefg" (7), size 3, overlap 1 → step 2: windows 0-3, 2-5, 4-7. + // The last window reaches the end, so there is no trailing "g" chunk. + got, err := chunkText("abcdefg", 3, 1, "chars") + if err != nil { + t.Fatalf("chunkText error: %v", err) + } + want := []string{"abc", "cde", "efg"} + if strings.Join(got, "|") != strings.Join(want, "|") { + t.Errorf("chunks = %v, want %v", got, want) + } +} + +func TestChunkText_NoOverlapExactFit(t *testing.T) { + // 6 chars, size 3, overlap 0 → [abc, def], no trailing empty chunk. + got, err := chunkText("abcdef", 3, 0, "chars") + if err != nil { + t.Fatalf("error: %v", err) + } + if len(got) != 2 || got[0] != "abc" || got[1] != "def" { + t.Errorf("chunks = %v, want [abc def]", got) + } +} + +func TestChunkText_Words(t *testing.T) { + got, err := chunkText("the quick brown fox jumps", 2, 0, "words") + if err != nil { + t.Fatalf("error: %v", err) + } + want := []string{"the quick", "brown fox", "jumps"} + if strings.Join(got, "|") != strings.Join(want, "|") { + t.Errorf("chunks = %v, want %v", got, want) + } +} + +func TestChunkText_ShorterThanSize(t *testing.T) { + got, err := chunkText("hi", 100, 10, "chars") + if err != nil { + t.Fatalf("error: %v", err) + } + if len(got) != 1 || got[0] != "hi" { + t.Errorf("chunks = %v, want [hi]", got) + } +} + +func TestChunkText_Unicode(t *testing.T) { + // Multibyte runes must not be split mid-character. + got, err := chunkText("héllo wörld", 3, 0, "chars") + if err != nil { + t.Fatalf("error: %v", err) + } + // 11 runes / size 3 → 4 chunks; first is "hél". + if got[0] != "hél" { + t.Errorf("first chunk = %q, want %q", got[0], "hél") + } +} + +func TestChunkText_Errors(t *testing.T) { + cases := []struct { + name string + text string + size, overlap int + unit string + }{ + {"empty text", " ", 10, 0, "chars"}, + {"zero size", "abc", 0, 0, "chars"}, + {"negative overlap", "abc", 10, -1, "chars"}, + {"overlap >= size", "abc", 5, 5, "chars"}, + {"unknown unit", "abc", 5, 0, "tokens"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if _, err := chunkText(tc.text, tc.size, tc.overlap, tc.unit); err == nil { + t.Errorf("expected error for %s", tc.name) + } + }) + } +} + +func TestTextChunkConnector_Output(t *testing.T) { + out, err := (&TextChunkConnector{}).Execute(context.Background(), map[string]any{ + "text": "one two three four five", + "chunk_size": 2, + "chunk_overlap": 0, + "unit": "words", + }) + if err != nil { + t.Fatalf("Execute error: %v", err) + } + chunks := out["chunks"].([]string) + if out["count"].(int) != len(chunks) || len(chunks) != 3 { + t.Errorf("count/chunks = %v / %v", out["count"], chunks) + } +} diff --git a/packages/engine/internal/connector/connector.go b/packages/engine/internal/connector/connector.go index bd84cb0..08d29bb 100644 --- a/packages/engine/internal/connector/connector.go +++ b/packages/engine/internal/connector/connector.go @@ -33,6 +33,7 @@ func NewRegistry() *Registry { r.Register("postgres/query", &PostgresQueryConnector{}) r.Register("kb/upsert", &KBUpsertConnector{}) r.Register("kb/query", &KBQueryConnector{}) + r.Register("text/chunk", &TextChunkConnector{}) 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 index 4de94f0..f63f21c 100644 --- a/packages/engine/internal/connector/kb.go +++ b/packages/engine/internal/connector/kb.go @@ -92,7 +92,9 @@ func kbStrings(params map[string]any, singleKey, manyKey string) ([]string, erro } // kbMetadata normalizes the optional metadata param into JSON strings aligned -// with the rows. Returns present=false when no metadata param is set. +// with the rows. A single `metadata` object is broadcast to every row (handy +// for chunked ingest that shares a title/source); `metadatas` must match the +// row count one-to-one. 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"] @@ -103,25 +105,30 @@ func kbMetadata(params map[string]any, rows int) (jsons []string, present bool, 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) + b, mErr := json.Marshal(single) + if mErr != nil { + return nil, false, fmt.Errorf("marshaling metadata: %w", mErr) } - 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, rows) + for i := range out { + out[i] = string(b) + } + return out, true, nil } - out := make([]string, len(objs)) - for i, o := range objs { + arr, ok := many.([]any) + if !ok { + return nil, false, fmt.Errorf("metadatas must be an array of objects, got %T", many) + } + if len(arr) != rows { + return nil, false, fmt.Errorf("metadatas count %d does not match content count %d", len(arr), rows) + } + out := make([]string, len(arr)) + for i, o := range arr { b, mErr := json.Marshal(o) if mErr != nil { - return nil, false, fmt.Errorf("marshaling metadata[%d]: %w", i, mErr) + return nil, false, fmt.Errorf("marshaling metadatas[%d]: %w", i, mErr) } out[i] = string(b) } diff --git a/packages/engine/internal/connector/kb_test.go b/packages/engine/internal/connector/kb_test.go index fb46419..9f7e802 100644 --- a/packages/engine/internal/connector/kb_test.go +++ b/packages/engine/internal/connector/kb_test.go @@ -98,6 +98,28 @@ func TestPrepareUpsert_BatchWithMetadataAndConflict(t *testing.T) { } } +func TestPrepareUpsert_SingleMetadataBroadcast(t *testing.T) { + // One `metadata` object applies to every row of a batch (chunked ingest). + _, args, err := prepareUpsert(map[string]any{ + "table": "kb_documents", + "contents": []any{"a", "b", "c"}, + "vectors": []any{"[1]", "[2]", "[3]"}, + "metadata": map[string]any{"source": "doc-1"}, + }) + if err != nil { + t.Fatalf("prepareUpsert error: %v", err) + } + metas := args[2].([]string) + if len(metas) != 3 { + t.Fatalf("metadata args = %d, want 3 (broadcast)", len(metas)) + } + for i, m := range metas { + if !strings.Contains(m, `"source":"doc-1"`) { + t.Errorf("metas[%d] = %q, want broadcast source", i, m) + } + } +} + func TestPrepareUpsert_Errors(t *testing.T) { cases := []struct { name string diff --git a/packages/site/src/content/docs/rag-guide.md b/packages/site/src/content/docs/rag-guide.md index 9a46c8a..da46b8e 100644 --- a/packages/site/src/content/docs/rag-guide.md +++ b/packages/site/src/content/docs/rag-guide.md @@ -82,10 +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 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)`). +`rag-ingest` splits the document into overlapping chunks with `text/chunk`, +embeds them all in one `ai/embed` call, and stores them in one `kb/upsert`. A +single `metadata` object is broadcast to every chunk. Re-ingesting identical +content is a no-op (the schema dedupes on `md5(content)`). ## Asking questions @@ -111,6 +111,7 @@ outputs with `--output json` or `-v`). to match (Titan v2 defaults to 1024). Cohere Bedrock models are a follow-up. - **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). +- **Chunking is fixed-size.** `text/chunk` splits by a sliding window over + characters or words with overlap. Semantic/recursive (separator-aware or + token-accurate) chunking 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 25c2cd1..cda2bba 100644 --- a/packages/site/src/content/docs/workflow-reference/connectors.md +++ b/packages/site/src/content/docs/workflow-reference/connectors.md @@ -385,6 +385,50 @@ Nearest-neighbour search over a pgvector table for a query embedding. A thin hel top_k: 5 ``` +## text/chunk + +Splits a long document into overlapping chunks for embedding. Feeds the chunk → embed → store pipeline: pass `output.chunks` to `ai/embed`'s `input` and to `kb/upsert`'s `contents`. Pure text processing — no credential needed. + +**Params:** + +| Param | Type | Required | Description | +|---|---|---|---| +| `text` | string | Yes | The document text to split. | +| `chunk_size` | integer | No | Target chunk size, in `unit`s. Default `1000`. | +| `chunk_overlap` | integer | No | Overlap between consecutive chunks, in `unit`s. Must be `< chunk_size`. Default `0`. | +| `unit` | string | No | `chars` (default, Unicode-aware) or `words` (whitespace-separated). | + +**Output:** `chunks` (list of strings) and `count` (number). + +**Example — chunk → embed → upsert:** + +```yaml +- name: chunk + action: text/chunk + params: + text: "{{ inputs.content }}" + chunk_size: 1200 + chunk_overlap: 150 + +- name: embed + action: ai/embed + credential: my-openai + depends_on: [chunk] + params: + model: text-embedding-3-small + input: "{{ steps['chunk'].output.chunks }}" + +- name: store + action: kb/upsert + credential: kb-db + depends_on: [chunk, embed] + params: + table: kb_documents + contents: "{{ steps['chunk'].output.chunks }}" + vectors: "{{ steps['embed'].output.vectors }}" + metadata: { source: "{{ inputs.source }}" } # broadcast to every chunk +``` + ## email/send Sends an email via SMTP. Supports plaintext and HTML content. diff --git a/packages/site/src/content/examples/rag-ingest.yaml b/packages/site/src/content/examples/rag-ingest.yaml index 5072732..0b2e126 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 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). + Ingest a document into a pgvector knowledge base: split it into overlapping + chunks with text/chunk, embed them in one batch with ai/embed, and store them + with kb/upsert. Pair with rag-ask.yaml. Requires the schema in + rag-kb-schema.sql. inputs: title: @@ -16,31 +16,43 @@ inputs: default: "" content: type: string - description: The document text to embed and store + description: The document text to chunk, embed, and store steps: - # 1. Produce an embedding. output.vector is a pgvector text literal that - # kb/upsert stores directly. + # 1. Split the document into overlapping chunks (defaults are chars; use + # unit: words for word counts). + - name: chunk + action: text/chunk + params: + text: "{{ inputs.content }}" + chunk_size: 1200 + chunk_overlap: 150 + + # 2. Embed all chunks in one call. output.vectors is an array of pgvector + # literals aligned with the chunks. - name: embed action: ai/embed credential: openai - timeout: "30s" + depends_on: + - chunk + timeout: "60s" params: model: text-embedding-3-small - input: "{{ inputs.content }}" + input: "{{ steps['chunk'].output.chunks }}" - # 2. Store the document, embedding, and metadata. conflict_target makes - # re-ingesting identical content a no-op (dedupe_key is md5(content)). + # 3. Store every chunk + embedding in one batch. The single `metadata` object + # is applied to every chunk. conflict_target dedupes on md5(content). - name: store action: kb/upsert credential: kb-db depends_on: + - chunk - embed - timeout: "15s" + timeout: "30s" params: table: kb_documents - content: "{{ inputs.content }}" - vector: "{{ steps['embed'].output.vector }}" + contents: "{{ steps['chunk'].output.chunks }}" + vectors: "{{ steps['embed'].output.vectors }}" metadata: title: "{{ inputs.title }}" source: "{{ inputs.source }}"