fix(ingestion): clamp unconfigured embedding limit; document chunker title parity (comment-only) (#17539)

## 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)
This commit is contained in:
Jack
2026-07-29 21:10:19 +08:00
committed by GitHub
parent 90f46b0b4d
commit 0cb4039be9
3 changed files with 92 additions and 28 deletions

View File

@@ -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

View File

@@ -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) {

View File

@@ -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