chunking.test.ts

chunking.test.tspackages/search/src/chunking.test.tsraw ← src6,009 B← Back to the app
1import { describe, expect, it } from "vitest";
2 
3import { chunkText, DEFAULT_CHUNK_OPTIONS, estimateTokens } from "./chunking.js";
4 
5describe("estimateTokens", () => {
6  it("estimates ~4 chars/token via words/0.75", () => {
7    expect(estimateTokens("")).toBe(0);
8    expect(estimateTokens("one two three four")).toBe(Math.ceil(4 / 0.75));
9  });
10});
11 
12describe("chunkText", () => {
13  it("returns [] for empty/whitespace input", () => {
14    expect(chunkText("")).toEqual([]);
15    expect(chunkText("   \n\n  ")).toEqual([]);
16  });
17 
18  it("keeps a short text as a single chunk", () => {
19    const chunks = chunkText("Hello world");
20    expect(chunks).toHaveLength(1);
21    expect(chunks[0]!.text).toBe("Hello world");
22    expect(chunks[0]!.chunkIndex).toBe(0);
23    expect(chunks[0]!.charStart).toBe(0);
24    expect(chunks[0]!.charEnd).toBe(11);
25  });
26 
27  it("splits on paragraph boundaries and keeps char offsets exact", () => {
28    const a = "alpha paragraph here";
29    const b = "beta paragraph here";
30    const text = a + "\n\n" + b;
31    // maxChars 30 < 40 (a + sep + b) so the pieces stay separate
32    const chunks = chunkText(text, { maxChars: 30, overlapChars: 0 });
33    expect(chunks).toHaveLength(2);
34    // offsets point into the original text
35    expect(text.slice(chunks[0]!.charStart, chunks[0]!.charEnd)).toBe(a);
36    expect(text.slice(chunks[1]!.charStart, chunks[1]!.charEnd)).toBe(b);
37  });
38 
39  it("merges two paragraphs that fit in one chunk (byte-exact source span)", () => {
40    const a = "alpha paragraph here";
41    const b = "beta paragraph here";
42    const text = a + "\n\n" + b;
43    const chunks = chunkText(text, { maxChars: 50, overlapChars: 0 });
44    expect(chunks).toHaveLength(1);
45    expect(chunks[0]!.charStart).toBe(0);
46    expect(chunks[0]!.charEnd).toBe(text.length);
47    // the merged text is exactly the source span (original gap preserved)
48    expect(chunks[0]!.text).toBe(text);
49    expect(text.slice(chunks[0]!.charStart, chunks[0]!.charEnd)).toBe(chunks[0]!.text);
50  });
51 
52  it("hard-splits oversized paragraphs on word boundaries", () => {
53    const words = Array.from({ length: 100 }, (_, i) => `w${i}`).join(" ");
54    const chunks = chunkText(words, { maxChars: 100, overlapChars: 0 });
55    expect(chunks.length).toBeGreaterThan(1);
56    const originalWords = words.split(" ");
57    let cursor = 0;
58    for (const c of chunks) {
59      expect(c.text.length).toBeLessThanOrEqual(100);
60      // each piece is a contiguous run of whole words from the original
61      const pieceWords = c.text.split(/\s+/).filter(Boolean);
62      const run = originalWords.slice(cursor, cursor + pieceWords.length).join(" ");
63      expect(c.text).toBe(run);
64      cursor += pieceWords.length;
65    }
66    expect(cursor).toBe(100); // every word accounted for, in order
67  });
68 
69  it("keeps small fences intact in a single chunk", () => {
70    const code = "```ts\n" + "line ".repeat(80).trim() + "\n```";
71    const text = "before\n\n" + code + "\nafter";
72    const chunks = chunkText(text, { maxChars: 500, overlapChars: 0 });
73    // the fence is never split across chunks
74    const fenceChunks = chunks.filter((c) => c.text.includes("```ts"));
75    expect(fenceChunks).toHaveLength(1);
76    expect(fenceChunks[0]!.text).toContain("line line line");
77    expect(fenceChunks[0]!.text).toContain("\n```");
78  });
79 
80  it("re-brackets an oversized fence that must be hard-split", () => {
81    const code = "```ts\n" + "x ".repeat(200).trim() + "\n```";
82    const chunks = chunkText(code, { maxChars: 100, overlapChars: 0 });
83    expect(chunks.length).toBeGreaterThan(1);
84    // first chunk opens the fence
85    expect(chunks[0]!.text.startsWith("```ts")).toBe(true);
86    // every chunk re-opens it; only the last closes it
87    for (const c of chunks) {
88      expect(c.text.startsWith("```ts")).toBe(true);
89    }
90    const last = chunks[chunks.length - 1]!;
91    expect(last.text.trimEnd().endsWith("```")).toBe(true);
92    for (const c of chunks.slice(0, -1)) {
93      expect(c.text.trimEnd().endsWith("```")).toBe(false);
94    }
95  });
96 
97  it("blank lines inside a fence do not split", () => {
98    const code = "```\nalpha\n\nbeta\n```";
99    const chunks = chunkText("intro\n\n" + code, { maxChars: 500, overlapChars: 0 });
100    const fence = chunks.find((c) => c.text.includes("```"));
101    expect(fence).toBeDefined();
102    expect(fence!.text).toContain("alpha\n\nbeta");
103  });
104 
105  it("merges adjacent small pieces up to maxChars", () => {
106    const text = ["one", "two", "three", "four", "five"].map((w) => w + " " + w).join("\n\n");
107    const chunks = chunkText(text, { maxChars: 200, overlapChars: 0 });
108    expect(chunks).toHaveLength(1); // all small pieces merge
109    expect(chunks[0]!.text).toContain("one");
110    expect(chunks[0]!.text).toContain("five");
111  });
112 
113  it("overlaps consecutive chunks by the tail of the previous one", () => {
114    const para1 = "A".repeat(300) + " TAILMARK";
115    const para2 = "B".repeat(300);
116    const text = para1 + "\n\n" + para2;
117    const chunks = chunkText(text, { maxChars: 400, overlapChars: 60 });
118    expect(chunks).toHaveLength(2);
119    expect(chunks[1]!.text).toContain("TAILMARK");
120    // the overlap is a tail of chunk 0's own text
121    expect(chunks[0]!.text.endsWith(chunks[1]!.text.split("\n")[0]!)).toBe(true);
122    // overlap never exceeds budget
123    for (const c of chunks) expect(c.text.length).toBeLessThanOrEqual(400);
124  });
125 
126  it("clamps overlap to maxChars/4", () => {
127    const text = "x ".repeat(500);
128    const chunks = chunkText(text, { maxChars: 200, overlapChars: 5000 });
129    expect(chunks.length).toBeGreaterThan(0);
130  });
131 
132  it("is deterministic", () => {
133    const text = "para one\n\npara two is longer than the first one so it wraps around a bit\n\npara three";
134    const a = chunkText(text, { maxChars: 90, overlapChars: 20 });
135    const b = chunkText(text, { maxChars: 90, overlapChars: 20 });
136    expect(a).toEqual(b);
137  });
138 
139  it("defaults to 800/120", () => {
140    expect(DEFAULT_CHUNK_OPTIONS).toEqual({ maxChars: 800, overlapChars: 120 });
141  });
142});
143