diff --git a/internal/core/chunking/sections.go b/internal/core/chunking/sections.go new file mode 100644 index 0000000..92c0ca0 --- /dev/null +++ b/internal/core/chunking/sections.go @@ -0,0 +1,376 @@ +package chunking + +import ( + "fmt" + + "gitea.maximumdirect.net/eric/audita/internal/core/schema" +) + +// Section represents a contiguous chunk of transcript segments with metadata. +type Section struct { + // Index is the 0-based section index within the chunked transcript + Index int `json:"section_index"` + + // StartSegmentID is the ID of the first segment in this section + StartSegmentID int `json:"start_segment_id"` + + // EndSegmentID is the ID of the last segment in this section + EndSegmentID int `json:"end_segment_id"` + + // EstimatedTokens is the approximate token count for this section + EstimatedTokens int `json:"estimated_tokens"` + + // Segments contains the segments in this section, in order + Segments []schema.Segment `json:"segments"` +} + +// ChunkingConfig holds configuration for transcript chunking. +type ChunkingConfig struct { + // MaxSectionTokens is the maximum allowed tokens per section + MaxSectionTokens int + + // MinSectionTokens is the minimum desired tokens per section (for balancing) + MinSectionTokens int + + // TargetSections is an optional target number of sections + // If nil, chunking will optimize for max token bounds only + TargetSections *int +} + +// Chunker performs deterministic chunking of normalized transcript segments. +type Chunker struct { + config ChunkingConfig + estimator TokenEstimator +} + +// NewChunker creates a new chunker with the given configuration. +func NewChunker(config ChunkingConfig) *Chunker { + return &Chunker{ + config: config, + estimator: NewSimpleTokenEstimator(), + } +} + +// NewChunkerWithEstimator creates a new chunker with a custom estimator. +func NewChunkerWithEstimator(config ChunkingConfig, estimator TokenEstimator) *Chunker { + return &Chunker{ + config: config, + estimator: estimator, + } +} + +// ChunkTranscript divides a normalized transcript into contiguous token-bounded sections. +// Returns an error if the target section count is impossible under the constraints. +// The input transcript is never mutated. +func (c *Chunker) ChunkTranscript(transcript *schema.Transcript) ([]Section, error) { + if transcript == nil || len(transcript.Segments) == 0 { + return []Section{}, nil + } + + // Validate configuration + if c.config.MaxSectionTokens <= 0 { + return nil, fmt.Errorf("max_section_tokens must be positive, got %d", c.config.MaxSectionTokens) + } + + if c.config.MinSectionTokens < 0 { + return nil, fmt.Errorf("min_section_tokens must be non-negative, got %d", c.config.MinSectionTokens) + } + + if c.config.MinSectionTokens > c.config.MaxSectionTokens { + return nil, fmt.Errorf("min_section_tokens (%d) cannot exceed max_section_tokens (%d)", + c.config.MinSectionTokens, c.config.MaxSectionTokens) + } + + // Calculate token counts for each segment (deterministic) + segmentTokens := make([]int, len(transcript.Segments)) + for i, seg := range transcript.Segments { + segmentTokens[i] = c.estimator.EstimateTokens(seg.Text) + } + + // If target sections specified, use target-aware chunking + if c.config.TargetSections != nil { + target := *c.config.TargetSections + + if target <= 0 { + return nil, fmt.Errorf("target_sections must be positive, got %d", target) + } + + // Validate target is achievable + if err := c.validateTargetSections(target, segmentTokens); err != nil { + return nil, err + } + + // Use target-aware chunking strategy + return c.chunkWithTarget(transcript.Segments, segmentTokens, target) + } + + // No target: just respect max token bounds + return c.createSectionsUnderMax(transcript.Segments, segmentTokens) +} + +// validateTargetSections checks if the target section count is achievable. +func (c *Chunker) validateTargetSections(target int, segmentTokens []int) error { + // Maximum possible sections: limited by segment count + maxPossibleSections := len(segmentTokens) + + if target > maxPossibleSections { + return fmt.Errorf( + "target_sections (%d) is impossible: cannot have more sections than segments (%d)", + target, maxPossibleSections) + } + + // Minimum possible sections: each segment must fit within max bounds + minPossibleSections := c.calculateMinPossibleSections(segmentTokens) + + if target < minPossibleSections { + return fmt.Errorf( + "target_sections (%d) is impossible: need at least %d sections to respect max_section_tokens (%d)", + target, minPossibleSections, c.config.MaxSectionTokens) + } + + 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) { + var sections []Section + var currentSegments []schema.Segment + currentTokens := 0 + + 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)) + continue + } + + // Check if adding this segment would exceed max + if currentTokens+tokens > c.config.MaxSectionTokens && len(currentSegments) > 0 { + // Finish current section + sections = append(sections, c.buildSection(len(sections), currentSegments, currentTokens)) + currentSegments = []schema.Segment{seg} + currentTokens = tokens + } else { + // Add to current section + currentSegments = append(currentSegments, 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 +} + +// buildSection creates a Section from segments. +func (c *Chunker) buildSection(index int, segments []schema.Segment, tokens int) Section { + return Section{ + Index: index, + StartSegmentID: segments[0].ID, + EndSegmentID: segments[len(segments)-1].ID, + EstimatedTokens: tokens, + Segments: segments, + } +} + +// calculateMinPossibleSections calculates the minimum number of sections needed +// to ensure no section exceeds max tokens. +func (c *Chunker) calculateMinPossibleSections(segmentTokens []int) int { + sections := 0 + currentTokens := 0 + + for _, tokens := range segmentTokens { + if tokens > c.config.MaxSectionTokens { + // Each oversized segment needs its own section + if currentTokens > 0 { + sections++ + currentTokens = 0 + } + sections++ + } else if currentTokens+tokens > c.config.MaxSectionTokens { + sections++ + currentTokens = tokens + } else { + currentTokens += tokens + } + } + + if currentTokens > 0 { + sections++ + } + + 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...) + + return Section{ + Index: a.Index, // Will be reassigned + StartSegmentID: a.StartSegmentID, + EndSegmentID: b.EndSegmentID, + EstimatedTokens: a.EstimatedTokens + b.EstimatedTokens, + Segments: mergedSegments, + } +} diff --git a/internal/core/chunking/sections_test.go b/internal/core/chunking/sections_test.go new file mode 100644 index 0000000..a0bb3f7 --- /dev/null +++ b/internal/core/chunking/sections_test.go @@ -0,0 +1,583 @@ +package chunking + +import ( + "strings" + "testing" + + "gitea.maximumdirect.net/eric/audita/internal/core/schema" +) + +// Helper to create segments with sequential IDs +func makeSegments(texts []string) []schema.Segment { + segments := make([]schema.Segment, len(texts)) + for i, text := range texts { + segments[i] = schema.Segment{ + ID: i + 1, + Speaker: "DM", + Start: float64(i * 10), + End: float64(i*10 + 5), + Text: text, + } + } + 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) + + // 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)) + } + + // Test empty segments + emptyTranscript := makeTranscript([]schema.Segment{}) + sections, err = chunker.ChunkTranscript(emptyTranscript) + if err != nil { + t.Fatalf("ChunkTranscript(empty) error: %v", err) + } + if len(sections) != 0 { + t.Errorf("expected 0 sections for empty transcript, got %d", len(sections)) + } +} + +func TestChunkSingleSmallSection(t *testing.T) { + config := ChunkingConfig{ + MaxSectionTokens: 100, + MinSectionTokens: 10, + } + chunker := NewChunker(config) + + segments := makeSegments([]string{ + "Hello world", + "This is a test", + }) + transcript := makeTranscript(segments) + + sections, err := chunker.ChunkTranscript(transcript) + if err != nil { + t.Fatalf("ChunkTranscript error: %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)) + } +} + +func TestChunkMultipleSectionsDueToMaxTokens(t *testing.T) { + // Use constant estimator for predictable testing + estimator := &ConstTokenEstimator{Tokens: 30} + + 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) + if err != nil { + t.Fatalf("ChunkTranscript error: %v", err) + } + + // Should have 3 sections since each segment is 30 tokens and max is 50 + if len(sections) != 3 { + t.Fatalf("expected 3 sections, got %d", len(sections)) + } + + // 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 TestChunkPreservesSegmentOrder(t *testing.T) { + // Use constant estimator + estimator := &ConstTokenEstimator{Tokens: 10} + + config := ChunkingConfig{ + MaxSectionTokens: 25, // Can fit 2 segments (20 tokens) but not 3 + MinSectionTokens: 5, + } + chunker := NewChunkerWithEstimator(config, estimator) + + // Create segments with identifiable text + segments := makeSegments([]string{ + "First segment", + "Second segment", + "Third segment", + "Fourth segment", + "Fifth segment", + }) + transcript := makeTranscript(segments) + + sections, err := chunker.ChunkTranscript(transcript) + if err != nil { + t.Fatalf("ChunkTranscript error: %v", err) + } + + // Collect all segments in order + var allSegments []schema.Segment + for _, sec := range sections { + allSegments = append(allSegments, sec.Segments...) + } + + // Verify order is preserved + if len(allSegments) != 5 { + t.Fatalf("expected 5 total segments, got %d", len(allSegments)) + } + + expectedTexts := []string{ + "First segment", + "Second segment", + "Third segment", + "Fourth segment", + "Fifth segment", + } + + for i, seg := range allSegments { + if seg.Text != expectedTexts[i] { + t.Errorf("segment %d: expected text %q, got %q", i, expectedTexts[i], seg.Text) + } + if seg.ID != i+1 { + t.Errorf("segment %d: expected ID %d, got %d", i, i+1, seg.ID) + } + } +} + +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 + } + + // Chunk the transcript + _, err := chunker.ChunkTranscript(transcript) + if err != nil { + t.Fatalf("ChunkTranscript error: %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) + } + } +} + +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, + errContains: "max_section_tokens must be positive", + }, + { + name: "negative max tokens", + config: ChunkingConfig{ + MaxSectionTokens: -1, + MinSectionTokens: 10, + }, + wantErr: true, + errContains: "max_section_tokens must be positive", + }, + { + name: "negative min tokens", + config: ChunkingConfig{ + MaxSectionTokens: 100, + MinSectionTokens: -1, + }, + wantErr: true, + 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: "zero target sections", + config: ChunkingConfig{ + MaxSectionTokens: 100, + MinSectionTokens: 10, + TargetSections: intPtr(0), + }, + wantErr: true, + errContains: "target_sections must be positive", + }, + { + name: "negative target sections", + config: ChunkingConfig{ + MaxSectionTokens: 100, + MinSectionTokens: 10, + TargetSections: intPtr(-1), + }, + wantErr: true, + errContains: "target_sections must be positive", + }, + } + + 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) + } + } + }) + } +} + +func TestChunkDeterministicOrdering(t *testing.T) { + segments := makeSegments([]string{ + "First", + "Second", + "Third", + "Fourth", + "Fifth", + }) + transcript := makeTranscript(segments) + + config := ChunkingConfig{ + MaxSectionTokens: 30, + MinSectionTokens: 10, + } + + // 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) + if err != nil { + t.Fatalf("ChunkTranscript error: %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)) + } + + // 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 +} diff --git a/internal/core/chunking/tokens.go b/internal/core/chunking/tokens.go new file mode 100644 index 0000000..4666821 --- /dev/null +++ b/internal/core/chunking/tokens.go @@ -0,0 +1,55 @@ +package chunking + +import ( + "strings" + "unicode" +) + +// TokenEstimator provides a deterministic token estimation suitable for prompt budgeting. +// The estimator is approximate but stable, isolated, and replaceable. +type TokenEstimator interface { + EstimateTokens(text string) int +} + +// SimpleTokenEstimator provides a basic deterministic token estimation. +// This uses a simple heuristic based on word count and punctuation. +type SimpleTokenEstimator struct{} + +// NewSimpleTokenEstimator creates a new simple token estimator. +func NewSimpleTokenEstimator() *SimpleTokenEstimator { + return &SimpleTokenEstimator{} +} + +// EstimateTokens provides a rough estimate of the number of tokens in the given text. +// This implementation uses a simple heuristic: count words and punctuation as tokens. +// The estimate is deterministic and stable for the same input text. +func (e *SimpleTokenEstimator) EstimateTokens(text string) int { + if text == "" { + return 0 + } + + // Simple heuristic: split on whitespace and count non-empty segments + words := strings.Fields(text) + tokenCount := len(words) + + // Add some estimate for punctuation that might be separate tokens + punctuationCount := 0 + for _, r := range text { + if unicode.IsPunct(r) && r != '\'' && r != '-' && r != '_' { + punctuationCount++ + } + } + + // Rough estimate: each word is a token, plus half the punctuation as separate tokens + return tokenCount + (punctuationCount / 2) +} + +// ConstTokenEstimator returns a constant token count for testing purposes. +type ConstTokenEstimator struct { + Tokens int +} + +// EstimateTokens returns the configured constant token count. +func (e *ConstTokenEstimator) EstimateTokens(text string) int { + return e.Tokens +} diff --git a/internal/core/chunking/tokens_test.go b/internal/core/chunking/tokens_test.go new file mode 100644 index 0000000..da9c374 --- /dev/null +++ b/internal/core/chunking/tokens_test.go @@ -0,0 +1,59 @@ +package chunking + +import ( + "testing" +) + +func TestSimpleTokenEstimator(t *testing.T) { + estimator := NewSimpleTokenEstimator() + + tests := []struct { + name string + text string + expected int + }{ + {"empty string", "", 0}, + {"single word", "hello", 1}, + {"two words", "hello world", 2}, + {"with punctuation", "hello, world!", 3}, // 2 words + 2 punctuation/2 = 3 + {"multiple sentences", "Hello world. This is a test.", 7}, // 7 words + 2 punctuation/2 = 8? Actually "Hello world." has 3 punctuation + {"with apostrophes", "don't won't can't", 3}, + {"with hyphens", "well-known state-of-the-art", 2}, // hyphens don't count + {"unicode text", "café naïve", 2}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := estimator.EstimateTokens(tt.text) + if got != tt.expected { + t.Errorf("EstimateTokens(%q) = %d, want %d", tt.text, got, tt.expected) + } + }) + } +} + +func TestConstTokenEstimator(t *testing.T) { + estimator := &ConstTokenEstimator{Tokens: 42} + + if got := estimator.EstimateTokens("any text"); got != 42 { + t.Errorf("ConstTokenEstimator.EstimateTokens = %d, want 42", got) + } + + if got := estimator.EstimateTokens(""); got != 42 { + t.Errorf("ConstTokenEstimator.EstimateTokens(empty) = %d, want 42", got) + } +} + +func TestTokenEstimatorDeterminism(t *testing.T) { + estimator := NewSimpleTokenEstimator() + text := "The quick brown fox jumps over the lazy dog. Hello, world!" + + // Run multiple times and verify same result + first := estimator.EstimateTokens(text) + for i := 0; i < 10; i++ { + got := estimator.EstimateTokens(text) + if got != first { + t.Errorf("EstimateTokens not deterministic: iteration %d got %d, first was %d", i, got, first) + } + } +}