mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-10 14:27:00 -04:00
Remove LLMLingua: Kompress is the sole text compressor
LLMLingua was the original ML text compressor (BERT-based). Kompress (ModernBERT, trained on 330K structured tool outputs) replaced it with better compression quality and simpler architecture. Removed across 35 files: - Deleted headroom/transforms/llmlingua_compressor.py - Deleted tests/test_transforms/test_llmlingua_compressor.py - Deleted tests/test_proxy_llmlingua.py - Removed all enable_llmlingua config, _get_llmlingua methods, LLMLingua fallback paths, LLMLINGUA strategy enum values - Removed CLI flags, model configs, compression handler references - Simplified ContentRouter: Kompress is primary and only text compressor
This commit is contained in:
chopratejas
parent
7a79aa0792
commit
3290a3d582
35 changed files with 126 additions and 2461 deletions
3
.gitignore
vendored
3
.gitignore
vendored
|
|
@ -1,3 +1,6 @@
|
|||
# Swift SDK (separate repo)
|
||||
swift/
|
||||
|
||||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
|
|
|
|||
|
|
@ -44,13 +44,13 @@ try:
|
|||
except ImportError:
|
||||
HEADROOM_AVAILABLE = False
|
||||
|
||||
# LLMLingua imports (SOTA baseline)
|
||||
# Kompress imports (ML baseline)
|
||||
try:
|
||||
from headroom.transforms.llmlingua_compressor import LLMLinguaCompressor, LLMLinguaConfig
|
||||
from headroom.transforms.kompress_compressor import KompressCompressor, is_kompress_available
|
||||
|
||||
LLMLINGUA_AVAILABLE = True
|
||||
KOMPRESS_AVAILABLE = is_kompress_available()
|
||||
except ImportError:
|
||||
LLMLINGUA_AVAILABLE = False
|
||||
KOMPRESS_AVAILABLE = False
|
||||
|
||||
|
||||
@dataclass
|
||||
|
|
@ -427,27 +427,21 @@ Provide a structured summary that retains all critical details."""
|
|||
return summary, cost, latency
|
||||
|
||||
|
||||
def llmlingua_compress(data: list[dict]) -> tuple[str, dict]:
|
||||
def kompress_compress(data: list[dict]) -> tuple[str, dict]:
|
||||
"""
|
||||
Use LLMLingua-2 (Microsoft SOTA) for ML-based compression.
|
||||
Use Kompress (ModernBERT) for ML-based compression.
|
||||
Returns (compressed_text, metadata).
|
||||
"""
|
||||
if not LLMLINGUA_AVAILABLE:
|
||||
raise RuntimeError(
|
||||
"LLMLingua not available. Install with: pip install headroom-ai[llmlingua]"
|
||||
)
|
||||
if not KOMPRESS_AVAILABLE:
|
||||
raise RuntimeError("Kompress not available. Install with: pip install headroom-ai[ml]")
|
||||
|
||||
config = LLMLinguaConfig(
|
||||
target_compression_rate=0.3, # Keep ~30% of tokens
|
||||
min_tokens_for_compression=50,
|
||||
)
|
||||
compressor = LLMLinguaCompressor(config)
|
||||
compressor = KompressCompressor()
|
||||
|
||||
# Convert data to string for LLMLingua (it works on text, not structured data)
|
||||
# Convert data to string for Kompress (it works on text, not structured data)
|
||||
data_str = json.dumps(data, indent=2)
|
||||
|
||||
start = time.time()
|
||||
result = compressor.compress(data_str, content_type="json")
|
||||
result = compressor.compress(data_str)
|
||||
latency = (time.time() - start) * 1000
|
||||
|
||||
metadata = {
|
||||
|
|
@ -590,7 +584,7 @@ class BenchmarkConfig:
|
|||
max_truncate_items: int = 20
|
||||
max_headroom_items: int = 20
|
||||
run_summarization: bool = True # Can disable to save cost
|
||||
run_llmlingua: bool = True # Run LLMLingua-2 (SOTA baseline)
|
||||
run_kompress: bool = True # Run Kompress (ML baseline)
|
||||
|
||||
|
||||
def run_scenario_benchmark(
|
||||
|
|
@ -701,12 +695,12 @@ def run_scenario_benchmark(
|
|||
except Exception as e:
|
||||
print(f" Summarization failed: {e}")
|
||||
|
||||
# --- LLMLINGUA-2 (SOTA) ---
|
||||
if config.run_llmlingua:
|
||||
print("\n[3/4] Running LLMLingua-2 (Microsoft SOTA)...")
|
||||
if LLMLINGUA_AVAILABLE:
|
||||
# --- KOMPRESS (ML baseline) ---
|
||||
if config.run_kompress:
|
||||
print("\n[3/4] Running Kompress (ModernBERT ML baseline)...")
|
||||
if KOMPRESS_AVAILABLE:
|
||||
try:
|
||||
ll_compressed, ll_metadata = llmlingua_compress(scenario.data)
|
||||
ll_compressed, ll_metadata = kompress_compress(scenario.data)
|
||||
ll_tokens = count_tokens(ll_compressed)
|
||||
|
||||
ll_answers = []
|
||||
|
|
@ -729,7 +723,7 @@ def run_scenario_benchmark(
|
|||
|
||||
results.append(
|
||||
ApproachResult(
|
||||
approach="llmlingua-2",
|
||||
approach="kompress",
|
||||
scenario=scenario.name,
|
||||
tokens_original=original_tokens,
|
||||
tokens_after=ll_tokens,
|
||||
|
|
@ -747,9 +741,9 @@ def run_scenario_benchmark(
|
|||
print(f" Accuracy: {ll_accuracy:.1%}")
|
||||
print(f" Compression latency: {ll_metadata['latency_ms']:.1f}ms")
|
||||
except Exception as e:
|
||||
print(f" LLMLingua-2 failed: {e}")
|
||||
print(f" Kompress failed: {e}")
|
||||
else:
|
||||
print(" LLMLingua-2 not available. Install with: pip install headroom-ai[llmlingua]")
|
||||
print(" Kompress not available. Install with: pip install headroom-ai[ml]")
|
||||
|
||||
# --- HEADROOM ---
|
||||
print("\n[4/4] Running Headroom...")
|
||||
|
|
|
|||
|
|
@ -196,7 +196,6 @@ class CCRToolInjector:
|
|||
_detected_hashes: list[str] = field(default_factory=list)
|
||||
# Multiple marker patterns to match different compressors:
|
||||
# - SmartCrusher: [100 items compressed to 10. Retrieve more: hash=abc123]
|
||||
# - LLMLingua: [1000 items compressed to 300. Retrieve more: hash=abc123]
|
||||
# - Kompress: [100 lines compressed to 10. Retrieve more: hash=abc123]
|
||||
# - LogCompressor: [200 lines compressed to 20. Retrieve more: hash=abc123]
|
||||
# - SearchCompressor: [50 matches compressed to 5. Retrieve more: hash=abc123]
|
||||
|
|
|
|||
|
|
@ -21,20 +21,6 @@ from .main import main
|
|||
@click.option("--no-rate-limit", is_flag=True, help="Disable rate limiting")
|
||||
@click.option("--log-file", default=None, help="Path to JSONL log file")
|
||||
@click.option("--budget", type=float, default=None, help="Daily budget limit in USD")
|
||||
# LLMLingua ML-based compression (ON by default if installed)
|
||||
@click.option("--no-llmlingua", is_flag=True, help="Disable LLMLingua-2 ML-based compression")
|
||||
@click.option(
|
||||
"--llmlingua-device",
|
||||
type=click.Choice(["auto", "cuda", "cpu", "mps"]),
|
||||
default="auto",
|
||||
help="Device for LLMLingua model (default: auto)",
|
||||
)
|
||||
@click.option(
|
||||
"--llmlingua-rate",
|
||||
type=float,
|
||||
default=0.3,
|
||||
help="LLMLingua compression rate 0.0-1.0 (default: 0.3 = keep 30%)",
|
||||
)
|
||||
# Code-aware compression (ON by default if installed)
|
||||
@click.option("--no-code-aware", is_flag=True, help="Disable AST-based code compression")
|
||||
# Read lifecycle (ON by default: compresses stale/superseded Read outputs)
|
||||
|
|
@ -149,9 +135,6 @@ def proxy(
|
|||
no_rate_limit: bool,
|
||||
log_file: str | None,
|
||||
budget: float | None,
|
||||
no_llmlingua: bool,
|
||||
llmlingua_device: str,
|
||||
llmlingua_rate: float,
|
||||
no_code_aware: bool,
|
||||
no_read_lifecycle: bool,
|
||||
no_intelligent_context: bool,
|
||||
|
|
@ -225,10 +208,6 @@ def proxy(
|
|||
rate_limit_enabled=not no_rate_limit,
|
||||
log_file=log_file,
|
||||
budget_limit_usd=budget,
|
||||
# LLMLingua: ON by default (use --no-llmlingua to disable)
|
||||
llmlingua_enabled=not no_llmlingua,
|
||||
llmlingua_device=llmlingua_device,
|
||||
llmlingua_target_rate=llmlingua_rate,
|
||||
# Code-aware: ON by default (use --no-code-aware to disable)
|
||||
code_aware_enabled=not no_code_aware,
|
||||
# Read lifecycle: ON by default (use --no-read-lifecycle to disable)
|
||||
|
|
|
|||
|
|
@ -203,7 +203,7 @@ def _get_pipeline() -> Any:
|
|||
# Default pipeline: CacheAligner → ContentRouter → IntelligentContext
|
||||
# CacheAligner: stabilizes prefix for provider KV cache hits
|
||||
# ContentRouter: routes to the right compressor per content type
|
||||
# (SmartCrusher for JSON, CodeCompressor for code, LLMLingua for text)
|
||||
# (SmartCrusher for JSON, CodeCompressor for code, Kompress for text)
|
||||
# IntelligentContext: enforces token limits with score-based dropping
|
||||
_pipeline = TransformPipeline()
|
||||
logger.debug("Headroom compression pipeline initialized")
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@
|
|||
This module provides intelligent, automatic compression that:
|
||||
1. Detects content type using ML (Magika)
|
||||
2. Preserves structure (keys, signatures, templates)
|
||||
3. Compresses content with LLMLingua
|
||||
3. Compresses content with Kompress
|
||||
4. Enables retrieval via CCR
|
||||
|
||||
Quick Start:
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ Each handler knows how to extract structural information from a specific
|
|||
content type and create a StructureMask marking what should be preserved.
|
||||
|
||||
Handlers don't compress - they only identify structure. The actual
|
||||
compression is done by LLMLingua on the non-structural parts.
|
||||
compression is done by Kompress on the non-structural parts.
|
||||
"""
|
||||
|
||||
from headroom.compression.handlers.base import (
|
||||
|
|
|
|||
|
|
@ -181,7 +181,7 @@ class BaseStructureHandler(ABC):
|
|||
|
||||
Subclasses may override for more sophisticated tokenization.
|
||||
For mask purposes, character-level is often sufficient and
|
||||
aligns well with LLMLingua's token-level compression.
|
||||
aligns well with token-level compression.
|
||||
|
||||
Args:
|
||||
content: Content to tokenize.
|
||||
|
|
|
|||
|
|
@ -1,11 +1,11 @@
|
|||
"""Structure mask system for compression.
|
||||
|
||||
A StructureMask identifies which parts of content are "structural" (should be
|
||||
preserved) vs "compressible" (can be compressed by LLMLingua).
|
||||
preserved) vs "compressible" (can be compressed by Kompress).
|
||||
|
||||
This separates the concerns of:
|
||||
1. Structure detection (handlers) - What tokens are navigational?
|
||||
2. Content compression (LLMLingua) - What tokens can be removed?
|
||||
2. Content compression (Kompress) - What tokens can be removed?
|
||||
|
||||
The mask is content-agnostic - it's just a boolean array aligned to tokens.
|
||||
"""
|
||||
|
|
@ -21,7 +21,7 @@ class StructureMask:
|
|||
"""A mask identifying structural vs compressible tokens.
|
||||
|
||||
The mask is aligned to a token sequence. True means "preserve this token"
|
||||
(it's structural/navigational), False means "compressible" (LLMLingua can
|
||||
(it's structural/navigational), False means "compressible" (Kompress can
|
||||
potentially remove it).
|
||||
|
||||
Attributes:
|
||||
|
|
@ -207,7 +207,7 @@ def apply_mask_to_text(
|
|||
Args:
|
||||
text: Original text.
|
||||
mask: Structure mask aligned to tokens.
|
||||
compress_fn: Function to compress text (e.g., LLMLingua).
|
||||
compress_fn: Function to compress text (e.g., Kompress).
|
||||
tokenizer_decode: Optional function to decode tokens to text.
|
||||
If not provided, assumes tokens are strings and joins them.
|
||||
|
||||
|
|
@ -245,7 +245,7 @@ class EntropyScore:
|
|||
be preserved because:
|
||||
1. They're information-dense (can't be reconstructed)
|
||||
2. They're often identifiers (semantically important)
|
||||
3. LLMLingua may mangle them
|
||||
3. Token-level compressors may mangle them
|
||||
|
||||
This is a self-signal - no external classifier needed.
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@
|
|||
This is the main entry point for compression. It:
|
||||
1. Detects content type using Magika (ML)
|
||||
2. Extracts structure using appropriate handler
|
||||
3. Compresses non-structural content with LLMLingua
|
||||
3. Compresses non-structural content with Kompress
|
||||
4. Optionally stores original in CCR for retrieval
|
||||
|
||||
Usage:
|
||||
|
|
@ -51,7 +51,7 @@ class UniversalCompressorConfig:
|
|||
|
||||
Attributes:
|
||||
use_magika: Use ML-based detection (requires magika package).
|
||||
use_llmlingua: Use LLMLingua for content compression.
|
||||
use_kompress: Use Kompress for content compression.
|
||||
use_entropy_preservation: Preserve high-entropy tokens (UUIDs, etc.).
|
||||
entropy_threshold: Threshold for entropy-based preservation.
|
||||
min_content_length: Minimum content length to compress.
|
||||
|
|
@ -60,7 +60,7 @@ class UniversalCompressorConfig:
|
|||
"""
|
||||
|
||||
use_magika: bool = True
|
||||
use_llmlingua: bool = True
|
||||
use_kompress: bool = True
|
||||
use_entropy_preservation: bool = True
|
||||
entropy_threshold: float = 0.85
|
||||
min_content_length: int = 100
|
||||
|
|
@ -139,7 +139,7 @@ class UniversalCompressor:
|
|||
config: Compression configuration.
|
||||
handlers: Custom handlers for content types.
|
||||
compress_fn: Custom compression function. If None, uses
|
||||
LLMLingua when available, else simple truncation.
|
||||
Kompress when available, else simple truncation.
|
||||
"""
|
||||
self.config = config or UniversalCompressorConfig()
|
||||
|
||||
|
|
@ -165,18 +165,18 @@ class UniversalCompressor:
|
|||
def _get_default_compress_fn(self) -> Callable[[str], str]:
|
||||
"""Get default compression function.
|
||||
|
||||
Returns LLMLingua wrapper if available, else simple truncation.
|
||||
Returns Kompress wrapper if available, else simple truncation.
|
||||
"""
|
||||
if self.config.use_llmlingua:
|
||||
if self.config.use_kompress:
|
||||
try:
|
||||
return self._llmlingua_compress
|
||||
return self._kompress_compress
|
||||
except ImportError:
|
||||
logger.info("LLMLingua not available, using simple compression")
|
||||
logger.info("Kompress not available, using simple compression")
|
||||
|
||||
return self._simple_compress
|
||||
|
||||
def _llmlingua_compress(self, text: str) -> str:
|
||||
"""Compress using LLMLingua.
|
||||
def _kompress_compress(self, text: str) -> str:
|
||||
"""Compress using Kompress.
|
||||
|
||||
Args:
|
||||
text: Text to compress.
|
||||
|
|
@ -185,16 +185,15 @@ class UniversalCompressor:
|
|||
Compressed text.
|
||||
"""
|
||||
try:
|
||||
from headroom.transforms.llmlingua_compressor import compress_with_llmlingua
|
||||
from headroom.transforms.kompress_compressor import KompressCompressor
|
||||
|
||||
return compress_with_llmlingua(
|
||||
text,
|
||||
compression_rate=self.config.compression_ratio_target,
|
||||
)
|
||||
compressor = KompressCompressor()
|
||||
result = compressor.compress(text)
|
||||
return result.compressed
|
||||
except ImportError:
|
||||
return self._simple_compress(text)
|
||||
except Exception as e:
|
||||
logger.warning("LLMLingua compression failed: %s", e)
|
||||
logger.warning("Kompress compression failed: %s", e)
|
||||
return self._simple_compress(text)
|
||||
|
||||
def _simple_compress(self, text: str) -> str:
|
||||
|
|
|
|||
|
|
@ -619,7 +619,7 @@ class HeadroomConfig:
|
|||
prefix_freeze: PrefixFreezeConfig = field(default_factory=PrefixFreezeConfig)
|
||||
|
||||
# Content Router - intelligent content-type based compression
|
||||
# Routes content to appropriate compressor (LLMLingua for text, SmartCrusher for JSON,
|
||||
# Routes content to appropriate compressor (Kompress for text, SmartCrusher for JSON,
|
||||
# CodeCompressor for code, LogCompressor for logs, etc.)
|
||||
content_router_enabled: bool = True
|
||||
|
||||
|
|
|
|||
|
|
@ -353,7 +353,7 @@ def generate_factual_test_cases() -> list[BatchTestCase]:
|
|||
Context:
|
||||
The Headroom SDK is a context optimization layer for LLM applications. It was created
|
||||
by Anthropic in 2024. The main features include SmartCrusher for JSON compression,
|
||||
LLMLingua for text compression, and CCR (Compress-Cache-Retrieve) for reversible
|
||||
Kompress for text compression, and CCR (Compress-Cache-Retrieve) for reversible
|
||||
compression. The SDK supports Python 3.9+ and can save up to 70% of tokens on
|
||||
large JSON arrays.
|
||||
|
||||
|
|
@ -363,7 +363,7 @@ Question: What percentage of tokens can the SDK save on large JSON arrays?""",
|
|||
),
|
||||
ground_truth="70%",
|
||||
ground_truth_keywords=["70", "percent", "%"],
|
||||
context_facts=["70%", "SmartCrusher", "LLMLingua", "CCR", "Python 3.9"],
|
||||
context_facts=["70%", "SmartCrusher", "Kompress", "CCR", "Python 3.9"],
|
||||
),
|
||||
BatchTestCase(
|
||||
id="factual_002",
|
||||
|
|
|
|||
|
|
@ -4,7 +4,6 @@ This module evaluates whether HTMLExtractor preserves the information
|
|||
that LLMs need to answer questions about web content. We compare:
|
||||
1. LLM answers from original HTML
|
||||
2. LLM answers from HTMLExtractor output
|
||||
3. LLM answers from LLMLingua baseline (current fallback)
|
||||
|
||||
Uses LLM-as-judge to score answer quality on a 1-5 scale.
|
||||
"""
|
||||
|
|
@ -81,7 +80,7 @@ class HTMLEvalResult:
|
|||
# Answers from different methods
|
||||
answer_from_original: str
|
||||
answer_from_extracted: str
|
||||
answer_from_baseline: str | None = None # LLMLingua baseline
|
||||
answer_from_baseline: str | None = None # Baseline comparison
|
||||
|
||||
# Judge scores (1-5 scale)
|
||||
extracted_score: float = 0.0
|
||||
|
|
@ -211,7 +210,7 @@ class HTMLExtractionEvaluator:
|
|||
Args:
|
||||
answer_model: Model for generating answers from content.
|
||||
judge_model: Model for judging answer quality.
|
||||
compare_baseline: Whether to also test LLMLingua baseline.
|
||||
compare_baseline: Whether to also test Kompress baseline.
|
||||
provider: API provider ("openai", "anthropic", "litellm").
|
||||
"""
|
||||
self.answer_model = answer_model
|
||||
|
|
@ -221,7 +220,7 @@ class HTMLExtractionEvaluator:
|
|||
|
||||
# Lazy-loaded components
|
||||
self._extractor: HTMLExtractor | None = None
|
||||
self._llmlingua: Any = None
|
||||
self._kompress: Any = None
|
||||
self._judge_fn: Callable[[str, str, str], tuple[float, str]] | None = None
|
||||
self._answer_fn: Any = None
|
||||
|
||||
|
|
@ -235,16 +234,16 @@ class HTMLExtractionEvaluator:
|
|||
return self._extractor
|
||||
|
||||
@property
|
||||
def llmlingua(self) -> Any:
|
||||
"""Lazy-load LLMLingua compressor for baseline."""
|
||||
if self._llmlingua is None and self.compare_baseline:
|
||||
def kompress(self) -> Any:
|
||||
"""Lazy-load Kompress compressor for baseline."""
|
||||
if self._kompress is None and self.compare_baseline:
|
||||
try:
|
||||
from headroom.transforms.llmlingua_compressor import LLMLinguaCompressor
|
||||
from headroom.transforms.kompress_compressor import KompressCompressor
|
||||
|
||||
self._llmlingua = LLMLinguaCompressor()
|
||||
self._kompress = KompressCompressor()
|
||||
except ImportError:
|
||||
logger.warning("LLMLingua not available for baseline comparison")
|
||||
return self._llmlingua
|
||||
logger.warning("Kompress not available for baseline comparison")
|
||||
return self._kompress
|
||||
|
||||
@property
|
||||
def judge_fn(self) -> Callable[[str, str, str], tuple[float, str]]:
|
||||
|
|
@ -429,14 +428,14 @@ Answer concisely and factually based only on the content provided."""
|
|||
answer_from_extracted,
|
||||
)
|
||||
|
||||
# Optionally compare with LLMLingua baseline
|
||||
# Optionally compare with Kompress baseline
|
||||
baseline_answer = None
|
||||
baseline_score = None
|
||||
baseline_reasoning = None
|
||||
|
||||
if self.compare_baseline and self.llmlingua:
|
||||
if self.compare_baseline and self.kompress:
|
||||
try:
|
||||
baseline_result = self.llmlingua.compress(case.html)
|
||||
baseline_result = self.kompress.compress(case.html)
|
||||
baseline_content = baseline_result.compressed
|
||||
baseline_answer = self._get_answer(baseline_content, case.question)
|
||||
baseline_score, baseline_reasoning = self.judge_fn(
|
||||
|
|
|
|||
|
|
@ -12,7 +12,6 @@ Usage:
|
|||
# Or use environment variables to override at runtime:
|
||||
# HEADROOM_SENTENCE_TRANSFORMER=intfloat/e5-small-v2
|
||||
# HEADROOM_SIGLIP=google/siglip-base-patch16-224
|
||||
# HEADROOM_LLMLINGUA=microsoft/llmlingua-2-bert-base-multilingual-cased-meetingbank
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -31,7 +30,6 @@ class MLModelConfig:
|
|||
Environment variables can override any default:
|
||||
- HEADROOM_SENTENCE_TRANSFORMER
|
||||
- HEADROOM_SIGLIP
|
||||
- HEADROOM_LLMLINGUA
|
||||
- HEADROOM_SPACY
|
||||
- HEADROOM_TECHNIQUE_ROUTER
|
||||
|
||||
|
|
@ -47,10 +45,6 @@ class MLModelConfig:
|
|||
Default: google/siglip-base-patch16-224 (~400MB)
|
||||
Alternative: google/siglip-so400m-patch14-384 (larger, more accurate)
|
||||
|
||||
llmlingua: Model for ML-based prompt compression.
|
||||
Default: microsoft/llmlingua-2-bert-base-multilingual-cased-meetingbank (~350MB)
|
||||
Alternative: microsoft/llmlingua-2-xlm-roberta-large-meetingbank (~1GB, slightly more accurate)
|
||||
|
||||
spacy: Model for named entity recognition.
|
||||
Default: en_core_web_sm (~40MB)
|
||||
Alternative: en_core_web_md (~120MB, more accurate)
|
||||
|
|
@ -70,13 +64,6 @@ class MLModelConfig:
|
|||
default_factory=lambda: os.environ.get("HEADROOM_SIGLIP", "google/siglip-base-patch16-224")
|
||||
)
|
||||
|
||||
# Prompt Compression (LLMLingua-2)
|
||||
llmlingua: str = field(
|
||||
default_factory=lambda: os.environ.get(
|
||||
"HEADROOM_LLMLINGUA", "microsoft/llmlingua-2-bert-base-multilingual-cased-meetingbank"
|
||||
)
|
||||
)
|
||||
|
||||
# Named Entity Recognition (spaCy)
|
||||
spacy: str = field(default_factory=lambda: os.environ.get("HEADROOM_SPACY", "en_core_web_sm"))
|
||||
|
||||
|
|
@ -99,9 +86,6 @@ class MLModelConfig:
|
|||
"google/siglip-base-patch16-224": 400,
|
||||
"google/siglip-so400m-patch14-384": 900,
|
||||
"google/siglip-large-patch16-384": 1200,
|
||||
# LLMLingua
|
||||
"microsoft/llmlingua-2-xlm-roberta-large-meetingbank": 1000,
|
||||
"microsoft/llmlingua-2-bert-base-multilingual-cased-meetingbank": 350,
|
||||
# spaCy
|
||||
"en_core_web_sm": 40,
|
||||
"en_core_web_md": 120,
|
||||
|
|
@ -131,7 +115,6 @@ class MLModelConfig:
|
|||
return (
|
||||
self.get_memory_estimate(self.sentence_transformer)
|
||||
+ self.get_memory_estimate(self.siglip)
|
||||
+ self.get_memory_estimate(self.llmlingua)
|
||||
+ self.get_memory_estimate(self.spacy)
|
||||
+ self.get_memory_estimate(self.technique_router)
|
||||
)
|
||||
|
|
@ -160,8 +143,3 @@ def get_default_spacy_model() -> str:
|
|||
def get_default_siglip_model() -> str:
|
||||
"""Get the default SIGLIP model name."""
|
||||
return ML_MODEL_DEFAULTS.siglip
|
||||
|
||||
|
||||
def get_default_llmlingua_model() -> str:
|
||||
"""Get the default LLMLingua model name."""
|
||||
return ML_MODEL_DEFAULTS.llmlingua
|
||||
|
|
|
|||
|
|
@ -252,34 +252,6 @@ class MLModelRegistry:
|
|||
result: tuple[Any, Any] = instance._models[key]
|
||||
return result
|
||||
|
||||
# =========================================================================
|
||||
# LLMLingua (uses existing singleton pattern)
|
||||
# =========================================================================
|
||||
|
||||
@classmethod
|
||||
def get_llmlingua(cls, device: str | None = None, model_name: str | None = None) -> Any:
|
||||
"""Get the LLMLingua compressor.
|
||||
|
||||
Note: LLMLingua already has its own singleton in llmlingua_compressor.py.
|
||||
This method delegates to that implementation.
|
||||
|
||||
Args:
|
||||
device: Device to use. Auto-detected if None.
|
||||
model_name: Model name (default: microsoft/llmlingua-2-xlm-roberta-large-meetingbank).
|
||||
|
||||
Returns:
|
||||
PromptCompressor instance.
|
||||
"""
|
||||
from headroom.transforms.llmlingua_compressor import _get_llmlingua_compressor
|
||||
|
||||
if device is None:
|
||||
device = cls._detect_device()
|
||||
|
||||
if model_name is None:
|
||||
model_name = ML_MODEL_DEFAULTS.llmlingua
|
||||
|
||||
return _get_llmlingua_compressor(model_name=model_name, device=device)
|
||||
|
||||
# =========================================================================
|
||||
# Utility Methods
|
||||
# =========================================================================
|
||||
|
|
|
|||
|
|
@ -562,7 +562,7 @@ def _generate_recommendations(report: PerfReport) -> list[str]:
|
|||
if len(slow) > len(report.perf_records) * 0.2:
|
||||
recs.append(
|
||||
f"{len(slow)} requests took >500ms for optimization — "
|
||||
"consider disabling LLMLingua or reducing transform pipeline"
|
||||
"consider reducing transform pipeline"
|
||||
)
|
||||
|
||||
if report.router_records:
|
||||
|
|
|
|||
|
|
@ -93,7 +93,6 @@ from headroom.telemetry import get_telemetry_collector
|
|||
from headroom.telemetry.toin import get_toin
|
||||
from headroom.tokenizers import get_tokenizer
|
||||
from headroom.transforms import (
|
||||
_LLMLINGUA_AVAILABLE,
|
||||
CacheAligner,
|
||||
CodeAwareCompressor,
|
||||
CodeCompressorConfig,
|
||||
|
|
@ -127,10 +126,6 @@ def _get_image_compressor():
|
|||
return _image_compressor if _image_compressor else None
|
||||
|
||||
|
||||
# Conditionally import LLMLingua if available
|
||||
if _LLMLINGUA_AVAILABLE:
|
||||
from headroom.transforms import LLMLinguaCompressor, LLMLinguaConfig
|
||||
|
||||
# Try to import LiteLLM for pricing
|
||||
try:
|
||||
import litellm
|
||||
|
|
@ -648,11 +643,6 @@ class ProxyConfig:
|
|||
ccr_proactive_expansion: bool = True # Proactively expand based on query relevance
|
||||
ccr_max_proactive_expansions: int = 2 # Max contexts to proactively expand per turn
|
||||
|
||||
# LLMLingua ML-based compression (ON by default if installed)
|
||||
llmlingua_enabled: bool = True # Enable LLMLingua-2 for ML-based compression
|
||||
llmlingua_device: str = "auto" # Device: 'auto', 'cuda', 'cpu', 'mps'
|
||||
llmlingua_target_rate: float = 0.3 # Target compression rate (0.3 = keep 30%)
|
||||
|
||||
# Code-aware compression (ON by default if installed)
|
||||
code_aware_enabled: bool = True # Enable AST-based code compression
|
||||
|
||||
|
|
@ -1641,9 +1631,8 @@ class HeadroomProxy:
|
|||
|
||||
if config.smart_routing:
|
||||
# Smart routing: ContentRouter handles all content types intelligently
|
||||
# It lazy-loads compressors (including LLMLingua) only when needed
|
||||
# It lazy-loads compressors only when needed
|
||||
router_config = ContentRouterConfig(
|
||||
enable_llmlingua=config.llmlingua_enabled,
|
||||
enable_code_aware=config.code_aware_enabled,
|
||||
tool_profiles=config.tool_profiles,
|
||||
read_lifecycle=ReadLifecycleConfig(enabled=config.read_lifecycle),
|
||||
|
|
@ -1656,7 +1645,6 @@ class HeadroomProxy:
|
|||
ContentRouter(router_config),
|
||||
context_manager,
|
||||
]
|
||||
self._llmlingua_status = "lazy" if config.llmlingua_enabled else "disabled"
|
||||
self._code_aware_status = "lazy" if config.code_aware_enabled else "disabled"
|
||||
else:
|
||||
# Legacy mode: sequential pipeline
|
||||
|
|
@ -1675,8 +1663,6 @@ class HeadroomProxy:
|
|||
),
|
||||
context_manager,
|
||||
]
|
||||
# Add LLMLingua if enabled and available
|
||||
self._llmlingua_status = self._setup_llmlingua(config, transforms)
|
||||
# Add CodeAware if enabled and available
|
||||
self._code_aware_status = self._setup_code_aware(config, transforms)
|
||||
|
||||
|
|
@ -1893,38 +1879,6 @@ class HeadroomProxy:
|
|||
self._compression_caches[session_id] = CompressionCache()
|
||||
return self._compression_caches[session_id]
|
||||
|
||||
def _setup_llmlingua(self, config: ProxyConfig, transforms: list) -> str:
|
||||
"""Set up LLMLingua compression if enabled.
|
||||
|
||||
Args:
|
||||
config: Proxy configuration
|
||||
transforms: Transform list to append to
|
||||
|
||||
Returns:
|
||||
Status string for logging: 'enabled', 'disabled', 'available', 'unavailable'
|
||||
"""
|
||||
if config.llmlingua_enabled:
|
||||
if _LLMLINGUA_AVAILABLE:
|
||||
llmlingua_config = LLMLinguaConfig(
|
||||
device=config.llmlingua_device,
|
||||
target_compression_rate=config.llmlingua_target_rate,
|
||||
enable_ccr=config.ccr_inject_tool, # Link to CCR
|
||||
)
|
||||
# Insert before RollingWindow (which should be last)
|
||||
# LLMLingua works best on individual tool outputs before windowing
|
||||
transforms.insert(-1, LLMLinguaCompressor(llmlingua_config))
|
||||
return "enabled"
|
||||
else:
|
||||
logger.warning(
|
||||
"LLMLingua requested but not installed. "
|
||||
"Install with: pip install headroom-ai[llmlingua]"
|
||||
)
|
||||
return "unavailable"
|
||||
else:
|
||||
if _LLMLINGUA_AVAILABLE:
|
||||
return "available" # Available but not enabled - hint to user
|
||||
return "disabled"
|
||||
|
||||
def _setup_code_aware(self, config: ProxyConfig, transforms: list) -> str:
|
||||
"""Set up code-aware compression if enabled.
|
||||
|
||||
|
|
@ -2013,8 +1967,6 @@ class HeadroomProxy:
|
|||
# Update internal status from eager loading results
|
||||
if eager_status.get("kompress") == "enabled":
|
||||
self._kompress_status = "enabled"
|
||||
if eager_status.get("llmlingua") == "enabled":
|
||||
self._llmlingua_status = "enabled"
|
||||
if eager_status.get("code_aware") == "enabled":
|
||||
self._code_aware_status = "enabled"
|
||||
|
||||
|
|
@ -2024,18 +1976,6 @@ class HeadroomProxy:
|
|||
elif self.config.optimize:
|
||||
logger.info("Kompress: not installed (pip install headroom-ai[ml] for ML compression)")
|
||||
|
||||
if self._llmlingua_status == "enabled":
|
||||
logger.info(
|
||||
f"LLMLingua: ENABLED (device={self.config.llmlingua_device}, "
|
||||
f"rate={self.config.llmlingua_target_rate})"
|
||||
)
|
||||
elif self._kompress_status == "enabled":
|
||||
logger.info("LLMLingua: skipped (Kompress is active)")
|
||||
elif self._llmlingua_status == "lazy":
|
||||
logger.info("LLMLingua: LAZY (will load when prose content detected)")
|
||||
elif self._llmlingua_status == "disabled":
|
||||
logger.info("LLMLingua: DISABLED")
|
||||
|
||||
if self._code_aware_status == "enabled":
|
||||
logger.info("Code-Aware: ENABLED (AST-based compression)")
|
||||
if "tree_sitter" in eager_status:
|
||||
|
|
@ -7969,21 +7909,6 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
return app
|
||||
|
||||
|
||||
def _get_llmlingua_banner_status(config: ProxyConfig) -> str:
|
||||
"""Get LLMLingua status line for banner."""
|
||||
if config.llmlingua_enabled:
|
||||
if _LLMLINGUA_AVAILABLE:
|
||||
return (
|
||||
f"ENABLED (device={config.llmlingua_device}, rate={config.llmlingua_target_rate})"
|
||||
)
|
||||
else:
|
||||
return "NOT INSTALLED (pip install headroom-ai[llmlingua])"
|
||||
else:
|
||||
if _LLMLINGUA_AVAILABLE:
|
||||
return "DISABLED (remove --no-llmlingua to enable)"
|
||||
return "DISABLED"
|
||||
|
||||
|
||||
def _get_code_aware_banner_status(config: ProxyConfig) -> str:
|
||||
"""Get code-aware compression status line for banner."""
|
||||
if config.code_aware_enabled:
|
||||
|
|
@ -8016,7 +7941,6 @@ def run_server(
|
|||
config = config or ProxyConfig()
|
||||
app = create_app(config)
|
||||
|
||||
llmlingua_status = _get_llmlingua_banner_status(config)
|
||||
code_aware_status = _get_code_aware_banner_status(config)
|
||||
|
||||
# Format connection pool info
|
||||
|
|
@ -8055,7 +7979,6 @@ def run_server(
|
|||
║ Rate Limiting: {"ENABLED " if config.rate_limit_enabled else "DISABLED"} ({config.rate_limit_requests_per_minute} req/min, {config.rate_limit_tokens_per_minute:,} tok/min) ║
|
||||
║ Retry: {"ENABLED " if config.retry_enabled else "DISABLED"} (max {config.retry_max_attempts} attempts) ║
|
||||
║ Cost Tracking: {"ENABLED " if config.cost_tracking_enabled else "DISABLED"} (budget: {"$" + str(config.budget_limit_usd) + "/" + config.budget_period if config.budget_limit_usd else "unlimited"}) ║
|
||||
║ LLMLingua: {llmlingua_status:<52}║
|
||||
║ Code-Aware: {code_aware_status:<52}║
|
||||
║ HTTP/2: {http2_status:<52}║
|
||||
║ Conn Pool: {pool_info:<52}║
|
||||
|
|
@ -8269,30 +8192,6 @@ if __name__ == "__main__":
|
|||
help="Disable smart routing (use legacy sequential pipeline)",
|
||||
)
|
||||
|
||||
# LLMLingua ML-based compression
|
||||
parser.add_argument(
|
||||
"--llmlingua",
|
||||
action="store_true",
|
||||
help="Enable LLMLingua-2 ML-based compression (requires: pip install headroom-ai[llmlingua])",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-llmlingua",
|
||||
action="store_true",
|
||||
help="Disable LLMLingua compression",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--llmlingua-device",
|
||||
choices=["auto", "cuda", "cpu", "mps"],
|
||||
default="auto",
|
||||
help="Device for LLMLingua model (default: auto)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--llmlingua-rate",
|
||||
type=float,
|
||||
default=0.3,
|
||||
help="LLMLingua target compression rate, 0.0-1.0 (default: 0.3 = keep 30%%)",
|
||||
)
|
||||
|
||||
# Code-aware compression
|
||||
parser.add_argument(
|
||||
"--code-aware",
|
||||
|
|
@ -8310,7 +8209,6 @@ if __name__ == "__main__":
|
|||
# Environment variable defaults (HEADROOM_* prefix)
|
||||
# CLI args override env vars, env vars override ProxyConfig defaults
|
||||
env_smart_routing = _get_env_bool("HEADROOM_SMART_ROUTING", True)
|
||||
env_llmlingua = _get_env_bool("HEADROOM_LLMLINGUA_ENABLED", True)
|
||||
env_code_aware = _get_env_bool("HEADROOM_CODE_AWARE_ENABLED", True)
|
||||
env_optimize = _get_env_bool("HEADROOM_OPTIMIZE", True)
|
||||
env_cache = _get_env_bool("HEADROOM_CACHE_ENABLED", True)
|
||||
|
|
@ -8319,11 +8217,6 @@ if __name__ == "__main__":
|
|||
# Determine settings: CLI flags override env vars
|
||||
# --no-X explicitly disables, --X explicitly enables, neither uses env var
|
||||
smart_routing = env_smart_routing if not args.no_smart_routing else False
|
||||
llmlingua_enabled = (
|
||||
env_llmlingua
|
||||
if not (args.llmlingua or args.no_llmlingua)
|
||||
else (args.llmlingua or not args.no_llmlingua)
|
||||
)
|
||||
code_aware_enabled = (
|
||||
env_code_aware
|
||||
if not (args.code_aware or args.no_code_aware)
|
||||
|
|
@ -8364,9 +8257,6 @@ if __name__ == "__main__":
|
|||
else os.environ.get("HEADROOM_LOG_FILE"),
|
||||
log_full_messages=args.log_messages or _get_env_bool("HEADROOM_LOG_MESSAGES", False),
|
||||
smart_routing=smart_routing,
|
||||
llmlingua_enabled=llmlingua_enabled,
|
||||
llmlingua_device=_get_env_str("HEADROOM_LLMLINGUA_DEVICE", args.llmlingua_device),
|
||||
llmlingua_target_rate=_get_env_float("HEADROOM_LLMLINGUA_RATE", args.llmlingua_rate),
|
||||
code_aware_enabled=code_aware_enabled,
|
||||
# Connection pool settings
|
||||
max_connections=_get_env_int("HEADROOM_MAX_CONNECTIONS", args.max_connections),
|
||||
|
|
|
|||
|
|
@ -25,21 +25,6 @@ from .search_compressor import (
|
|||
from .smart_crusher import SmartCrusher, SmartCrusherConfig
|
||||
from .tool_crusher import ToolCrusher
|
||||
|
||||
# ML-based compression (optional dependency)
|
||||
try:
|
||||
from .llmlingua_compressor import ( # noqa: F401
|
||||
LLMLinguaCompressor,
|
||||
LLMLinguaConfig,
|
||||
LLMLinguaResult,
|
||||
compress_with_llmlingua,
|
||||
is_llmlingua_model_loaded,
|
||||
unload_llmlingua_model,
|
||||
)
|
||||
|
||||
_LLMLINGUA_AVAILABLE = True
|
||||
except ImportError:
|
||||
_LLMLINGUA_AVAILABLE = False
|
||||
|
||||
# HTML content extraction (optional dependency - requires trafilatura)
|
||||
try:
|
||||
from .html_extractor import ( # noqa: F401
|
||||
|
|
@ -122,25 +107,10 @@ __all__ = [
|
|||
"MessageScorer",
|
||||
"MessageScore",
|
||||
"EmbeddingProvider",
|
||||
# ML-based compression (optional)
|
||||
"_LLMLINGUA_AVAILABLE",
|
||||
# HTML extraction (optional)
|
||||
"_HTML_EXTRACTOR_AVAILABLE",
|
||||
]
|
||||
|
||||
# Conditionally add LLMLingua exports
|
||||
if _LLMLINGUA_AVAILABLE:
|
||||
__all__.extend(
|
||||
[
|
||||
"LLMLinguaCompressor",
|
||||
"LLMLinguaConfig",
|
||||
"LLMLinguaResult",
|
||||
"compress_with_llmlingua",
|
||||
"is_llmlingua_model_loaded",
|
||||
"unload_llmlingua_model",
|
||||
]
|
||||
)
|
||||
|
||||
# Conditionally add HTML extractor exports
|
||||
if _HTML_EXTRACTOR_AVAILABLE:
|
||||
__all__.extend(
|
||||
|
|
|
|||
|
|
@ -1,8 +1,8 @@
|
|||
"""Code-aware compressor using AST parsing for syntax-preserving compression.
|
||||
|
||||
This module provides AST-based compression for source code that guarantees
|
||||
valid syntax output. Unlike token-level compression (LLMLingua), this
|
||||
preserves structural elements while compressing function bodies.
|
||||
valid syntax output. Unlike token-level compression, this preserves
|
||||
structural elements while compressing function bodies.
|
||||
|
||||
Key Features:
|
||||
- Syntax validity guaranteed (output always parses)
|
||||
|
|
@ -329,7 +329,7 @@ class CodeCompressorConfig:
|
|||
compress_comments: Remove non-docstring comments.
|
||||
min_tokens_for_compression: Minimum tokens to trigger compression.
|
||||
language_hint: Explicit language (None = auto-detect).
|
||||
fallback_to_llmlingua: Use LLMLingua for unknown languages.
|
||||
fallback_to_kompress: Use Kompress for unknown languages.
|
||||
enable_ccr: Store originals for retrieval.
|
||||
ccr_ttl: TTL for CCR entries in seconds.
|
||||
"""
|
||||
|
|
@ -351,7 +351,7 @@ class CodeCompressorConfig:
|
|||
|
||||
# Language handling
|
||||
language_hint: str | None = None
|
||||
fallback_to_llmlingua: bool = True
|
||||
fallback_to_kompress: bool = True
|
||||
|
||||
# Semantic analysis (symbol importance scoring)
|
||||
semantic_analysis: bool = True
|
||||
|
|
@ -615,7 +615,7 @@ class CodeAwareCompressor(Transform):
|
|||
- Syntax validity guaranteed
|
||||
- Preserves imports, signatures, types
|
||||
- Better compression ratios for code (5-8x vs 3-5x)
|
||||
- Lower latency (~20-50ms vs 50-200ms for LLMLingua)
|
||||
- Lower latency (~20-50ms vs 50-200ms for token-level compressors)
|
||||
- Smaller memory footprint (~50MB vs ~1GB)
|
||||
- Thread-safe (no mutable instance state during compression)
|
||||
|
||||
|
|
@ -947,9 +947,9 @@ class CodeAwareCompressor(Transform):
|
|||
else:
|
||||
detected_lang, confidence = detect_language(code)
|
||||
|
||||
# If language unknown and fallback enabled, try LLMLingua
|
||||
# If language unknown and fallback enabled, try Kompress
|
||||
if detected_lang == CodeLanguage.UNKNOWN:
|
||||
if self.config.fallback_to_llmlingua:
|
||||
if self.config.fallback_to_kompress:
|
||||
return self._fallback_compress(code, original_tokens)
|
||||
else:
|
||||
return CodeCompressionResult(
|
||||
|
|
@ -966,7 +966,7 @@ class CodeAwareCompressor(Transform):
|
|||
# Check if tree-sitter is available
|
||||
if not _check_tree_sitter_available():
|
||||
logger.warning("tree-sitter not available. Install with: pip install headroom-ai[code]")
|
||||
if self.config.fallback_to_llmlingua:
|
||||
if self.config.fallback_to_kompress:
|
||||
return self._fallback_compress(code, original_tokens)
|
||||
return CodeCompressionResult(
|
||||
compressed=code,
|
||||
|
|
@ -1062,7 +1062,7 @@ class CodeAwareCompressor(Transform):
|
|||
|
||||
except Exception as e:
|
||||
logger.warning("AST compression failed: %s, falling back", e)
|
||||
if self.config.fallback_to_llmlingua:
|
||||
if self.config.fallback_to_kompress:
|
||||
return self._fallback_compress(code, original_tokens)
|
||||
return CodeCompressionResult(
|
||||
compressed=code,
|
||||
|
|
@ -1629,13 +1629,13 @@ class CodeAwareCompressor(Transform):
|
|||
return False
|
||||
|
||||
def _fallback_compress(self, code: str, original_tokens: int) -> CodeCompressionResult:
|
||||
"""Fall back to LLMLingua compression."""
|
||||
"""Fall back to Kompress compression."""
|
||||
try:
|
||||
from .llmlingua_compressor import LLMLinguaCompressor, _check_llmlingua_available
|
||||
from .kompress_compressor import KompressCompressor, is_kompress_available
|
||||
|
||||
if _check_llmlingua_available():
|
||||
compressor = LLMLinguaCompressor()
|
||||
result = compressor.compress(code, content_type="code")
|
||||
if is_kompress_available():
|
||||
compressor = KompressCompressor()
|
||||
result = compressor.compress(code)
|
||||
return CodeCompressionResult(
|
||||
compressed=result.compressed,
|
||||
original=code,
|
||||
|
|
@ -1644,7 +1644,7 @@ class CodeAwareCompressor(Transform):
|
|||
compression_ratio=result.compression_ratio,
|
||||
language=CodeLanguage.UNKNOWN,
|
||||
language_confidence=0.0,
|
||||
# LLMLingua does NOT guarantee syntax validity
|
||||
# Kompress does NOT guarantee syntax validity
|
||||
syntax_valid=False,
|
||||
)
|
||||
except ImportError:
|
||||
|
|
|
|||
|
|
@ -224,7 +224,7 @@ def _try_detect_html(content: str) -> DetectionResult | None:
|
|||
"""Try to detect HTML content.
|
||||
|
||||
HTML needs content extraction (removing scripts, styles, nav, etc.),
|
||||
not token-level compression like LLMLingua.
|
||||
not token-level compression like Kompress.
|
||||
"""
|
||||
# Check first 3000 chars for HTML indicators
|
||||
sample = content[:3000]
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ Supported Compressors:
|
|||
- SmartCrusher: JSON arrays
|
||||
- SearchCompressor: grep/ripgrep results
|
||||
- LogCompressor: Build/test output
|
||||
- LLMLinguaCompressor: Plain text (ML-based)
|
||||
- KompressCompressor: Plain text (ML-based)
|
||||
- Kompress: Plain text (ML-based, requires [ml] extra)
|
||||
|
||||
Routing Strategy:
|
||||
|
|
@ -251,7 +251,6 @@ class CompressionStrategy(Enum):
|
|||
SEARCH = "search"
|
||||
LOG = "log"
|
||||
KOMPRESS = "kompress"
|
||||
LLMLINGUA = "llmlingua"
|
||||
TEXT = "text"
|
||||
DIFF = "diff"
|
||||
HTML = "html"
|
||||
|
|
@ -360,12 +359,11 @@ class ContentRouterConfig:
|
|||
|
||||
Attributes:
|
||||
enable_code_aware: Enable AST-based code compression.
|
||||
enable_llmlingua: Enable ML-based text compression.
|
||||
enable_smart_crusher: Enable JSON array compression.
|
||||
enable_search_compressor: Enable search result compression.
|
||||
enable_log_compressor: Enable build/test log compression.
|
||||
enable_image_optimizer: Enable image token optimization.
|
||||
prefer_code_aware_for_code: Use CodeAware over LLMLingua for code.
|
||||
prefer_code_aware_for_code: Use CodeAware over Kompress for code.
|
||||
mixed_content_threshold: Min distinct types to consider "mixed".
|
||||
min_section_tokens: Minimum tokens for a section to compress.
|
||||
fallback_strategy: Strategy when no compressor matches.
|
||||
|
|
@ -376,8 +374,7 @@ class ContentRouterConfig:
|
|||
|
||||
# Enable/disable specific compressors
|
||||
enable_code_aware: bool = True
|
||||
enable_kompress: bool = True # Kompress: ModernBERT token compressor (preferred over LLMLingua)
|
||||
enable_llmlingua: bool = True
|
||||
enable_kompress: bool = True # Kompress: ModernBERT token compressor
|
||||
enable_smart_crusher: bool = True
|
||||
enable_search_compressor: bool = True
|
||||
enable_log_compressor: bool = True
|
||||
|
|
@ -647,7 +644,6 @@ class ContentRouter(Transform):
|
|||
self._diff_compressor: Any = None
|
||||
self._html_extractor: Any = None
|
||||
self._kompress: Any = None
|
||||
self._llmlingua: Any = None
|
||||
self._image_optimizer: Any = None
|
||||
|
||||
# TOIN integration for cross-strategy learning
|
||||
|
|
@ -813,7 +809,7 @@ class ContentRouter(Transform):
|
|||
strategy == CompressionStrategy.CODE_AWARE
|
||||
and not self.config.prefer_code_aware_for_code
|
||||
):
|
||||
strategy = CompressionStrategy.LLMLINGUA
|
||||
strategy = CompressionStrategy.KOMPRESS
|
||||
|
||||
return strategy
|
||||
|
||||
|
|
@ -960,9 +956,11 @@ class ContentRouter(Transform):
|
|||
result = compressor.compress(content, language=language, context=context)
|
||||
compressed, compressed_tokens = result.compressed, result.compressed_tokens
|
||||
if compressed is None:
|
||||
# Fallback to LLMLingua
|
||||
compressed, compressed_tokens = self._try_llmlingua(content, context, question)
|
||||
strategy = CompressionStrategy.LLMLINGUA # Update for TOIN
|
||||
# Fallback to Kompress
|
||||
compressed, compressed_tokens = self._try_ml_compressor(
|
||||
content, context, question
|
||||
)
|
||||
strategy = CompressionStrategy.KOMPRESS # Update for TOIN
|
||||
|
||||
elif strategy == CompressionStrategy.SMART_CRUSHER:
|
||||
# SmartCrusher handles its own TOIN recording
|
||||
|
|
@ -1013,12 +1011,9 @@ class ContentRouter(Transform):
|
|||
elif strategy == CompressionStrategy.KOMPRESS:
|
||||
compressed, compressed_tokens = self._try_ml_compressor(content, context, question)
|
||||
|
||||
elif strategy == CompressionStrategy.LLMLINGUA:
|
||||
compressed, compressed_tokens = self._try_ml_compressor(content, context, question)
|
||||
|
||||
elif strategy == CompressionStrategy.TEXT:
|
||||
# Prefer ML compressor (Kompress > LLMLingua) for text
|
||||
# Passes through unchanged if neither Kompress nor LLMLingua available
|
||||
# Prefer Kompress ML compressor for text
|
||||
# Passes through unchanged if Kompress not available
|
||||
compressed, compressed_tokens = self._try_ml_compressor(content, context, question)
|
||||
|
||||
except Exception as e:
|
||||
|
|
@ -1043,7 +1038,7 @@ class ContentRouter(Transform):
|
|||
def _try_ml_compressor(
|
||||
self, content: str, context: str, question: str | None = None
|
||||
) -> tuple[str, int]:
|
||||
"""ML-based compression: Kompress (primary), LLMLingua (fallback only).
|
||||
"""ML-based compression using Kompress.
|
||||
|
||||
Kompress (ModernBERT, trained on 330K structured tool outputs)
|
||||
auto-downloads from HuggingFace on first use. No heuristic fallback.
|
||||
|
|
@ -1090,19 +1085,6 @@ class ContentRouter(Transform):
|
|||
except Exception as e:
|
||||
logger.warning("Kompress failed: %s", e)
|
||||
|
||||
# Fallback: LLMLingua (only if Kompress not installed)
|
||||
if compressed is None and self.config.enable_llmlingua:
|
||||
compressor = self._get_llmlingua()
|
||||
if compressor:
|
||||
try:
|
||||
result = compressor.compress(
|
||||
text_to_compress, context=context, question=question
|
||||
)
|
||||
compressed = result.compressed
|
||||
compressed_tokens = result.compressed_tokens
|
||||
except Exception as e:
|
||||
logger.warning("LLMLingua failed: %s", e)
|
||||
|
||||
if compressed is None:
|
||||
return content, len(content.split())
|
||||
|
||||
|
|
@ -1113,9 +1095,6 @@ class ContentRouter(Transform):
|
|||
|
||||
return compressed, compressed_tokens or len(compressed.split())
|
||||
|
||||
# Backwards compatibility
|
||||
_try_llmlingua = _try_ml_compressor
|
||||
|
||||
def _strategy_from_detection_type(self, content_type: ContentType) -> CompressionStrategy:
|
||||
"""Get strategy from ContentType enum."""
|
||||
mapping = {
|
||||
|
|
@ -1140,7 +1119,6 @@ class ContentRouter(Transform):
|
|||
CompressionStrategy.HTML: ContentType.HTML,
|
||||
CompressionStrategy.TEXT: ContentType.PLAIN_TEXT,
|
||||
CompressionStrategy.KOMPRESS: ContentType.PLAIN_TEXT,
|
||||
CompressionStrategy.LLMLINGUA: ContentType.PLAIN_TEXT,
|
||||
CompressionStrategy.PASSTHROUGH: ContentType.PLAIN_TEXT,
|
||||
}
|
||||
return mapping.get(strategy, ContentType.PLAIN_TEXT)
|
||||
|
|
@ -1233,7 +1211,7 @@ class ContentRouter(Transform):
|
|||
"""
|
||||
status: dict[str, str] = {}
|
||||
|
||||
# 1. ML text compressor: Kompress or LLMLingua fallback
|
||||
# 1. ML text compressor: Kompress
|
||||
if self.config.enable_kompress:
|
||||
compressor = self._get_kompress()
|
||||
if compressor:
|
||||
|
|
@ -1241,20 +1219,6 @@ class ContentRouter(Transform):
|
|||
status["kompress"] = "enabled"
|
||||
else:
|
||||
status["kompress"] = "unavailable"
|
||||
if "kompress" not in status or status["kompress"] != "enabled":
|
||||
if self.config.enable_llmlingua:
|
||||
compressor = self._get_llmlingua()
|
||||
if compressor:
|
||||
try:
|
||||
from .llmlingua_compressor import _get_llmlingua_compressor
|
||||
|
||||
device = compressor._resolve_device()
|
||||
_get_llmlingua_compressor(compressor.config.model_name, device)
|
||||
logger.info("LLMLingua model pre-loaded at startup")
|
||||
status["llmlingua"] = "enabled"
|
||||
except Exception as e:
|
||||
logger.warning("Failed to pre-load LLMLingua model: %s", e)
|
||||
status["llmlingua"] = f"failed: {e}"
|
||||
|
||||
# 2. Magika content detector (avoids 100-200ms on first content detection)
|
||||
try:
|
||||
|
|
@ -1326,21 +1290,6 @@ class ContentRouter(Transform):
|
|||
logger.debug("Kompress dependencies not available")
|
||||
return self._kompress
|
||||
|
||||
def _get_llmlingua(self) -> Any:
|
||||
"""Get LLMLinguaCompressor (lazy load). Fallback if Kompress unavailable."""
|
||||
if self._llmlingua is None:
|
||||
try:
|
||||
from .llmlingua_compressor import (
|
||||
LLMLinguaCompressor,
|
||||
_check_llmlingua_available,
|
||||
)
|
||||
|
||||
if _check_llmlingua_available():
|
||||
self._llmlingua = LLMLinguaCompressor()
|
||||
except ImportError:
|
||||
logger.debug("LLMLinguaCompressor not available")
|
||||
return self._llmlingua
|
||||
|
||||
def _get_image_optimizer(self) -> Any:
|
||||
"""Get ImageCompressor (lazy load).
|
||||
|
||||
|
|
@ -1558,7 +1507,7 @@ class ContentRouter(Transform):
|
|||
"non_string": 0,
|
||||
"content_blocks": 0,
|
||||
}
|
||||
compressed_details: list[str] = [] # e.g. ["code_aware:0.72", "llmlingua:0.65"]
|
||||
compressed_details: list[str] = [] # e.g. ["code_aware:0.72", "kompress:0.65"]
|
||||
|
||||
# Check for analysis intent in the most recent user message
|
||||
analysis_intent = False
|
||||
|
|
|
|||
|
|
@ -501,7 +501,6 @@ class IntelligentContextManager(Transform):
|
|||
# Configure for aggressive compression in COMPRESS_FIRST context
|
||||
router_config = ContentRouterConfig(
|
||||
enable_code_aware=True,
|
||||
enable_llmlingua=True,
|
||||
enable_smart_crusher=True,
|
||||
enable_search_compressor=True,
|
||||
enable_log_compressor=True,
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
"""Kompress: ModernBERT token compressor for structured tool outputs.
|
||||
|
||||
Drop-in replacement for LLMLingua-2. Auto-downloads the model from
|
||||
HuggingFace (chopratejas/kompress-base) on first use.
|
||||
Auto-downloads the model from HuggingFace (chopratejas/kompress-base)
|
||||
on first use.
|
||||
|
||||
Requires the [ml] extra: pip install headroom-ai[ml]
|
||||
|
||||
|
|
@ -222,7 +222,6 @@ class KompressCompressor(Transform):
|
|||
"""Kompress: ModernBERT token compressor for structured tool outputs.
|
||||
|
||||
Auto-downloads chopratejas/kompress-base from HuggingFace on first use.
|
||||
Drop-in replacement for LLMLinguaCompressor with identical interface.
|
||||
"""
|
||||
|
||||
name: str = "kompress_compressor"
|
||||
|
|
@ -242,9 +241,9 @@ class KompressCompressor(Transform):
|
|||
|
||||
Args:
|
||||
content: Text to compress.
|
||||
context: Optional surrounding context (unused by model, kept for interface compat).
|
||||
context: Optional surrounding context (unused by model).
|
||||
content_type: Ignored — model decides importance per content type.
|
||||
question: Ignored — kept for LLMLingua interface compat.
|
||||
question: Ignored — reserved for future QA-aware compression.
|
||||
target_ratio: If None (default), model decides how much to keep using
|
||||
score threshold. If set (e.g. 0.3), forces that keep ratio.
|
||||
The proxy never sets this — only user-facing API does.
|
||||
|
|
|
|||
|
|
@ -1,652 +0,0 @@
|
|||
"""LLMLingua-2 compressor for ML-based prompt compression.
|
||||
|
||||
This module provides integration with LLMLingua-2, a BERT-based token classifier
|
||||
trained via GPT-4 distillation. It achieves superior compression (up to 20x)
|
||||
while maintaining high fidelity on tool outputs and structured content.
|
||||
|
||||
Key Features:
|
||||
- Token-level classification (keep/remove) using fine-tuned BERT
|
||||
- 3-6x faster than LLMLingua-1 with better results
|
||||
- Especially effective on tool outputs, code, and structured data
|
||||
- Reversible compression via CCR integration
|
||||
|
||||
Reference:
|
||||
LLMLingua-2: Data Distillation for Efficient and Faithful Task-Agnostic Prompt Compression
|
||||
https://arxiv.org/abs/2403.12968
|
||||
|
||||
Installation:
|
||||
pip install headroom-ai[llmlingua]
|
||||
|
||||
Usage:
|
||||
>>> from headroom.transforms import LLMLinguaCompressor
|
||||
>>> compressor = LLMLinguaCompressor()
|
||||
>>> result = compressor.compress(long_tool_output)
|
||||
>>> print(result.compressed) # Significantly reduced output
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import threading
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
from ..config import TransformResult
|
||||
from ..models.config import ML_MODEL_DEFAULTS
|
||||
from ..tokenizer import Tokenizer
|
||||
from .base import Transform
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Lazy import for optional dependency
|
||||
_llmlingua_available: bool | None = None
|
||||
_llmlingua_instance: Any = None
|
||||
_llmlingua_lock = threading.Lock() # Thread safety for model access
|
||||
|
||||
|
||||
def _check_llmlingua_available() -> bool:
|
||||
"""Check if llmlingua package is available."""
|
||||
global _llmlingua_available
|
||||
if _llmlingua_available is None:
|
||||
try:
|
||||
import llmlingua # noqa: F401
|
||||
|
||||
_llmlingua_available = True
|
||||
except ImportError:
|
||||
_llmlingua_available = False
|
||||
return _llmlingua_available
|
||||
|
||||
|
||||
def _get_llmlingua_compressor(model_name: str, device: str) -> Any:
|
||||
"""Get or create the LLMLingua compressor instance.
|
||||
|
||||
Uses lazy initialization and caches the instance to avoid repeated model loading.
|
||||
Thread-safe: uses lock to prevent race conditions during model initialization.
|
||||
|
||||
Args:
|
||||
model_name: HuggingFace model name for the compressor.
|
||||
device: Device to run the model on ('cuda', 'cpu', or 'auto').
|
||||
|
||||
Returns:
|
||||
PromptCompressor instance from llmlingua.
|
||||
|
||||
Raises:
|
||||
ImportError: If llmlingua is not installed.
|
||||
RuntimeError: If model loading fails.
|
||||
"""
|
||||
global _llmlingua_instance
|
||||
|
||||
if not _check_llmlingua_available():
|
||||
raise ImportError(
|
||||
"llmlingua is not installed. Install with: pip install headroom-ai[llmlingua]\n"
|
||||
"Note: This requires ~2GB of disk space and ~1GB RAM for the model."
|
||||
)
|
||||
|
||||
with _llmlingua_lock:
|
||||
# Double-check after acquiring lock
|
||||
if _llmlingua_instance is None or _llmlingua_instance._model_name != model_name:
|
||||
try:
|
||||
from llmlingua import PromptCompressor
|
||||
|
||||
logger.info(
|
||||
"Loading LLMLingua-2 model: %s on device: %s "
|
||||
"(this may take 10-30s on first run)",
|
||||
model_name,
|
||||
device,
|
||||
)
|
||||
_llmlingua_instance = PromptCompressor(
|
||||
model_name=model_name,
|
||||
device_map=device,
|
||||
use_llmlingua2=True, # Use LLMLingua-2 (BERT classifier)
|
||||
)
|
||||
# Store model name for later comparison
|
||||
_llmlingua_instance._model_name = model_name
|
||||
logger.info("LLMLingua-2 model loaded successfully")
|
||||
|
||||
except Exception as e:
|
||||
error_msg = str(e).lower()
|
||||
if "out of memory" in error_msg or "oom" in error_msg:
|
||||
raise RuntimeError(
|
||||
f"Out of memory loading LLMLingua model. Try:\n"
|
||||
f" 1. Use device='cpu' instead of 'cuda'\n"
|
||||
f" 2. Close other GPU applications\n"
|
||||
f" 3. Use a smaller model\n"
|
||||
f"Original error: {e}"
|
||||
) from e
|
||||
elif "not found" in error_msg or "404" in error_msg:
|
||||
raise RuntimeError(
|
||||
f"Model '{model_name}' not found on HuggingFace. Try:\n"
|
||||
f" 1. Check the model name is correct\n"
|
||||
f" 2. Use default: 'microsoft/llmlingua-2-xlm-roberta-large-meetingbank'\n"
|
||||
f"Original error: {e}"
|
||||
) from e
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"Failed to load LLMLingua model: {e}\n"
|
||||
f"Ensure you have sufficient disk space and memory."
|
||||
) from e
|
||||
|
||||
return _llmlingua_instance
|
||||
|
||||
|
||||
def unload_llmlingua_model() -> bool:
|
||||
"""Unload the LLMLingua model to free memory.
|
||||
|
||||
Use this when you're done with compression and want to reclaim GPU/CPU memory.
|
||||
The model will be reloaded automatically on the next compression call.
|
||||
|
||||
Returns:
|
||||
True if a model was unloaded, False if no model was loaded.
|
||||
|
||||
Example:
|
||||
>>> from headroom.transforms import LLMLinguaCompressor, unload_llmlingua_model
|
||||
>>> compressor = LLMLinguaCompressor()
|
||||
>>> result = compressor.compress(content) # Model loaded here
|
||||
>>> # ... do other work ...
|
||||
>>> unload_llmlingua_model() # Free ~1GB of memory
|
||||
"""
|
||||
global _llmlingua_instance
|
||||
|
||||
with _llmlingua_lock:
|
||||
if _llmlingua_instance is not None:
|
||||
model_name = getattr(_llmlingua_instance, "_model_name", "unknown")
|
||||
logger.info("Unloading LLMLingua model: %s", model_name)
|
||||
|
||||
# Clear the instance
|
||||
_llmlingua_instance = None
|
||||
|
||||
# Attempt to free GPU memory if torch is available
|
||||
try:
|
||||
import torch
|
||||
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
logger.debug("Cleared CUDA cache")
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def is_llmlingua_model_loaded() -> bool:
|
||||
"""Check if an LLMLingua model is currently loaded.
|
||||
|
||||
Returns:
|
||||
True if a model is loaded in memory, False otherwise.
|
||||
"""
|
||||
return _llmlingua_instance is not None
|
||||
|
||||
|
||||
@dataclass
|
||||
class LLMLinguaConfig:
|
||||
"""Configuration for LLMLingua-2 compression.
|
||||
|
||||
Attributes:
|
||||
model_name: HuggingFace model for the compressor. Default is the
|
||||
LLMLingua-2 xlm-roberta-large model fine-tuned for compression.
|
||||
device: Device to run on ('cuda', 'cpu', 'auto'). Auto will use CUDA if available.
|
||||
target_compression_rate: Target compression ratio (e.g., 0.3 = keep 30% of tokens).
|
||||
force_tokens: Tokens to always preserve (e.g., important keywords).
|
||||
drop_consecutive: Whether to drop consecutive punctuation/whitespace.
|
||||
min_tokens_for_compression: Minimum token count to trigger compression.
|
||||
Content below this threshold is passed through unchanged.
|
||||
enable_ccr: Whether to store originals in CCR for retrieval.
|
||||
ccr_ttl: TTL for CCR entries in seconds.
|
||||
|
||||
GOTCHA: Lower target_compression_rate = more aggressive compression.
|
||||
A rate of 0.2 means keeping only 20% of tokens.
|
||||
"""
|
||||
|
||||
# Model configuration
|
||||
model_name: str = field(default_factory=lambda: ML_MODEL_DEFAULTS.llmlingua)
|
||||
device: str = "auto"
|
||||
|
||||
# Compression parameters
|
||||
target_compression_rate: float = 0.3
|
||||
force_tokens: list[str] = field(default_factory=list)
|
||||
drop_consecutive: bool = True
|
||||
|
||||
# Thresholds
|
||||
min_tokens_for_compression: int = 100
|
||||
|
||||
# CCR integration
|
||||
enable_ccr: bool = True
|
||||
ccr_ttl: int = 300 # 5 minutes
|
||||
|
||||
# Content type specific settings
|
||||
code_compression_rate: float = 0.5 # Conservative for code
|
||||
json_compression_rate: float = 0.4 # Somewhat conservative for JSON
|
||||
text_compression_rate: float = 0.5 # Balanced for plain text (higher = more accurate)
|
||||
|
||||
|
||||
@dataclass
|
||||
class LLMLinguaResult:
|
||||
"""Result of LLMLingua-2 compression.
|
||||
|
||||
Attributes:
|
||||
compressed: Compressed content.
|
||||
original: Original content before compression.
|
||||
original_tokens: Token count of original content.
|
||||
compressed_tokens: Token count after compression.
|
||||
compression_ratio: Actual compression ratio achieved.
|
||||
cache_key: CCR cache key if stored.
|
||||
model_used: Model that performed the compression.
|
||||
tokens_saved: Number of tokens saved.
|
||||
"""
|
||||
|
||||
compressed: str
|
||||
original: str
|
||||
original_tokens: int
|
||||
compressed_tokens: int
|
||||
compression_ratio: float
|
||||
cache_key: str | None = None
|
||||
model_used: str | None = None
|
||||
|
||||
@property
|
||||
def tokens_saved(self) -> int:
|
||||
"""Number of tokens saved by compression."""
|
||||
return max(0, self.original_tokens - self.compressed_tokens)
|
||||
|
||||
@property
|
||||
def savings_percentage(self) -> float:
|
||||
"""Percentage of tokens saved."""
|
||||
if self.original_tokens == 0:
|
||||
return 0.0
|
||||
return (self.tokens_saved / self.original_tokens) * 100
|
||||
|
||||
|
||||
class LLMLinguaCompressor(Transform):
|
||||
"""LLMLingua-2 based prompt compressor.
|
||||
|
||||
Uses a BERT-based token classifier trained via GPT-4 distillation to
|
||||
identify and remove non-essential tokens while preserving semantic meaning.
|
||||
|
||||
Key advantages over statistical compression:
|
||||
- Learned token importance from LLM feedback
|
||||
- Better handling of context-dependent importance
|
||||
- More aggressive compression with less information loss
|
||||
- Especially effective on structured outputs (JSON, code, logs)
|
||||
|
||||
Example:
|
||||
>>> compressor = LLMLinguaCompressor()
|
||||
>>> result = compressor.compress(long_tool_output)
|
||||
>>> print(f"Saved {result.tokens_saved} tokens ({result.savings_percentage:.1f}%)")
|
||||
|
||||
>>> # Use as a Transform in pipeline
|
||||
>>> from headroom.transforms import TransformPipeline
|
||||
>>> pipeline = TransformPipeline([LLMLinguaCompressor()])
|
||||
>>> result = pipeline.apply(messages, tokenizer)
|
||||
"""
|
||||
|
||||
name: str = "llmlingua_compressor"
|
||||
|
||||
def __init__(self, config: LLMLinguaConfig | None = None):
|
||||
"""Initialize LLMLingua compressor.
|
||||
|
||||
Args:
|
||||
config: Compression configuration. If None, uses defaults.
|
||||
|
||||
Note:
|
||||
The underlying model is loaded lazily on first use to avoid
|
||||
startup overhead when the compressor isn't used.
|
||||
"""
|
||||
self.config = config or LLMLinguaConfig()
|
||||
self._compressor: Any = None # Lazy loaded
|
||||
|
||||
def compress(
|
||||
self,
|
||||
content: str,
|
||||
context: str = "",
|
||||
content_type: str | None = None,
|
||||
question: str | None = None,
|
||||
) -> LLMLinguaResult:
|
||||
"""Compress content using LLMLingua-2.
|
||||
|
||||
Args:
|
||||
content: Content to compress.
|
||||
context: Optional context for relevance-aware compression.
|
||||
content_type: Type of content ('code', 'json', 'text').
|
||||
If None, auto-detected.
|
||||
question: Optional question for QA-aware compression. When provided,
|
||||
LLMLingua preserves tokens relevant to answering this question.
|
||||
This significantly improves accuracy for QA tasks.
|
||||
|
||||
Returns:
|
||||
LLMLinguaResult with compressed content and metadata.
|
||||
|
||||
Raises:
|
||||
ImportError: If llmlingua is not installed.
|
||||
"""
|
||||
# Check availability
|
||||
if not _check_llmlingua_available():
|
||||
logger.warning(
|
||||
"LLMLingua not available. Install with: pip install headroom-ai[llmlingua]"
|
||||
)
|
||||
return LLMLinguaResult(
|
||||
compressed=content,
|
||||
original=content,
|
||||
original_tokens=len(content.split()), # Rough estimate
|
||||
compressed_tokens=len(content.split()),
|
||||
compression_ratio=1.0,
|
||||
)
|
||||
|
||||
# Estimate token count (rough)
|
||||
estimated_tokens = len(content.split())
|
||||
|
||||
# Skip compression for small content
|
||||
if estimated_tokens < self.config.min_tokens_for_compression:
|
||||
return LLMLinguaResult(
|
||||
compressed=content,
|
||||
original=content,
|
||||
original_tokens=estimated_tokens,
|
||||
compressed_tokens=estimated_tokens,
|
||||
compression_ratio=1.0,
|
||||
)
|
||||
|
||||
# Get compression rate based on content type
|
||||
compression_rate = self._get_compression_rate(content, content_type)
|
||||
|
||||
# Get or initialize compressor
|
||||
device = self._resolve_device()
|
||||
compressor = _get_llmlingua_compressor(self.config.model_name, device)
|
||||
|
||||
# Prepare force tokens
|
||||
force_tokens = list(self.config.force_tokens)
|
||||
|
||||
# Add context words as force tokens if provided
|
||||
if context:
|
||||
context_words = [w for w in context.split() if len(w) > 3]
|
||||
force_tokens.extend(context_words[:10]) # Limit to avoid overhead
|
||||
|
||||
# Perform compression
|
||||
try:
|
||||
# Build compress_prompt kwargs
|
||||
compress_kwargs: dict[str, Any] = {
|
||||
"context": [content], # LLMLingua expects a list of context strings
|
||||
"rate": compression_rate,
|
||||
"force_tokens": force_tokens if force_tokens else [],
|
||||
"drop_consecutive": self.config.drop_consecutive,
|
||||
}
|
||||
|
||||
# Add question for QA-aware token selection (LLMLingua-2 feature)
|
||||
# This enables relevance-aware compression where tokens relevant
|
||||
# to answering the question are preserved with higher probability
|
||||
if question:
|
||||
compress_kwargs["question"] = question
|
||||
|
||||
result = compressor.compress_prompt(**compress_kwargs)
|
||||
|
||||
compressed = result.get("compressed_prompt", content)
|
||||
original_tokens = result.get("origin_tokens", estimated_tokens)
|
||||
compressed_tokens = result.get("compressed_tokens", len(compressed.split()))
|
||||
|
||||
except Exception as e:
|
||||
logger.warning("LLMLingua compression failed: %s", e)
|
||||
return LLMLinguaResult(
|
||||
compressed=content,
|
||||
original=content,
|
||||
original_tokens=estimated_tokens,
|
||||
compressed_tokens=estimated_tokens,
|
||||
compression_ratio=1.0,
|
||||
)
|
||||
|
||||
# Calculate actual ratio
|
||||
ratio = compressed_tokens / max(original_tokens, 1)
|
||||
|
||||
# Store in CCR if enabled
|
||||
cache_key = None
|
||||
if self.config.enable_ccr and ratio < 0.8:
|
||||
cache_key = self._store_in_ccr(content, compressed, original_tokens)
|
||||
if cache_key:
|
||||
# Use standard CCR marker format for CCRToolInjector detection
|
||||
compressed += f"\n[{original_tokens} items compressed to {compressed_tokens}. Retrieve more: hash={cache_key}]"
|
||||
|
||||
return LLMLinguaResult(
|
||||
compressed=compressed,
|
||||
original=content,
|
||||
original_tokens=original_tokens,
|
||||
compressed_tokens=compressed_tokens,
|
||||
compression_ratio=ratio,
|
||||
cache_key=cache_key,
|
||||
model_used=self.config.model_name,
|
||||
)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
messages: list[dict[str, Any]],
|
||||
tokenizer: Tokenizer,
|
||||
**kwargs: Any,
|
||||
) -> TransformResult:
|
||||
"""Apply LLMLingua compression to messages.
|
||||
|
||||
This method implements the Transform interface for use in pipelines.
|
||||
It compresses tool outputs and long assistant/user messages.
|
||||
|
||||
Args:
|
||||
messages: List of message dicts to transform.
|
||||
tokenizer: Tokenizer for accurate token counting.
|
||||
**kwargs: Additional arguments (e.g., 'context' for relevance).
|
||||
|
||||
Returns:
|
||||
TransformResult with compressed messages and metadata.
|
||||
"""
|
||||
tokens_before = sum(tokenizer.count_text(str(m.get("content", ""))) for m in messages)
|
||||
context = kwargs.get("context", "")
|
||||
|
||||
transformed_messages = []
|
||||
transforms_applied = []
|
||||
warnings: list[str] = []
|
||||
|
||||
for message in messages:
|
||||
role = message.get("role", "")
|
||||
content = message.get("content", "")
|
||||
|
||||
# Skip non-string content (multimodal messages with images)
|
||||
if not isinstance(content, str):
|
||||
transformed_messages.append(message)
|
||||
continue
|
||||
|
||||
# Compress tool results (highest value compression)
|
||||
if role == "tool" and content:
|
||||
result = self.compress(content, context=context, content_type="json")
|
||||
if result.compression_ratio < 0.9:
|
||||
transformed_messages.append({**message, "content": result.compressed})
|
||||
transforms_applied.append(f"llmlingua:tool:{result.compression_ratio:.2f}")
|
||||
else:
|
||||
transformed_messages.append(message)
|
||||
|
||||
# Compress long assistant messages (tool outputs often embedded)
|
||||
elif role == "assistant" and len(content) > 500:
|
||||
result = self.compress(content, context=context)
|
||||
if result.compression_ratio < 0.9:
|
||||
transformed_messages.append({**message, "content": result.compressed})
|
||||
transforms_applied.append(f"llmlingua:assistant:{result.compression_ratio:.2f}")
|
||||
else:
|
||||
transformed_messages.append(message)
|
||||
|
||||
# Pass through other messages
|
||||
else:
|
||||
transformed_messages.append(message)
|
||||
|
||||
tokens_after = sum(
|
||||
tokenizer.count_text(str(m.get("content", ""))) for m in transformed_messages
|
||||
)
|
||||
|
||||
# Add warning if llmlingua not available
|
||||
if not _check_llmlingua_available():
|
||||
warnings.append(
|
||||
"LLMLingua not installed. Install with: pip install headroom-ai[llmlingua]"
|
||||
)
|
||||
|
||||
return TransformResult(
|
||||
messages=transformed_messages,
|
||||
tokens_before=tokens_before,
|
||||
tokens_after=tokens_after,
|
||||
transforms_applied=transforms_applied if transforms_applied else ["llmlingua:noop"],
|
||||
warnings=warnings,
|
||||
)
|
||||
|
||||
def should_apply(
|
||||
self,
|
||||
messages: list[dict[str, Any]],
|
||||
tokenizer: Tokenizer,
|
||||
**kwargs: Any,
|
||||
) -> bool:
|
||||
"""Check if LLMLingua compression should be applied.
|
||||
|
||||
Returns True if:
|
||||
- LLMLingua is available, AND
|
||||
- Total token count exceeds minimum threshold
|
||||
|
||||
Args:
|
||||
messages: Messages to check.
|
||||
tokenizer: Tokenizer for counting.
|
||||
**kwargs: Additional arguments.
|
||||
|
||||
Returns:
|
||||
True if compression should be applied.
|
||||
"""
|
||||
if not _check_llmlingua_available():
|
||||
return False
|
||||
|
||||
total_tokens = sum(tokenizer.count_text(str(m.get("content", ""))) for m in messages)
|
||||
return total_tokens >= self.config.min_tokens_for_compression
|
||||
|
||||
def _get_compression_rate(
|
||||
self,
|
||||
content: str,
|
||||
content_type: str | None,
|
||||
) -> float:
|
||||
"""Get appropriate compression rate based on content type.
|
||||
|
||||
Args:
|
||||
content: Content to analyze.
|
||||
content_type: Explicit content type or None for auto-detection.
|
||||
|
||||
Returns:
|
||||
Target compression rate for this content.
|
||||
"""
|
||||
if content_type == "code":
|
||||
return self.config.code_compression_rate
|
||||
elif content_type == "json":
|
||||
return self.config.json_compression_rate
|
||||
elif content_type == "text":
|
||||
return self.config.text_compression_rate
|
||||
|
||||
# Auto-detect content type
|
||||
if self._looks_like_json(content):
|
||||
return self.config.json_compression_rate
|
||||
elif self._looks_like_code(content):
|
||||
return self.config.code_compression_rate
|
||||
else:
|
||||
return self.config.text_compression_rate
|
||||
|
||||
def _looks_like_json(self, content: str) -> bool:
|
||||
"""Check if content appears to be JSON."""
|
||||
stripped = content.strip()
|
||||
return (stripped.startswith("{") and stripped.endswith("}")) or (
|
||||
stripped.startswith("[") and stripped.endswith("]")
|
||||
)
|
||||
|
||||
def _looks_like_code(self, content: str) -> bool:
|
||||
"""Check if content appears to be code."""
|
||||
code_indicators = [
|
||||
"def ",
|
||||
"class ",
|
||||
"function ",
|
||||
"import ",
|
||||
"from ",
|
||||
"const ",
|
||||
"let ",
|
||||
"var ",
|
||||
"public ",
|
||||
"private ",
|
||||
"async ",
|
||||
"await ",
|
||||
"return ",
|
||||
"if (",
|
||||
"for (",
|
||||
"while (",
|
||||
]
|
||||
return any(indicator in content for indicator in code_indicators)
|
||||
|
||||
def _resolve_device(self) -> str:
|
||||
"""Resolve 'auto' device to actual device."""
|
||||
if self.config.device != "auto":
|
||||
return self.config.device
|
||||
|
||||
try:
|
||||
import torch
|
||||
|
||||
if torch.cuda.is_available():
|
||||
return "cuda"
|
||||
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
||||
return "mps"
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
return "cpu"
|
||||
|
||||
def _store_in_ccr(
|
||||
self,
|
||||
original: str,
|
||||
compressed: str,
|
||||
original_tokens: int,
|
||||
) -> str | None:
|
||||
"""Store original content in CCR for later retrieval.
|
||||
|
||||
Args:
|
||||
original: Original content before compression.
|
||||
compressed: Compressed content.
|
||||
original_tokens: Token count of original.
|
||||
|
||||
Returns:
|
||||
Cache key if stored successfully, None otherwise.
|
||||
"""
|
||||
try:
|
||||
from ..cache.compression_store import get_compression_store
|
||||
|
||||
store = get_compression_store()
|
||||
return store.store(
|
||||
original,
|
||||
compressed,
|
||||
original_tokens=original_tokens,
|
||||
compressed_tokens=len(compressed.split()),
|
||||
compression_strategy="llmlingua2",
|
||||
)
|
||||
except ImportError:
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.debug("CCR storage failed: %s", e)
|
||||
return None
|
||||
|
||||
|
||||
def compress_with_llmlingua(
|
||||
content: str,
|
||||
compression_rate: float = 0.3,
|
||||
context: str = "",
|
||||
model_name: str | None = None,
|
||||
) -> str:
|
||||
"""Convenience function for one-off compression.
|
||||
|
||||
Args:
|
||||
content: Content to compress.
|
||||
compression_rate: Target compression rate (0.0-1.0).
|
||||
context: Optional context for relevance-aware compression.
|
||||
model_name: Optional model name override.
|
||||
|
||||
Returns:
|
||||
Compressed content string.
|
||||
|
||||
Example:
|
||||
>>> compressed = compress_with_llmlingua(long_output, compression_rate=0.2)
|
||||
"""
|
||||
config = LLMLinguaConfig(target_compression_rate=compression_rate)
|
||||
if model_name:
|
||||
config.model_name = model_name
|
||||
|
||||
compressor = LLMLinguaCompressor(config)
|
||||
result = compressor.compress(content, context=context)
|
||||
return result.compressed
|
||||
|
|
@ -40,7 +40,7 @@ class TransformPipeline:
|
|||
Transform order:
|
||||
1. Cache Aligner - normalize prefix for cache hits
|
||||
2. Content Router - intelligent content-aware compression (routes to appropriate
|
||||
compressor: LLMLingua for text, SmartCrusher for JSON, CodeCompressor for code, etc.)
|
||||
compressor: Kompress for text, SmartCrusher for JSON, CodeCompressor for code, etc.)
|
||||
3. SmartCrusher/ToolCrusher - fallback if ContentRouter disabled
|
||||
4. IntelligentContextManager/RollingWindow - enforce token limits
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
"""Protect workflow/custom XML tags from text compression.
|
||||
|
||||
LLM workflows use XML-style tags (<system-reminder>, <tool_call>, <thinking>)
|
||||
as structural markers. Text compressors (Kompress, LLMLingua) treat these as
|
||||
as structural markers. Text compressors (Kompress) treat these as
|
||||
droppable noise and silently remove them, breaking downstream tools.
|
||||
|
||||
This module detects custom tags (anything NOT standard HTML), replaces entire
|
||||
|
|
|
|||
|
|
@ -413,7 +413,7 @@ class TestAlternativeMarkerFormats:
|
|||
- TextCompressor: [N lines compressed to M. Retrieve more: hash=xxx]
|
||||
- LogCompressor: [N lines compressed to M. Retrieve more: hash=xxx]
|
||||
- SearchCompressor: [N matches compressed to M. Retrieve more: hash=xxx]
|
||||
- LLMLingua: [N items compressed to M. Retrieve more: hash=xxx]
|
||||
- Kompress: [N items compressed to M. Retrieve more: hash=xxx]
|
||||
|
||||
The CCRToolInjector should detect all these formats.
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -755,10 +755,10 @@ class TestJSONAPIResponseEval:
|
|||
|
||||
@pytest.fixture
|
||||
def compressor(self):
|
||||
"""Create compressor with simple compression (no LLMLingua for tests)."""
|
||||
"""Create compressor with simple compression (no Kompress for tests)."""
|
||||
config = UniversalCompressorConfig(
|
||||
use_magika=False, # Use fallback for consistent tests
|
||||
use_llmlingua=False,
|
||||
use_kompress=False,
|
||||
ccr_enabled=False,
|
||||
)
|
||||
return UniversalCompressor(config=config)
|
||||
|
|
@ -885,7 +885,7 @@ class TestCodeFileEval:
|
|||
"""Create compressor."""
|
||||
config = UniversalCompressorConfig(
|
||||
use_magika=False,
|
||||
use_llmlingua=False,
|
||||
use_kompress=False,
|
||||
ccr_enabled=False,
|
||||
)
|
||||
return UniversalCompressor(config=config)
|
||||
|
|
@ -969,7 +969,7 @@ class TestLogOutputEval:
|
|||
"""Create compressor."""
|
||||
config = UniversalCompressorConfig(
|
||||
use_magika=False,
|
||||
use_llmlingua=False,
|
||||
use_kompress=False,
|
||||
ccr_enabled=False,
|
||||
)
|
||||
return UniversalCompressor(config=config)
|
||||
|
|
@ -1031,7 +1031,7 @@ class TestMultiToolAgentScenario:
|
|||
"""Create compressor."""
|
||||
config = UniversalCompressorConfig(
|
||||
use_magika=False,
|
||||
use_llmlingua=False,
|
||||
use_kompress=False,
|
||||
ccr_enabled=False,
|
||||
)
|
||||
return UniversalCompressor(config=config)
|
||||
|
|
@ -1129,7 +1129,7 @@ class TestCompressionQualityMetrics:
|
|||
"""Create compressor."""
|
||||
config = UniversalCompressorConfig(
|
||||
use_magika=False,
|
||||
use_llmlingua=False,
|
||||
use_kompress=False,
|
||||
ccr_enabled=False,
|
||||
compression_ratio_target=0.3, # Target 70% reduction
|
||||
)
|
||||
|
|
|
|||
|
|
@ -301,7 +301,7 @@ class TestJSONDiscoverability:
|
|||
"""Create compressor."""
|
||||
config = UniversalCompressorConfig(
|
||||
use_magika=False,
|
||||
use_llmlingua=False,
|
||||
use_kompress=False,
|
||||
ccr_enabled=False,
|
||||
)
|
||||
return UniversalCompressor(config=config)
|
||||
|
|
@ -418,7 +418,7 @@ class TestCodeUnderstanding:
|
|||
"""Create compressor."""
|
||||
config = UniversalCompressorConfig(
|
||||
use_magika=False,
|
||||
use_llmlingua=False,
|
||||
use_kompress=False,
|
||||
ccr_enabled=False,
|
||||
)
|
||||
return UniversalCompressor(config=config)
|
||||
|
|
@ -528,7 +528,7 @@ class TestMultiContentAgent:
|
|||
"""Create compressor."""
|
||||
config = UniversalCompressorConfig(
|
||||
use_magika=False,
|
||||
use_llmlingua=False,
|
||||
use_kompress=False,
|
||||
ccr_enabled=False,
|
||||
)
|
||||
return UniversalCompressor(config=config)
|
||||
|
|
@ -591,7 +591,7 @@ class TestCompressionEfficacy:
|
|||
"""Create compressor."""
|
||||
config = UniversalCompressorConfig(
|
||||
use_magika=False,
|
||||
use_llmlingua=False,
|
||||
use_kompress=False,
|
||||
ccr_enabled=False,
|
||||
)
|
||||
return UniversalCompressor(config=config)
|
||||
|
|
|
|||
|
|
@ -22,7 +22,7 @@ class TestUniversalCompressorConfig:
|
|||
config = UniversalCompressorConfig()
|
||||
|
||||
assert config.use_magika is True
|
||||
assert config.use_llmlingua is True
|
||||
assert config.use_kompress is True
|
||||
assert config.use_entropy_preservation is True
|
||||
assert config.entropy_threshold == 0.85
|
||||
assert config.min_content_length == 100
|
||||
|
|
@ -99,7 +99,7 @@ class TestUniversalCompressor:
|
|||
"""Create compressor with fallback detector (no Magika required)."""
|
||||
config = UniversalCompressorConfig(
|
||||
use_magika=False, # Use fallback detector
|
||||
use_llmlingua=False, # Use simple compression
|
||||
use_kompress=False, # Use simple compression
|
||||
ccr_enabled=False, # Skip CCR
|
||||
)
|
||||
return UniversalCompressor(config=config)
|
||||
|
|
@ -215,7 +215,7 @@ class TestUniversalCompressorBatch:
|
|||
"""Create compressor with fallback detector."""
|
||||
config = UniversalCompressorConfig(
|
||||
use_magika=False,
|
||||
use_llmlingua=False,
|
||||
use_kompress=False,
|
||||
ccr_enabled=False,
|
||||
)
|
||||
return UniversalCompressor(config=config)
|
||||
|
|
@ -262,7 +262,7 @@ class TestStructurePreservation:
|
|||
"""Create compressor."""
|
||||
config = UniversalCompressorConfig(
|
||||
use_magika=False,
|
||||
use_llmlingua=False,
|
||||
use_kompress=False,
|
||||
ccr_enabled=False,
|
||||
)
|
||||
return UniversalCompressor(config=config)
|
||||
|
|
|
|||
|
|
@ -262,7 +262,7 @@ class TestHTMLExtractionWithLLM:
|
|||
|
||||
@pytest.mark.skipif(not os.environ.get("OPENAI_API_KEY"), reason="OPENAI_API_KEY not set")
|
||||
class TestHTMLvsBaseline:
|
||||
"""Tests comparing HTMLExtractor vs LLMLingua baseline."""
|
||||
"""Tests comparing HTMLExtractor vs Kompress baseline."""
|
||||
|
||||
@pytest.fixture
|
||||
def evaluator_with_baseline(self):
|
||||
|
|
@ -274,9 +274,9 @@ class TestHTMLvsBaseline:
|
|||
provider="openai",
|
||||
)
|
||||
|
||||
@pytest.mark.skipif(True, reason="LLMLingua requires GPU, skip in CI")
|
||||
@pytest.mark.skipif(True, reason="Kompress requires GPU, skip in CI")
|
||||
def test_extraction_beats_baseline(self, evaluator_with_baseline):
|
||||
"""Test that HTMLExtractor outperforms LLMLingua on HTML."""
|
||||
"""Test that HTMLExtractor outperforms Kompress on HTML."""
|
||||
cases = get_sample_eval_cases()[:2] # Just test 2 for speed
|
||||
|
||||
results = evaluator_with_baseline.evaluate(cases)
|
||||
|
|
@ -286,9 +286,9 @@ class TestHTMLvsBaseline:
|
|||
print(f"Avg extraction score: {results.avg_extraction_score}/5")
|
||||
print(f"Avg baseline score: {results.avg_baseline_score}/5")
|
||||
|
||||
# HTMLExtractor should beat LLMLingua on HTML content
|
||||
# HTMLExtractor should beat Kompress on HTML content
|
||||
assert results.avg_extraction_score >= results.avg_baseline_score, (
|
||||
"HTMLExtractor should perform at least as well as LLMLingua on HTML"
|
||||
"HTMLExtractor should perform at least as well as Kompress on HTML"
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,468 +0,0 @@
|
|||
"""Tests for LLMLingua opt-in mechanism in the proxy server.
|
||||
|
||||
These tests verify:
|
||||
- ProxyConfig llmlingua settings
|
||||
- LLMLingua transform integration in pipeline
|
||||
- Status detection and logging hints
|
||||
- CLI flag parsing
|
||||
- DevEx: helpful messages when llmlingua unavailable
|
||||
"""
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
# Skip if fastapi not available
|
||||
pytest.importorskip("fastapi")
|
||||
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from headroom.proxy.server import (
|
||||
HeadroomProxy,
|
||||
ProxyConfig,
|
||||
_get_llmlingua_banner_status,
|
||||
create_app,
|
||||
)
|
||||
from headroom.transforms import _LLMLINGUA_AVAILABLE
|
||||
|
||||
# =============================================================================
|
||||
# Test Fixtures
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def base_config():
|
||||
"""Base config with optimization disabled for simpler tests."""
|
||||
return ProxyConfig(
|
||||
optimize=False,
|
||||
cache_enabled=False,
|
||||
rate_limit_enabled=False,
|
||||
cost_tracking_enabled=False,
|
||||
llmlingua_enabled=False, # Explicitly disable for base config
|
||||
smart_routing=False, # Use legacy mode for these tests
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def llmlingua_config():
|
||||
"""Config with LLMLingua enabled (legacy mode for explicit status testing)."""
|
||||
return ProxyConfig(
|
||||
optimize=True,
|
||||
cache_enabled=False,
|
||||
rate_limit_enabled=False,
|
||||
cost_tracking_enabled=False,
|
||||
llmlingua_enabled=True,
|
||||
llmlingua_device="cpu",
|
||||
llmlingua_target_rate=0.4,
|
||||
smart_routing=False, # Use legacy mode for explicit status testing
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client(base_config):
|
||||
"""Create test client with base config."""
|
||||
app = create_app(base_config)
|
||||
with TestClient(app) as client:
|
||||
yield client
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TestProxyConfigLLMLingua
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestProxyConfigLLMLingua:
|
||||
"""Tests for LLMLingua settings in ProxyConfig."""
|
||||
|
||||
def test_default_llmlingua_enabled(self):
|
||||
"""LLMLingua is enabled by default (with smart routing)."""
|
||||
config = ProxyConfig()
|
||||
|
||||
# LLMLingua is now enabled by default with smart routing
|
||||
assert config.llmlingua_enabled is True
|
||||
assert config.llmlingua_device == "auto"
|
||||
assert config.llmlingua_target_rate == 0.3
|
||||
|
||||
def test_llmlingua_can_be_enabled(self):
|
||||
"""LLMLingua can be enabled via config."""
|
||||
config = ProxyConfig(
|
||||
llmlingua_enabled=True,
|
||||
llmlingua_device="cuda",
|
||||
llmlingua_target_rate=0.5,
|
||||
)
|
||||
|
||||
assert config.llmlingua_enabled is True
|
||||
assert config.llmlingua_device == "cuda"
|
||||
assert config.llmlingua_target_rate == 0.5
|
||||
|
||||
def test_llmlingua_device_options(self):
|
||||
"""LLMLingua device accepts valid options."""
|
||||
for device in ["auto", "cuda", "cpu", "mps"]:
|
||||
config = ProxyConfig(llmlingua_device=device)
|
||||
assert config.llmlingua_device == device
|
||||
|
||||
def test_llmlingua_target_rate_range(self):
|
||||
"""LLMLingua target rate accepts 0.0-1.0 range."""
|
||||
# Low rate (aggressive compression)
|
||||
config_low = ProxyConfig(llmlingua_target_rate=0.1)
|
||||
assert config_low.llmlingua_target_rate == 0.1
|
||||
|
||||
# High rate (conservative compression)
|
||||
config_high = ProxyConfig(llmlingua_target_rate=0.8)
|
||||
assert config_high.llmlingua_target_rate == 0.8
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TestLLMLinguaSetup
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestLLMLinguaSetup:
|
||||
"""Tests for LLMLingua setup in HeadroomProxy."""
|
||||
|
||||
def test_setup_returns_disabled_when_not_enabled(self, base_config):
|
||||
"""Setup returns 'disabled' when llmlingua not enabled and not available."""
|
||||
with patch("headroom.proxy.server._LLMLINGUA_AVAILABLE", False):
|
||||
proxy = HeadroomProxy(base_config)
|
||||
assert proxy._llmlingua_status == "disabled"
|
||||
|
||||
def test_setup_returns_available_when_installed_but_not_enabled(self, base_config):
|
||||
"""Setup returns 'available' when llmlingua installed but not enabled."""
|
||||
with patch("headroom.proxy.server._LLMLINGUA_AVAILABLE", True):
|
||||
proxy = HeadroomProxy(base_config)
|
||||
assert proxy._llmlingua_status == "available"
|
||||
|
||||
def test_setup_returns_enabled_when_enabled_and_available(self, llmlingua_config):
|
||||
"""Setup returns 'enabled' when llmlingua enabled and available."""
|
||||
mock_compressor = MagicMock()
|
||||
|
||||
with patch("headroom.proxy.server._LLMLINGUA_AVAILABLE", True):
|
||||
with patch("headroom.proxy.server.LLMLinguaCompressor", mock_compressor):
|
||||
with patch("headroom.proxy.server.LLMLinguaConfig"):
|
||||
proxy = HeadroomProxy(llmlingua_config)
|
||||
assert proxy._llmlingua_status == "enabled"
|
||||
|
||||
def test_setup_returns_unavailable_when_enabled_but_not_installed(self):
|
||||
"""Setup returns 'unavailable' when enabled but llmlingua not installed."""
|
||||
config = ProxyConfig(
|
||||
llmlingua_enabled=True,
|
||||
optimize=False,
|
||||
cache_enabled=False,
|
||||
rate_limit_enabled=False,
|
||||
smart_routing=False, # Use legacy mode for explicit status testing
|
||||
)
|
||||
|
||||
with patch("headroom.proxy.server._LLMLINGUA_AVAILABLE", False):
|
||||
proxy = HeadroomProxy(config)
|
||||
assert proxy._llmlingua_status == "unavailable"
|
||||
|
||||
def test_llmlingua_compressor_added_to_pipeline(self, llmlingua_config):
|
||||
"""LLMLinguaCompressor is added to pipeline when enabled."""
|
||||
mock_compressor_class = MagicMock()
|
||||
mock_compressor_instance = MagicMock()
|
||||
mock_compressor_class.return_value = mock_compressor_instance
|
||||
|
||||
with patch("headroom.proxy.server._LLMLINGUA_AVAILABLE", True):
|
||||
with patch("headroom.proxy.server.LLMLinguaCompressor", mock_compressor_class):
|
||||
with patch("headroom.proxy.server.LLMLinguaConfig") as mock_config:
|
||||
HeadroomProxy(llmlingua_config)
|
||||
|
||||
# Verify LLMLinguaCompressor was instantiated
|
||||
mock_compressor_class.assert_called_once()
|
||||
|
||||
# Verify config was passed with correct device and rate
|
||||
call_args = mock_config.call_args
|
||||
assert call_args.kwargs["device"] == "cpu"
|
||||
assert call_args.kwargs["target_compression_rate"] == 0.4
|
||||
|
||||
def test_llmlingua_not_added_when_disabled(self, base_config):
|
||||
"""LLMLinguaCompressor is NOT added when disabled."""
|
||||
mock_compressor_class = MagicMock()
|
||||
|
||||
with patch("headroom.proxy.server._LLMLINGUA_AVAILABLE", True):
|
||||
with patch("headroom.proxy.server.LLMLinguaCompressor", mock_compressor_class):
|
||||
HeadroomProxy(base_config)
|
||||
|
||||
# Should NOT be called when disabled
|
||||
mock_compressor_class.assert_not_called()
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TestBannerStatus
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestBannerStatus:
|
||||
"""Tests for banner status helper function."""
|
||||
|
||||
def test_banner_disabled_when_not_available(self):
|
||||
"""Banner shows DISABLED when llmlingua not available and not enabled."""
|
||||
config = ProxyConfig(llmlingua_enabled=False, smart_routing=False)
|
||||
|
||||
with patch("headroom.proxy.server._LLMLINGUA_AVAILABLE", False):
|
||||
status = _get_llmlingua_banner_status(config)
|
||||
assert status == "DISABLED"
|
||||
|
||||
def test_banner_available_hint_when_installed(self):
|
||||
"""Banner shows availability hint when installed but not enabled."""
|
||||
config = ProxyConfig(llmlingua_enabled=False, smart_routing=False)
|
||||
|
||||
with patch("headroom.proxy.server._LLMLINGUA_AVAILABLE", True):
|
||||
status = _get_llmlingua_banner_status(config)
|
||||
# Now shows DISABLED with hint to enable
|
||||
assert "DISABLED" in status
|
||||
assert "--no-llmlingua" in status # Hint to remove the flag to enable
|
||||
|
||||
def test_banner_enabled_when_active(self):
|
||||
"""Banner shows ENABLED with config when active."""
|
||||
config = ProxyConfig(
|
||||
llmlingua_enabled=True,
|
||||
llmlingua_device="cuda",
|
||||
llmlingua_target_rate=0.25,
|
||||
)
|
||||
|
||||
with patch("headroom.proxy.server._LLMLINGUA_AVAILABLE", True):
|
||||
status = _get_llmlingua_banner_status(config)
|
||||
assert "ENABLED" in status
|
||||
assert "cuda" in status
|
||||
assert "0.25" in status
|
||||
|
||||
def test_banner_shows_install_hint_when_requested_but_missing(self):
|
||||
"""Banner shows install hint when enabled but not installed."""
|
||||
config = ProxyConfig(llmlingua_enabled=True, smart_routing=False)
|
||||
|
||||
with patch("headroom.proxy.server._LLMLINGUA_AVAILABLE", False):
|
||||
status = _get_llmlingua_banner_status(config)
|
||||
assert "NOT INSTALLED" in status
|
||||
assert "pip install" in status
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TestHealthEndpointWithLLMLingua
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestHealthEndpointWithLLMLingua:
|
||||
"""Tests for health endpoint reflecting LLMLingua status."""
|
||||
|
||||
def test_health_returns_llmlingua_in_config(self, client):
|
||||
"""Health endpoint works regardless of LLMLingua status."""
|
||||
response = client.get("/health")
|
||||
assert response.status_code == 200
|
||||
|
||||
data = response.json()
|
||||
assert data["status"] == "healthy"
|
||||
assert "config" in data
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TestStatsEndpointWithLLMLingua
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestStatsEndpointWithLLMLingua:
|
||||
"""Tests for stats endpoint with LLMLingua integration."""
|
||||
|
||||
def test_stats_endpoint_works(self, client):
|
||||
"""Stats endpoint works with any LLMLingua configuration."""
|
||||
response = client.get("/stats")
|
||||
assert response.status_code == 200
|
||||
|
||||
data = response.json()
|
||||
assert "requests" in data
|
||||
assert "tokens" in data
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TestCLIArguments
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestCLIArguments:
|
||||
"""Tests for CLI argument parsing (without actually running server)."""
|
||||
|
||||
def test_llmlingua_flag_defaults(self):
|
||||
"""Default CLI values for LLMLingua settings."""
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--llmlingua", action="store_true")
|
||||
parser.add_argument("--llmlingua-device", default="auto")
|
||||
parser.add_argument("--llmlingua-rate", type=float, default=0.3)
|
||||
|
||||
args = parser.parse_args([])
|
||||
|
||||
assert args.llmlingua is False
|
||||
assert args.llmlingua_device == "auto"
|
||||
assert args.llmlingua_rate == 0.3
|
||||
|
||||
def test_llmlingua_flag_enabled(self):
|
||||
"""CLI --llmlingua flag enables LLMLingua."""
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--llmlingua", action="store_true")
|
||||
parser.add_argument("--llmlingua-device", default="auto")
|
||||
parser.add_argument("--llmlingua-rate", type=float, default=0.3)
|
||||
|
||||
args = parser.parse_args(["--llmlingua"])
|
||||
|
||||
assert args.llmlingua is True
|
||||
|
||||
def test_llmlingua_device_flag(self):
|
||||
"""CLI --llmlingua-device flag sets device."""
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--llmlingua-device", default="auto")
|
||||
|
||||
args = parser.parse_args(["--llmlingua-device", "cuda"])
|
||||
|
||||
assert args.llmlingua_device == "cuda"
|
||||
|
||||
def test_llmlingua_rate_flag(self):
|
||||
"""CLI --llmlingua-rate flag sets compression rate."""
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--llmlingua-rate", type=float, default=0.3)
|
||||
|
||||
args = parser.parse_args(["--llmlingua-rate", "0.5"])
|
||||
|
||||
assert args.llmlingua_rate == 0.5
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TestDevExMessages
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestDevExMessages:
|
||||
"""Tests for developer experience messages and hints."""
|
||||
|
||||
def test_warning_logged_when_enabled_but_unavailable(self, caplog):
|
||||
"""Warning is logged when llmlingua enabled but not installed."""
|
||||
import logging
|
||||
|
||||
config = ProxyConfig(
|
||||
llmlingua_enabled=True,
|
||||
optimize=False,
|
||||
cache_enabled=False,
|
||||
rate_limit_enabled=False,
|
||||
smart_routing=False, # Use legacy mode for explicit status testing
|
||||
)
|
||||
|
||||
with caplog.at_level(logging.WARNING):
|
||||
with patch("headroom.proxy.server._LLMLINGUA_AVAILABLE", False):
|
||||
proxy = HeadroomProxy(config)
|
||||
|
||||
# Should have logged a warning about missing llmlingua
|
||||
assert proxy._llmlingua_status == "unavailable"
|
||||
assert any("llmlingua" in r.message.lower() for r in caplog.records)
|
||||
assert any("pip install" in r.message for r in caplog.records)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TestIntegrationWithActualLLMLingua
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@pytest.mark.skipif(not _LLMLINGUA_AVAILABLE, reason="llmlingua not installed")
|
||||
class TestIntegrationWithActualLLMLingua:
|
||||
"""Integration tests that require actual llmlingua installation.
|
||||
|
||||
These tests verify the full integration path when llmlingua is installed.
|
||||
"""
|
||||
|
||||
def test_proxy_starts_with_llmlingua_enabled(self):
|
||||
"""Proxy starts successfully with LLMLingua enabled."""
|
||||
config = ProxyConfig(
|
||||
llmlingua_enabled=True,
|
||||
llmlingua_device="cpu", # CPU for CI/test environments
|
||||
llmlingua_target_rate=0.3,
|
||||
optimize=True,
|
||||
cache_enabled=False,
|
||||
rate_limit_enabled=False,
|
||||
smart_routing=False, # Use legacy mode for explicit status testing
|
||||
)
|
||||
|
||||
# Should not raise
|
||||
proxy = HeadroomProxy(config)
|
||||
|
||||
assert proxy._llmlingua_status == "enabled"
|
||||
|
||||
def test_app_creates_with_llmlingua(self):
|
||||
"""FastAPI app creates successfully with LLMLingua enabled."""
|
||||
config = ProxyConfig(
|
||||
llmlingua_enabled=True,
|
||||
llmlingua_device="cpu",
|
||||
optimize=True,
|
||||
cache_enabled=False,
|
||||
rate_limit_enabled=False,
|
||||
)
|
||||
|
||||
# Should not raise
|
||||
app = create_app(config)
|
||||
|
||||
assert app is not None
|
||||
|
||||
def test_health_endpoint_with_llmlingua_enabled(self):
|
||||
"""Health endpoint works with LLMLingua enabled."""
|
||||
config = ProxyConfig(
|
||||
llmlingua_enabled=True,
|
||||
llmlingua_device="cpu",
|
||||
optimize=True,
|
||||
cache_enabled=False,
|
||||
rate_limit_enabled=False,
|
||||
)
|
||||
|
||||
app = create_app(config)
|
||||
with TestClient(app) as client:
|
||||
response = client.get("/health")
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TestEdgeCases
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestEdgeCases:
|
||||
"""Edge cases for LLMLingua proxy integration."""
|
||||
|
||||
def test_multiple_proxy_instances_independent(self):
|
||||
"""Multiple proxy instances have independent LLMLingua status."""
|
||||
config_enabled = ProxyConfig(
|
||||
llmlingua_enabled=True,
|
||||
optimize=False,
|
||||
cache_enabled=False,
|
||||
rate_limit_enabled=False,
|
||||
smart_routing=False, # Use legacy mode for explicit status testing
|
||||
)
|
||||
config_disabled = ProxyConfig(
|
||||
llmlingua_enabled=False,
|
||||
optimize=False,
|
||||
cache_enabled=False,
|
||||
rate_limit_enabled=False,
|
||||
smart_routing=False, # Use legacy mode for explicit status testing
|
||||
)
|
||||
|
||||
with patch("headroom.proxy.server._LLMLINGUA_AVAILABLE", True):
|
||||
with patch("headroom.proxy.server.LLMLinguaCompressor"):
|
||||
with patch("headroom.proxy.server.LLMLinguaConfig"):
|
||||
proxy_enabled = HeadroomProxy(config_enabled)
|
||||
proxy_disabled = HeadroomProxy(config_disabled)
|
||||
|
||||
assert proxy_enabled._llmlingua_status == "enabled"
|
||||
assert proxy_disabled._llmlingua_status == "available"
|
||||
|
||||
def test_config_immutable_after_proxy_creation(self, base_config):
|
||||
"""Config values are captured at proxy creation time."""
|
||||
proxy = HeadroomProxy(base_config)
|
||||
|
||||
# Modifying config after creation doesn't affect proxy
|
||||
# (ProxyConfig is a dataclass, so this tests the pattern)
|
||||
original_status = proxy._llmlingua_status
|
||||
|
||||
# Status should remain unchanged
|
||||
assert proxy._llmlingua_status == original_status
|
||||
|
|
@ -499,14 +499,14 @@ class TestFallbackCompression:
|
|||
|
||||
# Should still return a result (fallback compression)
|
||||
assert result is not None
|
||||
# LLMLingua fallback does NOT guarantee syntax validity
|
||||
# If LLMLingua is unavailable, returns original (valid)
|
||||
# If LLMLingua IS available, syntax_valid=False (cannot guarantee)
|
||||
# Kompress fallback does NOT guarantee syntax validity
|
||||
# If Kompress is unavailable, returns original (valid)
|
||||
# If Kompress IS available, syntax_valid=False (cannot guarantee)
|
||||
|
||||
def test_fallback_preserves_structure(self, default_config):
|
||||
"""Fallback compression preserves basic structure when no compressor available.
|
||||
|
||||
When both tree-sitter and LLMLingua are unavailable, the fallback
|
||||
When both tree-sitter and Kompress are unavailable, the fallback
|
||||
returns the original code unchanged - preserving all structure.
|
||||
"""
|
||||
with (
|
||||
|
|
@ -515,7 +515,7 @@ class TestFallbackCompression:
|
|||
return_value=False,
|
||||
),
|
||||
patch(
|
||||
"headroom.transforms.llmlingua_compressor._check_llmlingua_available",
|
||||
"headroom.transforms.kompress_compressor.is_kompress_available",
|
||||
return_value=False,
|
||||
),
|
||||
):
|
||||
|
|
|
|||
|
|
@ -156,7 +156,6 @@ class TestContentRouterConfig:
|
|||
|
||||
assert config.enable_code_aware is True
|
||||
assert config.enable_kompress is True
|
||||
assert config.enable_llmlingua is True
|
||||
assert config.enable_smart_crusher is True
|
||||
assert config.enable_search_compressor is True
|
||||
assert config.enable_log_compressor is True
|
||||
|
|
@ -168,13 +167,11 @@ class TestContentRouterConfig:
|
|||
config = ContentRouterConfig(
|
||||
min_section_tokens=50,
|
||||
enable_code_aware=False,
|
||||
enable_llmlingua=False,
|
||||
fallback_strategy=CompressionStrategy.TEXT,
|
||||
)
|
||||
|
||||
assert config.min_section_tokens == 50
|
||||
assert config.enable_code_aware is False
|
||||
assert config.enable_llmlingua is False
|
||||
assert config.fallback_strategy == CompressionStrategy.TEXT
|
||||
|
||||
def test_all_strategies_in_enum(self):
|
||||
|
|
@ -184,7 +181,6 @@ class TestContentRouterConfig:
|
|||
"SMART_CRUSHER",
|
||||
"SEARCH",
|
||||
"LOG",
|
||||
"LLMLINGUA",
|
||||
"TEXT",
|
||||
"MIXED",
|
||||
"PASSTHROUGH",
|
||||
|
|
|
|||
|
|
@ -1,941 +0,0 @@
|
|||
"""Tests for LLMLingua-2 compressor integration.
|
||||
|
||||
Comprehensive tests covering:
|
||||
- LLMLinguaConfig: Configuration validation and defaults
|
||||
- LLMLinguaCompressor: Core compression functionality
|
||||
- Transform interface: apply(), should_apply() methods
|
||||
- Content type detection: JSON, code, plain text
|
||||
- CCR integration: Reversible compression storage
|
||||
- Edge cases: Empty content, unavailable dependency, fallbacks
|
||||
"""
|
||||
|
||||
import json
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom.transforms.llmlingua_compressor import (
|
||||
LLMLinguaCompressor,
|
||||
LLMLinguaConfig,
|
||||
LLMLinguaResult,
|
||||
compress_with_llmlingua,
|
||||
is_llmlingua_model_loaded,
|
||||
unload_llmlingua_model,
|
||||
)
|
||||
|
||||
# Try to import for availability check
|
||||
try:
|
||||
import llmlingua # noqa: F401
|
||||
|
||||
LLMLINGUA_INSTALLED = True
|
||||
except ImportError:
|
||||
LLMLINGUA_INSTALLED = False
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Test Fixtures
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def default_config():
|
||||
"""Default LLMLinguaConfig for testing."""
|
||||
return LLMLinguaConfig(
|
||||
min_tokens_for_compression=10, # Low threshold for tests
|
||||
enable_ccr=False, # Disable CCR for unit tests
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def compressor(default_config):
|
||||
"""LLMLinguaCompressor instance with default config."""
|
||||
return LLMLinguaCompressor(default_config)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_llmlingua():
|
||||
"""Mock the llmlingua module and PromptCompressor."""
|
||||
mock_compressor = MagicMock()
|
||||
mock_compressor._model_name = "test-model"
|
||||
|
||||
# Default compress_prompt return value
|
||||
mock_compressor.compress_prompt.return_value = {
|
||||
"compressed_prompt": "compressed content here",
|
||||
"origin_tokens": 100,
|
||||
"compressed_tokens": 30,
|
||||
}
|
||||
|
||||
with patch(
|
||||
"headroom.transforms.llmlingua_compressor._check_llmlingua_available",
|
||||
return_value=True,
|
||||
):
|
||||
with patch(
|
||||
"headroom.transforms.llmlingua_compressor._get_llmlingua_compressor",
|
||||
return_value=mock_compressor,
|
||||
):
|
||||
yield mock_compressor
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def tokenizer():
|
||||
"""Get a tokenizer for Transform interface tests."""
|
||||
from headroom.providers import OpenAIProvider
|
||||
from headroom.tokenizer import Tokenizer
|
||||
|
||||
provider = OpenAIProvider()
|
||||
token_counter = provider.get_token_counter("gpt-4o")
|
||||
return Tokenizer(token_counter, "gpt-4o")
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Test Data Generators
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def generate_long_text(n_words: int = 500) -> str:
|
||||
"""Generate long text content for compression testing."""
|
||||
words = ["the", "quick", "brown", "fox", "jumps", "over", "lazy", "dog"]
|
||||
return " ".join(words[i % len(words)] for i in range(n_words))
|
||||
|
||||
|
||||
def generate_long_json(n_items: int = 50) -> str:
|
||||
"""Generate long JSON content for compression testing."""
|
||||
items = [
|
||||
{
|
||||
"id": i,
|
||||
"name": f"Item {i}",
|
||||
"description": f"This is a detailed description for item number {i}",
|
||||
"value": i * 10,
|
||||
"active": i % 2 == 0,
|
||||
}
|
||||
for i in range(n_items)
|
||||
]
|
||||
return json.dumps(items)
|
||||
|
||||
|
||||
def generate_long_code(n_functions: int = 20) -> str:
|
||||
"""Generate Python code content for compression testing."""
|
||||
lines = ['"""Module with many functions."""', "", "import os", "from typing import Any", ""]
|
||||
for i in range(n_functions):
|
||||
lines.extend(
|
||||
[
|
||||
f"def function_{i}(arg: Any) -> str:",
|
||||
f' """Process argument {i}."""',
|
||||
" result = str(arg)",
|
||||
f' return f"Function {i}: {{result}}"',
|
||||
"",
|
||||
]
|
||||
)
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TestLLMLinguaConfig
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestLLMLinguaConfig:
|
||||
"""Tests for LLMLinguaConfig dataclass."""
|
||||
|
||||
def test_default_values(self):
|
||||
"""Default config values are sensible."""
|
||||
config = LLMLinguaConfig()
|
||||
|
||||
assert config.model_name == "microsoft/llmlingua-2-bert-base-multilingual-cased-meetingbank"
|
||||
assert config.device == "auto"
|
||||
assert config.target_compression_rate == 0.3
|
||||
assert config.min_tokens_for_compression == 100
|
||||
assert config.enable_ccr is True
|
||||
assert config.drop_consecutive is True
|
||||
|
||||
def test_custom_values(self):
|
||||
"""Custom config values are applied."""
|
||||
config = LLMLinguaConfig(
|
||||
model_name="custom/model",
|
||||
device="cuda",
|
||||
target_compression_rate=0.5,
|
||||
min_tokens_for_compression=50,
|
||||
force_tokens=["important", "keep"],
|
||||
)
|
||||
|
||||
assert config.model_name == "custom/model"
|
||||
assert config.device == "cuda"
|
||||
assert config.target_compression_rate == 0.5
|
||||
assert config.min_tokens_for_compression == 50
|
||||
assert "important" in config.force_tokens
|
||||
|
||||
def test_content_type_rates(self):
|
||||
"""Different content types have appropriate compression rates."""
|
||||
config = LLMLinguaConfig()
|
||||
|
||||
# Code and text are equally conservative for accuracy
|
||||
assert config.code_compression_rate >= config.text_compression_rate
|
||||
# JSON can be slightly more aggressive since structure is preserved
|
||||
assert config.json_compression_rate <= config.code_compression_rate
|
||||
assert config.json_compression_rate <= config.text_compression_rate
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TestLLMLinguaResult
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestLLMLinguaResult:
|
||||
"""Tests for LLMLinguaResult dataclass."""
|
||||
|
||||
def test_tokens_saved(self):
|
||||
"""tokens_saved property calculates correctly."""
|
||||
result = LLMLinguaResult(
|
||||
compressed="short",
|
||||
original="long content here",
|
||||
original_tokens=100,
|
||||
compressed_tokens=30,
|
||||
compression_ratio=0.3,
|
||||
)
|
||||
|
||||
assert result.tokens_saved == 70
|
||||
|
||||
def test_tokens_saved_no_negative(self):
|
||||
"""tokens_saved never returns negative."""
|
||||
result = LLMLinguaResult(
|
||||
compressed="expanded content",
|
||||
original="short",
|
||||
original_tokens=10,
|
||||
compressed_tokens=20, # Expanded (unusual case)
|
||||
compression_ratio=2.0,
|
||||
)
|
||||
|
||||
assert result.tokens_saved == 0
|
||||
|
||||
def test_savings_percentage(self):
|
||||
"""savings_percentage property calculates correctly."""
|
||||
result = LLMLinguaResult(
|
||||
compressed="short",
|
||||
original="long content",
|
||||
original_tokens=100,
|
||||
compressed_tokens=25,
|
||||
compression_ratio=0.25,
|
||||
)
|
||||
|
||||
assert result.savings_percentage == 75.0
|
||||
|
||||
def test_savings_percentage_zero_original(self):
|
||||
"""savings_percentage handles zero original tokens."""
|
||||
result = LLMLinguaResult(
|
||||
compressed="",
|
||||
original="",
|
||||
original_tokens=0,
|
||||
compressed_tokens=0,
|
||||
compression_ratio=1.0,
|
||||
)
|
||||
|
||||
assert result.savings_percentage == 0.0
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TestLLMLinguaCompressor
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestLLMLinguaCompressor:
|
||||
"""Tests for LLMLinguaCompressor core functionality."""
|
||||
|
||||
def test_init_with_default_config(self):
|
||||
"""Compressor initializes with default config."""
|
||||
compressor = LLMLinguaCompressor()
|
||||
|
||||
assert compressor.config is not None
|
||||
assert compressor.config.model_name is not None
|
||||
|
||||
def test_init_with_custom_config(self, default_config):
|
||||
"""Compressor initializes with custom config."""
|
||||
compressor = LLMLinguaCompressor(default_config)
|
||||
|
||||
assert compressor.config == default_config
|
||||
|
||||
def test_compress_returns_result_when_unavailable(self, compressor):
|
||||
"""Compress returns passthrough result when llmlingua unavailable."""
|
||||
with patch(
|
||||
"headroom.transforms.llmlingua_compressor._check_llmlingua_available",
|
||||
return_value=False,
|
||||
):
|
||||
content = generate_long_text(100)
|
||||
result = compressor.compress(content)
|
||||
|
||||
# Should return unchanged content
|
||||
assert result.compressed == content
|
||||
assert result.compression_ratio == 1.0
|
||||
|
||||
def test_compress_skips_small_content(self, compressor):
|
||||
"""Small content is not compressed."""
|
||||
small_content = "short text"
|
||||
result = compressor.compress(small_content)
|
||||
|
||||
assert result.compressed == small_content
|
||||
assert result.compression_ratio == 1.0
|
||||
|
||||
def test_compress_with_llmlingua(self, default_config, mock_llmlingua):
|
||||
"""Compression uses llmlingua when available."""
|
||||
compressor = LLMLinguaCompressor(default_config)
|
||||
content = generate_long_text(200)
|
||||
|
||||
result = compressor.compress(content)
|
||||
|
||||
# Should have called compress_prompt
|
||||
mock_llmlingua.compress_prompt.assert_called_once()
|
||||
assert result.compressed == "compressed content here"
|
||||
assert result.compression_ratio < 1.0
|
||||
|
||||
def test_compress_with_context(self, default_config, mock_llmlingua):
|
||||
"""Context words are used as force tokens."""
|
||||
compressor = LLMLinguaCompressor(default_config)
|
||||
content = generate_long_text(200)
|
||||
context = "important keywords here"
|
||||
|
||||
compressor.compress(content, context=context)
|
||||
|
||||
# Check force_tokens includes context words
|
||||
call_args = mock_llmlingua.compress_prompt.call_args
|
||||
force_tokens = call_args.kwargs.get("force_tokens", [])
|
||||
# Should include context words longer than 3 chars
|
||||
assert "important" in force_tokens or "keywords" in force_tokens
|
||||
|
||||
def test_compress_handles_exception(self, default_config, mock_llmlingua):
|
||||
"""Exceptions from llmlingua are handled gracefully."""
|
||||
mock_llmlingua.compress_prompt.side_effect = RuntimeError("Model error")
|
||||
|
||||
compressor = LLMLinguaCompressor(default_config)
|
||||
content = generate_long_text(200)
|
||||
|
||||
result = compressor.compress(content)
|
||||
|
||||
# Should return original content on error
|
||||
assert result.compressed == content
|
||||
assert result.compression_ratio == 1.0
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TestContentTypeDetection
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestContentTypeDetection:
|
||||
"""Tests for content type auto-detection."""
|
||||
|
||||
def test_detect_json_content(self, default_config, mock_llmlingua):
|
||||
"""JSON content is detected and uses JSON compression rate."""
|
||||
compressor = LLMLinguaCompressor(default_config)
|
||||
|
||||
rate = compressor._get_compression_rate(generate_long_json(50), None)
|
||||
|
||||
assert rate == default_config.json_compression_rate
|
||||
|
||||
def test_detect_code_content(self, default_config, mock_llmlingua):
|
||||
"""Code content is detected and uses code compression rate."""
|
||||
compressor = LLMLinguaCompressor(default_config)
|
||||
code = generate_long_code(20)
|
||||
|
||||
rate = compressor._get_compression_rate(code, None)
|
||||
|
||||
assert rate == default_config.code_compression_rate
|
||||
|
||||
def test_detect_plain_text(self, default_config, mock_llmlingua):
|
||||
"""Plain text uses text compression rate."""
|
||||
compressor = LLMLinguaCompressor(default_config)
|
||||
text = generate_long_text(200)
|
||||
|
||||
rate = compressor._get_compression_rate(text, None)
|
||||
|
||||
assert rate == default_config.text_compression_rate
|
||||
|
||||
def test_explicit_content_type(self, default_config, mock_llmlingua):
|
||||
"""Explicit content_type overrides detection."""
|
||||
compressor = LLMLinguaCompressor(default_config)
|
||||
# JSON-looking content but marked as text
|
||||
json_content = generate_long_json(50)
|
||||
|
||||
rate = compressor._get_compression_rate(json_content, content_type="text")
|
||||
|
||||
assert rate == default_config.text_compression_rate
|
||||
|
||||
def test_looks_like_json_detection(self, default_config):
|
||||
"""JSON detection works for arrays and objects."""
|
||||
compressor = LLMLinguaCompressor(default_config)
|
||||
|
||||
assert compressor._looks_like_json('[{"key": "value"}]')
|
||||
assert compressor._looks_like_json('{"key": "value"}')
|
||||
assert not compressor._looks_like_json("plain text")
|
||||
assert not compressor._looks_like_json("def function():")
|
||||
|
||||
def test_looks_like_code_detection(self, default_config):
|
||||
"""Code detection works for common patterns."""
|
||||
compressor = LLMLinguaCompressor(default_config)
|
||||
|
||||
assert compressor._looks_like_code("def function():")
|
||||
assert compressor._looks_like_code("class MyClass:")
|
||||
assert compressor._looks_like_code("import os")
|
||||
assert compressor._looks_like_code("function test() {")
|
||||
assert compressor._looks_like_code("const x = 5")
|
||||
assert not compressor._looks_like_code("plain text content")
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TestTransformInterface
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestTransformInterface:
|
||||
"""Tests for Transform interface (apply, should_apply)."""
|
||||
|
||||
def test_should_apply_returns_false_when_unavailable(self, compressor, tokenizer):
|
||||
"""should_apply returns False when llmlingua unavailable."""
|
||||
messages = [{"role": "user", "content": generate_long_text(200)}]
|
||||
|
||||
with patch(
|
||||
"headroom.transforms.llmlingua_compressor._check_llmlingua_available",
|
||||
return_value=False,
|
||||
):
|
||||
assert not compressor.should_apply(messages, tokenizer)
|
||||
|
||||
def test_should_apply_returns_false_for_small_content(self, default_config, tokenizer):
|
||||
"""should_apply returns False for small content."""
|
||||
config = LLMLinguaConfig(min_tokens_for_compression=1000)
|
||||
compressor = LLMLinguaCompressor(config)
|
||||
messages = [{"role": "user", "content": "small"}]
|
||||
|
||||
with patch(
|
||||
"headroom.transforms.llmlingua_compressor._check_llmlingua_available",
|
||||
return_value=True,
|
||||
):
|
||||
assert not compressor.should_apply(messages, tokenizer)
|
||||
|
||||
def test_should_apply_returns_true_for_large_content(self, default_config, tokenizer):
|
||||
"""should_apply returns True for large content."""
|
||||
compressor = LLMLinguaCompressor(default_config)
|
||||
messages = [{"role": "user", "content": generate_long_text(500)}]
|
||||
|
||||
with patch(
|
||||
"headroom.transforms.llmlingua_compressor._check_llmlingua_available",
|
||||
return_value=True,
|
||||
):
|
||||
assert compressor.should_apply(messages, tokenizer)
|
||||
|
||||
def test_apply_compresses_tool_messages(self, default_config, tokenizer, mock_llmlingua):
|
||||
"""apply() compresses tool message content."""
|
||||
compressor = LLMLinguaCompressor(default_config)
|
||||
tool_content = generate_long_json(100)
|
||||
messages = [
|
||||
{"role": "user", "content": "Get data"},
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": tool_content},
|
||||
]
|
||||
|
||||
result = compressor.apply(messages, tokenizer)
|
||||
|
||||
# Tool content should be compressed
|
||||
assert result.messages[1]["content"] != tool_content
|
||||
assert "compressed content here" in result.messages[1]["content"]
|
||||
assert len(result.transforms_applied) > 0
|
||||
|
||||
def test_apply_compresses_long_assistant_messages(
|
||||
self, default_config, tokenizer, mock_llmlingua
|
||||
):
|
||||
"""apply() compresses long assistant messages."""
|
||||
compressor = LLMLinguaCompressor(default_config)
|
||||
long_content = generate_long_text(1000)
|
||||
messages = [
|
||||
{"role": "user", "content": "Tell me a story"},
|
||||
{"role": "assistant", "content": long_content},
|
||||
]
|
||||
|
||||
result = compressor.apply(messages, tokenizer)
|
||||
|
||||
# Assistant content should be compressed (>500 chars)
|
||||
assert result.messages[1]["content"] != long_content
|
||||
|
||||
def test_apply_passes_through_short_messages(self, default_config, tokenizer, mock_llmlingua):
|
||||
"""apply() passes through short messages unchanged."""
|
||||
compressor = LLMLinguaCompressor(default_config)
|
||||
messages = [
|
||||
{"role": "user", "content": "Hello"},
|
||||
{"role": "assistant", "content": "Hi there!"},
|
||||
]
|
||||
|
||||
result = compressor.apply(messages, tokenizer)
|
||||
|
||||
# Short messages unchanged
|
||||
assert result.messages[0]["content"] == "Hello"
|
||||
assert result.messages[1]["content"] == "Hi there!"
|
||||
|
||||
def test_apply_tracks_transform_metadata(self, default_config, tokenizer, mock_llmlingua):
|
||||
"""apply() returns proper TransformResult metadata."""
|
||||
compressor = LLMLinguaCompressor(default_config)
|
||||
messages = [
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": generate_long_json(100)},
|
||||
]
|
||||
|
||||
result = compressor.apply(messages, tokenizer)
|
||||
|
||||
assert result.tokens_before > 0
|
||||
assert result.tokens_after > 0
|
||||
assert len(result.transforms_applied) > 0
|
||||
assert "llmlingua" in result.transforms_applied[0]
|
||||
|
||||
def test_apply_adds_warning_when_unavailable(self, default_config, tokenizer):
|
||||
"""apply() adds warning when llmlingua unavailable."""
|
||||
compressor = LLMLinguaCompressor(default_config)
|
||||
messages = [{"role": "user", "content": "test"}]
|
||||
|
||||
with patch(
|
||||
"headroom.transforms.llmlingua_compressor._check_llmlingua_available",
|
||||
return_value=False,
|
||||
):
|
||||
result = compressor.apply(messages, tokenizer)
|
||||
|
||||
assert len(result.warnings) > 0
|
||||
assert "llmlingua" in result.warnings[0].lower()
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TestDeviceResolution
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestDeviceResolution:
|
||||
"""Tests for device resolution logic."""
|
||||
|
||||
def test_resolve_explicit_device(self, default_config):
|
||||
"""Explicit device is returned unchanged."""
|
||||
config = LLMLinguaConfig(device="cuda")
|
||||
compressor = LLMLinguaCompressor(config)
|
||||
|
||||
assert compressor._resolve_device() == "cuda"
|
||||
|
||||
def test_resolve_auto_to_cpu_no_torch(self, default_config):
|
||||
"""Auto resolves to CPU when torch unavailable."""
|
||||
config = LLMLinguaConfig(device="auto")
|
||||
compressor = LLMLinguaCompressor(config)
|
||||
|
||||
with patch.dict("sys.modules", {"torch": None}):
|
||||
with patch(
|
||||
"headroom.transforms.llmlingua_compressor.LLMLinguaCompressor._resolve_device"
|
||||
) as mock_resolve:
|
||||
mock_resolve.return_value = "cpu"
|
||||
assert compressor._resolve_device() == "cpu"
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TestCCRIntegration
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestCCRIntegration:
|
||||
"""Tests for CCR (Compress-Cache-Retrieve) integration."""
|
||||
|
||||
def test_ccr_stores_original(self, mock_llmlingua):
|
||||
"""Compressed content is stored in CCR when enabled."""
|
||||
config = LLMLinguaConfig(
|
||||
enable_ccr=True,
|
||||
min_tokens_for_compression=10,
|
||||
)
|
||||
compressor = LLMLinguaCompressor(config)
|
||||
content = generate_long_text(200)
|
||||
|
||||
with patch(
|
||||
"headroom.transforms.llmlingua_compressor.LLMLinguaCompressor._store_in_ccr"
|
||||
) as mock_store:
|
||||
mock_store.return_value = "hash123"
|
||||
|
||||
result = compressor.compress(content)
|
||||
|
||||
mock_store.assert_called_once()
|
||||
assert result.cache_key == "hash123"
|
||||
|
||||
def test_ccr_skipped_when_disabled(self, mock_llmlingua):
|
||||
"""CCR is not used when disabled in config."""
|
||||
config = LLMLinguaConfig(
|
||||
enable_ccr=False,
|
||||
min_tokens_for_compression=10,
|
||||
)
|
||||
compressor = LLMLinguaCompressor(config)
|
||||
content = generate_long_text(200)
|
||||
|
||||
with patch(
|
||||
"headroom.transforms.llmlingua_compressor.LLMLinguaCompressor._store_in_ccr"
|
||||
) as mock_store:
|
||||
result = compressor.compress(content)
|
||||
|
||||
mock_store.assert_not_called()
|
||||
assert result.cache_key is None
|
||||
|
||||
def test_ccr_handles_storage_error(self, mock_llmlingua):
|
||||
"""CCR storage errors are handled gracefully."""
|
||||
config = LLMLinguaConfig(
|
||||
enable_ccr=True,
|
||||
min_tokens_for_compression=10,
|
||||
)
|
||||
compressor = LLMLinguaCompressor(config)
|
||||
content = generate_long_text(200)
|
||||
|
||||
with patch(
|
||||
"headroom.transforms.llmlingua_compressor.LLMLinguaCompressor._store_in_ccr"
|
||||
) as mock_store:
|
||||
# Return None to simulate storage failure (internal error handling)
|
||||
mock_store.return_value = None
|
||||
|
||||
# Should not raise
|
||||
result = compressor.compress(content)
|
||||
|
||||
# Storage failed, so cache_key should be None
|
||||
assert result.cache_key is None
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TestConvenienceFunction
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestConvenienceFunction:
|
||||
"""Tests for compress_with_llmlingua convenience function."""
|
||||
|
||||
def test_compress_with_llmlingua_basic(self, mock_llmlingua):
|
||||
"""compress_with_llmlingua works with default settings."""
|
||||
content = generate_long_text(200)
|
||||
|
||||
# Disable CCR for this test to avoid hash suffix
|
||||
with patch(
|
||||
"headroom.transforms.llmlingua_compressor.LLMLinguaCompressor._store_in_ccr"
|
||||
) as mock_store:
|
||||
mock_store.return_value = None
|
||||
result = compress_with_llmlingua(content)
|
||||
|
||||
# Should contain the compressed content
|
||||
assert "compressed content here" in result
|
||||
|
||||
def test_compress_with_llmlingua_custom_rate(self, mock_llmlingua):
|
||||
"""compress_with_llmlingua accepts custom compression rate."""
|
||||
content = generate_long_text(200)
|
||||
|
||||
compress_with_llmlingua(content, compression_rate=0.5)
|
||||
|
||||
# Verify compress_prompt was called
|
||||
mock_llmlingua.compress_prompt.assert_called()
|
||||
|
||||
def test_compress_with_llmlingua_with_context(self, mock_llmlingua):
|
||||
"""compress_with_llmlingua passes context."""
|
||||
content = generate_long_text(200)
|
||||
context = "important keywords"
|
||||
|
||||
compress_with_llmlingua(content, context=context)
|
||||
|
||||
call_args = mock_llmlingua.compress_prompt.call_args
|
||||
force_tokens = call_args.kwargs.get("force_tokens", [])
|
||||
# Context words should be in force_tokens
|
||||
assert any("important" in str(t) for t in force_tokens) or len(force_tokens) > 0
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TestEdgeCases
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestEdgeCases:
|
||||
"""Edge case tests for LLMLingua compressor."""
|
||||
|
||||
def test_empty_content(self, compressor):
|
||||
"""Empty content is handled gracefully."""
|
||||
result = compressor.compress("")
|
||||
|
||||
assert result.compressed == ""
|
||||
assert result.compression_ratio == 1.0
|
||||
|
||||
def test_whitespace_only_content(self, compressor):
|
||||
"""Whitespace-only content is handled gracefully."""
|
||||
result = compressor.compress(" \n\t\n ")
|
||||
|
||||
assert result.compression_ratio == 1.0
|
||||
|
||||
def test_unicode_content(self, default_config, mock_llmlingua):
|
||||
"""Unicode content is handled correctly."""
|
||||
mock_llmlingua.compress_prompt.return_value = {
|
||||
"compressed_prompt": "compressed \u4e2d\u6587 content",
|
||||
"origin_tokens": 100,
|
||||
"compressed_tokens": 30,
|
||||
}
|
||||
|
||||
compressor = LLMLinguaCompressor(default_config)
|
||||
content = "\u4e2d\u6587 \u65e5\u672c\u8a9e " * 100 # Chinese/Japanese text
|
||||
|
||||
result = compressor.compress(content)
|
||||
|
||||
assert "\u4e2d\u6587" in result.compressed
|
||||
|
||||
def test_very_long_content(self, default_config, mock_llmlingua):
|
||||
"""Very long content is compressed."""
|
||||
compressor = LLMLinguaCompressor(default_config)
|
||||
content = generate_long_text(10000)
|
||||
|
||||
compressor.compress(content)
|
||||
|
||||
mock_llmlingua.compress_prompt.assert_called_once()
|
||||
|
||||
def test_mixed_content_types(self, default_config, mock_llmlingua):
|
||||
"""Mixed content (JSON with text) is handled."""
|
||||
compressor = LLMLinguaCompressor(default_config)
|
||||
# JSON-like but with extra text
|
||||
content = 'Some preamble text\n{"key": "value"}\nMore text after'
|
||||
|
||||
# Should not crash
|
||||
result = compressor.compress(content)
|
||||
assert result is not None
|
||||
|
||||
def test_malformed_json_content(self, default_config, mock_llmlingua):
|
||||
"""Malformed JSON is treated as text."""
|
||||
compressor = LLMLinguaCompressor(default_config)
|
||||
content = "{malformed: json, missing quotes" * 50
|
||||
|
||||
rate = compressor._get_compression_rate(content, None)
|
||||
|
||||
# Should not detect as JSON
|
||||
assert rate == default_config.text_compression_rate
|
||||
|
||||
def test_force_tokens_list_handling(self, default_config, mock_llmlingua):
|
||||
"""Force tokens list is properly passed."""
|
||||
config = LLMLinguaConfig(
|
||||
force_tokens=["keep", "these", "tokens"],
|
||||
min_tokens_for_compression=10,
|
||||
)
|
||||
compressor = LLMLinguaCompressor(config)
|
||||
content = generate_long_text(200)
|
||||
|
||||
compressor.compress(content)
|
||||
|
||||
call_args = mock_llmlingua.compress_prompt.call_args
|
||||
force_tokens = call_args.kwargs.get("force_tokens", [])
|
||||
assert "keep" in force_tokens
|
||||
assert "these" in force_tokens
|
||||
assert "tokens" in force_tokens
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Integration Tests (only run if llmlingua is installed)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@pytest.mark.skipif(not LLMLINGUA_INSTALLED, reason="llmlingua not installed")
|
||||
class TestLLMLinguaIntegration:
|
||||
"""Integration tests that require actual llmlingua installation.
|
||||
|
||||
These tests verify the actual compression behavior and should be run
|
||||
in environments where llmlingua is installed.
|
||||
"""
|
||||
|
||||
def test_actual_compression(self):
|
||||
"""Test actual compression with real llmlingua."""
|
||||
config = LLMLinguaConfig(
|
||||
target_compression_rate=0.3,
|
||||
min_tokens_for_compression=50,
|
||||
enable_ccr=False,
|
||||
)
|
||||
compressor = LLMLinguaCompressor(config)
|
||||
content = generate_long_text(500)
|
||||
|
||||
result = compressor.compress(content)
|
||||
|
||||
# Should achieve actual compression
|
||||
assert result.compression_ratio < 1.0
|
||||
assert result.tokens_saved > 0
|
||||
assert len(result.compressed) < len(content)
|
||||
|
||||
def test_actual_json_compression(self):
|
||||
"""Test JSON content compression with real llmlingua."""
|
||||
config = LLMLinguaConfig(
|
||||
target_compression_rate=0.35,
|
||||
min_tokens_for_compression=50,
|
||||
enable_ccr=False,
|
||||
)
|
||||
compressor = LLMLinguaCompressor(config)
|
||||
content = generate_long_json(50)
|
||||
|
||||
result = compressor.compress(content, content_type="json")
|
||||
|
||||
assert result.compression_ratio < 1.0
|
||||
|
||||
def test_actual_code_compression(self):
|
||||
"""Test code content compression with real llmlingua."""
|
||||
config = LLMLinguaConfig(
|
||||
target_compression_rate=0.4,
|
||||
min_tokens_for_compression=50,
|
||||
enable_ccr=False,
|
||||
)
|
||||
compressor = LLMLinguaCompressor(config)
|
||||
content = generate_long_code(30)
|
||||
|
||||
result = compressor.compress(content, content_type="code")
|
||||
|
||||
assert result.compression_ratio < 1.0
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TestMemoryManagement
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestMemoryManagement:
|
||||
"""Tests for memory management functions (unload_llmlingua_model, is_llmlingua_model_loaded)."""
|
||||
|
||||
def test_is_model_loaded_returns_false_initially(self):
|
||||
"""is_llmlingua_model_loaded returns False when no model loaded."""
|
||||
# Ensure model is unloaded
|
||||
with patch(
|
||||
"headroom.transforms.llmlingua_compressor._llmlingua_instance",
|
||||
None,
|
||||
):
|
||||
assert is_llmlingua_model_loaded() is False
|
||||
|
||||
def test_is_model_loaded_returns_true_when_loaded(self):
|
||||
"""is_llmlingua_model_loaded returns True when model is loaded."""
|
||||
mock_instance = MagicMock()
|
||||
|
||||
with patch(
|
||||
"headroom.transforms.llmlingua_compressor._llmlingua_instance",
|
||||
mock_instance,
|
||||
):
|
||||
assert is_llmlingua_model_loaded() is True
|
||||
|
||||
def test_unload_returns_false_when_no_model(self):
|
||||
"""unload_llmlingua_model returns False when no model loaded."""
|
||||
import headroom.transforms.llmlingua_compressor as module
|
||||
|
||||
# Save original
|
||||
original = module._llmlingua_instance
|
||||
|
||||
try:
|
||||
module._llmlingua_instance = None
|
||||
result = unload_llmlingua_model()
|
||||
assert result is False
|
||||
finally:
|
||||
module._llmlingua_instance = original
|
||||
|
||||
def test_unload_clears_instance(self):
|
||||
"""unload_llmlingua_model clears the global instance."""
|
||||
import headroom.transforms.llmlingua_compressor as module
|
||||
|
||||
# Save original
|
||||
original = module._llmlingua_instance
|
||||
|
||||
try:
|
||||
# Set a mock instance
|
||||
mock_instance = MagicMock()
|
||||
mock_instance._model_name = "test-model"
|
||||
module._llmlingua_instance = mock_instance
|
||||
|
||||
# Unload
|
||||
result = unload_llmlingua_model()
|
||||
|
||||
assert result is True
|
||||
assert module._llmlingua_instance is None
|
||||
finally:
|
||||
module._llmlingua_instance = original
|
||||
|
||||
def test_unload_clears_cuda_cache(self):
|
||||
"""unload_llmlingua_model attempts to clear CUDA cache."""
|
||||
import headroom.transforms.llmlingua_compressor as module
|
||||
|
||||
original = module._llmlingua_instance
|
||||
|
||||
try:
|
||||
mock_instance = MagicMock()
|
||||
mock_instance._model_name = "test-model"
|
||||
module._llmlingua_instance = mock_instance
|
||||
|
||||
mock_torch = MagicMock()
|
||||
mock_torch.cuda.is_available.return_value = True
|
||||
|
||||
with patch.dict("sys.modules", {"torch": mock_torch}):
|
||||
with patch(
|
||||
"headroom.transforms.llmlingua_compressor.torch",
|
||||
mock_torch,
|
||||
create=True,
|
||||
):
|
||||
result = unload_llmlingua_model()
|
||||
|
||||
assert result is True
|
||||
finally:
|
||||
module._llmlingua_instance = original
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TestThreadSafety
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestThreadSafety:
|
||||
"""Tests for thread safety of model loading."""
|
||||
|
||||
def test_lock_exists(self):
|
||||
"""Verify thread lock is available."""
|
||||
import headroom.transforms.llmlingua_compressor as module
|
||||
|
||||
assert hasattr(module, "_llmlingua_lock")
|
||||
import threading
|
||||
|
||||
assert isinstance(module._llmlingua_lock, type(threading.Lock()))
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TestErrorMessages
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestErrorMessages:
|
||||
"""Tests for improved error messages."""
|
||||
|
||||
def test_import_error_message_includes_install_hint(self):
|
||||
"""ImportError includes installation instructions."""
|
||||
with patch(
|
||||
"headroom.transforms.llmlingua_compressor._check_llmlingua_available",
|
||||
return_value=False,
|
||||
):
|
||||
from headroom.transforms.llmlingua_compressor import _get_llmlingua_compressor
|
||||
|
||||
with pytest.raises(ImportError) as exc_info:
|
||||
_get_llmlingua_compressor("test-model", "cpu")
|
||||
|
||||
error_msg = str(exc_info.value)
|
||||
assert "pip install headroom-ai[llmlingua]" in error_msg
|
||||
assert "2GB" in error_msg or "disk space" in error_msg.lower()
|
||||
|
||||
def test_oom_error_provides_helpful_suggestions(self):
|
||||
"""Out of memory error provides helpful suggestions."""
|
||||
import headroom.transforms.llmlingua_compressor as module
|
||||
|
||||
# Save original state
|
||||
original_instance = module._llmlingua_instance
|
||||
original_available = module._llmlingua_available
|
||||
|
||||
try:
|
||||
module._llmlingua_instance = None
|
||||
module._llmlingua_available = True
|
||||
|
||||
# Create a mock that raises OOM when called
|
||||
mock_prompt_compressor_class = MagicMock()
|
||||
mock_prompt_compressor_class.side_effect = RuntimeError("CUDA out of memory")
|
||||
|
||||
with patch.dict("sys.modules", {"llmlingua": MagicMock()}):
|
||||
with patch(
|
||||
"llmlingua.PromptCompressor",
|
||||
mock_prompt_compressor_class,
|
||||
):
|
||||
from headroom.transforms.llmlingua_compressor import (
|
||||
_get_llmlingua_compressor,
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError) as exc_info:
|
||||
_get_llmlingua_compressor("test-model", "cuda")
|
||||
|
||||
error_msg = str(exc_info.value)
|
||||
# Should include helpful suggestions
|
||||
assert "cpu" in error_msg.lower() or "memory" in error_msg.lower()
|
||||
finally:
|
||||
module._llmlingua_instance = original_instance
|
||||
module._llmlingua_available = original_available
|
||||
Loading…
Add table
Add a link
Reference in a new issue