Skip to content

Commit 1fc0beb

Browse files
committed
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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WnKgCRfMvFcZAkoVbr8k79
1 parent 9a7e8a8 commit 1fc0beb

9 files changed

Lines changed: 322 additions & 35 deletions

File tree

.changeset/text-chunk-connector.md

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
"@mantle/engine": minor
3+
---
4+
5+
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.
Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
package connector
2+
3+
import (
4+
"context"
5+
"fmt"
6+
"strings"
7+
)
8+
9+
// TextChunkConnector implements text/chunk: split a long document into
10+
// overlapping chunks for embedding. It composes with ai/embed (pass the chunks
11+
// as `input`) and kb/upsert (pass them as `contents`): chunk → embed → upsert.
12+
type TextChunkConnector struct{}
13+
14+
const (
15+
chunkDefaultSize = 1000
16+
chunkDefaultOverlap = 0
17+
)
18+
19+
// chunkText splits text into fixed-size, optionally overlapping chunks. `unit`
20+
// is "chars" (default, Unicode-aware) or "words" (whitespace-separated). size
21+
// and overlap are counted in that unit. Pure — no I/O — so it is fully
22+
// unit-testable.
23+
func chunkText(text string, size, overlap int, unit string) ([]string, error) {
24+
if strings.TrimSpace(text) == "" {
25+
return nil, fmt.Errorf("text must not be empty")
26+
}
27+
if size <= 0 {
28+
return nil, fmt.Errorf("chunk_size must be positive, got %d", size)
29+
}
30+
if overlap < 0 {
31+
return nil, fmt.Errorf("chunk_overlap must be >= 0, got %d", overlap)
32+
}
33+
if overlap >= size {
34+
return nil, fmt.Errorf("chunk_overlap (%d) must be less than chunk_size (%d)", overlap, size)
35+
}
36+
step := size - overlap
37+
38+
switch unit {
39+
case "", "chars":
40+
runes := []rune(text)
41+
return sliceChunks(len(runes), size, step, func(a, b int) string {
42+
return string(runes[a:b])
43+
}), nil
44+
case "words":
45+
words := strings.Fields(text)
46+
return sliceChunks(len(words), size, step, func(a, b int) string {
47+
return strings.Join(words[a:b], " ")
48+
}), nil
49+
default:
50+
return nil, fmt.Errorf("unknown unit %q (use chars or words)", unit)
51+
}
52+
}
53+
54+
// sliceChunks walks [0,n) in windows of `size` advancing by `step`, emitting
55+
// each window via slice(start, end). It stops once a window reaches the end so
56+
// overlap never produces a trailing duplicate.
57+
func sliceChunks(n, size, step int, slice func(a, b int) string) []string {
58+
var chunks []string
59+
for start := 0; start < n; start += step {
60+
end := start + size
61+
if end >= n {
62+
end = n
63+
chunks = append(chunks, slice(start, end))
64+
break
65+
}
66+
chunks = append(chunks, slice(start, end))
67+
}
68+
return chunks
69+
}
70+
71+
func (c *TextChunkConnector) Execute(_ context.Context, params map[string]any) (map[string]any, error) {
72+
text, _ := params["text"].(string)
73+
74+
size := chunkDefaultSize
75+
if v, ok := extractInt(params["chunk_size"]); ok {
76+
size = v
77+
}
78+
overlap := chunkDefaultOverlap
79+
if v, ok := extractInt(params["chunk_overlap"]); ok {
80+
overlap = v
81+
}
82+
unit, _ := params["unit"].(string)
83+
84+
chunks, err := chunkText(text, size, overlap, unit)
85+
if err != nil {
86+
return nil, fmt.Errorf("text/chunk: %w", err)
87+
}
88+
89+
return map[string]any{
90+
"chunks": chunks,
91+
"count": len(chunks),
92+
}, nil
93+
}
Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
package connector
2+
3+
import (
4+
"context"
5+
"strings"
6+
"testing"
7+
)
8+
9+
func TestChunkText_Chars(t *testing.T) {
10+
// "abcdefg" (7), size 3, overlap 1 → step 2: windows 0-3, 2-5, 4-7.
11+
// The last window reaches the end, so there is no trailing "g" chunk.
12+
got, err := chunkText("abcdefg", 3, 1, "chars")
13+
if err != nil {
14+
t.Fatalf("chunkText error: %v", err)
15+
}
16+
want := []string{"abc", "cde", "efg"}
17+
if strings.Join(got, "|") != strings.Join(want, "|") {
18+
t.Errorf("chunks = %v, want %v", got, want)
19+
}
20+
}
21+
22+
func TestChunkText_NoOverlapExactFit(t *testing.T) {
23+
// 6 chars, size 3, overlap 0 → [abc, def], no trailing empty chunk.
24+
got, err := chunkText("abcdef", 3, 0, "chars")
25+
if err != nil {
26+
t.Fatalf("error: %v", err)
27+
}
28+
if len(got) != 2 || got[0] != "abc" || got[1] != "def" {
29+
t.Errorf("chunks = %v, want [abc def]", got)
30+
}
31+
}
32+
33+
func TestChunkText_Words(t *testing.T) {
34+
got, err := chunkText("the quick brown fox jumps", 2, 0, "words")
35+
if err != nil {
36+
t.Fatalf("error: %v", err)
37+
}
38+
want := []string{"the quick", "brown fox", "jumps"}
39+
if strings.Join(got, "|") != strings.Join(want, "|") {
40+
t.Errorf("chunks = %v, want %v", got, want)
41+
}
42+
}
43+
44+
func TestChunkText_ShorterThanSize(t *testing.T) {
45+
got, err := chunkText("hi", 100, 10, "chars")
46+
if err != nil {
47+
t.Fatalf("error: %v", err)
48+
}
49+
if len(got) != 1 || got[0] != "hi" {
50+
t.Errorf("chunks = %v, want [hi]", got)
51+
}
52+
}
53+
54+
func TestChunkText_Unicode(t *testing.T) {
55+
// Multibyte runes must not be split mid-character.
56+
got, err := chunkText("héllo wörld", 3, 0, "chars")
57+
if err != nil {
58+
t.Fatalf("error: %v", err)
59+
}
60+
// 11 runes / size 3 → 4 chunks; first is "hél".
61+
if got[0] != "hél" {
62+
t.Errorf("first chunk = %q, want %q", got[0], "hél")
63+
}
64+
}
65+
66+
func TestChunkText_Errors(t *testing.T) {
67+
cases := []struct {
68+
name string
69+
text string
70+
size, overlap int
71+
unit string
72+
}{
73+
{"empty text", " ", 10, 0, "chars"},
74+
{"zero size", "abc", 0, 0, "chars"},
75+
{"negative overlap", "abc", 10, -1, "chars"},
76+
{"overlap >= size", "abc", 5, 5, "chars"},
77+
{"unknown unit", "abc", 5, 0, "tokens"},
78+
}
79+
for _, tc := range cases {
80+
t.Run(tc.name, func(t *testing.T) {
81+
if _, err := chunkText(tc.text, tc.size, tc.overlap, tc.unit); err == nil {
82+
t.Errorf("expected error for %s", tc.name)
83+
}
84+
})
85+
}
86+
}
87+
88+
func TestTextChunkConnector_Output(t *testing.T) {
89+
out, err := (&TextChunkConnector{}).Execute(context.Background(), map[string]any{
90+
"text": "one two three four five",
91+
"chunk_size": 2,
92+
"chunk_overlap": 0,
93+
"unit": "words",
94+
})
95+
if err != nil {
96+
t.Fatalf("Execute error: %v", err)
97+
}
98+
chunks := out["chunks"].([]string)
99+
if out["count"].(int) != len(chunks) || len(chunks) != 3 {
100+
t.Errorf("count/chunks = %v / %v", out["count"], chunks)
101+
}
102+
}

