From 0cb4039be9c0691f89c391c5cc28ab40682a8163 Mon Sep 17 00:00:00 2001 From: Jack Date: Wed, 29 Jul 2026 21:10:19 +0800 Subject: [PATCH] fix(ingestion): clamp unconfigured embedding limit; document chunker title parity (comment-only) (#17539) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Fixes embedding inputs being emptied and documents chunker heading-detection parity. - `component/tokenizer.go`: `truncateForEmbedding` now mirrors Python `common/token_utils.py:183-185` `truncate(string, max_len)` (keep the first `max_len` tokens). For an unconfigured embedder reporting `maxTokens <= 0`, instead of mirroring Python's `truncate` (which returns `""` for `max_len <= 0` and would make the embeddings API reject the batch with "inputs cannot be empty"), Go **clamps the limit to a safe default** (`defaultEmbeddingTokenLimit = 8192`) so truncation stays active and never produces empty inputs. The previous code returned `""` for `maxTokens <= 10`, which emptied non-empty chunks into the batch and triggered the API error. - A **10-token safety margin** is reserved (mirroring Python's embedding path `rag/svr/task_executor.py` which uses `mdl.max_length - 10`) so Go does not over-send tokens to hard-capped embedding APIs; it is only applied when the limit `> 10` so small limits remain non-empty. - The clamp is applied **centrally inside `truncateForEmbedding`**, covering both Builtin and generic callers — no separate `model_service.go` change is required. - `component/chunker/title.go`: comment-only — documents the already-ported PDF-outline detection (omission 1.5 / Gap C) and the hardcoded `BULLET_PATTERN` fallback (omission 1.7 / Gap C), porting `common.py:_outline_similarity` and `rag/nlp` `BULLET_PATTERN`. - `tokenizer_unit_test.go`: update truncation expectations to the new positive-keeps-tokens / non-positive-clamps-to-default behaviour. ## Test plan - `tokenizer_unit_test.go` updated for the new truncation behaviour (`TestTruncateForEmbedding_SmallMaxTokens`, `TestTruncateForEmbedding_UnconfiguredClampsToDefault`). - `./build.sh --test ./internal/ingestion/component/` passes. 🤖 Generated with [CodeBuddy Code](https://cnb.cool/codebuddy) --- internal/ingestion/component/chunker/title.go | 22 ++++--- internal/ingestion/component/tokenizer.go | 41 +++++++++---- .../component/tokenizer_unit_test.go | 57 +++++++++++++++---- 3 files changed, 92 insertions(+), 28 deletions(-) diff --git a/internal/ingestion/component/chunker/title.go b/internal/ingestion/component/chunker/title.go index 40662c341f..7484925a5f 100644 --- a/internal/ingestion/component/chunker/title.go +++ b/internal/ingestion/component/chunker/title.go @@ -37,15 +37,21 @@ // a non-regex text record whose layout flags it as a // section/title/head and whose text passes not_title is promoted to // fallback_level = len(selected_group) + 1. Strategy (1) — -// PDF-outline detection — still requires deepdoc/parser binary -// access and remains the only parity gap. +// PDF-outline detection — is also ported: newLevelContext +// (title.go) reads the Parser-supplied file.outline via +// outlineFromInputs / resolveOutlineLevels / outlineSimilarity +// (port of common.py:_outline_similarity) before falling back to +// the regex/layout branch, wired into group.go / hierarchy.go +// (Chunker omission 1.5, Gap C closed). // -// - The Go port has NO hardcoded BULLET_PATTERN fallback. Heading -// detection relies entirely on the user-supplied `levels` param -// (which templates carry as comprehensive regex families) paired -// with the layout-hint fallback. A PDF without matching levels -// produces BODY_LEVEL-only records — Python's BULLET_PATTERN-based -// tree_merge / hierarchical_merge would still find structure. +// - The Go port SHIPS a hardcoded BULLET_PATTERN fallback +// (Chunker omission 1.7, Gap C closed): when the regex/layout +// branch yields BODY_LEVEL-only records, resolveTitleLevels +// (title.go) calls bulletsCategory over bulletPatterns (title.go, +// mirroring Python BULLET_PATTERN at rag/nlp/__init__.py:258) to +// recover structure from numbered/bulleted list entries, before +// the layout-title fallback. It never overrides an existing level +// assignment. // // - GROUP-TITLE and HIERARCHY-TITLE are separate Go files // (`group.go`, `hierarchy.go`); they share the resolve_levels diff --git a/internal/ingestion/component/tokenizer.go b/internal/ingestion/component/tokenizer.go index cf7c5644cc..8e8318b796 100644 --- a/internal/ingestion/component/tokenizer.go +++ b/internal/ingestion/component/tokenizer.go @@ -403,13 +403,14 @@ func (c *TokenizerComponent) embedChunks(ctx context.Context, tenantID, kbID, em texts := make([]string, 0, len(chunks)) pairs := make([]int, 0, len(chunks)) for i, ck := range chunks { - txt := concatFields(ck, c.param.Fields) - txt = htmlTableRE.ReplaceAllString(txt, " ") + raw := concatFields(ck, c.param.Fields) + txt := htmlTableRE.ReplaceAllString(raw, " ") txt = strings.TrimSpace(txt) + trunc := truncateForEmbedding(txt, embedder.MaxTokens()) if txt == "" { continue } - texts = append(texts, truncateForEmbedding(txt, embedder.MaxTokens())) + texts = append(texts, trunc) pairs = append(pairs, i) } if len(texts) == 0 { @@ -476,15 +477,35 @@ func (c *TokenizerComponent) embedChunks(ctx context.Context, tenantID, kbID, em return chunks, tokenCount, nil } -// truncateForEmbedding truncates text to fit within maxTokens, -// reserving 10 tokens for special/model overhead. Mirrors Python -// tokenizer.py truncation: when maxTokens <= 10, the text is -// truncated to empty +// defaultEmbeddingTokenLimit is the safe fallback used when an embedder reports +// no token limit (maxTokens <= 0). It both prevents empty embedding inputs and +// keeps truncation active for every path instead of passing the full text through. +const defaultEmbeddingTokenLimit = 8192 + +// truncateForEmbedding keeps the first maxTokens tokens of text so it fits the +// embedding model's limit. +// +// For a positive maxTokens it mirrors Python common/token_utils.py:183-185 +// `truncate(string, max_len)` (keep the first max_len tokens). +// +// An unconfigured embedder reports maxTokens <= 0. Rather than mirror Python's +// behaviour of returning "" (which would make the embeddings API reject the whole +// batch with "inputs cannot be empty"), Go clamps the limit to a safe default +// (defaultEmbeddingTokenLimit = 8192). This both prevents empty inputs AND keeps +// truncation active for every path (Builtin and generic) instead of silently +// passing the full, untruncated text when no limit is configured. func truncateForEmbedding(text string, maxTokens int) string { - if maxTokens <= 10 { - return "" + if maxTokens <= 0 { + maxTokens = defaultEmbeddingTokenLimit } - return tokenizer.TrimContentToTokenLimit(text, maxTokens-10) + // Keep a 10-token safety margin, mirroring Python's embedding path + // (rag/svr/task_executor.py uses `mdl.max_length - 10`). Only apply it + // when the limit is large enough; for small limits (<=10) keep the full + // value so the result stays non-empty instead of collapsing to "". + if maxTokens > 10 { + maxTokens -= 10 + } + return tokenizer.TrimContentToTokenLimit(text, maxTokens) } func mergeEmbeddingVectors(titleVec, contentVec []float64, titleWeight float64) ([]float64, error) { diff --git a/internal/ingestion/component/tokenizer_unit_test.go b/internal/ingestion/component/tokenizer_unit_test.go index 9e3bdcd65b..25dcc1f5f7 100644 --- a/internal/ingestion/component/tokenizer_unit_test.go +++ b/internal/ingestion/component/tokenizer_unit_test.go @@ -91,8 +91,8 @@ func (s *stubEmbedder) Encode(ctx context.Context, texts []string) ([]EmbeddingR } // newStubEmbedder returns a stub embedder for instance-level resolver injection. -// maxTokens defaults to 2048 so truncateForEmbedding (Diff 14: maxTokens <= 10 -> "") -// does not empty the content text; tests that exercise truncation set maxTokens explicitly. +// maxTokens defaults to 2048 so truncateForEmbedding truncates to the first +// 2048 tokens; tests that exercise truncation set maxTokens explicitly. func newStubEmbedder(dim int) *stubEmbedder { return &stubEmbedder{dim: dim, maxTokens: 2048} } @@ -376,21 +376,58 @@ func TestIsPhantomChunk(t *testing.T) { } } -// TestTruncateForEmbedding_SmallMaxTokens covers Tokenizer Diff-14: when -// maxTokens <= 10, Go must truncate to empty, matching Python's -// truncation-to-empty behaviour (common/token_utils.py:183-185). +// TestTruncateForEmbedding_SmallMaxTokens covers Tokenizer Diff-14. For any +// positive maxTokens, truncateForEmbedding keeps the first maxTokens tokens +// (non-empty) and the result is strictly shorter than the input. +// +// The unconfigured case (maxTokens <= 0) is covered separately by +// TestTruncateForEmbedding_UnconfiguredClampsToDefault: rather than mirroring +// Python's `truncate` (which returns "" for max_len <= 0 and would make the +// embeddings API reject the batch), Go clamps the limit to a safe default so +// every path still truncates instead of passing the full text through. func TestTruncateForEmbedding_SmallMaxTokens(t *testing.T) { - long := strings.Repeat("a", 100) - if got := truncateForEmbedding(long, 5); got != "" { - t.Errorf("truncateForEmbedding(maxTokens=5) = %q, want %q", got, "") + // Long enough to produce well over 50 tokens, so 5/10/50 are all clearly + // below the total and truncation is observable. + long := strings.Repeat("a", 2000) + if got := truncateForEmbedding(long, 5); got == "" { + t.Error("truncateForEmbedding(maxTokens=5) returned empty, want first 5 tokens") } - if got := truncateForEmbedding(long, 10); got != "" { - t.Errorf("truncateForEmbedding(maxTokens=10) = %q, want %q", got, "") + if got := truncateForEmbedding(long, 10); got == "" { + t.Error("truncateForEmbedding(maxTokens=10) returned empty, want first 10 tokens") } // Normal path: maxTokens > 10 should truncate (not return empty). if got := truncateForEmbedding(long, 50); got == "" { t.Error("truncateForEmbedding(maxTokens=50) returned empty, want truncated text") } + if got := truncateForEmbedding(long, 50); len(got) >= len(long) { + t.Errorf("truncateForEmbedding(maxTokens=50) len = %d, want strictly shorter than %d", len(got), len(long)) + } +} + +// TestTruncateForEmbedding_UnconfiguredClampsToDefault is the regression gate +// for the root cause behind embedding truncation: an embedder that reports no +// token limit (maxTokens <= 0) must NOT be passed through verbatim. Before the +// central clamp, the generic model_service branch left maxTokens = 0, which +// silently disabled truncation for the whole generic path. The fix clamps +// maxTokens <= 0 to defaultEmbeddingTokenLimit (8192) inside +// truncateForEmbedding itself, so every caller — Builtin and generic alike — +// keeps truncation active. +// +// For a clearly-over-limit input and maxTokens = 0, the result must be non-empty +// AND strictly shorter than the input: proof that it was truncated to the +// default, not returned verbatim. +func TestTruncateForEmbedding_UnconfiguredClampsToDefault(t *testing.T) { + // ~10000+ CL100K tokens, comfortably above the 8192 default so truncation + // is observable. + long := strings.Repeat("hello world ", 5000) + got := truncateForEmbedding(long, 0) + if got == "" { + t.Fatal("truncateForEmbedding(maxTokens=0) returned empty; clamp must keep truncation active") + } + if len(got) >= len(long) { + t.Errorf("truncateForEmbedding(maxTokens=0) len = %d, want strictly shorter than %d; clamp to default must truncate", + len(got), len(long)) + } } // TestEmbeddingBatchSizeEnvVar covers Tokenizer Omission-3: the batch size