diff --git a/.gitignore b/.gitignore index e9f580c5..f9dcd5a7 100644 --- a/.gitignore +++ b/.gitignore @@ -93,3 +93,6 @@ server/internal/httpapi/dashboard/dist/* portainer_mcp/ portainer-mcp* tools.yaml + +# Local load-test corpus + throwaway server instance (never committed) +/loadtests/ diff --git a/server/cmd/cix-server/main.go b/server/cmd/cix-server/main.go index a2b8e7c6..54c23cc8 100644 --- a/server/cmd/cix-server/main.go +++ b/server/cmd/cix-server/main.go @@ -437,6 +437,7 @@ func run() (restart bool, err error) { idx := indexer.New(database, vsHolder, embedSvc, logger) idx.SetEmbedIncludePath(cfg.EmbedIncludePath) + idx.SetMaxChunkTokens(cfg.MaxChunkTokens) // Record the active embedding model on every indexed project so the // dashboard can highlight stale vectors when the runtime provider / // model changes. Wire it as a live lookup so a runtime provider diff --git a/server/internal/chunker/chunker.go b/server/internal/chunker/chunker.go index 6391ba1b..5b147971 100644 --- a/server/internal/chunker/chunker.go +++ b/server/internal/chunker/chunker.go @@ -10,7 +10,10 @@ package chunker import ( + "errors" + "github.com/dvcdsys/code-index/server/internal/tokenizer" "log/slog" + "math" "path/filepath" "strings" "sync" @@ -457,15 +460,88 @@ type Reference struct { // falls back to sliding-window chunking for unsupported languages. The maxSize // parameter controls per-chunk character limit; pass 0 to use the default. func ChunkFile(filePath, content, language string, maxSize int) ([]Chunk, []Reference, error) { + return ChunkFileTokens(filePath, content, language, maxSize, nil, 0) +} + +// ChunkFileTokens is ChunkFile with a token budget. +// +// maxSize (bytes) has always been a stand-in for a token limit: the default +// 4500 is "1500 tokens x 3 bytes", a ratio that holds for dense ASCII code and +// falls apart everywhere else — Cyrillic or CJK comments cost two to three +// bytes per character, so a byte-sized chunk carries far fewer tokens than +// intended, while minified JavaScript packs far more. +// +// When budget is non-nil and reports exact counts, the size decision is made +// in tokens instead, and an over-budget chunk is cut on real token boundaries +// (via Budget.SplitPoints) rather than on a byte count. maxTokens <= 0 uses +// DefaultMaxChunkTokens. A nil or estimating budget keeps the byte path +// unchanged, so nothing about existing behaviour depends on a provider +// having a tokenizer. +func ChunkFileTokens(filePath, content, language string, maxSize int, budget tokenizer.Budget, maxTokens int) ([]Chunk, []Reference, error) { if maxSize <= 0 { maxSize = maxChunkSize } - chunks, refs, err := chunkWithTreesitter(filePath, content, language, maxSize) - if err != nil { - // Fallback: sliding window, no references. - return chunkFallback(filePath, content, language), nil, nil + if budget != nil && !budget.ExactCounts() { + // An estimate is what the byte path already is; do not pretend + // otherwise by routing through the token splitter. + budget = nil + } + if budget != nil { + if maxTokens <= 0 { + maxTokens = DefaultMaxChunkTokens + } + if lim := budget.MaxInputTokens(); lim > 0 && maxTokens > lim { + maxTokens = lim + } + } + // With a token budget the byte cap must not fire first: it is the very + // bias being removed (a byte limit cuts Cyrillic or CJK three times + // sooner than ASCII for the same token cost). Let the inner path emit + // whole semantic units and bound them in tokens afterwards. + innerMax := maxSize + if budget != nil { + innerMax = math.MaxInt32 + } + + chunks, refs, err := chunkWithTreesitter(filePath, content, language, innerMax) + if err != nil || len(chunks) == 0 { + // Tree-sitter declined (no grammar, parse failure, minified input) or + // produced nothing. Either way the fallback runs here, where the + // budget is known, rather than inside the tree-sitter path where it + // is not. + return chunkFallbackTokens(filePath, content, language, budget, maxTokens), nil, nil + } + return boundTokens(chunks, budget, maxTokens), refs, nil +} + +// boundTokens enforces the token budget over chunks from ANY path — the +// tree-sitter one, the bash regex extractor, or the sliding-window fallback. +// Applying it here rather than inside the tree-sitter path is deliberate: +// minified JavaScript and files with no grammar are exactly the inputs that +// reach the fallback, and they are also the ones most likely to blow the +// model's input limit. A budget that only covered the happy path would miss +// them. +func boundTokens(chunks []Chunk, budget tokenizer.Budget, maxTokens int) []Chunk { + if budget == nil || maxTokens <= 0 { + return chunks + } + out := make([]Chunk, 0, len(chunks)) + for _, c := range chunks { + // A byte-level BPE token can never cover less than one byte, so a + // chunk shorter than the budget provably fits and needs no counting + // at all. Most chunks are, so this skips the tokenizer on the hot + // path rather than paying for an answer already known. + if len(c.Content) <= maxTokens { + out = append(out, c) + continue + } + if budget.CountTokens(c.Content) > maxTokens { + out = append(out, splitChunkTokens(c, budget, maxTokens)...) + continue + } + out = append(out, c) } - return chunks, refs, nil + return out } // chunkFallback returns reasonable chunks for content that the tree-sitter @@ -476,6 +552,17 @@ func ChunkFile(filePath, content, language string, maxSize int) ([]Chunk, []Refe // `block` ones, which is much more useful for semantic search. If the // extractor returns nil (no symbols found), we fall through to the universal // sliding-window strategy so the file content is still indexed. +// errUseFallback tells the caller that the tree-sitter path declined this +// file and the fallback chunker must run instead. +// +// It exists because chunkWithTreesitter used to CALL chunkFallback itself and +// return the result as success. That made the caller's fallback branch +// unreachable — which is exactly how a token-aware fallback shipped as dead +// code: no-grammar, minified, parse-failure and empty-AST files all came back +// as byte windows wearing a success return, and no budget ever reached them. +// The decision belongs to whoever knows whether a token budget is in play. +var errUseFallback = errors.New("chunker: tree-sitter declined, use fallback") + func chunkFallback(filePath, content, language string) []Chunk { if language == "bash" { if c := bashRegexChunks(filePath, content); len(c) > 0 { @@ -485,6 +572,54 @@ func chunkFallback(filePath, content, language string) []Chunk { return chunkSlidingWindow(filePath, content, language) } +// chunkFallbackTokens is chunkFallback with a token budget: the sliding window +// walks token boundaries instead of a fixed byte count. +// +// The byte window is 4000 bytes regardless of what those bytes contain, so a +// file of Cyrillic or CJK prose — two to three bytes per character — produced +// windows worth a third of the intended tokens, and this is the path such +// files take, because "no grammar" and "not ASCII" go together often enough to +// matter. boundTokens alone could not fix it: it splits chunks that are too +// large and has no way to grow ones that are too small. +func chunkFallbackTokens(filePath, content, language string, budget tokenizer.Budget, maxTokens int) []Chunk { + if budget == nil || maxTokens <= 0 { + return chunkFallback(filePath, content, language) + } + if language == "bash" { + if c := bashRegexChunks(filePath, content); len(c) > 0 { + return boundTokens(c, budget, maxTokens) + } + } + if len(content) == 0 { + return nil + } + // One chunk, then let the token splitter cut it — same code path, same + // guarantees (pieces are substrings, none over budget, line numbers + // tracked) as every other over-budget chunk in this package. + // + // Note this drops the byte window's 500-byte overlap. That overlap existed + // so a match spanning a window boundary would still be found in one of the + // two windows, and nothing replaces it here: pieces are contiguous. The + // trade is deliberate for now — overlap has to be expressed in tokens to + // coexist with a token budget (size pieces at budget minus overlap, then + // extend each start back into the previous piece), which is a change worth + // making on its own rather than smuggling into a correctness fix. The + // tree-sitter path has never overlapped, so this makes the two paths + // consistent rather than making the fallback worse than its neighbours. + whole := Chunk{ + Content: content, + ChunkType: "block", + FilePath: filePath, + StartLine: 1, + EndLine: countNewlines(content) + 1, + Language: language, + } + if budget.CountTokens(content) <= maxTokens { + return []Chunk{whole} + } + return splitChunkTokens(whole, budget, maxTokens) +} + // --------------------------------------------------------------------------- // Tree-sitter path // --------------------------------------------------------------------------- @@ -529,11 +664,11 @@ func chunkWithTreesitter(filePath, content, language string, maxSize int) ([]Chu registryMu.RUnlock() if !ok { - return chunkFallback(filePath, content, language), nil, nil + return nil, nil, errUseFallback } if nodeKinds == nil { // Grammar exists but we don't have node definitions → sliding window. - return chunkFallback(filePath, content, language), nil, nil + return nil, nil, errUseFallback } if looksMinified(filePath, content, language) { // Minified/bundled sources are the parser's pathological case: a @@ -541,7 +676,7 @@ func chunkWithTreesitter(filePath, content, language string, maxSize int) ([]Chu // instance to its memory cap, and forces a pool recycle — all to // produce AST chunks with near-zero semantic-search value. Skip // straight to the sliding window. - return chunkFallback(filePath, content, language), nil, nil + return nil, nil, errUseFallback } // Build flat target → kind map. @@ -559,10 +694,10 @@ func chunkWithTreesitter(filePath, content, language string, maxSize int) ([]Chu // fall back to sliding window so the file is still indexed. slog.Warn("chunker: wasm parse failed, falling back to sliding window", "path", filePath, "language", language, "err", err) - return chunkFallback(filePath, content, language), nil, nil + return nil, nil, errUseFallback } if len(nodes) == 0 { - return chunkFallback(filePath, content, language), nil, nil + return nil, nil, errUseFallback } tree := buildFlatTree(nodes) @@ -604,7 +739,7 @@ func chunkWithTreesitter(filePath, content, language string, maxSize int) ([]Chu } if len(finalChunks) == 0 { - return chunkFallback(filePath, content, language), nil, nil + return nil, nil, errUseFallback } return finalChunks, refs, nil } @@ -1060,3 +1195,113 @@ func sortRanges(ranges [][2]int) { } } } + +// DefaultMaxChunkTokens is the token equivalent of maxChunkSize. The byte +// default was written as 1500*3 — a 1500-token target at three bytes each — +// so the token target is that same 1500, now expressed in the unit that +// actually matters. +// +// Exported so config.go can use it as the CIX_MAX_CHUNK_TOKENS default rather +// than repeating the number: two copies of a chunk-size default drifting apart +// is the exact failure this change removes for maxChunkSize. +const DefaultMaxChunkTokens = 1500 + +// splitChunkTokens cuts an over-budget chunk into pieces of <= maxTokens. +// +// It keeps splitChunk's attribution rule: only the first piece inherits +// SymbolName/ChunkType, the rest become `block`, so one long function does not +// produce N symbol rows all claiming to be it. +// +// The cut positions come from Budget.SplitPoints over the WHOLE content rather +// than from summing per-line counts. Per-line summing is off by the separators +// — joining lines reinserts newlines, and a newline plus the next line's +// indentation forms its own pre-token — so a budget of 1500 produced chunks of +// up to 1546 tokens on real files, roughly one extra token per line boundary. +// SplitPoints is exact by construction, so the bound actually holds. +// +// Each exact cut is then pulled BACK to the nearest line start, because a chunk +// that begins mid-line reads badly in search results and its line range lies. +// Moving a cut backwards only ever shrinks the piece before it, so the budget +// survives the adjustment. A line longer than the whole budget has no earlier +// boundary to snap to; there the exact cut stands and the line is split +// internally — which is the case the byte splitter could not handle at all. +func splitChunkTokens(chunk Chunk, budget tokenizer.Budget, maxTokens int) []Chunk { + content := chunk.Content + var out []Chunk + pos, line := 0, chunk.StartLine + var cuts []int // absolute offsets into content; nil means "recompute" + + for pos < len(content) { + if cuts == nil { + rel, _ := budget.SplitPoints(content[pos:], maxTokens) + if len(rel) == 0 { + out = append(out, mkPiece(chunk, content[pos:], line, len(out) == 0)) + break + } + cuts = make([]int, 0, len(rel)) + for _, r := range rel { + cuts = append(cuts, pos+r) + } + } + + at := cuts[0] + // Pull the cut back to a line start so a piece never begins mid-line: + // search results and the stored line range both lie otherwise. Moving + // backwards only shrinks this piece, so it stays inside the budget. A + // line wider than the whole budget has no earlier boundary — there the + // exact cut stands and the line is split internally, which is the case + // the byte splitter could not handle at all. + snapped := at + if nl := strings.LastIndexByte(content[pos:at], '\n'); nl >= 0 { + snapped = pos + nl + 1 + } + if snapped <= pos { + snapped = at + } + + piece := content[pos:snapped] + out = append(out, mkPiece(chunk, piece, line, len(out) == 0)) + line += strings.Count(piece, "\n") + pos = snapped + + if snapped == at { + // The boundary landed where the tokenizer put it, so the cuts + // after it are still valid and can be consumed without another + // pass. This is what keeps a 66 KB single line — where the + // newline snap never fires — from costing one full scan per + // piece. + cuts = cuts[1:] + if len(cuts) == 0 { + cuts = nil + } + continue + } + // Snapping moved the boundary: every later cut was measured from a + // start that no longer exists, and reusing them would hand the next + // piece the tokens this one gave up. Recompute. + cuts = nil + } + return out +} + +// mkPiece builds one output chunk, preserving splitChunk's attribution rule: +// only the first piece keeps SymbolName/ChunkType, so a long function does not +// produce N symbol rows all claiming to be it. +func mkPiece(src Chunk, content string, startLine int, first bool) Chunk { + c := Chunk{ + Content: content, + FilePath: src.FilePath, + StartLine: startLine, + EndLine: startLine + strings.Count(strings.TrimSuffix(content, "\n"), "\n"), + Language: src.Language, + ParentName: src.ParentName, + } + if first { + c.ChunkType = src.ChunkType + c.SymbolName = src.SymbolName + c.SymbolSignature = src.SymbolSignature + } else { + c.ChunkType = "block" + } + return c +} diff --git a/server/internal/chunker/chunker_tokens_test.go b/server/internal/chunker/chunker_tokens_test.go new file mode 100644 index 00000000..552ddeb5 --- /dev/null +++ b/server/internal/chunker/chunker_tokens_test.go @@ -0,0 +1,299 @@ +package chunker + +import ( + "strings" + "testing" + + "github.com/dvcdsys/code-index/server/internal/tokenizer" +) + +// fakeBudget is deliberately ADVERSARIAL: a newline costs a token, exactly +// like it does in the real tokenizer, where a line break plus the next line's +// indentation forms its own pre-token. +// +// The first version of this double counted whitespace-separated words and let +// newlines be free. That made the sum of per-line counts equal the count of +// the joined text — which is precisely the assumption the implementation got +// wrong, so the double agreed with the bug and the tests passed while real +// files came out 3% over budget. A test double that cannot express the +// failure mode cannot catch it. +// +// Counting rule: one token per word start, one per newline. Sum over pieces +// therefore does NOT equal the count of the concatenation unless the pieces +// are substrings — which is the property the splitter must have. +type fakeBudget struct{ maxInput int } + +func (f fakeBudget) MaxInputTokens() int { return f.maxInput } +func (f fakeBudget) ExactCounts() bool { return true } + +func (f fakeBudget) CountTokens(s string) int { + n, inWord := 0, false + for i := 0; i < len(s); i++ { + switch c := s[i]; { + case c == '\n': + n++ + inWord = false + case c == ' ' || c == '\t' || c == '\r': + inWord = false + default: + if !inWord { + inWord = true + n++ + } + } + } + return n +} + +// SplitPoints cuts before the token that would overflow the budget, so every +// piece it produces costs at most budget under CountTokens above. +func (f fakeBudget) SplitPoints(s string, budget int) ([]int, int) { + var offsets []int + total, since := 0, 0 + inWord := false + cut := func(at int) { + offsets = append(offsets, at) + since = 1 + } + for i := 0; i < len(s); i++ { + switch c := s[i]; { + case c == '\n': + total++ + since++ + if since > budget { + cut(i) + } + inWord = false + case c == ' ' || c == '\t' || c == '\r': + inWord = false + default: + if !inWord { + inWord = true + total++ + since++ + if since > budget { + cut(i) + } + } + } + } + return offsets, total +} + +var _ tokenizer.Budget = fakeBudget{} + +// TestTokenBudgetBoundsEveryChunk is the property the integration exists for: +// with a budget in hand, no emitted chunk may exceed it. +func TestTokenBudgetBoundsEveryChunk(t *testing.T) { + var sb strings.Builder + sb.WriteString("func run() {\n") + for i := 0; i < 300; i++ { + sb.WriteString("\tdo something with several words on this line\n") + } + sb.WriteString("}\n") + + b := fakeBudget{maxInput: 4096} + chunks, _, err := ChunkFileTokens("x.go", sb.String(), "go", 0, b, 50) + if err != nil { + t.Fatalf("chunk: %v", err) + } + if len(chunks) < 2 { + t.Fatalf("expected the body to be split, got %d chunk(s)", len(chunks)) + } + for i, c := range chunks { + if n := b.CountTokens(c.Content); n > 50 { + t.Errorf("chunk %d is %d tokens, over budget 50", i, n) + } + } +} + +// TestLongSingleLineIsSplit covers the hole the byte splitter had: its loop +// requires more than one line, so a minified file arrived at the embedder as +// one enormous chunk. On the reference corpus that produced a 65 KB chunk +// whose vector was an average of byte windows. +func TestLongSingleLineIsSplit(t *testing.T) { + line := strings.TrimSpace(strings.Repeat("token ", 5000)) + b := fakeBudget{maxInput: 4096} + + chunks, _, err := ChunkFileTokens("min.js", line, "javascript", 0, b, 100) + if err != nil { + t.Fatalf("chunk: %v", err) + } + for i, c := range chunks { + if n := b.CountTokens(c.Content); n > 100 { + t.Errorf("chunk %d is %d tokens, over budget 100 — long line not cut", i, n) + } + } + if len(chunks) < 2 { + t.Fatalf("expected the single line to be cut, got %d chunk(s)", len(chunks)) + } +} + +// TestBudgetCappedByModelContext — a chunk target above the model's own input +// window is meaningless; the smaller of the two must win. +func TestBudgetCappedByModelContext(t *testing.T) { + src := strings.TrimSpace(strings.Repeat("word ", 400)) + b := fakeBudget{maxInput: 40} + + chunks, _, err := ChunkFileTokens("x.txt", src, "text", 0, b, 10000) + if err != nil { + t.Fatalf("chunk: %v", err) + } + for i, c := range chunks { + if n := b.CountTokens(c.Content); n > 40 { + t.Errorf("chunk %d is %d tokens, over the model's %d-token window", i, n, 40) + } + } +} + +// TestEstimatingBudgetKeepsBytePath — a provider without a real tokenizer must +// not be routed through the token splitter: its numbers are the same guess the +// byte path already makes, and pretending otherwise hides that from the caller. +func TestEstimatingBudgetKeepsBytePath(t *testing.T) { + src := strings.Repeat("x := 1\n", 2000) + got, _, err := ChunkFileTokens("x.go", src, "go", 0, estimatingBudget{}, 10) + if err != nil { + t.Fatalf("chunk: %v", err) + } + want, _, err := ChunkFile("x.go", src, "go", 0) + if err != nil { + t.Fatalf("chunk: %v", err) + } + if len(got) != len(want) { + t.Errorf("estimating budget changed chunking: %d chunks vs %d on the byte path", + len(got), len(want)) + } +} + +type estimatingBudget struct{ fakeBudget } + +func (estimatingBudget) ExactCounts() bool { return false } + +// TestNilBudgetUnchanged pins that the default path is untouched. +func TestNilBudgetUnchanged(t *testing.T) { + src := strings.Repeat("func f() { return 1 }\n", 500) + a, _, err := ChunkFileTokens("x.go", src, "go", 0, nil, 0) + if err != nil { + t.Fatalf("chunk: %v", err) + } + b, _, err := ChunkFile("x.go", src, "go", 0) + if err != nil { + t.Fatalf("chunk: %v", err) + } + if len(a) != len(b) { + t.Errorf("nil budget diverged from ChunkFile: %d vs %d chunks", len(a), len(b)) + } +} + +// TestTokenSplitPreservesContent pins the invariant that makes the token +// splitter exact: its pieces are SUBSTRINGS of the chunk it was given, so +// concatenating them reproduces it byte for byte — nothing lost, nothing +// duplicated. +// +// The first implementation instead re-joined lines it had counted separately, +// and the newlines it reinserted cost tokens the running total never saw. A +// 1500-token budget produced 1546-token chunks on real files. Slicing the +// original removes that class of error rather than compensating for it. +// +// Asserted on splitChunkTokens directly, not through ChunkFileTokens: the +// sliding-window fallback deliberately overlaps its windows for recall, so +// whole-pipeline output is not expected to concatenate back. +func TestTokenSplitPreservesContent(t *testing.T) { + var sb strings.Builder + for i := 0; i < 200; i++ { + sb.WriteString("some line with a handful of words in it\n") + } + src := Chunk{ + Content: sb.String(), + FilePath: "x.txt", + StartLine: 1, + EndLine: 200, + ChunkType: "function", + SymbolName: strPtr("run"), + } + b := fakeBudget{maxInput: 4096} + + pieces := splitChunkTokens(src, b, 40) + var rebuilt strings.Builder + for i, c := range pieces { + rebuilt.WriteString(c.Content) + if n := b.CountTokens(c.Content); n > 40 { + t.Errorf("piece %d is %d tokens, over budget 40", i, n) + } + } + if rebuilt.String() != src.Content { + t.Errorf("concatenated pieces differ from the source (%d bytes vs %d)", + rebuilt.Len(), len(src.Content)) + } + if pieces[0].SymbolName == nil || *pieces[0].SymbolName != "run" || pieces[0].ChunkType != "function" { + t.Error("first piece must inherit the symbol") + } + for i, c := range pieces[1:] { + if c.SymbolName != nil || c.ChunkType != "block" { + t.Errorf("piece %d must be an anonymous block, got type %q", i+1, c.ChunkType) + } + } +} + +// TestTokenSplitLineNumbers — a piece that starts mid-file must report the +// line it actually starts on, or `cix search` sends the reader to the wrong +// place. +func TestTokenSplitLineNumbers(t *testing.T) { + var sb strings.Builder + for i := 0; i < 100; i++ { + sb.WriteString("word word word word word\n") + } + b := fakeBudget{maxInput: 4096} + + chunks := splitChunkTokens(Chunk{Content: sb.String(), FilePath: "x.txt", StartLine: 1}, b, 20) + line := 1 + for i, c := range chunks { + if c.StartLine != line { + t.Errorf("chunk %d starts at line %d, expected %d", i, c.StartLine, line) + } + line += strings.Count(c.Content, "\n") + } +} + +func strPtr(s string) *string { return &s } + +// TestFallbackFillsTheBudget covers the path a file with no grammar takes. +// The byte-sized sliding window cut every 4000 bytes regardless of content, so +// multi-byte text produced windows worth a fraction of the intended tokens — +// and boundTokens could not repair that, since it only splits chunks that are +// too big and cannot merge ones that are too small. +func TestFallbackFillsTheBudget(t *testing.T) { + // Two-bytes-per-character text, well past one byte window. + src := strings.Repeat("привіт світ це коментар українською\n", 400) + b := fakeBudget{maxInput: 4096} + // The budget must exceed one byte window's token worth, or the test + // passes on arithmetic: a 4000-byte window of this text is ~358 fake + // tokens, so at budget 200 boundTokens splits every window into 200+158 + // and both halves clear half-budget without the fallback ever being + // token-aware. At 800 the raw window is BELOW half the budget, so an + // under-filled chunk can only come from a byte-sized window. + const budget = 800 + + chunks, _, err := ChunkFileTokens("notes.unknownlang", src, "unknownlang", 0, b, budget) + if err != nil { + t.Fatalf("chunk: %v", err) + } + if len(chunks) < 2 { + t.Fatalf("expected several chunks, got %d", len(chunks)) + } + var under int + for i, c := range chunks { + n := b.CountTokens(c.Content) + if n > budget { + t.Errorf("chunk %d is %d tokens, over budget %d", i, n, budget) + } + // The last chunk is a remainder and may legitimately be short. + if i < len(chunks)-1 && n < budget/2 { + under++ + } + } + if under > 0 { + t.Errorf("%d of %d chunks are under half the budget — the window is still byte-sized", + under, len(chunks)) + } +} diff --git a/server/internal/chunker/corpus_property_test.go b/server/internal/chunker/corpus_property_test.go new file mode 100644 index 00000000..f414af6a --- /dev/null +++ b/server/internal/chunker/corpus_property_test.go @@ -0,0 +1,208 @@ +package chunker + +import ( + "math/rand" + "os" + "path/filepath" + "sort" + "strings" + "testing" + + "github.com/dvcdsys/code-index/server/internal/tokenizer" + "github.com/dvcdsys/code-index/server/internal/tokenizer/bpecount" +) + +// Property test over a real corpus. +// +// Chunk splitting under a token budget has no external oracle: unlike the +// tokenizer, which can be checked against HuggingFace's implementation and +// against Voyage's own billing, "where should a chunk be cut" is our decision +// and there is nothing to compare it to. What can be checked is that the +// properties we chose actually hold on inputs we did not write — and the +// bugs this file exists to catch were all found by real files rather than by +// hand-written cases: +// +// - a 65 KB single line (a Zig integer literal) that the byte splitter left +// whole, at 65,553 tokens against a 32K context; +// - minified JavaScript, which has no grammar and therefore reaches the +// sliding-window fallback rather than the tree-sitter path; +// - per-line token counting, which was 3% under the truth because joining +// lines reinserts newlines that cost tokens. +// +// The corpus is not in the repository — it is a local fixture of cloned +// repositories, tens of gigabytes. Point the test at one: +// +// CIX_TEST_CORPUS_DIR=…/loadtests/data/repos/repos \ +// CIX_TEST_TOKENIZER=…/voyage-code-3.tokenizer.json \ +// go test ./internal/chunker/ -run Corpus +// +// Without those it skips, so a clean checkout and CI stay green. + +type realBudget struct{ tk *bpecount.Counter } + +func (realBudget) MaxInputTokens() int { return 32000 } +func (realBudget) ExactCounts() bool { return true } +func (b realBudget) CountTokens(s string) int { return b.tk.Count(s) } +func (b realBudget) SplitPoints(s string, n int) ([]int, int) { return b.tk.SplitPoints(s, n) } + +var _ tokenizer.Budget = realBudget{} + +var extLang = map[string]string{ + ".go": "go", ".py": "python", ".ts": "typescript", ".tsx": "tsx", + ".js": "javascript", ".jsx": "javascript", ".java": "java", ".rs": "rust", + ".c": "c", ".h": "c", ".cpp": "cpp", ".rb": "ruby", ".php": "php", + ".kt": "kotlin", ".swift": "swift", ".ex": "elixir", ".zig": "zig", + ".lua": "lua", ".sh": "bash", ".md": "markdown", ".json": "json", +} + +// sampleCorpus walks the fixture and returns up to n files, deterministically +// shuffled so a failure is reproducible. +func sampleCorpus(t *testing.T, root string, n int) []string { + t.Helper() + var files []string + err := filepath.WalkDir(root, func(path string, d os.DirEntry, err error) error { + if err != nil { + return nil //nolint:nilerr // an unreadable entry is not this test's problem + } + if d.IsDir() { + if d.Name() == ".git" { + return filepath.SkipDir + } + return nil + } + if _, ok := extLang[strings.ToLower(filepath.Ext(path))]; ok { + files = append(files, path) + } + return nil + }) + if err != nil { + t.Fatalf("walk corpus: %v", err) + } + sort.Strings(files) + rng := rand.New(rand.NewSource(20260818)) + rng.Shuffle(len(files), func(i, j int) { files[i], files[j] = files[j], files[i] }) + if len(files) > n { + files = files[:n] + } + return files +} + +func corpusBudget(t *testing.T) (realBudget, string) { + t.Helper() + dir := os.Getenv("CIX_TEST_CORPUS_DIR") + tok := os.Getenv("CIX_TEST_TOKENIZER") + if dir == "" || tok == "" { + t.Skip("set CIX_TEST_CORPUS_DIR and CIX_TEST_TOKENIZER to run corpus property tests") + } + tk, err := bpecount.Load(tok) + if err != nil { + t.Fatalf("load tokenizer: %v", err) + } + return realBudget{tk}, dir +} + +// TestCorpusChunksRespectBudget is the property that matters to the API: no +// chunk may cost more tokens than the budget, whatever path produced it. +func TestCorpusChunksRespectBudget(t *testing.T) { + b, dir := corpusBudget(t) + const budget = 1500 + + files := sampleCorpus(t, dir, 400) + if len(files) == 0 { + t.Skip("corpus contains no recognised source files") + } + + var checked, chunks, worst int + for _, f := range files { + src, err := os.ReadFile(f) + if err != nil || len(src) == 0 { + continue + } + lang := extLang[strings.ToLower(filepath.Ext(f))] + got, _, err := ChunkFileTokens(f, string(src), lang, 0, b, budget) + if err != nil { + t.Errorf("%s: %v", f, err) + continue + } + checked++ + chunks += len(got) + for i, c := range got { + n := b.CountTokens(c.Content) + if n > worst { + worst = n + } + if n > budget { + t.Errorf("%s chunk %d: %d tokens, over budget %d", f, i, n, budget) + } + } + } + t.Logf("%d files, %d chunks, largest %d tokens (budget %d)", checked, chunks, worst, budget) +} + +// TestCorpusSplitPreservesContent asserts the splitter loses and duplicates +// nothing, on chunks taken from real files rather than constructed ones. Run +// against splitChunkTokens directly: the sliding-window fallback overlaps its +// windows by design, so whole-pipeline output is not expected to concatenate +// back to the source. +func TestCorpusSplitPreservesContent(t *testing.T) { + b, dir := corpusBudget(t) + const budget = 300 + + var split int + for _, f := range sampleCorpus(t, dir, 200) { + src, err := os.ReadFile(f) + if err != nil || len(src) == 0 { + continue + } + whole := Chunk{ + Content: string(src), + FilePath: f, + StartLine: 1, + ChunkType: "file", + } + if b.CountTokens(whole.Content) <= budget { + continue + } + pieces := splitChunkTokens(whole, b, budget) + split++ + + var rebuilt strings.Builder + for i, p := range pieces { + rebuilt.WriteString(p.Content) + if n := b.CountTokens(p.Content); n > budget { + t.Errorf("%s piece %d: %d tokens, over budget %d", f, i, n, budget) + } + } + if rebuilt.String() != whole.Content { + t.Errorf("%s: pieces do not reconstruct the file (%d bytes vs %d)", + f, rebuilt.Len(), len(whole.Content)) + } + } + t.Logf("%d files exceeded the budget and were split", split) +} + +// TestCorpusLineNumbers — a chunk's StartLine must point at the line its text +// actually begins on, or search results send the reader to the wrong place. +func TestCorpusLineNumbers(t *testing.T) { + b, dir := corpusBudget(t) + const budget = 300 + + for _, f := range sampleCorpus(t, dir, 150) { + src, err := os.ReadFile(f) + if err != nil || len(src) == 0 { + continue + } + whole := Chunk{Content: string(src), FilePath: f, StartLine: 1} + if b.CountTokens(whole.Content) <= budget { + continue + } + line := 1 + for i, p := range splitChunkTokens(whole, b, budget) { + if p.StartLine != line { + t.Errorf("%s piece %d starts at line %d, expected %d", f, i, p.StartLine, line) + break + } + line += strings.Count(p.Content, "\n") + } + } +} diff --git a/server/internal/config/config.go b/server/internal/config/config.go index 4743a947..2a346c59 100644 --- a/server/internal/config/config.go +++ b/server/internal/config/config.go @@ -5,6 +5,7 @@ package config import ( "fmt" + "github.com/dvcdsys/code-index/server/internal/chunker" "net" "os" "path/filepath" @@ -352,7 +353,7 @@ func Load() (*Config, error) { } c.ChunkMaxConcurrent = chunkConc - maxChunk, err := getenvInt("CIX_MAX_CHUNK_TOKENS", 1500) + maxChunk, err := getenvInt("CIX_MAX_CHUNK_TOKENS", chunker.DefaultMaxChunkTokens) if err != nil { return nil, err } diff --git a/server/internal/embeddings/provider/voyage/factory.go b/server/internal/embeddings/provider/voyage/factory.go index 64d12c57..926e4869 100644 --- a/server/internal/embeddings/provider/voyage/factory.go +++ b/server/internal/embeddings/provider/voyage/factory.go @@ -33,6 +33,13 @@ func (factory) SchemaJSON() []byte { Description: "int8 is dequantized to float32 on the server side.", }, {Name: "truncation", Label: "Truncate over-length input", Kind: "bool", Default: true}, + { + Name: "tokenizer_path", Label: "Tokenizer file", Kind: "string", + Description: "Absolute path to the model's tokenizer.json (huggingface.co/voyageai/). " + + "Set it and token counts become exact: batches pack to the real limit instead of a " + + "byte guess that overestimates ~2x, and over-long inputs split on token boundaries " + + "instead of byte windows. Empty falls back to the estimate.", + }, {Name: "api_key_env", Label: "API key env var", Kind: "secret-env", Required: true, Default: defaultAPIKeyEnv}, }, } diff --git a/server/internal/embeddings/provider/voyage/voyage.go b/server/internal/embeddings/provider/voyage/voyage.go index 04318f85..c9e73375 100644 --- a/server/internal/embeddings/provider/voyage/voyage.go +++ b/server/internal/embeddings/provider/voyage/voyage.go @@ -36,6 +36,7 @@ import ( "golang.org/x/time/rate" "github.com/dvcdsys/code-index/server/internal/embeddings/provider" + "github.com/dvcdsys/code-index/server/internal/tokenizer/bpecount" ) // voyageBatchTooLargeRegex matches Voyage's per-batch token-limit @@ -159,6 +160,15 @@ type Config struct { // all in-flight + recent requests). 0 = no throttling. RateLimitTPM int `json:"rate_limit_tpm,omitempty"` + // TokenizerPath points at the model's tokenizer.json (the file + // Voyage publishes at huggingface.co/voyageai/). When set and + // loadable, token counts become EXACT and the per-batch cap rises to + // exactTokensPerBatch — the 40K of headroom the byte heuristic needed + // is headroom against the heuristic, not against Voyage. When empty or + // unreadable the provider logs once and falls back to estimateTokens, + // so a missing file degrades throughput, never correctness. + TokenizerPath string `json:"tokenizer_path,omitempty"` + // MaxInputsPerRequest overrides defaultMaxBatchSize. 0 = use // the default (128, safe for voyage-code-*). Operators running // only voyage-3* may bump this to 1000 for fewer round-trips. @@ -189,7 +199,13 @@ func (c *Config) maxBatchSize() int { return defaultMaxBatchSize } -// maxTokensPerBatch returns the effective per-POST token cap. +// maxTokensPerBatch returns the cap implied by config alone — the operator's +// override, or the conservative byte-heuristic default. +// +// Callers on the hot path want (*Provider).maxTokensPerBatch instead, which +// also knows whether a tokenizer is loaded. Two same-named methods one on +// Config and one on Provider is how the batch log came to report 80K while +// packing used 115K, so this one is only for the Provider method to build on. func (c *Config) maxTokensPerBatch() int { if c.MaxTokensPerRequest > 0 { return c.MaxTokensPerRequest @@ -266,6 +282,10 @@ type Provider struct { // budget is a sliding minute and bursting saves nothing. reqLimiter *rate.Limiter + // counter is the model's real tokenizer, or nil when no tokenizer.json + // was configured or it failed to load. Safe for concurrent use. + counter *bpecount.Counter + // tokenLimiter caps tokens-per-minute when cfg.RateLimitTPM > 0. // Burst is set to maxTokensPerBatch so a single full-budget POST // can pass even when the bucket is otherwise empty (we'd just @@ -294,6 +314,19 @@ func New(cfg Config, secrets provider.SecretLookup, logger *slog.Logger) *Provid secrets: secrets, http: &http.Client{Timeout: 60 * time.Second}, } + if cfg.TokenizerPath != "" { + c, err := bpecount.Load(cfg.TokenizerPath) + if err != nil { + // Not fatal: the byte heuristic still works. Loud because the + // operator asked for exact counts and is not getting them. + logger.Warn("voyage: tokenizer load failed, falling back to byte estimate", + "path", cfg.TokenizerPath, "err", err) + } else { + p.counter = c + logger.Info("voyage: exact token counting enabled", "path", cfg.TokenizerPath) + } + } + // Convert RPM/TPM to per-second token-bucket rates. burst on the // request bucket is 1 (one request worth of "credit"); burst on // the token bucket equals one full POST so we don't deadlock a @@ -302,7 +335,7 @@ func New(cfg Config, secrets provider.SecretLookup, logger *slog.Logger) *Provid p.reqLimiter = rate.NewLimiter(rate.Limit(float64(cfg.RateLimitRPM)/60.0), 1) } if cfg.RateLimitTPM > 0 { - p.tokenLimiter = rate.NewLimiter(rate.Limit(float64(cfg.RateLimitTPM)/60.0), cfg.maxTokensPerBatch()) + p.tokenLimiter = rate.NewLimiter(rate.Limit(float64(cfg.RateLimitTPM)/60.0), p.maxTokensPerBatch()) } return p } @@ -402,6 +435,38 @@ func (p *Provider) EmbedDocuments(ctx context.Context, texts []string) ([][]floa // such chunk, but oversize chunks are rare on well-chunked // indexes — the indexer should already be cutting at function / // class boundaries. +// splitForInput cuts one input down to what the model can read. +// +// With a tokenizer, the question "does this fit" has an exact answer, so the +// byte cap is not consulted at all: an input under the model's context window +// goes through whole, however many bytes it is, and one over it is cut on real +// token boundaries. Without a tokenizer we are back to guessing, and the byte +// cap is the guess. +// +// This matters because the chunker now sizes in tokens. Its bound and the +// provider's byte cap are different units: at CIX_MAX_CHUNK_TOKENS=20000 — +// legal, well inside the 32K window — chunks of 40-80 KB are ordinary, and +// every one of them used to be byte-windowed here and have its window vectors +// averaged into a single vector representing neither half. The averaging path +// now only runs where it is genuinely needed: no tokenizer, no exact answer. +func (p *Provider) splitForInput(text string, maxBytes int) []string { + if p.counter == nil { + return splitOversizeInput(text, maxBytes) + } + limit := p.MaxInputTokens() + offsets, total := p.counter.SplitPoints(text, limit) + if total <= limit || len(offsets) == 0 { + return []string{text} + } + out := make([]string, 0, len(offsets)+1) + prev := 0 + for _, off := range offsets { + out = append(out, text[prev:off]) + prev = off + } + return append(out, text[prev:]) +} + func (p *Provider) embedAndAverage(ctx context.Context, texts []string, inputType string) ([][]float32, error) { maxIn := p.cfg.maxInputBytes() @@ -411,7 +476,7 @@ func (p *Provider) embedAndAverage(ctx context.Context, texts []string, inputTyp var expanded []string totalSplits := 0 for i, t := range texts { - windows := splitOversizeInput(t, maxIn) + windows := p.splitForInput(t, maxIn) spans[i] = span{start: len(expanded), length: len(windows)} expanded = append(expanded, windows...) if len(windows) > 1 { @@ -419,23 +484,35 @@ func (p *Provider) embedAndAverage(ctx context.Context, texts []string, inputTyp } } if totalSplits > 0 { - p.logger.Info("voyage: oversize inputs split into byte-windows", - "original_inputs", len(texts), - "total_windows", len(expanded), - "split_windows", totalSplits, - "max_input_bytes", maxIn, - ) + // Report the unit the split actually used: with a tokenizer the cut is + // on token boundaries against the model's context, and logging a byte + // cap there sends whoever reads this to the wrong knob. + if p.counter != nil { + p.logger.Info("voyage: oversize inputs split on token boundaries", + "original_inputs", len(texts), + "total_windows", len(expanded), + "split_windows", totalSplits, + "max_input_tokens", p.MaxInputTokens(), + ) + } else { + p.logger.Info("voyage: oversize inputs split into byte-windows", + "original_inputs", len(texts), + "total_windows", len(expanded), + "split_windows", totalSplits, + "max_input_bytes", maxIn, + ) + } } // Phase 2: batch + POST as before, on the expanded slice. - batches := planBatches(expanded, p.cfg.maxBatchSize(), p.cfg.maxTokensPerBatch()) + batches := planBatches(expanded, p.cfg.maxBatchSize(), p.maxTokensPerBatch(), p.CountTokens) if len(batches) > 1 { p.logger.Info("voyage: splitting batch", "model", p.cfg.Model, "total_inputs", len(expanded), "sub_batches", len(batches), "limit_inputs", p.cfg.maxBatchSize(), - "limit_tokens", p.cfg.maxTokensPerBatch(), + "limit_tokens", p.maxTokensPerBatch(), ) } allVecs := make([][]float32, 0, len(expanded)) @@ -561,7 +638,10 @@ func (p *Provider) embedWithAdaptiveSplit(ctx context.Context, texts []string, i // operator can override them via the admin form when their tier or // chosen model allows a higher cap (e.g. voyage-3-large at 1000 // inputs/POST instead of 128). -func planBatches(texts []string, maxInputs, maxTokens int) [][]string { +func planBatches(texts []string, maxInputs, maxTokens int, count func(string) int) [][]string { + if count == nil { + count = estimateTokens + } if len(texts) == 0 { return nil } @@ -569,7 +649,7 @@ func planBatches(texts []string, maxInputs, maxTokens int) [][]string { var current []string currentTokens := 0 for _, t := range texts { - est := estimateTokens(t) + est := count(t) // Close the current batch when adding this text would exceed // either limit (and the batch already has something to send). if len(current) > 0 && (len(current) >= maxInputs || currentTokens+est > maxTokens) { @@ -586,9 +666,12 @@ func planBatches(texts []string, maxInputs, maxTokens int) [][]string { return batches } -// estimateTokens returns a conservative upper bound on the token cost -// of one text, in Voyage's tokenizer. Uses byte-length divided by a -// chars-per-token heuristic; see bytesPerToken doc for rationale. +// estimateTokens is the FALLBACK used only when no tokenizer.json is +// loaded. Measured against Voyage's own usage.total_tokens on 20k real +// chunks it overestimates by 1.94x on average — which wastes round-trips +// — while still undercounting 0.5% of chunks, worst case -41%. That is +// the wrong error in both directions, and it is why loading the real +// tokenizer is worth the 7 MB: see Provider.countTokens. func estimateTokens(s string) int { return len(s) / bytesPerToken } @@ -784,3 +867,90 @@ func dequantize(raw json.RawMessage, dtype string) ([]float32, error) { func (p *Provider) apiKey() (string, bool) { return provider.ResolveAPIKey(p.secrets, p.cfg.APIKeyEnv) } + +// ---------- tokenizer.Budget ---------- +// +// Implemented on Provider so the chunker can be handed the live provider and +// stay ignorant of which model is active: only the provider knows whether +// tokens come from a real BPE table, from llama-server's /tokenize, or from a +// byte guess. + +// exactTokensPerBatch is the per-POST cap once counts are exact. +// +// The 80K default exists to survive the byte heuristic's ~43% undercount +// against Voyage's 120K hard limit. With the real tokenizer the count is the +// count — measured against usage.total_tokens it is never below what Voyage +// bills — so the headroom collapses to a margin for Voyage-side accounting +// drift rather than for our own error. +const exactTokensPerBatch = 115_000 + +// modelContextTokens is the per-input context window, per model. It is a table +// rather than a constant because the factory's own enum offers voyage-code-2, +// whose window is 16K — half of what the rest of the list takes. Treating that +// as 32K would let the chunker build inputs the model cannot read, and with +// truncation enabled Voyage would silently drop the tail. +// +// Unknown models fall back to the conservative 16K: undershooting costs an +// unnecessary split, overshooting costs silent data loss. +var modelContextTokens = map[string]int{ + "voyage-code-3": 32_000, + "voyage-3-large": 32_000, + "voyage-3": 32_000, + "voyage-3-lite": 32_000, + "voyage-code-2": 16_000, +} + +const fallbackContextTokens = 16_000 + +// maxTokensPerBatch is the provider-level cap: an explicit operator override +// wins, then the exact-counting cap, then the conservative byte-heuristic one. +func (p *Provider) maxTokensPerBatch() int { + if p.cfg.MaxTokensPerRequest > 0 { + return p.cfg.maxTokensPerBatch() + } + if p.counter != nil { + return exactTokensPerBatch + } + return p.cfg.maxTokensPerBatch() +} + +// MaxInputTokens reports the model's context window for a single input. +func (p *Provider) MaxInputTokens() int { + if n, ok := modelContextTokens[p.cfg.Model]; ok { + return n + } + return fallbackContextTokens +} + +// ExactCounts reports whether CountTokens/SplitPoints are exact rather than +// estimated. False means no tokenizer.json was loaded. +func (p *Provider) ExactCounts() bool { return p.counter != nil } + +// CountTokens returns the token cost of s. Allocation-free on the exact path; +// this is the hot one — it runs for every chunk that gets embedded. +func (p *Provider) CountTokens(s string) int { + if p.counter != nil { + return p.counter.Count(s) + } + return estimateTokens(s) +} + +// SplitPoints returns byte offsets at which s must be cut so no piece exceeds +// budget tokens, and s's total token count. +// +// Exact when a tokenizer is loaded: cuts land on pre-token boundaries, where +// BPE merges never reach across, so the pieces provably add up to the whole. +// Without a tokenizer it degrades to rune-aligned byte windows — the old +// behaviour, kept only so a caller that ignores ExactCounts still gets +// something it can send. Check ExactCounts before trusting these. +func (p *Provider) SplitPoints(s string, budget int) ([]int, int) { + if p.counter == nil { + // No tokenizer: there are no token boundaries to report. Returning + // byte windows here would be the old behaviour wearing the new + // interface's clothes, and callers check ExactCounts() precisely so + // they can avoid it. splitForInput still byte-windows internally + // where that is genuinely all we have. + return nil, estimateTokens(s) + } + return p.counter.SplitPoints(s, budget) +} diff --git a/server/internal/embeddings/provider/voyage/voyage_test.go b/server/internal/embeddings/provider/voyage/voyage_test.go index 542e0a3c..4eb4dadf 100644 --- a/server/internal/embeddings/provider/voyage/voyage_test.go +++ b/server/internal/embeddings/provider/voyage/voyage_test.go @@ -5,9 +5,12 @@ import ( "encoding/base64" "encoding/json" "fmt" + "github.com/dvcdsys/code-index/server/internal/tokenizer" + "github.com/dvcdsys/code-index/server/internal/tokenizer/bpecount" "io" "net/http" "net/http/httptest" + "os" "strings" "sync/atomic" "testing" @@ -187,7 +190,7 @@ func TestPlanBatches_SplitsByTokenBudget(t *testing.T) { small := "tiny" texts := []string{big, small, small, small, small, small} - batches := planBatches(texts, defaultMaxBatchSize, defaultMaxTokensPerBatch) + batches := planBatches(texts, defaultMaxBatchSize, defaultMaxTokensPerBatch, nil) if len(batches) < 2 { t.Fatalf("expected at least 2 batches, got %d", len(batches)) } @@ -213,7 +216,7 @@ func TestPlanBatches_RespectsCountCap(t *testing.T) { for i := range texts { texts[i] = "chunk" } - batches := planBatches(texts, defaultMaxBatchSize, defaultMaxTokensPerBatch) + batches := planBatches(texts, defaultMaxBatchSize, defaultMaxTokensPerBatch, nil) if len(batches) != 2 { t.Fatalf("expected 2 batches (128 + 72), got %d", len(batches)) } @@ -652,3 +655,78 @@ func TestInt8Dequantize_Base64(t *testing.T) { t.Errorf("base64 int8 dequantized values out of range: %v", v) } } + +// TestProviderSatisfiesBudget pins the provider to the interface the chunker +// consumes. A compile-time assertion rather than a runtime test: the whole +// point of the interface is that the chunker never imports this package. +func TestProviderSatisfiesBudget(t *testing.T) { + var _ tokenizer.Budget = (*Provider)(nil) +} + +// TestFallbackWithoutTokenizer covers the degraded path: no tokenizer.json +// means estimates, the conservative batch cap, and ExactCounts()==false so a +// caller can widen its margins instead of trusting the number. +func TestFallbackWithoutTokenizer(t *testing.T) { + p := &Provider{cfg: Config{Model: "voyage-code-3"}} + if p.ExactCounts() { + t.Error("ExactCounts must be false without a tokenizer") + } + if got := p.maxTokensPerBatch(); got != defaultMaxTokensPerBatch { + t.Errorf("batch cap = %d, want the conservative %d", got, defaultMaxTokensPerBatch) + } + if got, want := p.CountTokens("hello world"), len("hello world")/bytesPerToken; got != want { + t.Errorf("CountTokens = %d, want the byte estimate %d", got, want) + } +} + +// TestOperatorOverrideWinsOverExactCap — an explicit MaxTokensPerRequest is +// the operator's call and must not be silently raised by exact counting. +func TestOperatorOverrideWinsOverExactCap(t *testing.T) { + p := &Provider{cfg: Config{Model: "voyage-code-3", MaxTokensPerRequest: 42_000}} + if got := p.maxTokensPerBatch(); got != 42_000 { + t.Errorf("batch cap = %d, want the operator's 42000", got) + } +} + +// TestSplitForInputUsesTokenBoundaries exercises the provider-level split with +// a tokenizer loaded — the branch that keeps a large-but-legal chunk out of the +// byte-window-and-average path. It needs the real tokenizer.json, so it skips +// on a clean checkout like the other fixture-backed tests. +func TestSplitForInputUsesTokenBoundaries(t *testing.T) { + const tokPath = "../../../../../loadtests/bench/voyage-code-3.tokenizer.json" + if _, err := os.Stat(tokPath); err != nil { + t.Skip("tokenizer.json not present") + } + c, err := bpecount.Load(tokPath) + if err != nil { + t.Fatalf("load tokenizer: %v", err) + } + p := &Provider{cfg: Config{Model: "voyage-code-3"}, counter: c} + + // Comfortably over the old 30 KB byte cap, comfortably under the model's + // 32K-token window: byte-windowing would split and average this, token + // counting must pass it through whole. + big := strings.Repeat("func handler(w http.ResponseWriter) { defer r.Body.Close() }\n", 700) + if n := p.CountTokens(big); n >= p.MaxInputTokens() { + t.Fatalf("fixture is %d tokens, needs to be under %d", n, p.MaxInputTokens()) + } + if got := p.splitForInput(big, 30_000); len(got) != 1 { + t.Errorf("input of %d bytes / %d tokens split into %d windows; a token-sized "+ + "input must pass through whole", len(big), p.CountTokens(big), len(got)) + } + + // Past the window: must split, and every piece must fit. + huge := strings.Repeat("x := compute(alpha, beta, gamma) // annotate the result\n", 40_000) + pieces := p.splitForInput(huge, 30_000) + if len(pieces) < 2 { + t.Fatalf("input of %d tokens was not split", p.CountTokens(huge)) + } + for i, piece := range pieces { + if n := p.CountTokens(piece); n > p.MaxInputTokens() { + t.Errorf("piece %d is %d tokens, over the %d-token window", i, n, p.MaxInputTokens()) + } + } + if strings.Join(pieces, "") != huge { + t.Error("pieces do not reconstruct the input") + } +} diff --git a/server/internal/embeddings/service.go b/server/internal/embeddings/service.go index cbb12c70..b1b18c3d 100644 --- a/server/internal/embeddings/service.go +++ b/server/internal/embeddings/service.go @@ -20,6 +20,7 @@ import ( // provider purely by kind string — these imports are the wiring. _ "github.com/dvcdsys/code-index/server/internal/embeddings/provider/openai" _ "github.com/dvcdsys/code-index/server/internal/embeddings/provider/voyage" + "github.com/dvcdsys/code-index/server/internal/tokenizer" ) // Service is the public embeddings API used by handlers and the indexer. @@ -457,6 +458,31 @@ func (s *Service) Status() provider.Status { return st } +// TokenBudget returns the active provider as a token budget when it can +// count tokens, and nil otherwise. The chunker uses it to size chunks in the +// model's own unit; nil keeps it on the byte heuristic. +// +// Snapshotted under the read lock like CurrentKind, because a provider swap +// mid-file would otherwise mix two models' limits inside one chunk set. +func (s *Service) TokenBudget() tokenizer.Budget { + // A typed-nil *Service still satisfies the capability interface the + // indexer asserts on, so the guard is not decoration: without it the + // first indexed file panics inside RLock. + if s == nil || s.disabled { + return nil + } + s.mu.RLock() + defer s.mu.RUnlock() + if s.current == nil { + return nil + } + b, ok := s.current.(tokenizer.Budget) + if !ok { + return nil + } + return b +} + // CurrentKind reports the kind of the active provider, or "" when // disabled / not yet built. Used by /status and admin endpoints. func (s *Service) CurrentKind() string { diff --git a/server/internal/indexer/indexer.go b/server/internal/indexer/indexer.go index 03a17ad5..4608f01d 100644 --- a/server/internal/indexer/indexer.go +++ b/server/internal/indexer/indexer.go @@ -23,6 +23,7 @@ import ( "github.com/dvcdsys/code-index/server/internal/embeddings" "github.com/dvcdsys/code-index/server/internal/langdetect" "github.com/dvcdsys/code-index/server/internal/symbolindex" + "github.com/dvcdsys/code-index/server/internal/tokenizer" "github.com/dvcdsys/code-index/server/internal/vectorstore" ) @@ -112,6 +113,17 @@ type TokenAwareEmbedder interface { TokenizeAndEmbed(ctx context.Context, texts []string) ([][]float32, error) } +// TokenBudgetSource is the capability of telling the chunker what a chunk +// costs in the active model's tokens. Named rather than asserted inline so a +// rename of TokenBudget is a compile error somewhere instead of a silent +// return to byte-sized chunking everywhere. +// +// *embeddings.Service satisfies it; test fakes generally do not, and get the +// byte path. +type TokenBudgetSource interface { + TokenBudget() tokenizer.Budget +} + // Service owns sessions and wires dependencies for the three-phase protocol. type Service struct { db *sql.DB @@ -141,6 +153,10 @@ type Service struct { // reindexed under the new format. embedIncludePath bool + // maxChunkTokens is the per-chunk token target (CIX_MAX_CHUNK_TOKENS). + // 0 means the chunker's own default. + maxChunkTokens int + // embeddingModel is the active embedding model identifier persisted on // projects.indexed_with_model at FinishIndexing. Set via // SetEmbeddingModel from main; empty string keeps the column NULL so @@ -208,6 +224,12 @@ func (s *Service) SetEmbedIncludePath(v bool) { s.embedIncludePath = v } +// SetMaxChunkTokens sets the per-chunk token target used when the active +// embedding provider can count tokens exactly. +func (s *Service) SetMaxChunkTokens(n int) { + s.maxChunkTokens = n +} + // SetEmbeddingModel records the model identifier the indexer will write to // projects.indexed_with_model at FinishIndexing. Called from main once the // runtime config is resolved; empty string disables the write (the column @@ -665,6 +687,7 @@ func (s *Service) ProcessFilesStreaming( // is CPU-local and cheap, so it stays sequential to keep progress-event // order; the expensive embed work is parallelised in stage 2. prep := make([]*preparedFile, 0, len(files)) + budgetSrc, _ := s.emb.(TokenBudgetSource) for fi, fp := range files { // file_started — emit even for files we'll skip below, so the client // counter advances monotonically and rendering stays aligned with N. @@ -705,7 +728,14 @@ func (s *Service) ProcessFilesStreaming( language = "text" } - chunks, refs, err := chunker.ChunkFile(fp.Path, fp.Content, language, 0) + // The budget is re-read per file: a provider swap between files is + // legitimate, mixing two models' limits inside one file's chunks is + // not. The type assertion itself is hoisted out of the loop. + var budget tokenizer.Budget + if budgetSrc != nil { + budget = budgetSrc.TokenBudget() + } + chunks, refs, err := chunker.ChunkFileTokens(fp.Path, fp.Content, language, 0, budget, s.maxChunkTokens) if err != nil { s.logger.Warn("indexer: chunk file failed", "path", fp.Path, "err", err) progressSend(progress, ProgressEvent{ diff --git a/server/internal/tokenizer/bpecount/bpecount.go b/server/internal/tokenizer/bpecount/bpecount.go new file mode 100644 index 00000000..64c7387f --- /dev/null +++ b/server/internal/tokenizer/bpecount/bpecount.go @@ -0,0 +1,634 @@ +// Package bpecount is a minimal, count-only byte-level BPE tokenizer for +// GPT-2/Qwen2-style tokenizer.json files (voyage-code-3, Qwen2, GPT-4o…). +// +// It reproduces the HuggingFace pipeline: +// +// normalizer = NFC +// pretokenize = Split(Qwen2 regex, Isolated) + ByteLevel(add_prefix_space=false) +// model = BPE (greedy lowest-rank merge) +// +// The Split regex contains `\s+(?!\S)`, a negative lookahead Go's RE2 +// cannot express, so the splitter is hand-rolled rather than compiled. +// Only a COUNT is produced — no ids, no offsets. +package bpecount + +import ( + "container/heap" + "encoding/json" + "fmt" + "os" + "strings" + "sync" + "unicode" + "unicode/utf8" + + "golang.org/x/text/unicode/norm" +) + +// Counter holds the merge table and a memo of pre-token → token count. +type Counter struct { + merges map[string]int32 // "left right" -> rank + + mu sync.RWMutex + memo map[string]int +} + +type tokJSON struct { + Model struct { + Type string `json:"type"` + Merges json.RawMessage `json:"merges"` + } `json:"model"` + Normalizer struct { + Type string `json:"type"` + } `json:"normalizer"` + PreTokenizer struct { + Type string `json:"type"` + PreTokenizers []struct { + Type string `json:"type"` + Pattern struct { + Regex string `json:"Regex"` + } `json:"pattern"` + } `json:"pretokenizers"` + } `json:"pre_tokenizer"` +} + +// qwen2SplitPattern is the pre-tokenizer regex this package implements by +// hand. It is compared, not compiled: the point of the hand-rolled splitter is +// that Go's RE2 cannot express the `\s+(?!\S)` lookahead in it. +// +// The comparison is the load-time guard against a plausible and silent +// failure: GPT-2 and o200k tokenizer.json files parse fine, declare +// model.type "BPE", and would produce confidently wrong counts against this +// splitter — GPT-2 has no Split stage at all, o200k has a different pattern. +// Refusing them keeps ExactCounts() false and the caller on its estimate, +// which is wrong but knows it is. +const qwen2SplitPattern = `(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+` + +// checkPipeline refuses a tokenizer whose normalizer or pre-tokenizer is not +// the one this package reimplements. Counting a different pipeline with this +// splitter does not fail loudly — it returns plausible numbers that are wrong, +// and the caller then packs batches and sizes chunks against them. +func checkPipeline(tj tokJSON) error { + if tj.Normalizer.Type != "NFC" { + return fmt.Errorf("bpecount: normalizer is %q, this package implements NFC", + tj.Normalizer.Type) + } + if tj.PreTokenizer.Type != "Sequence" || len(tj.PreTokenizer.PreTokenizers) < 2 { + return fmt.Errorf("bpecount: pre_tokenizer is %q, expected Sequence[Split, ByteLevel]", + tj.PreTokenizer.Type) + } + split, byteLevel := tj.PreTokenizer.PreTokenizers[0], tj.PreTokenizer.PreTokenizers[1] + if split.Type != "Split" || byteLevel.Type != "ByteLevel" { + return fmt.Errorf("bpecount: pre_tokenizer is Sequence[%s, %s], expected Sequence[Split, ByteLevel]", + split.Type, byteLevel.Type) + } + if split.Pattern.Regex != qwen2SplitPattern { + return fmt.Errorf("bpecount: Split pattern is not the one implemented here " + + "(a GPT-2 or o200k tokenizer would count wrong rather than fail)") + } + return nil +} + +// Load reads a tokenizer.json and keeps only what a count needs: the merges. +func Load(path string) (*Counter, error) { + b, err := os.ReadFile(path) + if err != nil { + return nil, err + } + return LoadBytes(b) +} + +func LoadBytes(b []byte) (*Counter, error) { + var tj tokJSON + if err := json.Unmarshal(b, &tj); err != nil { + return nil, err + } + if tj.Model.Type != "BPE" { + return nil, fmt.Errorf("bpecount: unsupported model type %q", tj.Model.Type) + } + if err := checkPipeline(tj); err != nil { + return nil, err + } + // merges is either ["a b", ...] (v1) or [["a","b"], ...] (v2). + var flat []string + m := make(map[string]int32) + if err := json.Unmarshal(tj.Model.Merges, &flat); err == nil { + for i, s := range flat { + m[s] = int32(i) + } + } else { + var pairs [][]string + if err := json.Unmarshal(tj.Model.Merges, &pairs); err != nil { + return nil, fmt.Errorf("bpecount: merges: %w", err) + } + for i, p := range pairs { + if len(p) == 2 { + m[p[0]+" "+p[1]] = int32(i) + } + } + } + return &Counter{merges: m, memo: make(map[string]int, 1<<16)}, nil +} + +// ---------- byte-level alphabet (GPT-2 bytes_to_unicode) ---------- + +var byteRune [256]rune + +func init() { + for b := 0; b < 256; b++ { + r := rune(b) + switch { + case r == 0xad: + r = 0x143 + case r <= 0x20: + r += 0x100 + case r >= 0x7f && r <= 0xa0: + r += 0xa2 + } + byteRune[b] = r + } +} + +// ---------- hand-rolled splitter ---------- +// +// Qwen2 pattern, alternation tried left to right (Perl leftmost-first): +// +// (?i:'s|'t|'re|'ve|'m|'ll|'d) +// [^\r\n\p{L}\p{N}]?\p{L}+ +// \p{N} +// ?[^\s\p{L}\p{N}]+[\r\n]* +// \s*[\r\n]+ +// \s+(?!\S) +// \s+ +// +// Every rune is covered by some branch, so Split(Isolated) yields no gaps. + +func isL(r rune) bool { return unicode.IsLetter(r) } +func isN(r rune) bool { return unicode.IsNumber(r) } +func isWS(r rune) bool { return unicode.IsSpace(r) } +func isNL(r rune) bool { return r == '\r' || r == '\n' } + +var contractions = []string{"s", "t", "re", "ve", "m", "ll", "d"} + +// nextToken returns the byte length of the pre-token starting at s[0]. +func nextToken(s string) int { + r0, w0 := decode(s, 0) + + // A: contraction + // + // Only the two bytes after the apostrophe can matter (the longest + // contraction is "re"/"ve"/"ll"), so lowercase just those. Lowercasing the + // whole remaining string here allocated a copy of the suffix for every + // apostrophe in the file: a 512 KB source with 5,000 quotes moved over a + // gigabyte through the allocator, on the indexing hot path. + if r0 == '\'' && len(s) > w0 { + tail := s[w0:] + if len(tail) > 2 { + tail = tail[:2] + } + low := strings.ToLower(tail) + for _, c := range contractions { + if strings.HasPrefix(low, c) { + return w0 + len(c) + } + } + } + + // B: [^\r\n\p{L}\p{N}]? \p{L}+ + { + i := 0 + if !isNL(r0) && !isL(r0) && !isN(r0) { + i = w0 + } + j := i + for j < len(s) { + r, w := decode(s, j) + if !isL(r) { + break + } + j += w + } + if j > i { // at least one letter followed + return j + } + } + + // C: single \p{N} + if isN(r0) { + return w0 + } + + // D: " ?" [^\s\p{L}\p{N}]+ [\r\n]* + { + i := 0 + if r0 == ' ' { + i = w0 + } + j := i + for j < len(s) { + r, w := decode(s, j) + if isWS(r) || isL(r) || isN(r) { + break + } + j += w + } + if j > i { + for j < len(s) { + r, w := decode(s, j) + if !isNL(r) { + break + } + j += w + } + return j + } + } + + // E/F/G: whitespace run. + if isWS(r0) { + // maximal whitespace run + end := 0 + lastNL := -1 + for end < len(s) { + r, w := decode(s, end) + if !isWS(r) { + break + } + if isNL(r) { + lastNL = end + w + } + end += w + } + // E: \s*[\r\n]+ — run truncated after its LAST \r or \n. + if lastNL >= 0 { + return lastNL + } + // F: \s+(?!\S) — whole run at EOF, else run minus its last rune. + if end == len(s) { + return end + } + _, lw := decodeLast(s[:end]) + if end-lw > 0 { + return end - lw + } + // G: \s+ (single whitespace rune followed by a non-space) + return end + } + + // Unreachable for well-formed input; make progress anyway. + return w0 +} + +func decode(s string, i int) (rune, int) { + if s[i] < utf8.RuneSelf { + return rune(s[i]), 1 + } + return utf8.DecodeRuneInString(s[i:]) +} + +func decodeLast(s string) (rune, int) { + return utf8.DecodeLastRuneInString(s) +} + +// ---------- BPE ---------- + +func (c *Counter) bpeLen(piece string) int { + c.mu.RLock() + n, ok := c.memo[piece] + c.mu.RUnlock() + if ok { + return n + } + n = c.bpe(piece) + c.mu.Lock() + if len(c.memo) < 1<<20 { + c.memo[piece] = n + } + c.mu.Unlock() + return n +} + +// node is one symbol in the doubly-linked list the merge loop walks. +type node struct { + prev, next int + s string + alive bool +} + +// cand is a candidate merge sitting in the priority queue. +type cand struct { + rank int32 + l, r int + // len of the two symbols when the candidate was pushed; a stale entry + // (one side already merged into something longer) is detected by comparing. + ll, rl int +} + +type candHeap []cand + +func (h candHeap) Len() int { return len(h) } +func (h candHeap) Less(i, j int) bool { + if h[i].rank != h[j].rank { + return h[i].rank < h[j].rank + } + return h[i].l < h[j].l // ties: leftmost first, matching HF +} +func (h candHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] } +func (h *candHeap) Push(x any) { *h = append(*h, x.(cand)) } +func (h *candHeap) Pop() any { + old := *h + n := len(old) + v := old[n-1] + *h = old[:n-1] + return v +} + +// bpe applies the greedy lowest-rank-first merge with a linked list + heap, +// so a pathological single pre-token (a 30 KB run of '-' or one enormous +// identifier) stays near-linear instead of the O(n^2) rescan a naive loop does. +func (c *Counter) bpe(piece string) int { + nodes := make([]node, 0, len(piece)) + for _, r := range piece { + i := len(nodes) + nodes = append(nodes, node{prev: i - 1, next: i + 1, s: string(r), alive: true}) + } + n := len(nodes) + if n < 2 { + return n + } + nodes[n-1].next = -1 + + h := make(candHeap, 0, n) + push := func(l, r int) { + if l < 0 || r < 0 || r >= n { + return + } + if rk, ok := c.merges[nodes[l].s+" "+nodes[r].s]; ok { + h = append(h, cand{rank: rk, l: l, r: r, ll: len(nodes[l].s), rl: len(nodes[r].s)}) + } + } + for i := 0; i+1 < n; i++ { + push(i, i+1) + } + heap.Init(&h) + + live := n + for h.Len() > 0 { + cd := heap.Pop(&h).(cand) + l, r := cd.l, cd.r + // Reject stale entries: either side merged away or grew since push. + if !nodes[l].alive || !nodes[r].alive || nodes[l].next != r || + len(nodes[l].s) != cd.ll || len(nodes[r].s) != cd.rl { + continue + } + nodes[l].s += nodes[r].s + nodes[r].alive = false + nodes[l].next = nodes[r].next + if nodes[r].next >= 0 { + nodes[nodes[r].next].prev = l + } + live-- + if live == 1 { + return 1 + } + if p := nodes[l].prev; p >= 0 { + if rk, ok := c.merges[nodes[p].s+" "+nodes[l].s]; ok { + heap.Push(&h, cand{rank: rk, l: p, r: l, ll: len(nodes[p].s), rl: len(nodes[l].s)}) + } + } + if nx := nodes[l].next; nx >= 0 { + if rk, ok := c.merges[nodes[l].s+" "+nodes[nx].s]; ok { + heap.Push(&h, cand{rank: rk, l: l, r: nx, ll: len(nodes[l].s), rl: len(nodes[nx].s)}) + } + } + } + return live +} + +// Count returns the number of tokens voyage/HF would produce for text. +func (c *Counter) Count(text string) int { + if text == "" { + return 0 + } + s := norm.NFC.String(text) + total := 0 + var sb strings.Builder + for len(s) > 0 { + n := nextToken(s) + if n <= 0 { + n = 1 + } + piece := s[:n] + s = s[n:] + // ByteLevel map + sb.Reset() + sb.Grow(len(piece) * 2) + for i := 0; i < len(piece); i++ { + sb.WriteRune(byteRune[piece[i]]) + } + total += c.bpeLen(sb.String()) + } + return total +} + +// SplitPoints returns the byte offsets at which text must be cut so that no +// piece exceeds budget tokens, plus the total token count of the whole text. +// +// Offsets index the CALLER's string, which is not the same thing as indexing +// the normalised copy counting works on. NFC can both shrink the text (a +// decomposed "e"+U+0301 becomes one code point) and grow it (the composition +// exclusions at U+0958..U+095F decompose under NFC), so an offset taken from +// the normalised copy can land mid-rune or past the end of the original — the +// latter panics the caller's slice expression. Already-normalised input, which +// is nearly all source code, takes the fast path where the two coincide. +// +// The cuts are exact, not estimated, and they need no search. BPE merges never +// cross a pre-token boundary in this pipeline (Split runs with Isolated +// behaviour, and ByteLevel+BPE are applied per pre-token), so a text's token +// count is the SUM of its pre-tokens' counts. Cutting on a pre-token boundary +// therefore leaves both sides tokenising exactly as they did inside the whole: +// the parts always add up to the total, with no drift to correct for. +// +// A single pre-token larger than budget cannot be honoured on a boundary (its +// merges DO interact internally), so splitInside falls back to a binary search +// on bytes within that one pre-token. That path is for base64 blobs and +// minified lines with no whitespace; on a 45-repo corpus it fires for 5 chunks +// in 1.9M. +// +// Offsets are cut points only: nil means the text already fits. +func (c *Counter) SplitPoints(text string, budget int) (offsets []int, total int) { + if text == "" { + return nil, 0 + } + if budget <= 0 { + return nil, c.Count(text) + } + if norm.NFC.IsNormalString(text) { + return c.splitNormalized(text, budget) + } + return c.splitDenormalized(text, budget) +} + +// splitDenormalized handles input that is not already NFC. +// +// It cuts one piece at a time: the first cut is computed on the normalised +// form, mapped back to a raw offset on a normalisation boundary, and the +// remainder is then processed from its real start. Recomputing per piece +// rather than translating a whole offsets slice is deliberate — mapping a cut +// backwards to a boundary shrinks the piece before it and grows the one after, +// so offsets computed against the old start would no longer hold. The cost is +// one pass per piece, paid only by input that is not already normalised, which +// source code essentially never is. +func (c *Counter) splitDenormalized(text string, budget int) ([]int, int) { + total := c.Count(text) + var offsets []int + base := 0 + for base < len(text) { + rest := text[base:] + cuts, _ := c.splitNormalized(norm.NFC.String(rest), budget) + if len(cuts) == 0 { + break + } + at := rawOffsetOf(rest, cuts[0]) + if at <= 0 || at >= len(rest) { + break + } + offsets = append(offsets, base+at) + base += at + } + return offsets, total +} + +// rawOffsetOf maps a byte offset in NFC(raw) back to a byte offset in raw, +// rounding DOWN to a normalisation boundary — rounding down can only shrink +// the piece that ends there, so the budget survives the rounding. +func rawOffsetOf(raw string, normOff int) int { + rawPos, normPos, lastRaw := 0, 0, 0 + for rawPos < len(raw) { + n := norm.NFC.NextBoundaryInString(raw[rawPos:], true) + if n <= 0 { + break + } + segNorm := len(norm.NFC.String(raw[rawPos : rawPos+n])) + if normPos+segNorm > normOff { + return lastRaw + } + normPos += segNorm + rawPos += n + lastRaw = rawPos + } + return lastRaw +} + +// splitNormalized is SplitPoints for input already known to be NFC, where a +// normalised offset IS a raw offset. +func (c *Counter) splitNormalized(s string, budget int) (offsets []int, total int) { + acc := 0 // tokens accumulated in the current piece + pos := 0 // byte offset into s + var sb strings.Builder + for pos < len(s) { + n := nextToken(s[pos:]) + if n <= 0 { + n = 1 + } + piece := s[pos : pos+n] + + sb.Reset() + sb.Grow(len(piece) * 2) + for i := 0; i < len(piece); i++ { + sb.WriteRune(byteRune[piece[i]]) + } + tk := c.bpeLen(sb.String()) + + switch { + case tk > budget: + // Does not fit even alone. Close the current piece, then cut + // inside this pre-token. + if acc > 0 { + offsets = append(offsets, pos) + acc = 0 + } + inner := c.splitInside(piece, budget) + for _, off := range inner { + offsets = append(offsets, pos+off) + } + // The tail after the last inner cut opens the next piece and + // must be CHARGED for: leaving acc at zero let the following + // pre-tokens add a full budget on top of it, producing pieces of + // up to twice the budget. + tailStart := 0 + if len(inner) > 0 { + tailStart = inner[len(inner)-1] + } + acc = c.Count(piece[tailStart:]) + case acc+tk > budget: + offsets = append(offsets, pos) + acc = tk + default: + acc += tk + } + total += tk + pos += n + } + return offsets, total +} + +// splitInside cuts one over-budget pre-token by binary search on its bytes. +// Inside a pre-token counts are not additive, so every candidate cut is +// re-counted — but the search converges in a handful of probes because +// bytes-per-token is near-constant within a homogeneous run. +func (c *Counter) splitInside(piece string, budget int) []int { + // Candidate cut positions are rune starts, enumerated once. The search + // then runs over INDICES into that list rather than over byte offsets. + // + // The byte-offset version of this loop deadlocked: it aligned a midpoint + // to a rune start by decrementing, and when alignment pulled the midpoint + // below lo, the next lo = mid+1 did not advance, so the (lo, hi) pair + // repeated forever. Any run of multi-byte runes long enough to exceed the + // budget reached it — a box-drawing comment separator is enough, and that + // hung the indexing worker with no error and no progress. Searching over + // rune indices removes the failure rather than guarding it: every + // candidate is a valid boundary by construction, so no alignment step + // exists to misbehave. + starts := make([]int, 0, len(piece)/2+2) + for i := 0; i < len(piece); { + starts = append(starts, i) + _, w := utf8.DecodeRuneInString(piece[i:]) + if w <= 0 { + w = 1 + } + i += w + } + starts = append(starts, len(piece)) + + var cuts []int + si := 0 + for si < len(starts)-1 { + if c.Count(piece[starts[si]:]) <= budget { + break + } + // Largest j > si whose prefix still fits. + lo, hi, best := si+1, len(starts)-1, -1 + for lo <= hi { + mid := (lo + hi) / 2 + if c.Count(piece[starts[si]:starts[mid]]) <= budget { + best = mid + lo = mid + 1 + } else { + hi = mid - 1 + } + } + if best < 0 { + // Even one rune exceeds the budget. Emit it anyway: refusing to + // advance is the deadlock this rewrite exists to remove, and a + // budget smaller than a single token is the caller's problem. + best = si + 1 + } + if starts[best] >= len(piece) { + break + } + cuts = append(cuts, starts[best]) + si = best + } + return cuts +} diff --git a/server/internal/tokenizer/bpecount/bpecount_test.go b/server/internal/tokenizer/bpecount/bpecount_test.go new file mode 100644 index 00000000..d41a6dfc --- /dev/null +++ b/server/internal/tokenizer/bpecount/bpecount_test.go @@ -0,0 +1,374 @@ +package bpecount + +import ( + "encoding/json" + "os" + "strings" + "testing" + "time" + "unicode/utf8" +) + +// tokenizerPath is the real voyage-code-3 tokenizer.json. The tests that need +// it skip when it is absent so a checkout without the 7 MB file still builds +// and tests clean. +const tokenizerPath = "../../../../loadtests/bench/voyage-code-3.tokenizer.json" + +func load(t *testing.T) *Counter { + t.Helper() + if _, err := os.Stat(tokenizerPath); err != nil { + t.Skip("tokenizer.json not present") + } + c, err := Load(tokenizerPath) + if err != nil { + t.Fatalf("load: %v", err) + } + return c +} + +// TestCountMatchesReference pins the counts that were verified against +// Voyage's own usage.total_tokens and the HuggingFace Rust tokenizer. The tab +// cases are the ones a RE2 rewrite of the pre-tokenizer regex gets wrong: the +// `\s+(?!\S)` lookahead cannot be expressed, and a naive rewrite absorbs a +// leading tab into the punctuation branch that may only absorb a space. +func TestCountMatchesReference(t *testing.T) { + c := load(t) + for _, tc := range []struct { + in string + want int + }{ + {"func main() {\n\tfmt.Println(\"hi\")\n}\n", 10}, + {"\t\t\"a\"", 4}, + {"\t\t\t\"end\": {", 6}, + {"a\t\t-b", 4}, + {"class A:\n def g(self):\n x = 1\n", 14}, + {"hello world", 2}, + {"#ifdef USE_THREADS", 3}, + {"", 0}, + } { + if got := c.Count(tc.in); got != tc.want { + t.Errorf("Count(%q) = %d, want %d", tc.in, got, tc.want) + } + } +} + +// TestNFCNormalisation covers the second gap in the ollama tokenizer: the +// tokenizer.json declares an NFC normalizer, and skipping it makes decomposed +// input cost an extra token. +func TestNFCNormalisation(t *testing.T) { + c := load(t) + nfc := "caf\u00e9" // é as one code point + nfd := "cafe\u0301" // e + combining acute + if a, b := c.Count(nfc), c.Count(nfd); a != b { + t.Errorf("NFC %d != NFD %d — normaliser not applied", a, b) + } +} + +// TestSplitPointsAreExact is the property the splitter exists for: because BPE +// merges never cross a pre-token boundary, the pieces must add up to the whole +// and none may exceed the budget. +func TestSplitPointsAreExact(t *testing.T) { + c := load(t) + src := "" + for i := 0; i < 400; i++ { + src += "func handler(w http.ResponseWriter, r *http.Request) {\n\tdefer r.Body.Close()\n}\n" + } + const budget = 500 + + offsets, total := c.SplitPoints(src, budget) + if total != c.Count(src) { + t.Fatalf("SplitPoints total %d != Count %d", total, c.Count(src)) + } + if len(offsets) == 0 { + t.Fatalf("expected cuts for %d tokens at budget %d", total, budget) + } + + sum, prev := 0, 0 + for _, off := range append(offsets, len(src)) { + n := c.Count(src[prev:off]) + if n > budget { + t.Errorf("piece [%d:%d] is %d tokens, over budget %d", prev, off, n, budget) + } + sum += n + prev = off + } + if sum != total { + t.Errorf("pieces sum to %d, whole is %d — merges leaked across a cut", sum, total) + } +} + +// TestSplitPointsFitsAlready — a text under budget must not be cut. +func TestSplitPointsFitsAlready(t *testing.T) { + c := load(t) + offsets, total := c.SplitPoints("package main\n", 1000) + if offsets != nil { + t.Errorf("expected no cuts, got %v", offsets) + } + if total == 0 { + t.Error("total should be counted even when no cut is needed") + } +} + +// TestSplitInsidePreToken covers the one case boundaries cannot serve: a +// single pre-token bigger than the budget (base64 blobs, minified lines). +func TestSplitPointsOversizePreToken(t *testing.T) { + c := load(t) + // One unbroken run of a single character class. "aB3" repeated would NOT + // do: the pre-tokenizer breaks letters from digits, so it yields 2-byte + // pre-tokens that never exceed the budget and the splitInside path this + // test exists for is never entered. + blob := strings.Repeat("a", 4000) + const budget = 100 + offsets, _ := c.SplitPoints(blob, budget) + if len(offsets) == 0 { + t.Fatal("expected the blob to be cut") + } + prev := 0 + for _, off := range append(offsets, len(blob)) { + if n := c.Count(blob[prev:off]); n > budget { + t.Errorf("piece [%d:%d] is %d tokens, over budget %d", prev, off, n, budget) + } + prev = off + } +} + +// --- CI coverage without the 7 MB file --- +// +// The golden-count tests above need the real voyage-code-3 tokenizer.json, +// which is not in the repo, so they skip on a clean checkout. The mechanics — +// pre-token splitting, merge application, the additivity SplitPoints relies on +// — do not need that vocabulary. A hand-built merge table exercises them, so +// CI still fails if the splitter or the merge loop regresses. +func syntheticCounter(t *testing.T) *Counter { + t.Helper() + // Merges are ranked: "a b" collapses first, then "ab c". + c, err := LoadBytes(syntheticJSON("a b", "ab c", "f u", "fu n")) + if err != nil { + t.Fatalf("LoadBytes: %v", err) + } + return c +} + +// TestSyntheticMergesApply — with only "a b" known, "abc" costs one merge plus +// the leftover byte; the second merge then folds that leftover in. +func TestSyntheticMergesApply(t *testing.T) { + c := syntheticCounter(t) + if got, want := c.Count("abc"), 1; got != want { + t.Errorf(`Count("abc") = %d, want %d (a+b -> ab, ab+c -> abc)`, got, want) + } + if got, want := c.Count("ab"), 1; got != want { + t.Errorf(`Count("ab") = %d, want %d`, got, want) + } + if got, want := c.Count("acb"), 3; got != want { + t.Errorf(`Count("acb") = %d, want %d (no merge applies)`, got, want) + } +} + +// TestSyntheticAdditivity is the property the whole splitter rests on: BPE +// never merges across a pre-token boundary, so counts add up. If a future +// change made merges span boundaries, cuts would silently produce over-budget +// pieces — this catches it without needing the real vocabulary. +func TestSyntheticAdditivity(t *testing.T) { + c := syntheticCounter(t) + const text = "abc abc\n\tabc fun fun" + whole := c.Count(text) + + offsets, total := c.SplitPoints(text, 3) + if total != whole { + t.Fatalf("SplitPoints total %d != Count %d", total, whole) + } + sum, prev := 0, 0 + for _, off := range append(offsets, len(text)) { + n := c.Count(text[prev:off]) + if n > 3 { + t.Errorf("piece %q is %d tokens, over budget 3", text[prev:off], n) + } + sum += n + prev = off + } + if sum != whole { + t.Errorf("pieces sum to %d, whole is %d", sum, whole) + } +} + +// TestPreTokenBoundaries pins the hand-rolled splitter against the branches of +// the Qwen2 pattern that a RE2 rewrite gets wrong — whitespace runs, and a tab +// that must NOT be absorbed into the punctuation branch. +func TestPreTokenBoundaries(t *testing.T) { + for _, tc := range []struct { + in string + want []string + }{ + {"a b", []string{"a", " b"}}, + {"a b", []string{"a", " ", " b"}}, + // A tab is a legal single-character prefix for the letter branch + // ([^\r\n\p{L}\p{N}]?\p{L}+), so it attaches to what follows. + {"x\n\ty", []string{"x", "\n", "\ty"}}, + // The lookahead branch \s+(?!\S) matches a whitespace run only when + // nothing non-space follows it. The first tab qualifies (a tab + // follows); the second does not (a quote follows) and falls through + // to plain \s+. Hence two separate pre-tokens, not one run — this is + // precisely what a RE2 rewrite of the pattern gets wrong. + {"\t\t\"a\"", []string{"\t", "\t", "\"a", "\""}}, + {"it's", []string{"it", "'s"}}, + {"a1", []string{"a", "1"}}, + } { + var got []string + s := tc.in + for len(s) > 0 { + n := nextToken(s) + if n <= 0 { + t.Fatalf("nextToken(%q) returned %d", s, n) + } + got = append(got, s[:n]) + s = s[n:] + } + if len(got) != len(tc.want) { + t.Errorf("split(%q) = %q, want %q", tc.in, got, tc.want) + continue + } + for i := range got { + if got[i] != tc.want[i] { + t.Errorf("split(%q) = %q, want %q", tc.in, got, tc.want) + break + } + } + } +} + +// TestNFDPiecesRespectBudget covers input that is not already NFC, where the +// normalised copy counting works on has different byte offsets from the +// caller's string. Before the fix, offsets computed against the normalised +// copy were returned as-is: decomposed text made them land mid-rune, and the +// composition exclusions at U+0958..U+095F (which DEcompose under NFC, making +// the normalised form longer) pushed them past the end of the original, so the +// caller's slice expression panicked. Every other fixture in this file is +// ASCII or already-NFC and could not see it. +func TestNFDPiecesRespectBudget(t *testing.T) { + c := load(t) + for _, raw := range []string{ + strings.Repeat("// café comment here\n", 200), + strings.Repeat("x क़ख़ग़ ", 400), + strings.Repeat("Ώ ", 900), + } { + offs, total := c.SplitPoints(raw, 50) + prev := 0 + for _, off := range append(offs, len(raw)) { + if off > len(raw) || off < prev { + t.Fatalf("bad offset %d (len %d, prev %d)", off, len(raw), prev) + } + if n := c.Count(raw[prev:off]); n > 50 { + t.Errorf("piece [%d:%d] is %d tokens, over budget 50", prev, off, n) + } + prev = off + } + if total != c.Count(raw) { + t.Errorf("total %d != Count %d", total, c.Count(raw)) + } + } +} + +// TestSplitInsideChargesTail — an over-budget pre-token used to leave its tail +// uncounted, letting the next piece stack a full budget on top of it. +func TestSplitInsideChargesTail(t *testing.T) { + c, err := LoadBytes(syntheticJSON("a b")) + if err != nil { + t.Fatal(err) + } + in := strings.Repeat("x", 10) + " " + strings.Repeat("y", 4) + offs, _ := c.SplitPoints(in, 5) + prev := 0 + for _, off := range append(offs, len(in)) { + if n := c.Count(in[prev:off]); n > 5 { + t.Errorf("piece %q is %d tokens, over budget 5", in[prev:off], n) + } + prev = off + } +} + +// syntheticJSON builds a tokenizer.json with a hand-picked merge table and the +// real pipeline sections, so LoadBytes's compatibility check sees what it +// expects. Tests that only exercise merging still have to declare the pipeline +// they are pretending to be — which is the point of the check. +func syntheticJSON(merges ...string) []byte { + doc := map[string]any{ + "model": map[string]any{"type": "BPE", "merges": merges}, + "normalizer": map[string]any{"type": "NFC"}, + "pre_tokenizer": map[string]any{ + "type": "Sequence", + "pretokenizers": []any{ + map[string]any{"type": "Split", "pattern": map[string]any{"Regex": qwen2SplitPattern}}, + map[string]any{"type": "ByteLevel"}, + }, + }, + } + b, err := json.Marshal(doc) + if err != nil { + panic(err) + } + return b +} + +// TestRejectsForeignPipeline — a GPT-2 or o200k tokenizer.json parses cleanly +// and declares BPE, but its pre-tokenizer is not the one implemented here. It +// must be refused rather than counted wrongly. +func TestRejectsForeignPipeline(t *testing.T) { + for name, doc := range map[string]string{ + "gpt2 (no Split stage)": `{"model":{"type":"BPE","merges":["a b"]}, + "normalizer":null,"pre_tokenizer":{"type":"ByteLevel"}}`, + "o200k (different pattern)": `{"model":{"type":"BPE","merges":["a b"]}, + "normalizer":{"type":"NFC"},"pre_tokenizer":{"type":"Sequence","pretokenizers":[ + {"type":"Split","pattern":{"Regex":"[^\\r\\n\\p{L}\\p{N}]?[\\p{L}]+"}}, + {"type":"ByteLevel"}]}}`, + } { + if _, err := LoadBytes([]byte(doc)); err == nil { + t.Errorf("%s: expected a load error, got none", name) + } + } +} + +// TestMultibyteRunsTerminate covers runs of multi-byte runes long enough to +// exceed the budget as a single pre-token — a box-drawing comment separator, +// an arrow run, a run of combining marks. +// +// The byte-offset binary search this replaced aligned its midpoint to a rune +// start by DECREMENTING, so when alignment pulled the midpoint below lo, the +// next lo = mid+1 did not advance and the search spun forever. It took no +// error path and produced no output: the indexing worker simply stopped. All +// three inputs below hung at budgets 5 and 50. +func TestMultibyteRunsTerminate(t *testing.T) { + c := load(t) + inputs := map[string]string{ + "box drawing separator": "// " + strings.Repeat("\u2500", 400) + "\n", + "arrow run": strings.Repeat("\u2192", 400), + "combining marks": strings.Repeat("\u0301", 50), + "composition exclusion": strings.Repeat("\u0958", 2000), + } + for _, budget := range []int{5, 50} { + for name, in := range inputs { + done := make(chan struct{}) + var offs []int + go func(s string, b int) { + offs, _ = c.SplitPoints(s, b) + close(done) + }(in, budget) + + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatalf("%s at budget %d: SplitPoints did not return", name, budget) + } + + prev := 0 + for _, off := range append(offs, len(in)) { + if off > len(in) || off < prev { + t.Fatalf("%s: offset %d out of range (len %d, prev %d)", name, off, len(in), prev) + } + if !utf8.ValidString(in[prev:off]) { + t.Errorf("%s: piece [%d:%d] is not valid UTF-8 — cut mid-rune", name, prev, off) + } + prev = off + } + } + } +} diff --git a/server/internal/tokenizer/budget.go b/server/internal/tokenizer/budget.go new file mode 100644 index 00000000..90d6406c --- /dev/null +++ b/server/internal/tokenizer/budget.go @@ -0,0 +1,49 @@ +// Package tokenizer carries the model-token knowledge the chunker needs and +// the embedding providers own. +// +// The chunker decides WHERE to cut; only the provider knows WHAT the model +// counts. Keeping the two apart is what lets a model change without teaching +// the chunker about byte-level BPE, SentencePiece, or llama-server's +// /tokenize endpoint. +package tokenizer + +// Budget is the whole surface the chunker sees. One interface rather than a +// mandatory one plus an optional splitter: a single constructor argument, and +// callers that need both never have to type-assert. +// +// Cost note, because the two methods look interchangeable and are not: +// CountTokens and SplitPoints do the same single left-to-right pass over the +// same memo, so neither is algorithmically cheaper. What differs is +// allocation. CountTokens returns an int and allocates nothing; SplitPoints +// must build a slice of offsets — for a 60 KB input that is on the order of +// 15k entries. CountTokens runs on every chunk (1.9M of them on the reference +// corpus) while SplitPoints runs only on inputs that exceed the model's +// context (5 of that same 1.9M). Call CountTokens on the hot path and reach +// for SplitPoints only once a text is known not to fit. +// +// A third shortcut avoids both: byte-level BPE cannot emit a token covering +// less than one byte, so len(text) <= budget PROVES the text fits, with no +// tokenisation at all. Use it before calling anything here. +type Budget interface { + // MaxInputTokens is the model's context window for a single input. + MaxInputTokens() int + + // ExactCounts reports whether CountTokens and SplitPoints are exact. + // False means the provider has no tokenizer and is estimating from + // byte length: counts may be wrong in both directions and split points + // are byte windows, not token boundaries. Callers that need a + // guarantee must widen their safety margin when this is false — + // silently trusting an estimate is what the byte-window splitter used + // to do, and it produced averaged vectors nobody could see was wrong. + ExactCounts() bool + + // CountTokens returns the number of tokens the model will charge for. + CountTokens(s string) int + + // SplitPoints returns byte offsets at which s must be cut so no piece + // exceeds budget tokens, plus the total token count of s. Offsets, not + // substrings, so the caller keeps ownership of the metadata that hangs + // off those positions — line numbers, symbol names, byte ranges. + // A nil offsets slice means s already fits. + SplitPoints(s string, budget int) (offsets []int, total int) +}