From 2f86e5b13d4fe43945ad80684ee865ac757abe50 Mon Sep 17 00:00:00 2001 From: wizardchen Date: Wed, 6 May 2026 20:19:33 +0800 Subject: [PATCH] feat(chunker): coalesce tiny adjacent chunks in heading splitter Documents with many short headings (FAQ-style, quick refs) used to fail the heading-tier validator with "too many tiny chunks" and silently fall back to legacy splitting, defeating the purpose of heading-aware mode. Merge physically adjacent sections whose combined size still fits within ChunkSize, deriving a shared breadcrumb via commonHeadingPrefix so the merged chunk's ContextHeader stays meaningful. Tests cover the merge path, position-invariant preservation after merge, ChunkSize ceiling, and the breadcrumb prefix helper. Existing tests that relied on tiny fixtures were grown so each section stays distinct. --- .../chunker/heading_splitter.go | 81 +++++++++ .../chunker/heading_splitter_test.go | 163 +++++++++++++++--- 2 files changed, 221 insertions(+), 23 deletions(-) diff --git a/internal/infrastructure/chunker/heading_splitter.go b/internal/infrastructure/chunker/heading_splitter.go index 4a9e13416..4ea46dd71 100644 --- a/internal/infrastructure/chunker/heading_splitter.go +++ b/internal/infrastructure/chunker/heading_splitter.go @@ -107,9 +107,90 @@ func splitByHeadingsImpl(text string, cfg SplitterConfig) []Chunk { } } + return coalesceTinyChunks(out, cfg.ChunkSize) +} + +// coalesceTinyChunks merges adjacent small chunks under their shared heading +// context so that documents whose primary sections are mostly short (FAQs, +// install logs, change-lists) don't trip the validator's "too many tiny +// chunks" rule and fall through all the way to legacy. The merged breadcrumb +// is the line-prefix shared by both inputs; the original sub-headings remain +// visible because heading_splitter includes the heading line in each +// section's Content. +// +// Safety: +// - We only merge when cur.End == next.Start. That preserves the +// End-Start == len([]rune(Content)) invariant that document +// reconstruction relies on, and naturally skips legacy sub-chunks (which +// may overlap due to ChunkOverlap). +// - We stop accumulating once the running chunk reaches the merge target +// (≈ ChunkSize/2) so we don't aggressively pack chunks beyond what the +// validator considers comfortable. +func coalesceTinyChunks(in []Chunk, chunkSize int) []Chunk { + if len(in) <= 1 || chunkSize <= 0 { + return in + } + target := chunkSize / 2 + if target < 200 { + target = 200 + } + + out := make([]Chunk, 0, len(in)) + cur := in[0] + curLen := utf8.RuneCountInString(cur.Content) + + for i := 1; i < len(in); i++ { + next := in[i] + nextLen := utf8.RuneCountInString(next.Content) + // Adjacent + still-small + would not blow the size budget → merge. + if cur.End == next.Start && curLen < target && curLen+nextLen <= chunkSize { + cur.Content += next.Content + cur.ContextHeader = commonHeadingPrefix(cur.ContextHeader, next.ContextHeader) + cur.End = next.End + curLen += nextLen + continue + } + out = append(out, cur) + cur = next + curLen = nextLen + } + out = append(out, cur) + + // Re-sequence — downstream code (knowledge.go) expects Seq to be a dense + // 0..N-1 range over the returned slice. + for i := range out { + out[i].Seq = i + } return out } +// commonHeadingPrefix returns the longest line-aligned prefix shared by two +// breadcrumb strings. Heading hierarchies are emitted as +// "# Top\n## Section\n### Sub", so a line-by-line comparison is sufficient +// and avoids partial-line truncation that would corrupt the breadcrumb. +func commonHeadingPrefix(a, b string) string { + if a == b { + return a + } + la := strings.Split(a, "\n") + lb := strings.Split(b, "\n") + n := len(la) + if len(lb) < n { + n = len(lb) + } + common := 0 + for i := 0; i < n; i++ { + if la[i] != lb[i] { + break + } + common = i + 1 + } + if common == 0 { + return "" + } + return strings.Join(la[:common], "\n") +} + // findHeadingBoundaries returns one boundary at offset 0 plus one per // Markdown heading at level <= primaryLevel that sits outside fenced code // blocks. Heading detection is line-oriented — a heading must occupy a diff --git a/internal/infrastructure/chunker/heading_splitter_test.go b/internal/infrastructure/chunker/heading_splitter_test.go index f8257b306..1efcbd24c 100644 --- a/internal/infrastructure/chunker/heading_splitter_test.go +++ b/internal/infrastructure/chunker/heading_splitter_test.go @@ -6,18 +6,13 @@ import ( ) func TestSplitByHeadings_BasicSections(t *testing.T) { - doc := `# Top -preamble. - -## Section A -content of A. - -## Section B -content of B. - -## Section C -content of C.` - cfg := SplitterConfig{ChunkSize: 200, ChunkOverlap: 0} + // Each section is intentionally larger than the merge-target (≈ + // ChunkSize/2) so the post-split coalesce pass leaves them as distinct + // chunks. We're testing per-section emission + breadcrumb here, not + // merging. + body := strings.Repeat("Lorem ipsum dolor sit amet consectetur adipiscing elit. ", 4) + doc := "# Top\n" + body + "\n\n## Section A\n" + body + "\n\n## Section B\n" + body + "\n\n## Section C\n" + body + cfg := SplitterConfig{ChunkSize: 300, ChunkOverlap: 0} chunks := splitByHeadingsImpl(doc, cfg) if len(chunks) < 3 { t.Fatalf("expected ≥3 chunks (one per section), got %d", len(chunks)) @@ -36,12 +31,12 @@ content of C.` found := false for _, c := range chunks { - if strings.Contains(c.Content, "Section B") && strings.Contains(c.Content, "content of B") { + if strings.Contains(c.Content, "## Section B") && strings.Contains(c.Content, "Lorem ipsum") { found = true } } if !found { - t.Error("no chunk contains Section B with its content") + t.Error("no chunk contains Section B with its body") } } @@ -72,15 +67,12 @@ func TestSplitByHeadings_LargeSectionRecursesIntoLegacy(t *testing.T) { } func TestSplitByHeadings_BreadcrumbReflectsLatestPath(t *testing.T) { - doc := `# Chapter 1 -intro - -## Section A -text A - -## Section B -text B` - cfg := SplitterConfig{ChunkSize: 200, ChunkOverlap: 0} + // Sized so each section stays its own chunk after the tiny-section + // coalesce pass — we're verifying breadcrumb assignment per section, + // not the merge behavior. + body := strings.Repeat("Lorem ipsum dolor sit amet consectetur adipiscing elit. ", 4) + doc := "# Chapter 1\n" + body + "\n\n## Section A\n" + body + "\n\n## Section B\n" + body + cfg := SplitterConfig{ChunkSize: 300, ChunkOverlap: 0} chunks := splitByHeadingsImpl(doc, cfg) if len(chunks) < 3 { t.Fatalf("expected ≥3 chunks, got %d", len(chunks)) @@ -158,6 +150,131 @@ content of C here.` } } +// TestSplitByHeadings_CoalescesTinyAdjacentSections covers the FAQ / +// install-log case where a parent heading hosts many short sub-sections. +// Without merging, each `##` becomes its own <50-char chunk and the +// validator rejects the tier with "too many tiny chunks". After merging, +// they collapse into a small number of properly-sized chunks while still +// surfacing the shared parent breadcrumb. +func TestSplitByHeadings_CoalescesTinyAdjacentSections(t *testing.T) { + doc := `# Install Log + +## Docker镜像 +使用 daocloud 部署 v0.3.1。 + +## 前端老版本 +浏览器缓存了旧前端资源。 + +## 登录报错 +ERROR: column missing. + +## 解析失败 +embedding 表缺列。` + cfg := SplitterConfig{ChunkSize: 500, ChunkOverlap: 0} + chunks := splitByHeadingsImpl(doc, cfg) + if len(chunks) == 0 { + t.Fatal("expected at least one chunk") + } + if len(chunks) >= 5 { + t.Errorf("expected coalesce to produce <5 chunks, got %d", len(chunks)) + } + // Every merged chunk must carry the shared parent in its breadcrumb so + // retrieval can still answer "what document is this from". + for i, c := range chunks { + if !strings.Contains(c.ContextHeader, "# Install Log") { + t.Errorf("chunk %d missing parent H1 in breadcrumb: %q", i, c.ContextHeader) + } + } + // All four sub-section headings must remain visible somewhere in the + // merged content (heading_splitter keeps the heading line as part of + // each section's Content). + for _, h := range []string{"## Docker镜像", "## 前端老版本", "## 登录报错", "## 解析失败"} { + seen := false + for _, c := range chunks { + if strings.Contains(c.Content, h) { + seen = true + break + } + } + if !seen { + t.Errorf("merged chunks should still contain heading %q somewhere", h) + } + } +} + +// TestSplitByHeadings_CoalescePreservesPositionInvariant guards the +// End-Start == len([]rune(Content)) invariant after merging. Adjacent +// chunks (cur.End == next.Start) must concatenate cleanly; the merge must +// refuse to combine non-adjacent chunks (e.g. legacy sub-chunks from an +// oversized section that overlap). +func TestSplitByHeadings_CoalescePreservesPositionInvariant(t *testing.T) { + doc := `# Top + +## A +short A. + +## B +short B. + +## C +short C.` + cfg := SplitterConfig{ChunkSize: 500, ChunkOverlap: 0} + chunks := splitByHeadingsImpl(doc, cfg) + docRunes := []rune(doc) + for i, c := range chunks { + contentRuneLen := len([]rune(c.Content)) + if c.End-c.Start != contentRuneLen { + t.Errorf("chunk %d: End-Start(%d) != content_runes(%d) after merge", + i, c.End-c.Start, contentRuneLen) + } + if c.Start >= 0 && c.End <= len(docRunes) { + if string(docRunes[c.Start:c.End]) != c.Content { + t.Errorf("chunk %d: source[Start:End] != Content after merge", i) + } + } + } +} + +// TestSplitByHeadings_CoalesceRespectsChunkSize ensures the merge target +// stays within the ChunkSize budget — the validator caps oversize chunks +// at 2x and we should never approach that line via merging. +func TestSplitByHeadings_CoalesceRespectsChunkSize(t *testing.T) { + const sections = 30 + var sb strings.Builder + sb.WriteString("# Doc\n") + for i := 0; i < sections; i++ { + sb.WriteString("\n## Section ") + sb.WriteString(strings.Repeat("X", 1)) // unique-ish heading + sb.WriteString("\nshort body line.\n") + } + cfg := SplitterConfig{ChunkSize: 200, ChunkOverlap: 0} + chunks := splitByHeadingsImpl(sb.String(), cfg) + for i, c := range chunks { + if l := len([]rune(c.Content)); l > cfg.ChunkSize { + t.Errorf("chunk %d exceeds ChunkSize: %d > %d", i, l, cfg.ChunkSize) + } + } +} + +// TestCommonHeadingPrefix exercises the breadcrumb-prefix helper directly. +func TestCommonHeadingPrefix(t *testing.T) { + cases := []struct { + a, b, want string + }{ + {"# Top\n## A", "# Top\n## B", "# Top"}, + {"# Top", "# Top", "# Top"}, + {"# X", "# Y", ""}, + {"# Top\n## A\n### x", "# Top\n## A\n### y", "# Top\n## A"}, + {"", "# Top", ""}, + } + for _, tc := range cases { + got := commonHeadingPrefix(tc.a, tc.b) + if got != tc.want { + t.Errorf("commonHeadingPrefix(%q, %q) = %q, want %q", tc.a, tc.b, got, tc.want) + } + } +} + // TestSplitByHeadings_NoBreadcrumbDuplication ensures the section's own // heading line does not appear twice in the chunk content (once as part of // the breadcrumb, once as the section's first line).