diff --git a/internal/compress/algorithm.go b/internal/compress/algorithm.go index 234aeb712..d7178898a 100644 --- a/internal/compress/algorithm.go +++ b/internal/compress/algorithm.go @@ -81,10 +81,13 @@ func Compress(s string, intensity Intensity) (string, Stats) { var out strings.Builder dropSet := dropLists[intensity] for _, seg := range segs { + if seg.Raw == "" { + continue + } if !seg.Safe { stats.PassThroughSegments++ stats.SensitiveKeywordsHit = append(stats.SensitiveKeywordsHit, detectSensitive(seg.Text)) - out.WriteString(seg.Text) + out.WriteString(seg.Raw) continue } stats.CompressedSegments++ @@ -101,7 +104,11 @@ func Compress(s string, intensity Intensity) (string, Stats) { postDrops := applyDropList(postDict, dropSet, &stats) stats.BytesSavedByDrops += len(postDict) - len(postDrops) + // Re-attach the segment's original leading and trailing whitespace so + // the re-joined output preserves the input's sentence separation. + out.WriteString(segmentLeadingWS(seg.Raw)) out.WriteString(postDrops) + out.WriteString(segmentTrailingWS(seg.Raw)) } compressed := out.String() @@ -140,6 +147,35 @@ func applyDropList(s string, dropSet map[string]bool, stats *Stats) string { return normalizeWhitespace(out) } +// segmentLeadingWS returns the leading whitespace of s (spaces, tabs, +// newlines). An empty string is returned when s starts with non-space. +func segmentLeadingWS(s string) string { + i := 0 + for i < len(s) { + switch s[i] { + case ' ', '\t', '\n', '\r': + i++ + default: + return s[:i] + } + } + return s +} + +// segmentTrailingWS returns the trailing whitespace of s. +func segmentTrailingWS(s string) string { + i := len(s) + for i > 0 { + switch s[i-1] { + case ' ', '\t', '\n', '\r': + i-- + default: + return s[i:] + } + } + return s +} + func classifyDrop(phrase string, stats *Stats) { lower := strings.ToLower(phrase) switch { diff --git a/internal/compress/safety.go b/internal/compress/safety.go index eedaedfed..b3527712d 100644 --- a/internal/compress/safety.go +++ b/internal/compress/safety.go @@ -112,6 +112,10 @@ func IsSafeSegment(s string) bool { // SplitSafeSegments splits s into segments that are safe to compress and // segments that must be preserved verbatim. Returns a slice of (text, safe) // pairs. Adjacent safe/unsafe segments are merged by their flag. +// +// Each segment keeps its original, untrimmed text (Raw) so re-joins preserve +// the whitespace between sentences; Text holds the trimmed form used for +// classification and compression. func SplitSafeSegments(s string) []segment { if s == "" { return nil @@ -120,18 +124,22 @@ func SplitSafeSegments(s string) []segment { segs := splitSentences(s) out := make([]segment, 0, len(segs)) for _, seg := range segs { - seg = strings.TrimSpace(seg) - if seg == "" { + trimmed := strings.TrimSpace(seg) + if trimmed == "" { + // Whitespace-only segment (e.g. blank lines between sentences): + // keep the raw text and pass it through untouched. + out = append(out, segment{Text: "", Raw: seg, Safe: false}) continue } - out = append(out, segment{Text: seg, Safe: IsSafeSegment(seg)}) + out = append(out, segment{Text: trimmed, Raw: seg, Safe: IsSafeSegment(trimmed)}) } return out } type segment struct { - Text string - Safe bool // true if safe to compress; false = preserve verbatim + Text string // trimmed text used for classification and compression + Raw string // original (untrimmed) text, preserved for the re-join + Safe bool // true if safe to compress; false = preserve verbatim } // splitSentences splits s on sentence-ending punctuation (`.`, `!`, `?`, diff --git a/internal/compress/whitespace_regression_test.go b/internal/compress/whitespace_regression_test.go new file mode 100644 index 000000000..7da7471a9 --- /dev/null +++ b/internal/compress/whitespace_regression_test.go @@ -0,0 +1,65 @@ +package compress_test + +import ( + "strings" + "testing" + + "github.com/GrayCodeAI/tok/internal/compress" +) + +// Regression tests for the whitespace-preservation fix: compressed output must +// keep the sentence separation (newlines/spaces) of the input instead of +// concatenating trimmed sentences. + +func TestCompressPreservesSentenceWhitespace(t *testing.T) { + cases := []struct { + name string + in string + want string // substring that must survive from both sentences + }{ + {"single newline", "First sentence is here.\nSecond sentence is here.", "here."}, + {"double newline paragraph break", "Paragraph one ends with facts.\n\nParagraph two begins with facts.", "facts."}, + {"trailing newline", "Sentence with facts here.\n\n", "facts"}, + {"leading newline", "\n\nSentence with facts here.", "facts"}, + {"space separated", "First sentence has facts. Second sentence has facts.", "facts"}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + out, _ := compress.Compress(tc.in, compress.Full) + if !strings.Contains(out, "\n\n") && strings.Contains(tc.in, "\n\n") { + t.Errorf("paragraph break lost: input %q -> output %q", tc.in, out) + } + if !strings.Contains(out, "\n") && strings.Contains(tc.in, "\n") { + t.Errorf("newline lost: input %q -> output %q", tc.in, out) + } + // Counts may shift, but the separator must not vanish entirely. + if !strings.Contains(out, tc.want) { + t.Errorf("content missing after compression: %q", out) + } + }) + } +} + +func TestCompressPreservesBlankLinesInsidePassThroughSegments(t *testing.T) { + in := "SECRET_PASSWORD=abc\n\n\nKeep this line too, it has facts." + out, _ := compress.Compress(in, compress.Full) + // The sensitive line is passed through verbatim, including its trailing + // blank lines, and the next line keeps its own separation. + if !strings.Contains(out, "Keep this line too") { + t.Errorf("content lost: %q", out) + } + if !strings.Contains(out, "SECRET_PASSWORD=abc\n\n") { + t.Errorf("pass-through whitespace not preserved: %q", out) + } +} + +func TestCompressMultipleSentencesKeepSpacesBetween(t *testing.T) { + in := "Alpha has useful facts. Beta has useful facts." + out, _ := compress.Compress(in, compress.Full) + // Even when articles are dropped inside sentences, a space must separate + // the two sentences (no "facts.Beta" concatenation). + if strings.Contains(out, "facts.Beta") || strings.Contains(out, "sentence.sentence") { + t.Errorf("sentences concatenated without separator: %q", out) + } +} diff --git a/internal/filter/budget.go b/internal/filter/budget.go index e9840a321..e7b33e366 100644 --- a/internal/filter/budget.go +++ b/internal/filter/budget.go @@ -101,17 +101,31 @@ func (f *BudgetEnforcer) enforceBudget(input string, mode Mode) string { } // TruncateToBudget performs word-level truncation to fit within budget. +// It never returns an empty string for non-empty input: when even a single +// word exceeds the budget, a truncated fragment of the first word is +// returned so callers can distinguish "truncated" from "nothing to do". func (f *BudgetEnforcer) TruncateToBudget(input string, budget int) string { words := strings.Fields(input) if len(words) == 0 { return input } + if budget <= 0 { + return "" + } + if EstimateTokens(input) <= budget { + return input + } - // Binary search for the right number of words + // Binary search for the right number of words. Uses the fast heuristic + // estimator: the precise BPE counter would re-encode a freshly joined + // prefix on every iteration (cache-missing each time) — O(n log n) BPE + // passes and cache pollution for large inputs. The heuristic is within a + // few tokens of BPE for wordy text, and the result is trimmed back below + // with the precise counter below. lo, hi := 0, len(words) for lo < hi { mid := (lo + hi + 1) / 2 - if EstimateTokens(strings.Join(words[:mid], " ")) <= budget { + if EstimateTokensFast(strings.Join(words[:mid], " ")) <= budget { lo = mid } else { hi = mid - 1 @@ -119,9 +133,39 @@ func (f *BudgetEnforcer) TruncateToBudget(input string, budget int) string { } if lo == 0 { + // Even the first word alone exceeds the budget. Keep a fragment of + // it (rune-safe) so the output is never empty. + return firstWordFragment(words[0], budget) + } + + out := strings.Join(words[:lo], " ") + // Precise trim-back: the heuristic can land a few tokens over budget. + for lo > 0 && EstimateTokens(out) > budget { + lo-- + out = strings.Join(words[:lo], " ") + } + if lo == 0 { + return firstWordFragment(words[0], budget) + } + return out +} + +// firstWordFragment returns a leading, rune-safe fragment of word sized to +// fit roughly within budget tokens. It guarantees a non-empty result for +// non-empty input. +func firstWordFragment(word string, budget int) string { + if budget <= 0 { return "" } - return strings.Join(words[:lo], " ") + runes := []rune(word) + max := budget * 4 // heuristic: ~4 runes per token + if max < 1 { + max = 1 + } + if max > len(runes) { + max = len(runes) + } + return string(runes[:max]) } // scoredLine represents a line with its importance score diff --git a/internal/filter/budget_regression_test.go b/internal/filter/budget_regression_test.go new file mode 100644 index 000000000..0e742ac51 --- /dev/null +++ b/internal/filter/budget_regression_test.go @@ -0,0 +1,66 @@ +package filter + +import ( + "strings" + "testing" +) + +// Regression tests for budget enforcement edge cases. + +func TestTruncateToBudgetNeverReturnsEmptyForNonEmptyInput(t *testing.T) { + f := NewBudgetEnforcer(50) + // A single word that dominates the whole budget: the old binary search + // returned "" because it could not fit even one word. + in := strings.Repeat("supercalifragilistic", 2000) // ~72k chars, way over 50 tokens + out := f.TruncateToBudget(in, 50) + if out == "" { + t.Fatal("TruncateToBudget returned empty string for non-empty input") + } + if strings.Contains(out, " ") { + t.Fatalf("expected a fragment of the single word, got %q", out) + } + if len(out) > len(in) { + t.Fatalf("output longer than input") + } +} + +func TestTruncateToBudgetFitsBudget(t *testing.T) { + f := NewBudgetEnforcer(40) + in := strings.Join(strings.Fields("word one two three four five six seven eight nine ten eleven twelve thirteen fourteen fifteen"), " ") + in = strings.Repeat(in+" ", 100) + out := f.TruncateToBudget(in, 40) + if out == "" { + t.Fatal("expected non-empty truncation") + } + // Allow a small slack: the fast estimator drives the search, then the + // precise trim-back guarantees compliance. + if EstimateTokens(out) > 45 { + t.Fatalf("truncated output exceeds budget: %d tokens for budget 40", EstimateTokens(out)) + } + if !strings.HasPrefix(in, out) { + t.Fatalf("truncation must be a prefix of the input") + } +} + +func TestBudgetZeroIsNoOp(t *testing.T) { + // Budget 0 (or negative) disables enforcement entirely: input passes + // through untouched, even when far over any plausible budget. + f := NewBudgetEnforcer(0) + in := strings.Repeat("this is a fairly long sentence with useful facts. ", 500) + out, saved := f.Apply(in, ModeMinimal) + if out != in { + t.Fatalf("budget 0 must be a no-op, got %d/%d chars", len(out), len(in)) + } + if saved != 0 { + t.Fatalf("expected 0 tokens saved, got %d", saved) + } +} + +func TestTruncateToBudgetBudgetZero(t *testing.T) { + f := NewBudgetEnforcer(0) + // Guard: budget <= 0 returns "" (callers must not pass 0 as a real limit). + out := f.TruncateToBudget("some words here", 0) + if out != "" { + t.Fatalf("expected empty for zero budget, got %q", out) + } +} diff --git a/internal/filter/filter.go b/internal/filter/filter.go index e23558005..f4015a715 100644 --- a/internal/filter/filter.go +++ b/internal/filter/filter.go @@ -163,6 +163,13 @@ func EstimateTokens(text string) int { return core.EstimateTokens(text) } +// EstimateTokensFast provides a fast heuristic token count without BPE. +// Delegates to core.EstimateTokensFast; used by hot loops where an exact +// count is not required (e.g. binary-search budget fitting). +func EstimateTokensFast(text string) int { + return core.EstimateTokensFast(text) +} + // IsCode checks if the output looks like source code. func IsCode(output string) bool { codeIndicators := []string{ diff --git a/internal/filter/pipeline_process.go b/internal/filter/pipeline_process.go index 904f59f01..9b4c1ed67 100644 --- a/internal/filter/pipeline_process.go +++ b/internal/filter/pipeline_process.go @@ -72,7 +72,14 @@ func (p *PipelineCoordinator) Process(input string) (string, *PipelineStats) { if p.budgetEnforcer != nil && p.config.Budget > 0 { finalTokens := core.EstimateTokens(output) if finalTokens > p.config.Budget { + before := finalTokens output = p.budgetEnforcer.TruncateToBudget(output, p.config.Budget) + // Record the budget layer so Stats.HardTruncated() can tell a + // hard cut-off from structural compression. + saved := before - core.EstimateTokens(output) + if saved > 0 { + stats.AddLayerStatSafe(LayerBudget, LayerStat{TokensSaved: saved}) + } } } diff --git a/internal/filter/presets.go b/internal/filter/presets.go index 943efa32e..b80586854 100644 --- a/internal/filter/presets.go +++ b/internal/filter/presets.go @@ -55,7 +55,6 @@ func TierConfig(tier Tier, baseMode Mode) PipelineConfig { cfg.EnableEvaluator = true cfg.EnableH2O = true cfg.EnableAttentionSink = true - cfg.EnableMetaToken = true case TierExtract: cfg.EnableEntropy = true cfg.EnablePerplexity = true @@ -70,7 +69,6 @@ func TierConfig(tier Tier, baseMode Mode) PipelineConfig { cfg.EnableAttribution = true cfg.EnableH2O = true cfg.EnableAttentionSink = true - cfg.EnableMetaToken = true cfg.EnableSemanticChunk = true cfg.EnableLazyPruner = true cfg.EnableSemanticAnchor = true @@ -89,7 +87,6 @@ func TierConfig(tier Tier, baseMode Mode) PipelineConfig { cfg.EnableAttribution = true cfg.EnableH2O = true cfg.EnableAttentionSink = true - cfg.EnableMetaToken = true cfg.EnableSemanticChunk = true cfg.EnableSketchStore = true cfg.EnableLazyPruner = true @@ -100,7 +97,6 @@ func TierConfig(tier Tier, baseMode Mode) PipelineConfig { cfg.EnableAST = true cfg.EnableGoalDriven = true cfg.EnableH2O = true - cfg.EnableMetaToken = true case TierLog: cfg.EnableEntropy = true cfg.EnablePerplexity = true diff --git a/internal/filter/tier_config.go b/internal/filter/tier_config.go index fe7a60362..84a5d2aa1 100644 --- a/internal/filter/tier_config.go +++ b/internal/filter/tier_config.go @@ -223,7 +223,6 @@ func enableTierLayers(cfg *PipelineConfig, tier AutoTier) { cfg.EnableAttribution = true cfg.EnableH2O = true cfg.EnableAttentionSink = true - cfg.EnableMetaToken = true cfg.EnableSemanticChunk = true cfg.EnableSketchStore = true cfg.EnableLazyPruner = true diff --git a/options.go b/options.go index 3716b5abd..6fab3944b 100644 --- a/options.go +++ b/options.go @@ -45,10 +45,23 @@ func WithMode(m Mode) Option { } // WithBudget sets a hard token limit for the output. +// A value <= 0 disables the budget cap entirely: no truncation is applied +// (structural compression layers still run). Passing WithBudget(0) therefore +// means "no budget", not "zero-token output". func WithBudget(tokens int) Option { return optFunc(func(c *config) { c.budget = tokens }) } +// WithMetaToken enables the lossless meta-token compression layer for repeated +// token sequences. It is OFF by default: the layer replaces repeated spans +// with [META:...] placeholders whose originals live only inside the pipeline +// instance, so the string returned by the stateless Compress() API would be a +// data-losing digest for repetitive input. Enable it only when the consumer +// can round-trip through a live pipeline that decompresses. +func WithMetaToken() Option { + return optFunc(func(c *config) { c.metaToken = true }) +} + // WithQuery provides intent context for goal-driven filtering. func WithQuery(intent string) Option { return optFunc(func(c *config) { c.query = intent }) @@ -146,6 +159,7 @@ type config struct { codeLang string symbolProvider SymbolProvider customFilters *CustomFilter + metaToken bool perplexityEnabled bool perplexityScorer PerplexityScorer @@ -176,5 +190,6 @@ func (c *config) toPipelineConfig() filter.PipelineConfig { cfg.Budget = c.budget cfg.QueryIntent = c.query cfg.EnableLLMLinguaProse = c.llmLinguaProse + cfg.EnableMetaToken = c.metaToken return cfg } diff --git a/stats.go b/stats.go index 354764a2c..9bff8d803 100644 --- a/stats.go +++ b/stats.go @@ -19,6 +19,19 @@ type LayerStat struct { DurationMs int64 } +// HardTruncated reports whether the final budget enforcer did the bulk of the +// token reduction — i.e. the output is mostly the input cut short at the +// budget boundary — rather than the structural compression layers (dedupe, +// compaction, gist, ...). A hard-truncated output is NOT a rewrite of the +// content and must not be treated as a summary by callers. +func (s Stats) HardTruncated() bool { + if s.TokensSaved <= 0 { + return false + } + budgetSaved := s.Layers[filter.LayerBudget].TokensSaved + return budgetSaved > 0 && budgetSaved*2 > s.TokensSaved +} + func newStats(ps *filter.PipelineStats) Stats { if ps == nil { return Stats{} diff --git a/stats_regression_test.go b/stats_regression_test.go new file mode 100644 index 000000000..1387e47ea --- /dev/null +++ b/stats_regression_test.go @@ -0,0 +1,78 @@ +package tok + +import ( + "fmt" + "strings" + "testing" +) + +func TestStatsHardTruncated(t *testing.T) { + tests := []struct { + name string + layers map[string]LayerStat + saved int + want bool + }{ + { + name: "budget layer dominant -> hard truncation", + layers: map[string]LayerStat{"10_budget": {TokensSaved: 800}}, + saved: 1000, + want: true, + }, + { + name: "budget layer absent -> not hard truncation", + layers: map[string]LayerStat{"11_compaction": {TokensSaved: 800}}, + saved: 1000, + want: false, + }, + { + name: "structural layers dominant -> not hard truncation", + layers: map[string]LayerStat{"10_budget": {TokensSaved: 100}, "11_compaction": {TokensSaved: 900}}, + saved: 1000, + want: false, + }, + { + name: "nothing saved -> not hard truncation", + layers: map[string]LayerStat{}, + saved: 0, + want: false, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + s := Stats{TokensSaved: tc.saved, Layers: tc.layers} + if got := s.HardTruncated(); got != tc.want { + t.Errorf("HardTruncated() = %v, want %v", got, tc.want) + } + }) + } +} + +func TestCompressWithBudgetZeroDisablesTruncation(t *testing.T) { + // WithBudget(0) must mean "no budget cap": the budget-enforcer layer must + // be inert (record no savings), while a real budget activates it. + var sb strings.Builder + for i := 0; i < 500; i++ { + fmt.Fprintf(&sb, "Distinct sentence number %d adds factual informative content here. ", i) + } + in := sb.String() + + out, stats := Compress(in, WithBudget(0), WithMode(ModeMinimal)) + if out == "" { + t.Fatal("WithBudget(0) must not produce empty output") + } + if saved := stats.Layers["10_budget"].TokensSaved; saved != 0 { + t.Fatalf("budget enforcer must be inert with budget 0, saved %d tokens", saved) + } + + out2, stats2 := Compress(in, WithBudget(60), WithMode(ModeMinimal)) + if out2 == "" { + t.Fatal("budgeted compress must not produce empty output") + } + if stats2.Layers["10_budget"].TokensSaved <= 0 { + t.Fatalf("budget enforcer must cut with a real budget, layer savings = %d", stats2.Layers["10_budget"].TokensSaved) + } + if stats2.FinalTokens > 60 { + t.Fatalf("budgeted output exceeds budget: %d tokens", stats2.FinalTokens) + } +}