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