algorithm.go

// Package compress implements the prompt-compression algorithm
// ported natively to Go. Source: JuliusBrussee/prompt-compression (MIT).
//
// Algorithm overview:
//
//  1. Split input into segments (sentences, preserving code/URLs).
//  2. Per segment, run the auto-clarity check: segments containing
//     security/destructive keywords are passed through verbatim.
//  3. For safe segments, apply the dictionary (eng/v1) first
//     (longest-match-first phrase substitutions).
//  4. Then apply the drop-list for the chosen intensity level
//     (Lite/Full/Ultra).
//  5. Re-join segments, preserving original whitespace.
package compress

import (
	"strings"
)

// CompressResult reports what the compression did to the input.
type CompressResult struct {
	// Original is the input verbatim.
	Original string
	// Compressed is the rewritten text.
	Compressed string
	// Stats for the operation.
	Stats Stats
}

// Stats reports per-stage counts.
type Stats struct {
	Intensity             Intensity
	OriginalBytes         int
	CompressedBytes       int
	OriginalWords         int
	CompressedWords       int
	BytesSaved            int
	PercentOff            float64
	DroppedArticles       int
	DroppedFiller         int
	DroppedHedge          int
	DroppedConjunctions   int
	DictionaryHits        int
	PassThroughSegments   int
	CompressedSegments    int
	BytesSavedByDict      int
	BytesSavedByDrops     int
	RefusedDueToSensitive bool
	// Per-segment counts (informational; first N at most).
	SensitiveKeywordsHit []string
}

// Compress applies the algorithm to s at the given intensity.
//
// Behavior:
//   - Empty input returns (input, empty result) untouched.
//   - CJK-heavy text is passed through (compression rules assume Latin grammar).
//   - Sensitive segments (security/destructive keywords) are preserved
//     verbatim regardless of intensity.
//   - Multi-word drop-list entries ("of course", "feel free") are matched
//     as whole phrases.
//   - Dictionary substitutions are case-insensitive and preserve the
//     case of the first character of each match.
func Compress(s string, intensity Intensity) (string, Stats) {
	stats := Stats{Intensity: intensity, OriginalBytes: len(s), OriginalWords: wordCount(s)}
	if s == "" {
		return s, stats
	}
	if isCJK(s) {
		// CJK has no articles/filler; pass through.
		stats.CompressedBytes = len(s)
		stats.CompressedWords = stats.OriginalWords
		stats.PercentOff = 0
		return s, stats
	}

	// Split into segments, classify safe vs. pass-through.
	segs := SplitSafeSegments(s)

	// Build the output and per-stage counters.
	var out strings.Builder
	dropSet := dropLists[intensity]
	for _, seg := range segs {
		if !seg.Safe {
			stats.PassThroughSegments++
			stats.SensitiveKeywordsHit = append(stats.SensitiveKeywordsHit, detectSensitive(seg.Text))
			out.WriteString(seg.Text)
			continue
		}
		stats.CompressedSegments++

		// 1. Dictionary substitutions
		preDict := seg.Text
		postDict := SubstituteDict(seg.Text)
		stats.BytesSavedByDict += len(preDict) - len(postDict)
		if postDict != preDict {
			stats.DictionaryHits++
		}

		// 2. Drop-list application
		postDrops := applyDropList(postDict, dropSet, &stats)
		stats.BytesSavedByDrops += len(postDict) - len(postDrops)

		out.WriteString(postDrops)
	}

	compressed := out.String()
	stats.CompressedBytes = len(compressed)
	stats.CompressedWords = wordCount(compressed)
	stats.BytesSaved = stats.OriginalBytes - stats.CompressedBytes
	if stats.OriginalBytes > 0 {
		stats.PercentOff = float64(stats.BytesSaved) / float64(stats.OriginalBytes) * 100
	}
	return compressed, stats
}

// applyDropList removes all words/phrases in dropSet from s.
// Counts are tallied per category (article, filler, hedge, conjunction).
// Whitespace is normalized but not destroyed: a doubled space becomes
// a single space, leading/trailing spaces on a segment are preserved.
func applyDropList(s string, dropSet map[string]bool, stats *Stats) string {
	if len(dropSet) == 0 {
		return s
	}
	// Sort by length desc so multi-word phrases match first.
	phrases := make([]string, 0, len(dropSet))
	for k := range dropSet {
		phrases = append(phrases, k)
	}
	sortByLengthDesc(phrases)

	out := s
	for _, p := range phrases {
		before := out
		out = removePhraseCI(out, p)
		if out != before {
			classifyDrop(p, stats)
		}
	}
	return normalizeWhitespace(out)
}

func classifyDrop(phrase string, stats *Stats) {
	lower := strings.ToLower(phrase)
	switch {
	case lower == "a" || lower == "an" || lower == "the":
		stats.DroppedArticles++
	case isHedge(lower):
		stats.DroppedHedge++
	case isConjunction(lower):
		stats.DroppedConjunctions++
	default:
		stats.DroppedFiller++
	}
}

func isHedge(p string) bool {
	switch p {
	case "perhaps", "maybe", "i think", "i believe", "in my opinion",
		"it seems", "it appears", "somewhat", "kind of", "sort of":
		return true
	}
	return false
}

func isConjunction(p string) bool {
	switch p {
	case "however", "moreover", "furthermore", "nevertheless", "nonetheless",
		"therefore", "thus", "hence", "accordingly", "consequently",
		"additionally", "alternatively", "likewise", "similarly",
		"subsequently", "previously":
		return true
	}
	return false
}

// removePhraseCI removes all case-insensitive occurrences of needle from s.
// Removes the trailing whitespace if present (to avoid "word  word").
// Does NOT remove leading whitespace (to preserve word boundaries).
func removePhraseCI(s, needle string) string {
	if needle == "" {
		return s
	}
	lowerS := toLowerASCII(s)
	lowerN := toLowerASCII(needle)
	if len(lowerN) > len(lowerS) {
		return s
	}
	var out []byte
	i := 0
	for i < len(s) {
		if i+len(lowerN) <= len(lowerS) && lowerS[i:i+len(lowerN)] == lowerN {
			// Word boundary checks
			if i > 0 && isWordByte(s[i-1]) {
				out = append(out, s[i])
				i++
				continue
			}
			if i+len(lowerN) < len(s) && isWordByte(s[i+len(lowerN)]) {
				out = append(out, s[i])
				i++
				continue
			}
			// Skip the phrase
			i += len(lowerN)
			// Also skip one trailing space if present
			if i < len(s) && s[i] == ' ' {
				i++
			}
		} else {
			out = append(out, s[i])
			i++
		}
	}
	return string(out)
}

// normalizeWhitespace collapses runs of spaces and trims segment
// boundaries, but preserves newlines.
func normalizeWhitespace(s string) string {
	var out []byte
	prevSpace := false
	for i := 0; i < len(s); i++ {
		c := s[i]
		if c == ' ' || c == '\t' {
			if !prevSpace {
				out = append(out, ' ')
			}
			prevSpace = true
			continue
		}
		prevSpace = false
		out = append(out, c)
	}
	return string(out)
}

func wordCount(s string) int {
	return len(strings.Fields(s))
}

// detectSensitive returns the first safety keyword found in s (lower-case).
// Returns empty if none.
func detectSensitive(s string) string {
	lower := strings.ToLower(s)
	for _, kw := range safetyKeywords {
		if strings.Contains(lower, kw) {
			return kw
		}
	}
	return ""
}