From df96f9fdf688f053754b9c71ba0c032225d51252 Mon Sep 17 00:00:00 2001 From: Eric Rakestraw Date: Tue, 12 May 2026 15:42:32 -0500 Subject: [PATCH] Rewrite and simplify the section chunking algorithm --- docs/architecture.md | 6 +- internal/core/chunking/sections.go | 324 ++++----- internal/core/chunking/sections_test.go | 857 +++++++++++------------- 3 files changed, 530 insertions(+), 657 deletions(-) diff --git a/docs/architecture.md b/docs/architecture.md index 5d52a71..41cf9a6 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -268,12 +268,14 @@ Chunking (`internal/core/chunking`) currently provides: - deterministic heuristic token estimation; - contiguous sectioning with section metadata; - max/min section token validation; -- optional `target_sections` handling with target-aware merge/split logic; +- optional `target_sections` override for section-count planning; - summary and detailed summary generation. Current behavior details: - if a single segment exceeds max tokens, it is emitted as its own section (not hard-failed); -- section balancing is deterministic but heuristic. +- default section count is planned from `ceil(total_tokens / max_section_tokens)`; +- section sizing targets `ceil(total_tokens / section_count)` with a deterministic forward pass; +- sections remain contiguous and ordered, and segments are never split. ## Implemented proposal/replacement infrastructure `internal/framework/proposals` provides deterministic proposal composition logic: diff --git a/internal/core/chunking/sections.go b/internal/core/chunking/sections.go index 92c0ca0..b1f4606 100644 --- a/internal/core/chunking/sections.go +++ b/internal/core/chunking/sections.go @@ -29,11 +29,12 @@ type ChunkingConfig struct { // MaxSectionTokens is the maximum allowed tokens per section MaxSectionTokens int - // MinSectionTokens is the minimum desired tokens per section (for balancing) + // MinSectionTokens is a validated soft lower-bound setting retained for + // configuration/reporting compatibility. MinSectionTokens int // TargetSections is an optional target number of sections - // If nil, chunking will optimize for max token bounds only + // If nil, section count is derived from total/max token budgeting. TargetSections *int } @@ -59,8 +60,18 @@ func NewChunkerWithEstimator(config ChunkingConfig, estimator TokenEstimator) *C } } -// ChunkTranscript divides a normalized transcript into contiguous token-bounded sections. -// Returns an error if the target section count is impossible under the constraints. +// ChunkTranscript divides a normalized transcript into contiguous token-bounded +// sections using a deterministic balanced forward pass. +// +// Behavior: +// - preserve segment order and never split segments; +// - estimate per-segment tokens once, then compute total; +// - derive desired section count from ceil(total/max_section_tokens), unless +// target_sections is explicitly set; +// - prefer section sizes near ceil(total/section_count) while never exceeding +// max_section_tokens unless a section consists of a single oversized segment. +// +// Returns an error if explicit target_sections is impossible under constraints. // The input transcript is never mutated. func (c *Chunker) ChunkTranscript(transcript *schema.Transcript) ([]Section, error) { if transcript == nil || len(transcript.Segments) == 0 { @@ -81,31 +92,51 @@ func (c *Chunker) ChunkTranscript(transcript *schema.Transcript) ([]Section, err c.config.MinSectionTokens, c.config.MaxSectionTokens) } - // Calculate token counts for each segment (deterministic) + // Calculate token counts for each segment (deterministic). segmentTokens := make([]int, len(transcript.Segments)) + totalTokens := 0 for i, seg := range transcript.Segments { segmentTokens[i] = c.estimator.EstimateTokens(seg.Text) + totalTokens += segmentTokens[i] } - // If target sections specified, use target-aware chunking + minPossibleSections := c.calculateMinPossibleSections(segmentTokens) + + var desiredSections int + useExplicitTarget := false if c.config.TargetSections != nil { - target := *c.config.TargetSections - - if target <= 0 { - return nil, fmt.Errorf("target_sections must be positive, got %d", target) + desiredSections = *c.config.TargetSections + useExplicitTarget = true + if desiredSections <= 0 { + return nil, fmt.Errorf("target_sections must be positive, got %d", desiredSections) } - - // Validate target is achievable - if err := c.validateTargetSections(target, segmentTokens); err != nil { + if err := c.validateTargetSections(desiredSections, segmentTokens); err != nil { return nil, err } - - // Use target-aware chunking strategy - return c.chunkWithTarget(transcript.Segments, segmentTokens, target) + } else { + desiredSections = ceilDiv(totalTokens, c.config.MaxSectionTokens) + if desiredSections < minPossibleSections { + desiredSections = minPossibleSections + } + if desiredSections < 1 { + desiredSections = 1 + } + if desiredSections > len(transcript.Segments) { + desiredSections = len(transcript.Segments) + } } - // No target: just respect max token bounds - return c.createSectionsUnderMax(transcript.Segments, segmentTokens) + targetTokensPerSection := ceilDiv(totalTokens, desiredSections) + if useExplicitTarget { + return c.buildSectionsWithExplicitTarget( + transcript.Segments, + segmentTokens, + desiredSections, + targetTokensPerSection, + ) + } + + return c.buildSectionsBalanced(transcript.Segments, segmentTokens, targetTokensPerSection), nil } // validateTargetSections checks if the target section count is achievable. @@ -131,154 +162,8 @@ func (c *Chunker) validateTargetSections(target int, segmentTokens []int) error return nil } -// chunkWithTarget creates sections aiming for a specific target count. -func (c *Chunker) chunkWithTarget(segments []schema.Segment, segmentTokens []int, target int) ([]Section, error) { - // Strategy: - // 1. First, create sections respecting max bounds (greedy packing) - // 2. If we have more sections than target, try to merge adjacent sections - // 3. If we have fewer sections than target, split sections where possible - - // Step 1: Create initial sections respecting max bounds - sections, err := c.createSectionsUnderMax(segments, segmentTokens) - if err != nil { - return nil, err - } - - // Step 2: Adjust to reach target - if len(sections) > target { - // Need to merge sections - sections = c.mergeSectionsToTarget(sections, target) - } else if len(sections) < target { - // Need to split sections - split the largest sections first - sections = c.splitSectionsToTarget(sections, target, segmentTokens) - } - - // Verify we achieved the target - if len(sections) != target { - return nil, fmt.Errorf( - "target_sections (%d) is impossible under current constraints (got %d sections)", - target, len(sections)) - } - - // Reindex sections - for i := range sections { - sections[i].Index = i - } - - return sections, nil -} - -// mergeSectionsToTarget merges adjacent sections to reduce count to target. -func (c *Chunker) mergeSectionsToTarget(sections []Section, target int) []Section { - for len(sections) > target { - // Find the best pair to merge (smallest combined section that stays under max) - bestIdx := -1 - bestTokens := -1 - - for i := 0; i < len(sections)-1; i++ { - combinedTokens := sections[i].EstimatedTokens + sections[i+1].EstimatedTokens - - // Can only merge if combined is under max - if combinedTokens <= c.config.MaxSectionTokens { - // Prefer smaller combined sections to keep sections balanced - if bestIdx == -1 || combinedTokens < bestTokens { - bestIdx = i - bestTokens = combinedTokens - } - } - } - - // No valid merge found - if bestIdx == -1 { - break - } - - // Merge sections[bestIdx] and sections[bestIdx+1] - merged := c.mergeTwoSections(sections[bestIdx], sections[bestIdx+1]) - - // Replace the two sections with the merged one - newSections := make([]Section, 0, len(sections)-1) - newSections = append(newSections, sections[:bestIdx]...) - newSections = append(newSections, merged) - if bestIdx+2 < len(sections) { - newSections = append(newSections, sections[bestIdx+2:]...) - } - sections = newSections - } - - return sections -} - -// splitSectionsToTarget splits sections to increase count to target. -func (c *Chunker) splitSectionsToTarget(sections []Section, target int, segmentTokens []int) []Section { - // We can only split sections that have multiple segments - // Split the largest sections first - - for len(sections) < target { - // Find the best section to split (largest section with multiple segments) - bestIdx := -1 - bestTokenCount := -1 - - for i, sec := range sections { - if len(sec.Segments) > 1 { - if sec.EstimatedTokens > bestTokenCount { - bestIdx = i - bestTokenCount = sec.EstimatedTokens - } - } - } - - // No splittable section found - if bestIdx == -1 { - break - } - - // Split the section into two - sec := sections[bestIdx] - splitPoint := len(sec.Segments) / 2 - - // Ensure each half has at least one segment - if splitPoint < 1 || splitPoint >= len(sec.Segments) { - break - } - - // Calculate token counts for each half - leftTokens := 0 - for i := 0; i < splitPoint; i++ { - leftTokens += c.estimator.EstimateTokens(sec.Segments[i].Text) - } - rightTokens := sec.EstimatedTokens - leftTokens - - // Create two new sections - leftSection := Section{ - StartSegmentID: sec.Segments[0].ID, - EndSegmentID: sec.Segments[splitPoint-1].ID, - EstimatedTokens: leftTokens, - Segments: sec.Segments[:splitPoint], - } - - rightSection := Section{ - StartSegmentID: sec.Segments[splitPoint].ID, - EndSegmentID: sec.Segments[len(sec.Segments)-1].ID, - EstimatedTokens: rightTokens, - Segments: sec.Segments[splitPoint:], - } - - // Replace old section with two new ones - newSections := make([]Section, 0, len(sections)+1) - newSections = append(newSections, sections[:bestIdx]...) - newSections = append(newSections, leftSection, rightSection) - if bestIdx+1 < len(sections) { - newSections = append(newSections, sections[bestIdx+1:]...) - } - sections = newSections - } - - return sections -} - -// createSectionsUnderMax creates initial sections respecting max token bounds. -func (c *Chunker) createSectionsUnderMax(segments []schema.Segment, segmentTokens []int) ([]Section, error) { +// buildSectionsBalanced creates sections with a deterministic single-pass policy. +func (c *Chunker) buildSectionsBalanced(segments []schema.Segment, segmentTokens []int, targetTokensPerSection int) []Section { var sections []Section var currentSegments []schema.Segment currentTokens := 0 @@ -286,38 +171,38 @@ func (c *Chunker) createSectionsUnderMax(segments []schema.Segment, segmentToken for i, seg := range segments { tokens := segmentTokens[i] - // If a single segment exceeds max tokens, it must be its own section - if tokens > c.config.MaxSectionTokens { - // Flush current section if any - if len(currentSegments) > 0 { - sections = append(sections, c.buildSection(len(sections), currentSegments, currentTokens)) - currentSegments = nil - currentTokens = 0 - } - // Add the oversized segment as its own section - sections = append(sections, c.buildSection(len(sections), []schema.Segment{seg}, tokens)) + // Empty section: always accept the next segment, including oversized. + if len(currentSegments) == 0 { + currentSegments = append(currentSegments, seg) + currentTokens = tokens continue } - // Check if adding this segment would exceed max - if currentTokens+tokens > c.config.MaxSectionTokens && len(currentSegments) > 0 { - // Finish current section + // If adding next segment would exceed max, close current section. + if currentTokens+tokens > c.config.MaxSectionTokens { sections = append(sections, c.buildSection(len(sections), currentSegments, currentTokens)) currentSegments = []schema.Segment{seg} currentTokens = tokens - } else { - // Add to current section + continue + } + + // Prefer staying near target tokens per section. + if targetTokensPerSection == 0 || currentTokens < targetTokensPerSection { currentSegments = append(currentSegments, seg) currentTokens += tokens + continue } + + sections = append(sections, c.buildSection(len(sections), currentSegments, currentTokens)) + currentSegments = []schema.Segment{seg} + currentTokens = tokens } - // Don't forget the last section if len(currentSegments) > 0 { sections = append(sections, c.buildSection(len(sections), currentSegments, currentTokens)) } - return sections, nil + return sections } // buildSection creates a Section from segments. @@ -360,17 +245,68 @@ func (c *Chunker) calculateMinPossibleSections(segmentTokens []int) int { return sections } -// mergeTwoSections combines two adjacent sections. -func (c *Chunker) mergeTwoSections(a, b Section) Section { - mergedSegments := make([]schema.Segment, 0, len(a.Segments)+len(b.Segments)) - mergedSegments = append(mergedSegments, a.Segments...) - mergedSegments = append(mergedSegments, b.Segments...) +// buildSectionsWithExplicitTarget builds exactly desiredSections when feasible. +func (c *Chunker) buildSectionsWithExplicitTarget( + segments []schema.Segment, + segmentTokens []int, + desiredSections int, + targetTokensPerSection int, +) ([]Section, error) { + n := len(segments) + cursor := 0 + sections := make([]Section, 0, desiredSections) - return Section{ - Index: a.Index, // Will be reassigned - StartSegmentID: a.StartSegmentID, - EndSegmentID: b.EndSegmentID, - EstimatedTokens: a.EstimatedTokens + b.EstimatedTokens, - Segments: mergedSegments, + for sectionIdx := 0; sectionIdx < desiredSections; sectionIdx++ { + if cursor >= n { + break + } + + remainingSectionsAfter := desiredSections - sectionIdx - 1 + currentSegments := []schema.Segment{segments[cursor]} + currentTokens := segmentTokens[cursor] + cursor++ + + for cursor < n { + remainingSegments := n - cursor + + // Reserve one segment per future section to avoid empty sections. + if remainingSegments == remainingSectionsAfter { + break + } + + nextTokens := segmentTokens[cursor] + if currentTokens+nextTokens > c.config.MaxSectionTokens { + break + } + if targetTokensPerSection == 0 || currentTokens < targetTokensPerSection { + currentSegments = append(currentSegments, segments[cursor]) + currentTokens += nextTokens + cursor++ + continue + } + break + } + + sections = append(sections, c.buildSection(len(sections), currentSegments, currentTokens)) } + + if cursor != n || len(sections) != desiredSections { + return nil, fmt.Errorf( + "target_sections (%d) is impossible under current constraints (got %d sections)", + desiredSections, + len(sections), + ) + } + + return sections, nil +} + +func ceilDiv(numerator int, denominator int) int { + if denominator <= 0 { + return 0 + } + if numerator <= 0 { + return 0 + } + return (numerator + denominator - 1) / denominator } diff --git a/internal/core/chunking/sections_test.go b/internal/core/chunking/sections_test.go index a0bb3f7..b3eb8ae 100644 --- a/internal/core/chunking/sections_test.go +++ b/internal/core/chunking/sections_test.go @@ -7,7 +7,21 @@ import ( "gitea.maximumdirect.net/eric/audita/internal/core/schema" ) -// Helper to create segments with sequential IDs +// mapTokenEstimator provides deterministic per-segment token counts for tests. +type mapTokenEstimator struct { + byText map[string]int +} + +func (e *mapTokenEstimator) EstimateTokens(text string) int { + if e.byText == nil { + return 0 + } + if tokens, ok := e.byText[text]; ok { + return tokens + } + return 0 +} + func makeSegments(texts []string) []schema.Segment { segments := make([]schema.Segment, len(texts)) for i, text := range texts { @@ -22,425 +36,431 @@ func makeSegments(texts []string) []schema.Segment { return segments } -// Helper to create a transcript func makeTranscript(segments []schema.Segment) *schema.Transcript { return &schema.Transcript{Segments: segments} } -func TestChunkEmptyTranscript(t *testing.T) { - config := ChunkingConfig{ - MaxSectionTokens: 100, - MinSectionTokens: 10, - } - chunker := NewChunker(config) +func intPtr(i int) *int { + return &i +} - // Test nil transcript - sections, err := chunker.ChunkTranscript(nil) - if err != nil { - t.Fatalf("ChunkTranscript(nil) error: %v", err) - } - if len(sections) != 0 { - t.Errorf("expected 0 sections for nil transcript, got %d", len(sections)) +func assertSegmentCoverageAndOrder(t *testing.T, input []schema.Segment, sections []Section) { + t.Helper() + + seen := make([]schema.Segment, 0, len(input)) + for _, sec := range sections { + seen = append(seen, sec.Segments...) } - // Test empty segments - emptyTranscript := makeTranscript([]schema.Segment{}) - sections, err = chunker.ChunkTranscript(emptyTranscript) - if err != nil { - t.Fatalf("ChunkTranscript(empty) error: %v", err) + if len(seen) != len(input) { + t.Fatalf("expected %d total segment occurrences, got %d", len(input), len(seen)) } - if len(sections) != 0 { - t.Errorf("expected 0 sections for empty transcript, got %d", len(sections)) + + for i := range input { + if seen[i].ID != input[i].ID { + t.Fatalf("segment order mismatch at index %d: got id=%d want id=%d", i, seen[i].ID, input[i].ID) + } } } -func TestChunkSingleSmallSection(t *testing.T) { - config := ChunkingConfig{ - MaxSectionTokens: 100, - MinSectionTokens: 10, +func assertSectionMetadataConsistent(t *testing.T, sections []Section) { + t.Helper() + + for i, sec := range sections { + if sec.Index != i { + t.Fatalf("section %d: expected index=%d got=%d", i, i, sec.Index) + } + if len(sec.Segments) == 0 { + t.Fatalf("section %d: section must not be empty", i) + } + if sec.StartSegmentID != sec.Segments[0].ID { + t.Fatalf("section %d: start_segment_id mismatch", i) + } + if sec.EndSegmentID != sec.Segments[len(sec.Segments)-1].ID { + t.Fatalf("section %d: end_segment_id mismatch", i) + } } - chunker := NewChunker(config) +} - segments := makeSegments([]string{ - "Hello world", - "This is a test", - }) - transcript := makeTranscript(segments) +func assertMaxBoundExceptSingletonOversized(t *testing.T, sections []Section, max int) { + t.Helper() - sections, err := chunker.ChunkTranscript(transcript) + for i, sec := range sections { + if sec.EstimatedTokens <= max { + continue + } + if len(sec.Segments) != 1 { + t.Fatalf("section %d exceeds max tokens (%d>%d) with %d segments", i, sec.EstimatedTokens, max, len(sec.Segments)) + } + } +} + +func imbalance(sections []Section) int { + if len(sections) == 0 { + return 0 + } + minTokens := sections[0].EstimatedTokens + maxTokens := sections[0].EstimatedTokens + for _, sec := range sections { + if sec.EstimatedTokens < minTokens { + minTokens = sec.EstimatedTokens + } + if sec.EstimatedTokens > maxTokens { + maxTokens = sec.EstimatedTokens + } + } + return maxTokens - minTokens +} + +func greedyMaxFillSections(segments []schema.Segment, tokens []int, max int) []Section { + sections := make([]Section, 0) + var current []schema.Segment + currentTokens := 0 + + for i, seg := range segments { + tok := tokens[i] + if len(current) == 0 { + current = append(current, seg) + currentTokens = tok + continue + } + if currentTokens+tok > max { + sections = append(sections, Section{ + Index: len(sections), + StartSegmentID: current[0].ID, + EndSegmentID: current[len(current)-1].ID, + EstimatedTokens: currentTokens, + Segments: append([]schema.Segment(nil), current...), + }) + current = []schema.Segment{seg} + currentTokens = tok + continue + } + current = append(current, seg) + currentTokens += tok + } + + if len(current) > 0 { + sections = append(sections, Section{ + Index: len(sections), + StartSegmentID: current[0].ID, + EndSegmentID: current[len(current)-1].ID, + EstimatedTokens: currentTokens, + Segments: append([]schema.Segment(nil), current...), + }) + } + + return sections +} + +func TestChunkEmptyTranscript(t *testing.T) { + chunker := NewChunker(ChunkingConfig{MaxSectionTokens: 100, MinSectionTokens: 10}) + + sections, err := chunker.ChunkTranscript(nil) if err != nil { - t.Fatalf("ChunkTranscript error: %v", err) + t.Fatalf("ChunkTranscript(nil): %v", err) + } + if len(sections) != 0 { + t.Fatalf("expected 0 sections for nil transcript, got %d", len(sections)) } + sections, err = chunker.ChunkTranscript(makeTranscript(nil)) + if err != nil { + t.Fatalf("ChunkTranscript(empty): %v", err) + } + if len(sections) != 0 { + t.Fatalf("expected 0 sections for empty transcript, got %d", len(sections)) + } +} + +func TestChunkSingleSegment(t *testing.T) { + segments := makeSegments([]string{"s1"}) + chunker := NewChunkerWithEstimator( + ChunkingConfig{MaxSectionTokens: 100, MinSectionTokens: 10}, + &mapTokenEstimator{byText: map[string]int{"s1": 7}}, + ) + + sections, err := chunker.ChunkTranscript(makeTranscript(segments)) + if err != nil { + t.Fatalf("ChunkTranscript: %v", err) + } if len(sections) != 1 { t.Fatalf("expected 1 section, got %d", len(sections)) } - - sec := sections[0] - if sec.Index != 0 { - t.Errorf("expected index 0, got %d", sec.Index) - } - if sec.StartSegmentID != 1 { - t.Errorf("expected start segment ID 1, got %d", sec.StartSegmentID) - } - if sec.EndSegmentID != 2 { - t.Errorf("expected end segment ID 2, got %d", sec.EndSegmentID) - } - if len(sec.Segments) != 2 { - t.Errorf("expected 2 segments, got %d", len(sec.Segments)) + if sections[0].EstimatedTokens != 7 { + t.Fatalf("expected estimated_tokens=7, got %d", sections[0].EstimatedTokens) } + assertSegmentCoverageAndOrder(t, segments, sections) + assertSectionMetadataConsistent(t, sections) } -func TestChunkMultipleSectionsDueToMaxTokens(t *testing.T) { - // Use constant estimator for predictable testing - estimator := &ConstTokenEstimator{Tokens: 30} +func TestChunkSingleOversizedSegment(t *testing.T) { + segments := makeSegments([]string{"big"}) + chunker := NewChunkerWithEstimator( + ChunkingConfig{MaxSectionTokens: 50, MinSectionTokens: 5}, + &mapTokenEstimator{byText: map[string]int{"big": 120}}, + ) - config := ChunkingConfig{ - MaxSectionTokens: 50, // Each section can hold at most 1 segment (30 tokens each) - MinSectionTokens: 10, - } - chunker := NewChunkerWithEstimator(config, estimator) - - segments := makeSegments([]string{ - "Segment one", - "Segment two", - "Segment three", - }) - transcript := makeTranscript(segments) - - sections, err := chunker.ChunkTranscript(transcript) + sections, err := chunker.ChunkTranscript(makeTranscript(segments)) if err != nil { - t.Fatalf("ChunkTranscript error: %v", err) + t.Fatalf("ChunkTranscript: %v", err) } + if len(sections) != 1 { + t.Fatalf("expected 1 section, got %d", len(sections)) + } + if sections[0].EstimatedTokens != 120 { + t.Fatalf("expected oversized singleton section, got %d", sections[0].EstimatedTokens) + } + assertMaxBoundExceptSingletonOversized(t, sections, 50) +} - // Should have 3 sections since each segment is 30 tokens and max is 50 +func TestChunkTotalBelowMaxSingleSection(t *testing.T) { + segments := makeSegments([]string{"a", "b", "c"}) + chunker := NewChunkerWithEstimator( + ChunkingConfig{MaxSectionTokens: 50, MinSectionTokens: 5}, + &ConstTokenEstimator{Tokens: 10}, + ) + + sections, err := chunker.ChunkTranscript(makeTranscript(segments)) + if err != nil { + t.Fatalf("ChunkTranscript: %v", err) + } + if len(sections) != 1 { + t.Fatalf("expected 1 section, got %d", len(sections)) + } + if sections[0].EstimatedTokens != 30 { + t.Fatalf("expected 30 section tokens, got %d", sections[0].EstimatedTokens) + } + assertSegmentCoverageAndOrder(t, segments, sections) +} + +func TestChunkTotalExactlyDivisibleByMax(t *testing.T) { + segments := makeSegments([]string{"a", "b", "c", "d"}) + chunker := NewChunkerWithEstimator( + ChunkingConfig{MaxSectionTokens: 10, MinSectionTokens: 1}, + &ConstTokenEstimator{Tokens: 5}, + ) + + sections, err := chunker.ChunkTranscript(makeTranscript(segments)) + if err != nil { + t.Fatalf("ChunkTranscript: %v", err) + } + if len(sections) != 2 { + t.Fatalf("expected 2 sections, got %d", len(sections)) + } + if sections[0].EstimatedTokens != 10 || sections[1].EstimatedTokens != 10 { + t.Fatalf("expected [10,10] tokens, got [%d,%d]", sections[0].EstimatedTokens, sections[1].EstimatedTokens) + } + assertSegmentCoverageAndOrder(t, segments, sections) + assertMaxBoundExceptSingletonOversized(t, sections, 10) +} + +func TestChunkTotalNotDivisibleByMax(t *testing.T) { + segments := makeSegments([]string{"a", "b", "c", "d", "e"}) + chunker := NewChunkerWithEstimator( + ChunkingConfig{MaxSectionTokens: 10, MinSectionTokens: 1}, + &ConstTokenEstimator{Tokens: 5}, + ) + + sections, err := chunker.ChunkTranscript(makeTranscript(segments)) + if err != nil { + t.Fatalf("ChunkTranscript: %v", err) + } if len(sections) != 3 { t.Fatalf("expected 3 sections, got %d", len(sections)) } + if sections[0].EstimatedTokens != 10 || sections[1].EstimatedTokens != 10 || sections[2].EstimatedTokens != 5 { + t.Fatalf("expected [10,10,5] tokens, got [%d,%d,%d]", sections[0].EstimatedTokens, sections[1].EstimatedTokens, sections[2].EstimatedTokens) + } + assertSegmentCoverageAndOrder(t, segments, sections) + assertMaxBoundExceptSingletonOversized(t, sections, 10) +} - // Verify section contents - for i, sec := range sections { - if sec.Index != i { - t.Errorf("section %d: expected index %d, got %d", i, i, sec.Index) - } - if len(sec.Segments) != 1 { - t.Errorf("section %d: expected 1 segment, got %d", i, len(sec.Segments)) - } - if sec.StartSegmentID != i+1 { - t.Errorf("section %d: expected start segment ID %d, got %d", i, i+1, sec.StartSegmentID) - } - if sec.EndSegmentID != i+1 { - t.Errorf("section %d: expected end segment ID %d, got %d", i, i+1, sec.EndSegmentID) - } +func TestChunkTargetSectionsPrecedenceAndSuccess(t *testing.T) { + segments := makeSegments([]string{"1", "2", "3", "4", "5", "6", "7", "8", "9", "10"}) + target := 3 + chunker := NewChunkerWithEstimator( + ChunkingConfig{MaxSectionTokens: 200, MinSectionTokens: 1, TargetSections: &target}, + &ConstTokenEstimator{Tokens: 10}, + ) + + sections, err := chunker.ChunkTranscript(makeTranscript(segments)) + if err != nil { + t.Fatalf("ChunkTranscript: %v", err) + } + if len(sections) != target { + t.Fatalf("expected %d sections from explicit target, got %d", target, len(sections)) + } + assertSegmentCoverageAndOrder(t, segments, sections) + assertSectionMetadataConsistent(t, sections) +} + +func TestChunkTargetSectionsImpossibleTooMany(t *testing.T) { + segments := makeSegments([]string{"1", "2", "3", "4", "5"}) + target := 10 + chunker := NewChunker(ChunkingConfig{MaxSectionTokens: 100, MinSectionTokens: 1, TargetSections: &target}) + + _, err := chunker.ChunkTranscript(makeTranscript(segments)) + if err == nil { + t.Fatal("expected error for impossible target_sections") + } + if !strings.Contains(err.Error(), "cannot have more sections than segments") { + t.Fatalf("unexpected error: %v", err) } } -func TestChunkPreservesSegmentOrder(t *testing.T) { - // Use constant estimator - estimator := &ConstTokenEstimator{Tokens: 10} +func TestChunkTargetSectionsImpossibleTooFew(t *testing.T) { + segments := makeSegments([]string{"1", "2", "3"}) + target := 1 + chunker := NewChunkerWithEstimator( + ChunkingConfig{MaxSectionTokens: 50, MinSectionTokens: 1, TargetSections: &target}, + &ConstTokenEstimator{Tokens: 30}, + ) - config := ChunkingConfig{ - MaxSectionTokens: 25, // Can fit 2 segments (20 tokens) but not 3 - MinSectionTokens: 5, + _, err := chunker.ChunkTranscript(makeTranscript(segments)) + if err == nil { + t.Fatal("expected error for impossible target_sections") } - chunker := NewChunkerWithEstimator(config, estimator) + if !strings.Contains(err.Error(), "need at least 3 sections") { + t.Fatalf("unexpected error: %v", err) + } +} - // Create segments with identifiable text - segments := makeSegments([]string{ - "First segment", - "Second segment", - "Third segment", - "Fourth segment", - "Fifth segment", - }) - transcript := makeTranscript(segments) +func TestChunkManySmallSegmentsBalanced(t *testing.T) { + texts := []string{"1", "2", "3", "4", "5", "6", "7", "8", "9", "10", "11"} + segments := makeSegments(texts) + chunker := NewChunkerWithEstimator( + ChunkingConfig{MaxSectionTokens: 50, MinSectionTokens: 1}, + &ConstTokenEstimator{Tokens: 10}, + ) - sections, err := chunker.ChunkTranscript(transcript) + sections, err := chunker.ChunkTranscript(makeTranscript(segments)) if err != nil { - t.Fatalf("ChunkTranscript error: %v", err) + t.Fatalf("ChunkTranscript: %v", err) } - - // Collect all segments in order - var allSegments []schema.Segment - for _, sec := range sections { - allSegments = append(allSegments, sec.Segments...) + if len(sections) != 3 { + t.Fatalf("expected 3 sections, got %d", len(sections)) } - - // Verify order is preserved - if len(allSegments) != 5 { - t.Fatalf("expected 5 total segments, got %d", len(allSegments)) + if sections[0].EstimatedTokens != 40 || sections[1].EstimatedTokens != 40 || sections[2].EstimatedTokens != 30 { + t.Fatalf("expected [40,40,30], got [%d,%d,%d]", sections[0].EstimatedTokens, sections[1].EstimatedTokens, sections[2].EstimatedTokens) } + assertSegmentCoverageAndOrder(t, segments, sections) + assertMaxBoundExceptSingletonOversized(t, sections, 50) +} - expectedTexts := []string{ - "First segment", - "Second segment", - "Third segment", - "Fourth segment", - "Fifth segment", +func TestChunkMixedLargeAndSmallSegments(t *testing.T) { + segments := makeSegments([]string{"big1", "s1", "s2", "s3", "big2", "s4"}) + estimator := &mapTokenEstimator{byText: map[string]int{ + "big1": 120, + "s1": 10, + "s2": 10, + "s3": 10, + "big2": 120, + "s4": 10, + }} + chunker := NewChunkerWithEstimator(ChunkingConfig{MaxSectionTokens: 100, MinSectionTokens: 1}, estimator) + + sections, err := chunker.ChunkTranscript(makeTranscript(segments)) + if err != nil { + t.Fatalf("ChunkTranscript: %v", err) } + if len(sections) != 4 { + t.Fatalf("expected 4 sections, got %d", len(sections)) + } + if len(sections[0].Segments) != 1 || sections[0].Segments[0].Text != "big1" { + t.Fatalf("expected first oversized segment in singleton section, got %+v", sections[0].Segments) + } + if len(sections[2].Segments) != 1 || sections[2].Segments[0].Text != "big2" { + t.Fatalf("expected second oversized segment in singleton section, got %+v", sections[2].Segments) + } + assertSegmentCoverageAndOrder(t, segments, sections) + assertMaxBoundExceptSingletonOversized(t, sections, 100) +} - for i, seg := range allSegments { - if seg.Text != expectedTexts[i] { - t.Errorf("segment %d: expected text %q, got %q", i, expectedTexts[i], seg.Text) +func TestChunkDeterministicOrdering(t *testing.T) { + segments := makeSegments([]string{"a", "b", "c", "d", "e", "f"}) + chunkerCfg := ChunkingConfig{MaxSectionTokens: 15, MinSectionTokens: 1} + estimator := &ConstTokenEstimator{Tokens: 5} + + var first []Section + for i := 0; i < 5; i++ { + chunker := NewChunkerWithEstimator(chunkerCfg, estimator) + sections, err := chunker.ChunkTranscript(makeTranscript(segments)) + if err != nil { + t.Fatalf("iteration %d: %v", i, err) } - if seg.ID != i+1 { - t.Errorf("segment %d: expected ID %d, got %d", i, i+1, seg.ID) + if i == 0 { + first = sections + continue + } + if len(sections) != len(first) { + t.Fatalf("iteration %d: section count mismatch (%d vs %d)", i, len(sections), len(first)) + } + for j := range sections { + if sections[j].Index != first[j].Index || + sections[j].StartSegmentID != first[j].StartSegmentID || + sections[j].EndSegmentID != first[j].EndSegmentID || + sections[j].EstimatedTokens != first[j].EstimatedTokens || + len(sections[j].Segments) != len(first[j].Segments) { + t.Fatalf("iteration %d section %d mismatch", i, j) + } } } } func TestChunkNoMutationOfInput(t *testing.T) { - config := ChunkingConfig{ - MaxSectionTokens: 50, - MinSectionTokens: 10, - } - chunker := NewChunker(config) - - originalSegments := makeSegments([]string{ - "Original text one", - "Original text two", - }) - transcript := makeTranscript(originalSegments) - - // Store original state - originalTexts := make([]string, len(originalSegments)) - for i, seg := range transcript.Segments { - originalTexts[i] = seg.Text + segments := makeSegments([]string{"original one", "original two"}) + transcript := makeTranscript(segments) + original := make([]string, len(transcript.Segments)) + for i := range transcript.Segments { + original[i] = transcript.Segments[i].Text } - // Chunk the transcript - _, err := chunker.ChunkTranscript(transcript) - if err != nil { - t.Fatalf("ChunkTranscript error: %v", err) + chunker := NewChunker(ChunkingConfig{MaxSectionTokens: 50, MinSectionTokens: 1}) + if _, err := chunker.ChunkTranscript(transcript); err != nil { + t.Fatalf("ChunkTranscript: %v", err) } - - // Verify original transcript is unchanged - for i, seg := range transcript.Segments { - if seg.Text != originalTexts[i] { - t.Errorf("segment %d was mutated: expected %q, got %q", i, originalTexts[i], seg.Text) + for i := range transcript.Segments { + if transcript.Segments[i].Text != original[i] { + t.Fatalf("segment %d mutated", i) } } } -func TestChunkTargetSectionsSuccess(t *testing.T) { - // Use constant estimator for predictable testing - estimator := &ConstTokenEstimator{Tokens: 10} - - targetSections := 3 - config := ChunkingConfig{ - MaxSectionTokens: 50, // Can fit up to 5 segments per section - MinSectionTokens: 5, - TargetSections: &targetSections, - } - chunker := NewChunkerWithEstimator(config, estimator) - - segments := makeSegments([]string{ - "One", "Two", "Three", "Four", "Five", - "Six", "Seven", "Eight", "Nine", "Ten", - }) - transcript := makeTranscript(segments) - - sections, err := chunker.ChunkTranscript(transcript) - if err != nil { - t.Fatalf("ChunkTranscript error: %v", err) - } - - if len(sections) != targetSections { - t.Fatalf("expected %d sections, got %d", targetSections, len(sections)) - } - - // Verify all segments are included - totalSegments := 0 - for _, sec := range sections { - totalSegments += len(sec.Segments) - } - if totalSegments != 10 { - t.Errorf("expected 10 total segments, got %d", totalSegments) - } -} - -func TestChunkTargetSectionsImpossibleTooFew(t *testing.T) { - estimator := &ConstTokenEstimator{Tokens: 30} - - targetSections := 1 // Impossible: need at least 3 sections for 3 segments at 30 tokens each with max 50 - config := ChunkingConfig{ - MaxSectionTokens: 50, - MinSectionTokens: 10, - TargetSections: &targetSections, - } - chunker := NewChunkerWithEstimator(config, estimator) - - segments := makeSegments([]string{ - "Segment one", - "Segment two", - "Segment three", - }) - transcript := makeTranscript(segments) - - _, err := chunker.ChunkTranscript(transcript) - if err == nil { - t.Fatal("expected error for impossible target_sections, got nil") - } - - expectedMsg := "target_sections (1) is impossible: need at least 3 sections" - if !strings.Contains(err.Error(), expectedMsg) { - t.Errorf("expected error containing %q, got %q", expectedMsg, err.Error()) - } -} - -func TestChunkTargetSectionsImpossibleTooMany(t *testing.T) { - targetSections := 10 // Impossible: cannot have more sections than segments (5) - config := ChunkingConfig{ - MaxSectionTokens: 100, - MinSectionTokens: 10, - TargetSections: &targetSections, - } - chunker := NewChunker(config) - - segments := makeSegments([]string{ - "One", "Two", "Three", "Four", "Five", - }) - transcript := makeTranscript(segments) - - _, err := chunker.ChunkTranscript(transcript) - if err == nil { - t.Fatal("expected error for impossible target_sections, got nil") - } - - expectedMsg := "target_sections (10) is impossible: cannot have more sections than segments" - if !strings.Contains(err.Error(), expectedMsg) { - t.Errorf("expected error containing %q, got %q", expectedMsg, err.Error()) - } -} - -func TestChunkMinTokenBalancing(t *testing.T) { - // Create segments with varying sizes - texts := []string{ - "Short", - "Also short", - "This one is a bit longer", - "Tiny", - "Another medium length segment here", - } - segments := makeSegments(texts) - transcript := makeTranscript(segments) - - config := ChunkingConfig{ - MaxSectionTokens: 50, - MinSectionTokens: 15, // Try to keep sections above this - } - chunker := NewChunker(config) - - sections, err := chunker.ChunkTranscript(transcript) - if err != nil { - t.Fatalf("ChunkTranscript error: %v", err) - } - - // Verify all segments are included - totalSegments := 0 - for _, sec := range sections { - totalSegments += len(sec.Segments) - } - if totalSegments != 5 { - t.Errorf("expected 5 total segments, got %d", totalSegments) - } -} - -func TestChunkWithOversizedSegment(t *testing.T) { - // Create an estimator that returns a very large value for one specific text - segments := makeSegments([]string{ - "Small", - strings.Repeat("very large ", 100), // This will have many tokens - "Also small", - }) - transcript := makeTranscript(segments) - - config := ChunkingConfig{ - MaxSectionTokens: 50, - MinSectionTokens: 5, - } - chunker := NewChunker(config) - - sections, err := chunker.ChunkTranscript(transcript) - if err != nil { - t.Fatalf("ChunkTranscript error: %v", err) - } - - // The oversized segment should be in its own section - foundLargeInOwnSection := false - for _, sec := range sections { - if len(sec.Segments) == 1 && strings.HasPrefix(sec.Segments[0].Text, "very large") { - foundLargeInOwnSection = true - break - } - } - - if !foundLargeInOwnSection { - t.Error("expected oversized segment to be in its own section") - } -} - func TestChunkConfigValidation(t *testing.T) { tests := []struct { name string config ChunkingConfig - wantErr bool errContains string }{ { - name: "zero max tokens", - config: ChunkingConfig{ - MaxSectionTokens: 0, - MinSectionTokens: 10, - }, - wantErr: true, + name: "zero max tokens", + config: ChunkingConfig{MaxSectionTokens: 0, MinSectionTokens: 1}, errContains: "max_section_tokens must be positive", }, { - name: "negative max tokens", - config: ChunkingConfig{ - MaxSectionTokens: -1, - MinSectionTokens: 10, - }, - wantErr: true, + name: "negative max tokens", + config: ChunkingConfig{MaxSectionTokens: -1, MinSectionTokens: 1}, errContains: "max_section_tokens must be positive", }, { - name: "negative min tokens", - config: ChunkingConfig{ - MaxSectionTokens: 100, - MinSectionTokens: -1, - }, - wantErr: true, + name: "negative min tokens", + config: ChunkingConfig{MaxSectionTokens: 10, MinSectionTokens: -1}, errContains: "min_section_tokens must be non-negative", }, { - name: "min exceeds max", - config: ChunkingConfig{ - MaxSectionTokens: 50, - MinSectionTokens: 100, - }, - wantErr: true, - errContains: "min_section_tokens (100) cannot exceed max_section_tokens (50)", + name: "min exceeds max", + config: ChunkingConfig{MaxSectionTokens: 10, MinSectionTokens: 11}, + errContains: "min_section_tokens (11) cannot exceed max_section_tokens (10)", }, { - name: "zero target sections", - config: ChunkingConfig{ - MaxSectionTokens: 100, - MinSectionTokens: 10, - TargetSections: intPtr(0), - }, - wantErr: true, + name: "zero target sections", + config: ChunkingConfig{MaxSectionTokens: 10, MinSectionTokens: 1, TargetSections: intPtr(0)}, errContains: "target_sections must be positive", }, { - name: "negative target sections", - config: ChunkingConfig{ - MaxSectionTokens: 100, - MinSectionTokens: 10, - TargetSections: intPtr(-1), - }, - wantErr: true, + name: "negative target sections", + config: ChunkingConfig{MaxSectionTokens: 10, MinSectionTokens: 1, TargetSections: intPtr(-1)}, errContains: "target_sections must be positive", }, } @@ -448,136 +468,51 @@ func TestChunkConfigValidation(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { chunker := NewChunker(tt.config) - transcript := makeTranscript(makeSegments([]string{"test"})) - - _, err := chunker.ChunkTranscript(transcript) - if tt.wantErr { - if err == nil { - t.Errorf("expected error containing %q, got nil", tt.errContains) - } else if !strings.Contains(err.Error(), tt.errContains) { - t.Errorf("expected error containing %q, got %q", tt.errContains, err.Error()) - } - } else { - if err != nil { - t.Errorf("unexpected error: %v", err) - } + _, err := chunker.ChunkTranscript(makeTranscript(makeSegments([]string{"x"}))) + if err == nil { + t.Fatalf("expected error containing %q", tt.errContains) + } + if !strings.Contains(err.Error(), tt.errContains) { + t.Fatalf("expected error containing %q, got %q", tt.errContains, err.Error()) } }) } } -func TestChunkDeterministicOrdering(t *testing.T) { - segments := makeSegments([]string{ - "First", - "Second", - "Third", - "Fourth", - "Fifth", - }) - transcript := makeTranscript(segments) - - config := ChunkingConfig{ - MaxSectionTokens: 30, - MinSectionTokens: 10, +func TestChunkBalancedAlgorithmBeatsGreedyMaxFillOnUnevenTranscript(t *testing.T) { + segments := makeSegments([]string{"s1", "s2", "s3", "s4", "s5", "s6"}) + tokenMap := map[string]int{ + "s1": 50, + "s2": 10, + "s3": 10, + "s4": 10, + "s5": 10, + "s6": 10, } + estimator := &mapTokenEstimator{byText: tokenMap} + chunker := NewChunkerWithEstimator(ChunkingConfig{MaxSectionTokens: 80, MinSectionTokens: 1}, estimator) - // Run chunking multiple times and verify deterministic results - var firstResult []Section - for i := 0; i < 5; i++ { - chunker := NewChunker(config) - sections, err := chunker.ChunkTranscript(transcript) - if err != nil { - t.Fatalf("ChunkTranscript error at iteration %d: %v", i, err) - } - - if i == 0 { - firstResult = sections - } else { - // Compare with first result - if len(sections) != len(firstResult) { - t.Errorf("iteration %d: section count mismatch: got %d, want %d", i, len(sections), len(firstResult)) - continue - } - for j, sec := range sections { - if sec.Index != firstResult[j].Index { - t.Errorf("iteration %d, section %d: index mismatch", i, j) - } - if sec.StartSegmentID != firstResult[j].StartSegmentID { - t.Errorf("iteration %d, section %d: start_segment_id mismatch", i, j) - } - if sec.EndSegmentID != firstResult[j].EndSegmentID { - t.Errorf("iteration %d, section %d: end_segment_id mismatch", i, j) - } - if sec.EstimatedTokens != firstResult[j].EstimatedTokens { - t.Errorf("iteration %d, section %d: estimated_tokens mismatch", i, j) - } - if len(sec.Segments) != len(firstResult[j].Segments) { - t.Errorf("iteration %d, section %d: segment count mismatch", i, j) - } - } - } - } -} - -func TestSectionMetadata(t *testing.T) { - estimator := &ConstTokenEstimator{Tokens: 5} - - config := ChunkingConfig{ - MaxSectionTokens: 15, // Can fit 3 segments of 5 tokens each - MinSectionTokens: 5, - } - chunker := NewChunkerWithEstimator(config, estimator) - - segments := makeSegments([]string{ - "One", "Two", "Three", "Four", "Five", - }) - transcript := makeTranscript(segments) - - sections, err := chunker.ChunkTranscript(transcript) + balancedSections, err := chunker.ChunkTranscript(makeTranscript(segments)) if err != nil { - t.Fatalf("ChunkTranscript error: %v", err) + t.Fatalf("ChunkTranscript: %v", err) } - // Should have 2 sections: [1,2,3] and [4,5] - if len(sections) != 2 { - t.Fatalf("expected 2 sections, got %d", len(sections)) + tokens := make([]int, 0, len(segments)) + for _, seg := range segments { + tokens = append(tokens, tokenMap[seg.Text]) + } + greedySections := greedyMaxFillSections(segments, tokens, 80) + + balancedImbalance := imbalance(balancedSections) + greedyImbalance := imbalance(greedySections) + if balancedImbalance >= greedyImbalance { + t.Fatalf( + "expected balanced chunking to improve over greedy max-fill; balanced=%d greedy=%d", + balancedImbalance, + greedyImbalance, + ) } - // First section - if sections[0].Index != 0 { - t.Errorf("section 0: expected index 0, got %d", sections[0].Index) - } - if sections[0].StartSegmentID != 1 { - t.Errorf("section 0: expected start segment ID 1, got %d", sections[0].StartSegmentID) - } - if sections[0].EndSegmentID != 3 { - t.Errorf("section 0: expected end segment ID 3, got %d", sections[0].EndSegmentID) - } - if sections[0].EstimatedTokens != 15 { - t.Errorf("section 0: expected 15 tokens, got %d", sections[0].EstimatedTokens) - } - if len(sections[0].Segments) != 3 { - t.Errorf("section 0: expected 3 segments, got %d", len(sections[0].Segments)) - } - - // Second section - if sections[1].Index != 1 { - t.Errorf("section 1: expected index 1, got %d", sections[1].Index) - } - if sections[1].StartSegmentID != 4 { - t.Errorf("section 1: expected start segment ID 4, got %d", sections[1].StartSegmentID) - } - if sections[1].EndSegmentID != 5 { - t.Errorf("section 1: expected end segment ID 5, got %d", sections[1].EndSegmentID) - } - if sections[1].EstimatedTokens != 10 { - t.Errorf("section 1: expected 10 tokens, got %d", sections[1].EstimatedTokens) - } - if len(sections[1].Segments) != 2 { - t.Errorf("section 1: expected 2 segments, got %d", len(sections[1].Segments)) - } -} - -func intPtr(i int) *int { - return &i + assertSegmentCoverageAndOrder(t, segments, balancedSections) + assertMaxBoundExceptSingletonOversized(t, balancedSections, 80) }