packages/engine/internal/connector/connector.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ func NewRegistry() *Registry {
3333
r.Register("postgres/query", &PostgresQueryConnector{})
3434
r.Register("kb/upsert", &KBUpsertConnector{})
3535
r.Register("kb/query", &KBQueryConnector{})
36+
r.Register("text/chunk", &TextChunkConnector{})
3637
r.Register("email/send", &EmailSendConnector{})
3738
r.Register("email/receive", &EmailReceiveConnector{})
3839
r.Register("email/move", &EmailMoveConnector{})

packages/engine/internal/connector/kb.go

Lines changed: 21 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -92,7 +92,9 @@ func kbStrings(params map[string]any, singleKey, manyKey string) ([]string, erro
9292
}
9393

9494
// kbMetadata normalizes the optional metadata param into JSON strings aligned
95-
// with the rows. Returns present=false when no metadata param is set.
95+
// with the rows. A single `metadata` object is broadcast to every row (handy
96+
// for chunked ingest that shares a title/source); `metadatas` must match the
97+
// row count one-to-one. Returns present=false when no metadata param is set.
9698
func kbMetadata(params map[string]any, rows int) (jsons []string, present bool, err error) {
9799
single, hasSingle := params["metadata"]
98100
many, hasMany := params["metadatas"]
@@ -103,25 +105,30 @@ func kbMetadata(params map[string]any, rows int) (jsons []string, present bool,
103105
return nil, false, fmt.Errorf("set only one of metadata or metadatas")
104106
}
105107

106-
var objs []any
107108
if hasSingle {
108-
objs = []any{single}
109-
} else {
110-
arr, ok := many.([]any)
111-
if !ok {
112-
return nil, false, fmt.Errorf("metadatas must be an array of objects, got %T", many)
109+
b, mErr := json.Marshal(single)
110+
if mErr != nil {
111+
return nil, false, fmt.Errorf("marshaling metadata: %w", mErr)
113112
}
114-
objs = arr
115-
}
116-
if len(objs) != rows {
117-
return nil, false, fmt.Errorf("metadata count %d does not match content count %d", len(objs), rows)
113+
out := make([]string, rows)
114+
for i := range out {
115+
out[i] = string(b)
116+
}
117+
return out, true, nil
118118
}
119119

120-
out := make([]string, len(objs))
121-
for i, o := range objs {
120+
arr, ok := many.([]any)
121+
if !ok {
122+
return nil, false, fmt.Errorf("metadatas must be an array of objects, got %T", many)
123+
}
124+
if len(arr) != rows {
125+
return nil, false, fmt.Errorf("metadatas count %d does not match content count %d", len(arr), rows)
126+
}
127+
out := make([]string, len(arr))
128+
for i, o := range arr {
122129
b, mErr := json.Marshal(o)
123130
if mErr != nil {
124-
return nil, false, fmt.Errorf("marshaling metadata[%d]: %w", i, mErr)
131+
return nil, false, fmt.Errorf("marshaling metadatas[%d]: %w", i, mErr)
125132
}
126133
out[i] = string(b)
127134
}

packages/engine/internal/connector/kb_test.go

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,28 @@ func TestPrepareUpsert_BatchWithMetadataAndConflict(t *testing.T) {
9898
}
9999
}
100100

101+
func TestPrepareUpsert_SingleMetadataBroadcast(t *testing.T) {
102+
// One `metadata` object applies to every row of a batch (chunked ingest).
103+
_, args, err := prepareUpsert(map[string]any{
104+
"table": "kb_documents",
105+
"contents": []any{"a", "b", "c"},
106+
"vectors": []any{"[1]", "[2]", "[3]"},
107+
"metadata": map[string]any{"source": "doc-1"},
108+
})
109+
if err != nil {
110+
t.Fatalf("prepareUpsert error: %v", err)
111+
}
112+
metas := args[2].([]string)
113+
if len(metas) != 3 {
114+
t.Fatalf("metadata args = %d, want 3 (broadcast)", len(metas))
115+
}
116+
for i, m := range metas {
117+
if !strings.Contains(m, `"source":"doc-1"`) {
118+
t.Errorf("metas[%d] = %q, want broadcast source", i, m)
119+
}
120+
}
121+
}
122+
101123
func TestPrepareUpsert_Errors(t *testing.T) {
102124
cases := []struct {
103125
name string

packages/site/src/content/docs/rag-guide.md

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -82,10 +82,10 @@ mantle run rag-ingest \
8282
--input content="Team A connects to Client C via the X integration; Team B uses the direct API."
8383
```
8484

85-
For a long document, split it into chunks and pass arrays to `ai/embed`
86-
(`input`) and `kb/upsert` (`contents` + `vectors`, plus `metadatas`) to store
87-
them all in one step. Re-ingesting identical content is a no-op (the schema
88-
dedupes on `md5(content)`).
85+
`rag-ingest` splits the document into overlapping chunks with `text/chunk`,
86+
embeds them all in one `ai/embed` call, and stores them in one `kb/upsert`. A
87+
single `metadata` object is broadcast to every chunk. Re-ingesting identical
88+
content is a no-op (the schema dedupes on `md5(content)`).
8989

9090
## Asking questions
9191

@@ -111,6 +111,7 @@ outputs with `--output json` or `-v`).
111111
to match (Titan v2 defaults to 1024). Cohere Bedrock models are a follow-up.
112112
- **You provide the vector store.** `kb/*` run against a pgvector database you
113113
create and point a credential at; they don't provision the extension or table.
114-
- **Chunking is manual.** Split long documents yourself and pass arrays to
115-
`ai/embed` / `kb/upsert`. A native chunking helper and metadata filtering on
116-
`kb/query` are tracked by [#153](https://github.com/dvflw/mantle/issues/153).
114+
- **Chunking is fixed-size.** `text/chunk` splits by a sliding window over
115+
characters or words with overlap. Semantic/recursive (separator-aware or
116+
token-accurate) chunking and metadata filtering on `kb/query` are tracked by
117+
[#153](https://github.com/dvflw/mantle/issues/153).

packages/site/src/content/docs/workflow-reference/connectors.md

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -385,6 +385,50 @@ Nearest-neighbour search over a pgvector table for a query embedding. A thin hel
385385
top_k: 5
386386
```
387387

388+
## text/chunk
389+
390+
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.
391+
392+
**Params:**
393+
394+
| Param | Type | Required | Description |
395+
|---|---|---|---|
396+
| `text` | string | Yes | The document text to split. |
397+
| `chunk_size` | integer | No | Target chunk size, in `unit`s. Default `1000`. |
398+
| `chunk_overlap` | integer | No | Overlap between consecutive chunks, in `unit`s. Must be `< chunk_size`. Default `0`. |
399+
| `unit` | string | No | `chars` (default, Unicode-aware) or `words` (whitespace-separated). |
400+
401+
**Output:** `chunks` (list of strings) and `count` (number).
402+
403+
**Example — chunk → embed → upsert:**
404+
405+
```yaml
406+
- name: chunk
407+
action: text/chunk
408+
params:
409+
text: "{{ inputs.content }}"
410+
chunk_size: 1200
411+
chunk_overlap: 150
412+
413+
- name: embed
414+
action: ai/embed
415+
credential: my-openai
416+
depends_on: [chunk]
417+
params:
418+
model: text-embedding-3-small
419+
input: "{{ steps['chunk'].output.chunks }}"
420+
421+
- name: store
422+
action: kb/upsert
423+
credential: kb-db
424+
depends_on: [chunk, embed]
425+
params:
426+
table: kb_documents
427+
contents: "{{ steps['chunk'].output.chunks }}"
428+
vectors: "{{ steps['embed'].output.vectors }}"
429+
metadata: { source: "{{ inputs.source }}" } # broadcast to every chunk
430+
```
431+
388432
## email/send
389433

390434
Sends an email via SMTP. Supports plaintext and HTML content.

0 commit comments

Comments
 (0)