mirror of
https://github.com/garrytan/gbrain.git
synced 2026-08-14 00:48:18 +00:00
fix(chunker): measured hard-split budgets + token-aware capByChars — the follow-up invited in #3477's merge review (#3564)
Co-Authored-By: YMYD <paul@ymyd.co.kr>
This commit is contained in:
Garry Tan
committed by
Sina Matian
Sina Matian
co-authored by
YMYD
YMYD
parent
273bd0e2be
commit
522cfb032a
+79
-87
@@ -20,7 +20,14 @@
|
||||
|
||||
import { chunkText as recursiveChunk } from './recursive.ts';
|
||||
import { buildQualifiedName } from './qualified-names.ts';
|
||||
import { CJK_SLUG_CHARS, CJK_RANGES_REGEX } from '../cjk.ts';
|
||||
import { estimateTokens, estimateEmbedTokens, estimateEmbedTokensCeiling, DEFAULT_MAX_CHUNK_TOKENS } from './token-estimate.ts';
|
||||
import { safeSplitIndex } from '../text-safe.ts';
|
||||
|
||||
// Both estimators moved to token-estimate.ts (#3477 follow-up) so
|
||||
// recursive.ts can share them without an import cycle. Re-exported here:
|
||||
// commands/sync.ts, commands/reindex-code.ts, and tests import them from
|
||||
// this module.
|
||||
export { estimateTokens, estimateEmbedTokens } from './token-estimate.ts';
|
||||
|
||||
// Embed the tree-sitter runtime + per-language grammars as files.
|
||||
// `with { type: 'file' }` returns a path (string) at runtime. Bun bundles
|
||||
@@ -559,7 +566,6 @@ export function parseWithTimeout(
|
||||
}
|
||||
|
||||
const DEFAULT_CHUNKER_TIMEOUT_MS = 30_000;
|
||||
const DEFAULT_MAX_CHUNK_TOKENS = 2000;
|
||||
|
||||
function resolveChunkerTimeoutMs(): number {
|
||||
const raw = process.env.GBRAIN_CHUNKER_TIMEOUT_MS;
|
||||
@@ -851,9 +857,19 @@ function capOversizedChunks(
|
||||
continue;
|
||||
}
|
||||
// Strip the structured header ("[Lang] path:N-M symbol\n\n") so the splitter
|
||||
// works on the raw body; buildChunk re-adds a header to each piece.
|
||||
const body = c.text.replace(/^\[[^\]]+\] [^\n]+\n\n/, '');
|
||||
for (const piece of splitToTokenBudget(body, cap, opts)) {
|
||||
// works on the raw body; buildChunk re-adds a header to each piece. The
|
||||
// re-added header costs tokens too — budget for it, or every piece split
|
||||
// to exactly `cap` re-emerges a header's-worth over it (measured: a 2,000
|
||||
// cap emitted 2,011-token fence chunks when the body alone was capped).
|
||||
// The reservation must be an UPPER bound on the header's contribution:
|
||||
// estimateEmbedTokens is super-additive across a mixed-script join, so the
|
||||
// header's standalone cl100k figure under-counts ~2.5x once the body
|
||||
// contains CJK and the weighted branch takes over (see
|
||||
// estimateEmbedTokensCeiling).
|
||||
const headerMatch = c.text.match(/^\[[^\]]+\] [^\n]+\n\n/);
|
||||
const body = headerMatch ? c.text.slice(headerMatch[0].length) : c.text;
|
||||
const bodyCap = Math.max(1, cap - (headerMatch ? estimateEmbedTokensCeiling(headerMatch[0]) : 0));
|
||||
for (const piece of splitToTokenBudget(body, bodyCap, opts)) {
|
||||
if (!piece.trim()) continue;
|
||||
out.push(buildChunk({
|
||||
body: piece,
|
||||
@@ -880,44 +896,68 @@ function splitToTokenBudget(text: string, cap: number, opts: CodeChunkOptions):
|
||||
chunkSize: opts.fallbackChunkSizeWords ?? 300,
|
||||
chunkOverlap: opts.fallbackOverlapWords ?? 50,
|
||||
}).map((p) => p.text);
|
||||
for (const piece of pieces) {
|
||||
if (estimateEmbedTokens(piece) <= cap) {
|
||||
out.push(piece);
|
||||
continue;
|
||||
// Hard-split budget is derived from each piece's own measured density
|
||||
// (chars per estimated token) scaled to the cap, not a fixed chars-per-
|
||||
// token guess: the previous 3.5 chars/token ASCII assumption undercuts
|
||||
// URL-dense JSON (~2.6 chars/token measured), leaving 2,070–2,095-token
|
||||
// slices past a 2,000 cap. Slices are re-measured and re-derived (density
|
||||
// varies within a piece), so the cap holds by construction.
|
||||
const hardSplit = (p: string): void => {
|
||||
const est = estimateEmbedTokens(p);
|
||||
if (est <= cap) {
|
||||
out.push(p);
|
||||
return;
|
||||
}
|
||||
// Hard-split slice size. Pure-ASCII pieces: ~3.5 chars/token is a
|
||||
// conservative cl100k estimate for source text. CJK-containing pieces:
|
||||
// the weighted estimate can reach 1 token/char, so budget 1 char/token
|
||||
// to keep every slice under cap by construction.
|
||||
const charBudget = Math.max(1, Math.floor(cap * (CJK_RANGES_REGEX.test(piece) ? 1 : 3.5)));
|
||||
for (let i = 0; i < piece.length; i += charBudget) out.push(piece.slice(i, i + charBudget));
|
||||
}
|
||||
const charBudget = Math.max(1, Math.floor((p.length * cap) / est));
|
||||
if (charBudget >= p.length) {
|
||||
out.push(p); // 1-char floor on a tiny cap — nothing left to split
|
||||
return;
|
||||
}
|
||||
// Even out the slice width instead of striding by charBudget and shedding
|
||||
// `p.length mod charBudget` as a standalone piece at EVERY recursion
|
||||
// level: buildChunk re-headers each remainder into its own embedding row,
|
||||
// so a 14.4K fence emitted 5 chunks of 50-86 chars (and, deeper in the
|
||||
// recursion, 3-char slivers) alongside its real content. Evening is free —
|
||||
// the piece count is ceil(length / charBudget) either way, so the same
|
||||
// content is spread over the same number of chunks — and width <=
|
||||
// charBudget by construction, so the token budget still holds.
|
||||
//
|
||||
// The width is re-derived from what REMAINS on every step rather than
|
||||
// fixed up front, because safeSplitIndex can back a cut off by up to two
|
||||
// units and a fixed width lets that drift accumulate into a tail runt
|
||||
// (measured on an all-astral blob: 4-unit chunks trailing 724-unit ones).
|
||||
let i = 0;
|
||||
while (i < p.length) {
|
||||
const remaining = p.length - i;
|
||||
const partsLeft = Math.ceil(remaining / charBudget);
|
||||
// `i === 0` cannot recurse on the whole piece — charBudget < p.length is
|
||||
// checked above, so partsLeft >= 2 on the first step. The guard keeps a
|
||||
// degenerate budget from looping instead of terminating.
|
||||
if (partsLeft <= 1) {
|
||||
if (i === 0) out.push(p);
|
||||
else hardSplit(p.slice(i));
|
||||
return;
|
||||
}
|
||||
// The budget is derived from measured density, so it has arbitrary
|
||||
// parity: a raw slice at `i + width` orphans a UTF-16 surrogate half.
|
||||
// safeSplitIndex backs the cut off a pair (#2011 — a lone surrogate is
|
||||
// rejected by Postgres inside a ::jsonb cast and aborts the whole batch).
|
||||
const end = safeSplitIndex(p, i + Math.ceil(remaining / partsLeft));
|
||||
if (end <= i) {
|
||||
// Degenerate width (a 1-char budget backing off a surrogate pair) —
|
||||
// emit rather than drop, and never re-enter on the same string.
|
||||
if (i === 0) out.push(p);
|
||||
else hardSplit(p.slice(i));
|
||||
return;
|
||||
}
|
||||
hardSplit(p.slice(i, end));
|
||||
i = end;
|
||||
}
|
||||
};
|
||||
for (const piece of pieces) hardSplit(piece);
|
||||
return out;
|
||||
}
|
||||
|
||||
const CJK_CHARS_G = new RegExp(`[${CJK_SLUG_CHARS}]`, 'g');
|
||||
|
||||
/**
|
||||
* Embedding-safe token estimate for the oversize cap. estimateTokens
|
||||
* (cl100k) matches embedding-family tokenizers closely on pure-ASCII source
|
||||
* (measured identical on English prose and JSON vs Qwen3-Embedding), but
|
||||
* UNDERCOUNTS mixed CJK+ASCII chunks — measured −31% on URL-dense Korean
|
||||
* text vs the Qwen3 embedding tokenizer, which is exactly the shape that
|
||||
* overflows strict embedding backends (#2826). For chunks containing CJK,
|
||||
* take the max of cl100k and a per-char-class overestimate (CJK 1.0/char,
|
||||
* other non-whitespace 0.75/char, whitespace 0.1/char). CJK-DOMINANT text
|
||||
* is unaffected too: cl100k already counts it above the weighted form, so
|
||||
* max() returns the same value as today. Only mixed-script chunks — the
|
||||
* measured divergence class — estimate higher.
|
||||
*/
|
||||
export function estimateEmbedTokens(text: string): number {
|
||||
const cjk = (text.match(CJK_CHARS_G) || []).length;
|
||||
if (cjk === 0) return estimateTokens(text);
|
||||
const ws = (text.match(/\s/g) || []).length;
|
||||
const weighted = Math.ceil(cjk + (text.length - cjk - ws) * 0.75 + ws * 0.1);
|
||||
return Math.max(estimateTokens(text), weighted);
|
||||
}
|
||||
|
||||
// ---------- Internals ----------
|
||||
|
||||
function fallbackChunks(
|
||||
@@ -1245,54 +1285,6 @@ function sanitize(name: string): string {
|
||||
return name.replace(/[\n\r\t]+/g, ' ').replace(/\s+/g, ' ').trim();
|
||||
}
|
||||
|
||||
// v0.19.0 (Layer 5): accurate token count via @dqbd/tiktoken cl100k_base,
|
||||
// the same encoder text-embedding-3-large uses. The old len/4 heuristic was
|
||||
// 2-3x off for code. Lazy-init so dev and compiled-binary both only pay
|
||||
// the init cost once. Falls back to the heuristic if the encoder fails
|
||||
// to load (vanishingly unlikely but keeps the chunker available).
|
||||
let tiktokenEncoder: { encode: (s: string) => Uint32Array; free: () => void } | null = null;
|
||||
let tiktokenInitialized = false;
|
||||
|
||||
// v0.20.0 Cathedral II Layer 8 (D1) — exported so commands/sync.ts can
|
||||
// estimate embed cost before a --all sync blows a surprise OpenAI bill.
|
||||
// Same cl100k_base tokenizer the embedding path actually uses, so cost
|
||||
// estimates match actual billing within tokenizer noise.
|
||||
export function estimateTokens(text: string): number {
|
||||
if (!text) return 0;
|
||||
if (!tiktokenInitialized) {
|
||||
try {
|
||||
// eslint-disable-next-line @typescript-eslint/no-var-requires
|
||||
const m = require('@dqbd/tiktoken');
|
||||
tiktokenEncoder = m.get_encoding('cl100k_base');
|
||||
} catch {
|
||||
tiktokenEncoder = null;
|
||||
}
|
||||
tiktokenInitialized = true;
|
||||
}
|
||||
if (tiktokenEncoder) {
|
||||
try {
|
||||
return tiktokenEncoder.encode(text).length;
|
||||
} catch {
|
||||
// Code legitimately contains tiktoken special-token strings (e.g. CLIP/GPT
|
||||
// tokenizers embed the literal "<|endoftext|>"). The default encode() uses
|
||||
// disallowed_special='all' and THROWS on those, crashing reindex-code on
|
||||
// valid source files. For a token COUNT we don't need special-token
|
||||
// semantics: re-encode treating them as ordinary text (never throws),
|
||||
// heuristic only if even that fails.
|
||||
try {
|
||||
return (
|
||||
tiktokenEncoder as unknown as {
|
||||
encode: (s: string, allowed: string[], disallowed: string[]) => Uint32Array;
|
||||
}
|
||||
).encode(text, [], []).length;
|
||||
} catch {
|
||||
return Math.max(1, Math.ceil(text.length / 4));
|
||||
}
|
||||
}
|
||||
}
|
||||
return Math.max(1, Math.ceil(text.length / 4));
|
||||
}
|
||||
|
||||
// v0.20.0 Cathedral II Layer 4: display name derived from the language
|
||||
// manifest. Single source of truth — adding a new language via
|
||||
// registerLanguage() automatically exposes its displayName to chunk
|
||||
|
||||
@@ -13,11 +13,16 @@
|
||||
* v0.32.7: maxChars hard cap (default 6000) sliding-window safety belt
|
||||
* guarantees no chunk overflows OpenAI's 8192-token embedding limit even
|
||||
* on pathological CJK / whitespace-less text.
|
||||
* #3477 follow-up: the belt also bounds ESTIMATED embedding tokens
|
||||
* (DEFAULT_MAX_CHUNK_TOKENS, shared with the code chunker's oversize cap) —
|
||||
* a char-only cap cannot bound tokens for CJK/dense text (#3037, #2826).
|
||||
*
|
||||
* Lossless invariant: non-overlapping portions reassemble to original.
|
||||
*/
|
||||
|
||||
import { countCJKAwareWords, CJK_SENTENCE_DELIMITERS, CJK_CLAUSE_DELIMITERS } from '../cjk.ts';
|
||||
import { estimateEmbedTokens, DEFAULT_MAX_CHUNK_TOKENS } from './token-estimate.ts';
|
||||
import { safeSplitIndex } from '../text-safe.ts';
|
||||
|
||||
/**
|
||||
* Markdown chunker version. Folded into the per-page chunker_version column
|
||||
@@ -109,29 +114,96 @@ export function chunkText(text: string, opts?: ChunkOptions): TextChunk[] {
|
||||
}
|
||||
|
||||
/**
|
||||
* Hard-cap a chunk's char length via a sliding window. Returns the input
|
||||
* unchanged when it's already ≤ maxChars.
|
||||
* Hard-cap a chunk via a sliding window — by char length AND by estimated
|
||||
* embedding tokens. Returns the input unchanged when it fits both budgets.
|
||||
*
|
||||
* Overlap is min(500, maxChars/10) so successive windows preserve semantic
|
||||
* The char budget (maxChars, default 6000) is the historical belt; the token
|
||||
* budget (DEFAULT_MAX_CHUNK_TOKENS, shared with the code chunker's oversize
|
||||
* cap) is the constraint embedders actually enforce. A char-only cap cannot
|
||||
* bound tokens: 6000 CJK-dense chars run 3-6k tokens, past strict embedder
|
||||
* contexts (nomic-embed-text 2048), so those chunks fail on every embed
|
||||
* sweep, silently, forever (#3037) — and URL-dense CJK markdown emits
|
||||
* over-limit chunks well under maxChars (#2826). When the text over-runs the
|
||||
* token budget, the window is derived from its own measured density —
|
||||
* floor(length × budget / estimate) — and every slice is re-checked (local
|
||||
* density can exceed the whole-text average), re-deriving on the slice until
|
||||
* each piece fits. ASCII prose is unaffected: 6000 chars measure ~1.5-1.7k
|
||||
* cl100k tokens, under the budget, so the window stays maxChars.
|
||||
*
|
||||
* Overlap is min(500, window/10) so successive windows preserve semantic
|
||||
* continuity across the cut.
|
||||
*
|
||||
* v0.32.7. BMP-only safe (does not split astral surrogate pairs in practice
|
||||
* because declared CJK ranges are all BMP; widening to astral Han support
|
||||
* is a v0.33+ follow-up that requires Array.from-style codepoint iteration).
|
||||
* v0.32.7. Surrogate-safe: the window is derived from measured density and so
|
||||
* has arbitrary parity, which a raw slice would use to cut an astral pair in
|
||||
* half — every boundary goes through safeSplitIndex. (The former "BMP-only
|
||||
* safe" note rested on maxChars=6000 and stride=5500 both being even;
|
||||
* deriving the window from density retired that guarantee.)
|
||||
*/
|
||||
function capByChars(text: string, maxChars: number): string[] {
|
||||
if (text.length <= maxChars) return text.length > 0 ? [text] : [];
|
||||
const overlap = Math.min(500, Math.floor(maxChars / 10));
|
||||
const stride = Math.max(1, maxChars - overlap);
|
||||
function capByChars(text: string, maxChars: number, knownEst?: number): string[] {
|
||||
if (text.length === 0) return [];
|
||||
const est = knownEst ?? probeEmbedTokens(text);
|
||||
const window = est <= DEFAULT_MAX_CHUNK_TOKENS
|
||||
? maxChars
|
||||
: Math.max(1, Math.min(maxChars, Math.floor((text.length * DEFAULT_MAX_CHUNK_TOKENS) / est)));
|
||||
if (text.length <= window) {
|
||||
// Emitting the text whole is the one path that skips the per-slice
|
||||
// re-check below, so a PROBED estimate has to be confirmed exactly first:
|
||||
// a sparse ASCII head can under-read a dense CJK tail.
|
||||
if (knownEst !== undefined || text.length <= DENSITY_PROBE_CHARS) return [text];
|
||||
const exact = estimateEmbedTokens(text);
|
||||
return exact <= DEFAULT_MAX_CHUNK_TOKENS ? [text] : capByChars(text, maxChars, exact);
|
||||
}
|
||||
// The stride keeps its nominal window-minus-overlap value. Evening the
|
||||
// windows out (as the header-budget hard split does) is WRONG here: that
|
||||
// splitter partitions, this one overlaps, so shrinking the stride to land
|
||||
// the last window flush against the end collapses successive windows into
|
||||
// near-duplicates — measured on scripts/test-weights.json, two 6,047-char
|
||||
// chunks differing by 47 chars. A short final window is the cheaper end of
|
||||
// that trade and is the behavior this loop has always had.
|
||||
const overlap = Math.min(500, Math.floor(window / 10));
|
||||
const stride = Math.max(1, window - overlap);
|
||||
const out: string[] = [];
|
||||
for (let i = 0; i < text.length; i += stride) {
|
||||
const slice = text.slice(i, i + maxChars).trim();
|
||||
if (slice.length > 0) out.push(slice);
|
||||
if (i + maxChars >= text.length) break;
|
||||
let i = 0;
|
||||
while (i < text.length) {
|
||||
const end = safeSplitIndex(text, Math.min(text.length, i + window));
|
||||
const slice = text.slice(i, end).trim();
|
||||
if (slice.length > 0) {
|
||||
const sliceEst = estimateEmbedTokens(slice);
|
||||
if (sliceEst > DEFAULT_MAX_CHUNK_TOKENS) {
|
||||
// Denser than the text average — re-derive locally, reusing the exact
|
||||
// figure just measured (it also guarantees window < slice.length, so
|
||||
// the recursion strictly shrinks).
|
||||
out.push(...capByChars(slice, maxChars, sliceEst));
|
||||
} else {
|
||||
out.push(slice);
|
||||
}
|
||||
}
|
||||
if (end >= text.length) break;
|
||||
const next = safeSplitIndex(text, Math.min(text.length, i + stride));
|
||||
i = next > i ? next : i + 1;
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
/**
|
||||
* Chars measured to derive the window. estimateEmbedTokens is SUPERLINEAR on
|
||||
* CJK — measured on this repo's encoder: 2K chars 11ms, 6K 99ms, 20K 1,138ms —
|
||||
* and capByChars runs on every chunk, so measuring the whole text up front
|
||||
* dominates the chunker (the 20K-char whitespace-less CJK cap test went from
|
||||
* an O(1) length compare to a 6.7s run, past bun's 5s per-test limit, on a
|
||||
* cold encoder). The window only needs an approximate density: every emitted
|
||||
* slice is re-measured exactly, denser-than-average slices recurse on that
|
||||
* exact figure, and the one path that emits without a re-check confirms
|
||||
* exactly first — so the cap holds regardless of what the probe reads.
|
||||
*/
|
||||
const DENSITY_PROBE_CHARS = 2000;
|
||||
|
||||
function probeEmbedTokens(text: string): number {
|
||||
if (text.length <= DENSITY_PROBE_CHARS) return estimateEmbedTokens(text);
|
||||
const head = text.slice(0, safeSplitIndex(text, DENSITY_PROBE_CHARS));
|
||||
return Math.ceil((estimateEmbedTokens(head) * text.length) / head.length);
|
||||
}
|
||||
|
||||
function recursiveSplit(text: string, level: number, target: number): string[] {
|
||||
if (level >= DELIMITERS.length) {
|
||||
// Level 4: split on whitespace
|
||||
|
||||
@@ -0,0 +1,121 @@
|
||||
/**
|
||||
* Embedding-token estimation — shared by both chunkers (#3477 follow-up).
|
||||
*
|
||||
* Moved out of code.ts so recursive.ts can measure with the same estimator:
|
||||
* code.ts imports recursive.ts, so recursive.ts could never import these from
|
||||
* code.ts without a cycle. The natural home suggested in #3477's merge review
|
||||
* was cjk.ts, but cjk.ts is a check:fuzz-purity bundle target and tiktoken's
|
||||
* loader pulls node:fs into the bundle — so the estimators live here, one
|
||||
* layer above cjk.ts (whose char classes they reuse) and below both chunkers.
|
||||
*/
|
||||
|
||||
import { CJK_SLUG_CHARS } from '../cjk.ts';
|
||||
|
||||
/**
|
||||
* Default hard budget for any emitted chunk's estimated embedding tokens.
|
||||
* Shared by capOversizedChunks (code.ts, #1675) and capByChars (recursive.ts).
|
||||
* 2000 keeps a margin under the smallest common strict embedder contexts
|
||||
* (nomic-embed-text 2048 — #3037; llama-server -ub 2048 — #2826).
|
||||
*/
|
||||
export const DEFAULT_MAX_CHUNK_TOKENS = 2000;
|
||||
|
||||
// v0.19.0 (Layer 5): accurate token count via @dqbd/tiktoken cl100k_base,
|
||||
// the same encoder text-embedding-3-large uses. The old len/4 heuristic was
|
||||
// 2-3x off for code. Lazy-init so dev and compiled-binary both only pay
|
||||
// the init cost once. Falls back to the heuristic if the encoder fails
|
||||
// to load (vanishingly unlikely but keeps the chunker available).
|
||||
let tiktokenEncoder: { encode: (s: string) => Uint32Array; free: () => void } | null = null;
|
||||
let tiktokenInitialized = false;
|
||||
|
||||
// v0.20.0 Cathedral II Layer 8 (D1) — re-exported from code.ts so
|
||||
// commands/sync.ts can estimate embed cost before a --all sync blows a
|
||||
// surprise OpenAI bill. Same cl100k_base tokenizer the embedding path
|
||||
// actually uses, so cost estimates match actual billing within tokenizer
|
||||
// noise.
|
||||
export function estimateTokens(text: string): number {
|
||||
if (!text) return 0;
|
||||
if (!tiktokenInitialized) {
|
||||
try {
|
||||
// eslint-disable-next-line @typescript-eslint/no-var-requires
|
||||
const m = require('@dqbd/tiktoken');
|
||||
tiktokenEncoder = m.get_encoding('cl100k_base');
|
||||
} catch {
|
||||
tiktokenEncoder = null;
|
||||
}
|
||||
tiktokenInitialized = true;
|
||||
}
|
||||
if (tiktokenEncoder) {
|
||||
try {
|
||||
return tiktokenEncoder.encode(text).length;
|
||||
} catch {
|
||||
// Code legitimately contains tiktoken special-token strings (e.g. CLIP/GPT
|
||||
// tokenizers embed the literal "<|endoftext|>"). The default encode() uses
|
||||
// disallowed_special='all' and THROWS on those, crashing reindex-code on
|
||||
// valid source files. For a token COUNT we don't need special-token
|
||||
// semantics: re-encode treating them as ordinary text (never throws),
|
||||
// heuristic only if even that fails.
|
||||
try {
|
||||
return (
|
||||
tiktokenEncoder as unknown as {
|
||||
encode: (s: string, allowed: string[], disallowed: string[]) => Uint32Array;
|
||||
}
|
||||
).encode(text, [], []).length;
|
||||
} catch {
|
||||
return Math.max(1, Math.ceil(text.length / 4));
|
||||
}
|
||||
}
|
||||
}
|
||||
return Math.max(1, Math.ceil(text.length / 4));
|
||||
}
|
||||
|
||||
const CJK_CHARS_G = new RegExp(`[${CJK_SLUG_CHARS}]`, 'g');
|
||||
|
||||
/**
|
||||
* Embedding-safe token estimate for the oversize caps. estimateTokens
|
||||
* (cl100k) matches embedding-family tokenizers closely on pure-ASCII source
|
||||
* (measured identical on English prose and JSON vs Qwen3-Embedding), but
|
||||
* UNDERCOUNTS mixed CJK+ASCII chunks — measured −31% on URL-dense Korean
|
||||
* text vs the Qwen3 embedding tokenizer, which is exactly the shape that
|
||||
* overflows strict embedding backends (#2826). For chunks containing CJK,
|
||||
* take the max of cl100k and a per-char-class overestimate (CJK 1.0/char,
|
||||
* other non-whitespace 0.75/char, whitespace 0.1/char). CJK-DOMINANT text
|
||||
* is unaffected too: cl100k already counts it above the weighted form, so
|
||||
* max() returns the same value as today. Only mixed-script chunks — the
|
||||
* measured divergence class — estimate higher.
|
||||
*/
|
||||
export function estimateEmbedTokens(text: string): number {
|
||||
const cjk = (text.match(CJK_CHARS_G) || []).length;
|
||||
if (cjk === 0) return estimateTokens(text);
|
||||
return Math.max(estimateTokens(text), weightedTokens(text, cjk));
|
||||
}
|
||||
|
||||
/** The per-char-class overestimate half of estimateEmbedTokens. Linear (two
|
||||
* regex scans), unlike the cl100k encoder — see estimateEmbedTokensCeiling. */
|
||||
function weightedTokens(text: string, cjk: number): number {
|
||||
const ws = (text.match(/\s/g) || []).length;
|
||||
return Math.ceil(cjk + (text.length - cjk - ws) * 0.75 + ws * 0.1);
|
||||
}
|
||||
|
||||
/**
|
||||
* Upper bound on what `text` contributes to `estimateEmbedTokens(text + rest)`
|
||||
* for ANY `rest` — i.e. the figure to RESERVE when a fragment will be glued
|
||||
* onto a body whose script mix is not yet known.
|
||||
*
|
||||
* estimateEmbedTokens is super-additive across a mixed-script join. It only
|
||||
* switches to the weighted branch when the text it is handed contains CJK, so
|
||||
* a pure-ASCII fragment measured ALONE costs cl100k (a 59-char structured
|
||||
* chunk header = 17 tokens) while the SAME fragment inside a chunk whose body
|
||||
* contains CJK costs ~0.75/char on the weighted branch (~42 tokens). Reserving
|
||||
* the standalone figure under-counts ~2.5x, and capOversizedChunks then emits
|
||||
* body-capped pieces that re-emerge over the cap once the header is re-added
|
||||
* (measured: 2,006- and 2,023-token chunks on src/core/migrate.ts against a
|
||||
* 2,000 cap, where the pre-#3564 chunker emitted none).
|
||||
*
|
||||
* Taking max(cl100k, weighted) unconditionally is a true bound because the
|
||||
* weighted form is additive per char class, so weighted(a + b) <=
|
||||
* weighted(a) + weighted(b), and cl100k does not gain tokens across the
|
||||
* header's trailing blank line.
|
||||
*/
|
||||
export function estimateEmbedTokensCeiling(text: string): number {
|
||||
return Math.max(estimateTokens(text), weightedTokens(text, (text.match(CJK_CHARS_G) || []).length));
|
||||
}
|
||||
@@ -0,0 +1,216 @@
|
||||
/**
|
||||
* #3477 follow-up — the two items flagged in its merge review:
|
||||
*
|
||||
* (1) splitToTokenBudget's hard-split budget is derived from each piece's
|
||||
* own measured density (chars per estimated token) instead of a fixed
|
||||
* 3.5 chars/token guess. URL-dense ASCII JSON runs ~2.6 chars/token, so
|
||||
* the old budget let 2,070–2,299-token slices past a 2,000 cap.
|
||||
*
|
||||
* (3) estimateTokens/estimateEmbedTokens moved below both chunkers
|
||||
* (token-estimate.ts — cjk.ts itself is a check:fuzz-purity target and
|
||||
* tiktoken's loader pulls node:fs; code.ts imports recursive.ts, so
|
||||
* recursive.ts could never reuse them without a cycle), letting
|
||||
* capByChars bound estimated embedding tokens too —
|
||||
* the fix for the #3037 shape (CJK-dense chunks under maxChars=6000
|
||||
* but over the embedder context, permanently unembeddable, silently)
|
||||
* and #2826's markdown reproduction (URL-dense Korean at defaults
|
||||
* emitting ~4,200-char / ~2,200-token chunks).
|
||||
*/
|
||||
|
||||
import { describe, test, expect } from 'bun:test';
|
||||
import {
|
||||
estimateTokens as estimateTokensViaCode,
|
||||
estimateEmbedTokens as estimateEmbedTokensViaCode,
|
||||
chunkCodeText,
|
||||
} from '../../src/core/chunkers/code.ts';
|
||||
import { estimateTokens, estimateEmbedTokens, DEFAULT_MAX_CHUNK_TOKENS } from '../../src/core/chunkers/token-estimate.ts';
|
||||
import { chunkText } from '../../src/core/chunkers/recursive.ts';
|
||||
|
||||
/** A string survives a UTF-8 round trip only if it is well-formed UTF-16 —
|
||||
* i.e. no orphaned surrogate half. Postgres rejects a lone surrogate inside a
|
||||
* `::jsonb` cast and aborts the whole batch (#2011). */
|
||||
function isWellFormedUtf16(s: string): boolean {
|
||||
return Buffer.from(s, 'utf8').toString('utf8') === s;
|
||||
}
|
||||
|
||||
/** ASCII structured header + CJK-dense body: the shape where the re-added
|
||||
* header's token cost is SUPER-additive (see the header-reservation test). */
|
||||
function mixedScriptTypeScript(lines: number): string {
|
||||
const body = Array.from({ length: lines }, (_, i) =>
|
||||
` // 설정 항목 ${i}: 환경 변수와 기본값을 병합해 최종 구성을 만든다 (참조 config/${i})\n` +
|
||||
` const option_${i} = resolveOption('key_${i}', defaults.key_${i}, { 우선순위: ${i} });`,
|
||||
).join('\n');
|
||||
return `export function loadEverything(defaults: Defaults) {\n${body}\n return { ok: true };\n}\n`;
|
||||
}
|
||||
|
||||
/** URL-dense ASCII JSON — ~2.6 chars/token, the (1) leak shape. */
|
||||
function urlDenseAsciiJson(targetChars: number): string {
|
||||
const entries: string[] = [];
|
||||
let i = 0;
|
||||
let len = 0;
|
||||
while (len < targetChars) {
|
||||
const hex = ((i * 48271) % 65521).toString(16) + ((i * 69621) % 233280).toString(16) + ((i * 16807) % 104729).toString(16);
|
||||
const row =
|
||||
` "row_${i}": { "href": "https://api.example.com/v3/resources/${hex}?sig=ab${i}cd&expires=17${i}&scope=read%2Fwrite", "etag": "W/\\"x${i}y\\"", "n": ${i} }`;
|
||||
entries.push(row);
|
||||
len += row.length;
|
||||
i++;
|
||||
}
|
||||
return `{\n${entries.join(',\n')}\n}`;
|
||||
}
|
||||
|
||||
/** URL-dense Korean rollup lines — #2826's markdown reproduction shape. */
|
||||
function urlDenseKoreanMarkdown(lines: number): string {
|
||||
return Array.from({ length: lines }, (_, i) =>
|
||||
`- 항목 ${i}: 검증용 한국어 설명 문장이 이어집니다 · 링크: https://docs.example.com/pages/${String(i).padStart(32, '0')}?v=abcdef0123456789&ref=sample`,
|
||||
).join('\n');
|
||||
}
|
||||
|
||||
describe('estimator home (cjk.ts) — the (3) move', () => {
|
||||
test('code.ts re-exports are the same functions (import sites unchanged)', () => {
|
||||
expect(estimateTokensViaCode).toBe(estimateTokens);
|
||||
expect(estimateEmbedTokensViaCode).toBe(estimateEmbedTokens);
|
||||
});
|
||||
});
|
||||
|
||||
describe('splitToTokenBudget — measured hard-split budget, the (1) leak', () => {
|
||||
test('URL-dense ASCII json fence stays under the default cap — headers included, no slack (previously 2,379-token max)', async () => {
|
||||
const src = urlDenseAsciiJson(14_400);
|
||||
const chunks = await chunkCodeText(src, 'fence.json');
|
||||
expect(chunks.length).toBeGreaterThan(1);
|
||||
for (const c of chunks) {
|
||||
// STRICT: the emitted chunk (structured header + body) fits the cap.
|
||||
// The splitter reserves the header's tokens from the body budget, so
|
||||
// no "body capped, header pushed it over" residue survives.
|
||||
expect(estimateEmbedTokens(c.text)).toBeLessThanOrEqual(DEFAULT_MAX_CHUNK_TOKENS);
|
||||
}
|
||||
// Content preserved — first and last rows survive the re-split.
|
||||
const joined = chunks.map((c) => c.text).join('\n');
|
||||
expect(joined).toContain('"row_0"');
|
||||
expect(joined).toContain('scope=read%2Fwrite');
|
||||
});
|
||||
});
|
||||
|
||||
describe('capByChars — token-aware belt, the (3) payoff', () => {
|
||||
test('URL-dense Korean markdown at defaults stays under the token budget (previously ~2,200-token chunks)', () => {
|
||||
const chunks = chunkText(urlDenseKoreanMarkdown(120));
|
||||
expect(chunks.length).toBeGreaterThan(0);
|
||||
for (const c of chunks) {
|
||||
expect(estimateEmbedTokens(c.text)).toBeLessThanOrEqual(DEFAULT_MAX_CHUNK_TOKENS);
|
||||
}
|
||||
});
|
||||
|
||||
test('low-density bilingual table (the #3037 shape) splits under the budget instead of shipping over-context chunks', () => {
|
||||
// #3037's failing chunk: mostly-ASCII with a CJK minority (their repro:
|
||||
// 6001 chars, 942 CJK). Density sits BELOW CJK_DENSITY_THRESHOLD, so the
|
||||
// word pipeline counts whitespace tokens and happily builds multi-
|
||||
// thousand-char chunks; the old belt only checked chars (6000), so these
|
||||
// shipped at token counts past strict embedder contexts.
|
||||
const table = Array.from({ length: 120 }, (_, i) =>
|
||||
`ITEM-${String(i).padStart(6, '0')} | 环境配置说明 段落${i} | https://wiki.example.com/pages/${String(i).padStart(20, '0')}?rev=${i}&lang=zh | flags=prod,readonly,audit`,
|
||||
).join('\n');
|
||||
const chunks = chunkText(table);
|
||||
expect(chunks.length).toBeGreaterThan(1);
|
||||
for (const c of chunks) {
|
||||
expect(estimateEmbedTokens(c.text)).toBeLessThanOrEqual(DEFAULT_MAX_CHUNK_TOKENS);
|
||||
}
|
||||
});
|
||||
|
||||
// DROPPED: a 'mixed-density input — every slice re-checked' test used to sit
|
||||
// here. It passed against the pre-cap behavior too, so it proved nothing.
|
||||
// Measured why: its dense run was whitespace-less CJK, which countCJKAwareWords
|
||||
// scores per character, so the word pipeline had already cut it to <=300-char
|
||||
// pieces before the belt was ever consulted (baseline: 20 dense chunks, max
|
||||
// 300 chars, cap never fires). The discriminating shape for the belt is the
|
||||
// LOW-density one — CJK below the density threshold, where the word pipeline
|
||||
// counts whitespace tokens and builds multi-thousand-char chunks — and that
|
||||
// is the #3037 bilingual-table test above.
|
||||
|
||||
test('ASCII prose under both budgets passes through untouched (single chunk, verbatim)', () => {
|
||||
const prose = 'plain english prose that fits comfortably inside every budget. '.repeat(20).trim();
|
||||
const chunks = chunkText(prose);
|
||||
expect(chunks.length).toBe(1);
|
||||
expect(chunks[0]!.text).toBe(prose);
|
||||
});
|
||||
});
|
||||
|
||||
describe('header reservation is an UPPER bound — estimateEmbedTokens is super-additive', () => {
|
||||
test('mixed-script source: the re-added ASCII header never pushes a piece over the cap', async () => {
|
||||
// estimateEmbedTokens takes max(cl100k, per-char-class weighted) and only
|
||||
// switches on the weighted branch when the text contains CJK. An ASCII
|
||||
// header measured ALONE therefore costs cl100k (a 59-char header = 17
|
||||
// tokens), but once it is glued onto a body containing CJK the whole
|
||||
// chunk measures on the weighted branch, where those same 59 chars cost
|
||||
// ~0.75/char (~42 tokens). Reserving the standalone figure under-counts
|
||||
// ~2.5x and the re-headered piece re-emerges over the cap: measured 2,006
|
||||
// and 2,023-token chunks on src/core/migrate.ts against a 2,000 cap,
|
||||
// where the pre-fix chunker emitted none (maxEst 1,439).
|
||||
const src = mixedScriptTypeScript(220);
|
||||
const chunks = await chunkCodeText(src, 'src/config/loader.ts');
|
||||
expect(chunks.length).toBeGreaterThan(1);
|
||||
for (const c of chunks) {
|
||||
expect(estimateEmbedTokens(c.text)).toBeLessThanOrEqual(DEFAULT_MAX_CHUNK_TOKENS);
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe('astral surrogate pairs survive the derived-window splits', () => {
|
||||
// The hard-split budget and the capByChars window are now DERIVED from
|
||||
// measured density, so they land on arbitrary parity. A raw .slice() at
|
||||
// such an offset orphans a UTF-16 surrogate half. The repo already ships
|
||||
// src/core/text-safe.ts:safeSplitIndex for exactly this: #2011 —
|
||||
// `extract --stale` died at ~1,550 pages because excerpt() raw-sliced a
|
||||
// window boundary through an emoji, and a lone surrogate is rejected by
|
||||
// Postgres inside a ::jsonb cast, aborting the WHOLE batch.
|
||||
const ASTRAL = '\u{20000}\u{20001}\u{20002}\u{1F600}';
|
||||
|
||||
test('markdown chunker emits well-formed UTF-16 on an astral-only document', () => {
|
||||
const chunks = chunkText(ASTRAL.repeat(1500));
|
||||
expect(chunks.length).toBeGreaterThan(1);
|
||||
const corrupted = chunks.filter((c) => !isWellFormedUtf16(c.text));
|
||||
expect(corrupted.length).toBe(0);
|
||||
});
|
||||
|
||||
test('code chunker emits well-formed UTF-16 on a whitespace-less astral blob', async () => {
|
||||
const chunks = await chunkCodeText(JSON.stringify({ k: ASTRAL.repeat(3000) }), 'blob.json');
|
||||
expect(chunks.length).toBeGreaterThan(1);
|
||||
const corrupted = chunks.filter((c) => !isWellFormedUtf16(c.text));
|
||||
expect(corrupted.length).toBe(0);
|
||||
});
|
||||
|
||||
test('astral content is preserved across the split, not dropped', () => {
|
||||
const doc = ASTRAL.repeat(1500);
|
||||
const rejoined = chunkText(doc).map((c) => c.text).join('');
|
||||
expect(rejoined).toContain(ASTRAL.repeat(4));
|
||||
});
|
||||
});
|
||||
|
||||
describe('hard split leaves no runt chunks', () => {
|
||||
// A fixed `i += charBudget` stride sheds `length mod charBudget` chars as a
|
||||
// standalone piece at EVERY recursion level, and buildChunk re-headers each
|
||||
// one into its own embedding row — near-empty fragments that still match
|
||||
// queries. Evening the slice width out costs nothing: the piece COUNT is
|
||||
// ceil(length / charBudget) either way, so the same content is redistributed
|
||||
// over the same number of chunks; only the smallest piece changes (measured
|
||||
// on a 45.6K blob: min 589 -> 1900 chars, same 24 pieces).
|
||||
const SLIVER_CHARS = 200;
|
||||
|
||||
test('URL-dense ASCII json fence sheds no sliver (pre-fix: 3-56 char fragments)', async () => {
|
||||
const chunks = await chunkCodeText(urlDenseAsciiJson(14_400), 'fence.json');
|
||||
expect(chunks.length).toBeGreaterThan(1);
|
||||
const bodies = chunks.map((c) => c.text.replace(/^\[[^\]]+\] [^\n]+\n\n/, ''));
|
||||
expect(bodies.filter((b) => b.length < SLIVER_CHARS)).toEqual([]);
|
||||
});
|
||||
|
||||
test('whitespace-less blob — the pure hard-split path — sheds no sliver', async () => {
|
||||
// No whitespace to break on, so recursiveChunk cannot help and every
|
||||
// boundary comes from the hard splitter: the shape where the remainder
|
||||
// stride was most visible (measured 32/54/56-char chunks pre-fix).
|
||||
const blob = JSON.stringify({ d: 'https://example.com/a/b/c?q=1&r=2#frag-'.repeat(1200) });
|
||||
const chunks = await chunkCodeText(blob, 'blob.json');
|
||||
expect(chunks.length).toBeGreaterThan(1);
|
||||
const bodies = chunks.map((c) => c.text.replace(/^\[[^\]]+\] [^\n]+\n\n/, ''));
|
||||
expect(bodies.filter((b) => b.length < SLIVER_CHARS)).toEqual([]);
|
||||
for (const c of chunks) expect(estimateEmbedTokens(c.text)).toBeLessThanOrEqual(DEFAULT_MAX_CHUNK_TOKENS);
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user