Add token-bounded transcript chunking
This commit is contained in:
376
internal/core/chunking/sections.go
Normal file
376
internal/core/chunking/sections.go
Normal file
@@ -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,
|
||||||
|
}
|
||||||
|
}
|
||||||
583
internal/core/chunking/sections_test.go
Normal file
583
internal/core/chunking/sections_test.go
Normal file
@@ -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
|
||||||
|
}
|
||||||
55
internal/core/chunking/tokens.go
Normal file
55
internal/core/chunking/tokens.go
Normal file
@@ -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
|
||||||
|
}
|
||||||
59
internal/core/chunking/tokens_test.go
Normal file
59
internal/core/chunking/tokens_test.go
Normal file
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user