audita/internal/core/normalization/normalize.go

package normalization

import (
	"sort"

	"gitea.maximumdirect.net/eric/audita/internal/core/schema"
)

// NormalizationConfig holds configuration parameters for transcript normalization
type NormalizationConfig struct {
	MaxSegmentGap      float64 // Maximum gap between segments to consider for merging (seconds)
	EllipsisGap        float64 // Gap threshold above which to insert ellipsis (seconds)
	MaxSegmentDuration float64 // Maximum duration for a merged segment (seconds)
	MaxSegmentTokens   int     // Maximum token estimate for a merged segment
}

// NormalizationSummary records the results and statistics of normalization
type NormalizationSummary struct {
	InputSegmentCount  int `json:"input_segment_count"`
	OutputSegmentCount int `json:"output_segment_count"`
	MergesPerformed    int `json:"merges_performed"`
	IDsReassigned      int `json:"ids_reassigned"`
	SkippedMerges      struct {
		DifferentSpeakers  int `json:"different_speakers"`
		GapTooLarge        int `json:"gap_too_large"`
		DurationExceeded   int `json:"duration_exceeded"`
		TokenLimitExceeded int `json:"token_limit_exceeded"`
	} `json:"skipped_merges"`
}

// NormalizeTranscript performs deterministic normalization on a transcript
type NormalizeTranscript struct {
	config    NormalizationConfig
	estimator *SimpleTokenEstimator
}

// NewNormalizer creates a new transcript normalizer with the given configuration
func NewNormalizer(config NormalizationConfig) *NormalizeTranscript {
	return &NormalizeTranscript{
		config:    config,
		estimator: &SimpleTokenEstimator{},
	}
}

// Normalize performs deterministic normalization on the given transcript
func (n *NormalizeTranscript) Normalize(transcript *schema.Transcript) (*schema.Transcript, *NormalizationSummary) {
	summary := &NormalizationSummary{
		InputSegmentCount: len(transcript.Segments),
	}

	if len(transcript.Segments) == 0 {
		return transcript, summary
	}

	// Sort segments chronologically
	sortedSegments := make([]schema.Segment, len(transcript.Segments))
	copy(sortedSegments, transcript.Segments)
	sort.Slice(sortedSegments, func(i, j int) bool {
		return sortedSegments[i].Start < sortedSegments[j].Start
	})

	// Merge same-speaker adjacent segments
	var normalizedSegments []schema.Segment
	var currentSegment schema.Segment

	for i, segment := range sortedSegments {
		if i == 0 {
			// Initialize with first segment
			currentSegment = segment
			continue
		}

		// Check if we should merge with current segment
		gap := segment.Start - currentSegment.End
		merge, reason := shouldMergeWithReason(n.config, n.estimator, &currentSegment, &segment, gap)
		if merge {
			summary.MergesPerformed++
			currentSegment = mergeSegments(&currentSegment, &segment, gap, n.config.EllipsisGap)
		} else {
			// Track the reason for not merging
			switch reason {
			case "different_speakers":
				summary.SkippedMerges.DifferentSpeakers++
			case "gap_too_large":
				summary.SkippedMerges.GapTooLarge++
			case "duration_exceeded":
				summary.SkippedMerges.DurationExceeded++
			case "token_limit_exceeded":
				summary.SkippedMerges.TokenLimitExceeded++
			}
			// Finalize current segment and start new one
			normalizedSegments = append(normalizedSegments, currentSegment)
			currentSegment = segment
		}
	}

	// Add the last segment
	if len(currentSegment.Text) > 0 {
		normalizedSegments = append(normalizedSegments, currentSegment)
	}

	// Reassign sequential IDs starting at 1
	for i, segment := range normalizedSegments {
		if segment.ID != i+1 {
			summary.IDsReassigned++
		}
		normalizedSegments[i].ID = i + 1
	}

	summary.OutputSegmentCount = len(normalizedSegments)

	return &schema.Transcript{Segments: normalizedSegments}, summary
}

func shouldMerge(config NormalizationConfig, estimator *SimpleTokenEstimator,
	current, next *schema.Segment, gap float64) bool {

	// Don't merge if different speakers
	if current.Speaker != next.Speaker {
		return false
	}

	// Don't merge if gap is too large
	if gap > config.MaxSegmentGap {
		return false
	}

	// Calculate what the merged segment would look like
	mergedText := current.Text
	if gap >= config.EllipsisGap {
		mergedText += "... "
	} else {
		mergedText += " "
	}
	mergedText += next.Text

	mergedDuration := next.End - current.Start

	// Don't merge if duration would be exceeded
	if mergedDuration > config.MaxSegmentDuration {
		return false
	}

	// Don't merge if token estimate would be exceeded
	tokenEstimate := estimator.EstimateTokens(mergedText)
	if tokenEstimate > config.MaxSegmentTokens {
		return false
	}

	return true
}

func shouldMergeWithReason(config NormalizationConfig, estimator *SimpleTokenEstimator,
	current, next *schema.Segment, gap float64) (bool, string) {

	if current.Speaker != next.Speaker {
		return false, "different_speakers"
	}
	if gap > config.MaxSegmentGap {
		return false, "gap_too_large"
	}

	mergedText := current.Text
	if gap >= config.EllipsisGap {
		mergedText += "... "
	} else {
		mergedText += " "
	}
	mergedText += next.Text

	mergedDuration := next.End - current.Start
	if mergedDuration > config.MaxSegmentDuration {
		return false, "duration_exceeded"
	}

	tokenEstimate := estimator.EstimateTokens(mergedText)
	if tokenEstimate > config.MaxSegmentTokens {
		return false, "token_limit_exceeded"
	}

	return true, ""
}

func mergeSegments(current, next *schema.Segment, gap, ellipsisGap float64) schema.Segment {
	mergedText := current.Text
	if gap >= ellipsisGap {
		mergedText += "... "
	} else {
		mergedText += " "
	}
	mergedText += next.Text

	// Merge categories
	categories := make([]string, 0, len(current.Categories)+len(next.Categories))
	categories = append(categories, current.Categories...)
	categories = append(categories, next.Categories...)

	return schema.Segment{
		ID:         current.ID, // Will be reassigned later
		Speaker:    current.Speaker,
		Start:      current.Start,
		End:        next.End,
		Text:       mergedText,
		Categories: categories,
	}
}