compress_test.go

package tok_test

import (
	"strings"
	"testing"

	"github.com/GrayCodeAI/tok"
)

func TestPromptCompress_Lite(t *testing.T) {
	in := "Sure, I can help you with that. Of course, the answer is yes."
	out, stats := tok.PromptCompress(in, tok.IntensityLite)
	if strings.Contains(out, "Sure,") {
		t.Errorf("expected 'Sure,' to be dropped at Lite, got %q", out)
	}
	if !strings.Contains(out, "help you with that") {
		t.Errorf("expected help text preserved, got %q", out)
	}
	if stats.Intensity != tok.IntensityLite {
		t.Errorf("expected intensity=Lite, got %v", stats.Intensity)
	}
}

func TestPromptCompress_Full(t *testing.T) {
	in := "The quick brown fox jumps over the lazy dog."
	out, _ := tok.PromptCompress(in, tok.IntensityFull)
	// "the" should be dropped
	if strings.Contains(out, " the ") && !strings.Contains(out, "over") {
		t.Errorf("expected 'the' to be dropped at Full, got %q", out)
	}
}

func TestPromptCompress_Ultra(t *testing.T) {
	in := "However, the system is basically quite slow. Therefore, we need to optimize it."
	out, stats := tok.PromptCompress(in, tok.IntensityUltra)
	if strings.Contains(out, "However,") {
		t.Errorf("expected 'However,' to be dropped at Ultra, got %q", out)
	}
	if strings.Contains(out, "Therefore,") {
		t.Errorf("expected 'Therefore,' to be dropped at Ultra, got %q", out)
	}
	if stats.DroppedConjunctions == 0 {
		t.Error("expected DroppedConjunctions > 0 for Ultra")
	}
}

func TestPromptCompress_SensitivePassThrough(t *testing.T) {
	in := "Be careful with rm -rf /tmp. The file is large."
	out, stats := tok.PromptCompress(in, tok.IntensityFull)
	// The segment containing "rm -rf" must be preserved verbatim.
	if !strings.Contains(out, "rm -rf /tmp.") {
		t.Errorf("expected 'rm -rf /tmp.' to be preserved, got %q", out)
	}
	if stats.PassThroughSegments == 0 {
		t.Error("expected PassThroughSegments > 0")
	}
}

func TestPromptCompress_Empty(t *testing.T) {
	out, stats := tok.PromptCompress("", tok.IntensityFull)
	if out != "" {
		t.Errorf("expected empty output, got %q", out)
	}
	if stats.OriginalBytes != 0 {
		t.Errorf("expected zero bytes, got %d", stats.OriginalBytes)
	}
}

func TestPromptCompress_Dictionary(t *testing.T) {
	in := "In order to install, you need to make use of the installer."
	out, _ := tok.PromptCompress(in, tok.IntensityFull)
	if strings.Contains(out, "In order to") {
		t.Errorf("expected 'In order to' to be replaced, got %q", out)
	}
	if strings.Contains(out, "make use of") {
		t.Errorf("expected 'make use of' to be replaced, got %q", out)
	}
}

func TestPromptCompress_DoesNotAffectTopLevelCompress(t *testing.T) {
	// Regression: adding prompt-compression API must not change the main Compress() output
	// for a typical input. We test that both APIs can be called and return
	// different shapes (PromptCompress returns PromptStats, main returns Stats).
	in := "The quick brown fox jumps over the lazy dog."
	mainOut, mainStats := tok.Compress(in, tok.Aggressive)
	if mainOut == "" {
		t.Error("expected non-empty main Compress output")
	}
	_ = mainStats.OriginalTokens

	promptCompressOut, _ := tok.PromptCompress(in, tok.IntensityFull)
	// PromptCompress output may equal main output (if both drop "The") or differ —
	// the point is they don't crash each other.
	if promptCompressOut == "" {
		t.Error("expected non-empty prompt-compression output")
	}
}