mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
Add LLMLingua-2 opt-in support to proxy server
Integrate Microsoft's LLMLingua-2 ML-based compression as an opt-in feature for the proxy server, with excellent developer experience. Features: - New CLI flags: --llmlingua, --llmlingua-device, --llmlingua-rate - ProxyConfig options: llmlingua_enabled, llmlingua_device, llmlingua_target_rate - Smart startup hints when llmlingua is available but not enabled - Helpful error messages when enabled but not installed - LLMLinguaCompressor inserted before RollingWindow in pipeline Why opt-in: - Heavy dependencies (~2GB torch, transformers) - 10-30s cold start for model loading - ~1GB RAM when loaded - Default proxy stays lightweight (<5ms overhead) Tests: - 26 new tests in test_proxy_llmlingua.py covering config, setup, banner status, CLI args, DevEx messages, and edge cases Documentation: - Updated README.md with proxy integration section - Updated docs/proxy.md with LLMLingua CLI options - Updated docs/transforms.md with LLMLinguaCompressor reference - Updated docs/ARCHITECTURE.md with pipeline and file structure - Updated CHANGELOG.md with new feature
This commit is contained in:
chopratejas
parent
c2117ab70c
commit
45633b69ab
11 changed files with 2518 additions and 13 deletions
|
|
@ -10,6 +10,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
|||
### Added
|
||||
- Production-ready proxy server with caching, rate limiting, and metrics
|
||||
- CLI command `headroom proxy` to start the proxy server
|
||||
- **LLMLingua-2 Integration** (opt-in ML-based compression)
|
||||
- `LLMLinguaCompressor` transform using Microsoft's LLMLingua-2 model
|
||||
- Content-aware compression rates (code: 0.4, JSON: 0.35, text: 0.3)
|
||||
- Memory management utilities: `unload_llmlingua_model()`, `is_llmlingua_model_loaded()`
|
||||
- Proxy integration via `--llmlingua` flag
|
||||
- Device selection: `--llmlingua-device` (auto/cuda/cpu/mps)
|
||||
- Custom compression rate: `--llmlingua-rate`
|
||||
- Helpful startup hints when llmlingua is available but not enabled
|
||||
- Install with: `pip install headroom-ai[llmlingua]`
|
||||
|
||||
## [0.2.0] - 2025-01-07
|
||||
|
||||
|
|
|
|||
137
README.md
137
README.md
|
|
@ -483,6 +483,143 @@ def compress_tool_output(content: str, context: str = "") -> str:
|
|||
|
||||
---
|
||||
|
||||
## ML-Based Compression with LLMLingua-2 (Optional)
|
||||
|
||||
For even more aggressive compression, Headroom integrates with **LLMLingua-2**, Microsoft's BERT-based token classifier trained via GPT-4 distillation. It achieves **up to 20x compression** while preserving semantic meaning.
|
||||
|
||||
### When to Use LLMLingua-2
|
||||
|
||||
| Approach | Best For | Compression | Speed |
|
||||
|----------|----------|-------------|-------|
|
||||
| **SmartCrusher** | JSON tool outputs | 70-90% | ~1ms |
|
||||
| **Text Utilities** | Search/logs | 50-90% | ~1ms |
|
||||
| **LLMLingua-2** | Any text, max compression | 80-95% | ~50-200ms |
|
||||
|
||||
LLMLingua-2 is ideal when you need maximum compression and can tolerate slightly higher latency (e.g., compressing large tool outputs before storage, offline processing).
|
||||
|
||||
### Installation
|
||||
|
||||
```bash
|
||||
# Adds ~2GB of model weights
|
||||
pip install "headroom-ai[llmlingua]"
|
||||
```
|
||||
|
||||
### Basic Usage
|
||||
|
||||
```python
|
||||
from headroom.transforms import LLMLinguaCompressor
|
||||
|
||||
# Create compressor (model loaded lazily on first use)
|
||||
compressor = LLMLinguaCompressor()
|
||||
|
||||
# Compress any text
|
||||
long_output = "The function processUserData takes a user object and validates..."
|
||||
result = compressor.compress(long_output)
|
||||
|
||||
print(f"Before: {result.original_tokens} tokens")
|
||||
print(f"After: {result.compressed_tokens} tokens")
|
||||
print(f"Saved: {result.savings_percentage:.1f}%")
|
||||
print(result.compressed)
|
||||
```
|
||||
|
||||
### Content-Aware Compression
|
||||
|
||||
LLMLingua-2 automatically adjusts compression based on content type:
|
||||
|
||||
```python
|
||||
from headroom.transforms import LLMLinguaCompressor, LLMLinguaConfig
|
||||
|
||||
# Conservative for code (keep 40% of tokens)
|
||||
config = LLMLinguaConfig(
|
||||
code_compression_rate=0.4, # More conservative
|
||||
json_compression_rate=0.35, # Moderate
|
||||
text_compression_rate=0.25, # Aggressive
|
||||
)
|
||||
|
||||
compressor = LLMLinguaCompressor(config)
|
||||
|
||||
# Auto-detects content type
|
||||
code_result = compressor.compress("def calculate(x): return x * 2")
|
||||
text_result = compressor.compress("This is a verbose explanation...")
|
||||
```
|
||||
|
||||
### Memory Management
|
||||
|
||||
The model uses ~1GB RAM. Unload it when done:
|
||||
|
||||
```python
|
||||
from headroom.transforms import (
|
||||
LLMLinguaCompressor,
|
||||
unload_llmlingua_model,
|
||||
is_llmlingua_model_loaded,
|
||||
)
|
||||
|
||||
compressor = LLMLinguaCompressor()
|
||||
result = compressor.compress(content) # Model loaded here
|
||||
|
||||
# Check if loaded
|
||||
print(is_llmlingua_model_loaded()) # True
|
||||
|
||||
# Free memory when done
|
||||
unload_llmlingua_model() # Frees ~1GB
|
||||
print(is_llmlingua_model_loaded()) # False
|
||||
|
||||
# Next compression will reload automatically
|
||||
```
|
||||
|
||||
### Use in Pipeline
|
||||
|
||||
```python
|
||||
from headroom.transforms import TransformPipeline, LLMLinguaCompressor, SmartCrusher
|
||||
|
||||
# Combine with other transforms
|
||||
pipeline = TransformPipeline([
|
||||
SmartCrusher(), # First: compress JSON
|
||||
LLMLinguaCompressor(), # Then: ML compression on remaining text
|
||||
])
|
||||
|
||||
result = pipeline.apply(messages, tokenizer)
|
||||
```
|
||||
|
||||
### Device Configuration
|
||||
|
||||
```python
|
||||
from headroom.transforms import LLMLinguaConfig, LLMLinguaCompressor
|
||||
|
||||
# Force CPU (slower but works everywhere)
|
||||
config = LLMLinguaConfig(device="cpu")
|
||||
|
||||
# Force GPU (faster but needs CUDA)
|
||||
config = LLMLinguaConfig(device="cuda")
|
||||
|
||||
# Auto-detect (default): uses CUDA > MPS > CPU
|
||||
config = LLMLinguaConfig(device="auto")
|
||||
|
||||
compressor = LLMLinguaCompressor(config)
|
||||
```
|
||||
|
||||
### Proxy Integration (Opt-In)
|
||||
|
||||
Enable LLMLingua in the proxy server for automatic ML compression of all requests:
|
||||
|
||||
```bash
|
||||
# Enable LLMLingua in proxy (requires: pip install headroom-ai[llmlingua,proxy])
|
||||
headroom proxy --llmlingua
|
||||
|
||||
# With custom settings
|
||||
headroom proxy --llmlingua --llmlingua-device cuda --llmlingua-rate 0.4
|
||||
|
||||
# The proxy shows LLMLingua status at startup:
|
||||
# LLMLingua: ENABLED (device=cuda, rate=0.4)
|
||||
#
|
||||
# If llmlingua is installed but not enabled, you'll see a helpful hint:
|
||||
# LLMLingua: available (enable with --llmlingua for ML compression)
|
||||
```
|
||||
|
||||
**Why opt-in?** LLMLingua adds ~2GB dependencies and 10-30s cold start. The default proxy is lightweight (~50MB) with <5ms overhead. Enable LLMLingua when you need maximum compression and can accept the tradeoffs.
|
||||
|
||||
---
|
||||
|
||||
## Metrics & Monitoring
|
||||
|
||||
### Prometheus Metrics (Proxy)
|
||||
|
|
|
|||
|
|
@ -193,7 +193,36 @@ analysis = {
|
|||
|
||||
---
|
||||
|
||||
#### Transform 4: Rolling Window
|
||||
#### Transform 4: LLMLingua Compressor (Optional)
|
||||
|
||||
**When to use:** Maximum compression needed and latency is acceptable.
|
||||
|
||||
```python
|
||||
# Opt-in ML-based compression using Microsoft's LLMLingua-2
|
||||
# BERT-based token classifier trained via GPT-4 distillation
|
||||
|
||||
# Before: Long tool output text
|
||||
"The function processUserData takes a user object and validates all fields..."
|
||||
|
||||
# After: Compressed while preserving semantic meaning
|
||||
"function processUserData validates user fields..."
|
||||
```
|
||||
|
||||
**Key characteristics:**
|
||||
- Uses `microsoft/llmlingua-2-xlm-roberta-large-meetingbank` model
|
||||
- Auto-detects content type (code, JSON, text) for optimal compression rates
|
||||
- Stores original in CCR for retrieval if needed
|
||||
- Adds 50-200ms latency per request
|
||||
- Requires ~1GB RAM when loaded
|
||||
|
||||
**Proxy integration (opt-in):**
|
||||
```bash
|
||||
headroom proxy --llmlingua --llmlingua-device cuda
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
#### Transform 5: Rolling Window
|
||||
|
||||
**Problem:** Even after compression, you might exceed the model's context limit.
|
||||
|
||||
|
|
@ -282,7 +311,12 @@ def apply(self, messages, ...):
|
|||
# - Compresses to 17 points (preserving spike)
|
||||
# - Factors out constant "host" field
|
||||
|
||||
# Transform 3: Rolling Window
|
||||
# Transform 3: LLMLingua (if enabled via --llmlingua)
|
||||
# - ML-based compression on remaining long text
|
||||
# - Auto-detects content type for optimal rate
|
||||
# - Stores original in CCR for retrieval
|
||||
|
||||
# Transform 4: Rolling Window
|
||||
# - Checks if we're under limit (we are)
|
||||
# - No drops needed
|
||||
|
||||
|
|
@ -670,12 +704,13 @@ headroom/
|
|||
│ └── anthropic.py # Anthropic-specific
|
||||
│
|
||||
├── transforms/
|
||||
│ ├── base.py # Transform protocol
|
||||
│ ├── pipeline.py # Orchestrates all transforms
|
||||
│ ├── cache_aligner.py # Date extraction for caching
|
||||
│ ├── tool_crusher.py # Naive compression (disabled)
|
||||
│ ├── smart_crusher.py # Statistical compression (default)
|
||||
│ └── rolling_window.py # Token limit enforcement
|
||||
│ ├── base.py # Transform protocol
|
||||
│ ├── pipeline.py # Orchestrates all transforms
|
||||
│ ├── cache_aligner.py # Date extraction for caching
|
||||
│ ├── tool_crusher.py # Naive compression (disabled)
|
||||
│ ├── smart_crusher.py # Statistical compression (default)
|
||||
│ ├── rolling_window.py # Token limit enforcement
|
||||
│ └── llmlingua_compressor.py # ML-based compression (opt-in)
|
||||
│
|
||||
├── cache/ # CCR Architecture
|
||||
│ ├── compression_store.py # Phase 1: Store original content
|
||||
|
|
|
|||
|
|
@ -21,6 +21,8 @@ headroom proxy \
|
|||
|
||||
## Command Line Options
|
||||
|
||||
### Core Options
|
||||
|
||||
| Option | Default | Description |
|
||||
|--------|---------|-------------|
|
||||
| `--host` | `127.0.0.1` | Host to bind to |
|
||||
|
|
@ -31,6 +33,27 @@ headroom proxy \
|
|||
| `--log-file` | None | Path to JSONL log file |
|
||||
| `--budget` | None | Daily budget limit in USD |
|
||||
|
||||
### LLMLingua Options (ML Compression)
|
||||
|
||||
| Option | Default | Description |
|
||||
|--------|---------|-------------|
|
||||
| `--llmlingua` | `false` | Enable LLMLingua-2 ML-based compression |
|
||||
| `--llmlingua-device` | `auto` | Device for model: `auto`, `cuda`, `cpu`, `mps` |
|
||||
| `--llmlingua-rate` | `0.3` | Target compression rate (0.3 = keep 30% of tokens) |
|
||||
|
||||
**Note:** LLMLingua requires additional dependencies: `pip install headroom-ai[llmlingua]`
|
||||
|
||||
```bash
|
||||
# Enable LLMLingua with GPU acceleration
|
||||
headroom proxy --llmlingua --llmlingua-device cuda
|
||||
|
||||
# More aggressive compression (keep only 20%)
|
||||
headroom proxy --llmlingua --llmlingua-rate 0.2
|
||||
|
||||
# Conservative compression for code (keep 50%)
|
||||
headroom proxy --llmlingua --llmlingua-rate 0.5
|
||||
```
|
||||
|
||||
## API Endpoints
|
||||
|
||||
### Health Check
|
||||
|
|
@ -104,6 +127,42 @@ client = OpenAI(
|
|||
|
||||
## Features
|
||||
|
||||
### LLMLingua ML Compression (Opt-In)
|
||||
|
||||
When enabled, the proxy uses Microsoft's LLMLingua-2 model for ML-based token compression:
|
||||
|
||||
```bash
|
||||
headroom proxy --llmlingua
|
||||
```
|
||||
|
||||
**How it works:**
|
||||
- LLMLinguaCompressor is added to the transform pipeline (before RollingWindow)
|
||||
- Automatically detects content type (JSON, code, text) and adjusts compression
|
||||
- Stores original content in CCR for retrieval if needed
|
||||
|
||||
**Startup feedback:**
|
||||
|
||||
```
|
||||
# When enabled and available:
|
||||
LLMLingua: ENABLED (device=cuda, rate=0.3)
|
||||
|
||||
# When installed but not enabled (helpful hint):
|
||||
LLMLingua: available (enable with --llmlingua for ML compression)
|
||||
|
||||
# When enabled but not installed:
|
||||
WARNING: LLMLingua requested but not installed. Install with: pip install headroom-ai[llmlingua]
|
||||
```
|
||||
|
||||
**Why opt-in?**
|
||||
| Concern | Default Proxy | With LLMLingua |
|
||||
|---------|---------------|----------------|
|
||||
| Dependencies | ~50MB | +2GB (torch, transformers) |
|
||||
| Cold start | <1s | 10-30s (model load) |
|
||||
| Memory | ~100MB | +1GB (model in RAM) |
|
||||
| Overhead | <5ms | 50-200ms per request |
|
||||
|
||||
Enable LLMLingua when maximum compression justifies the resource cost.
|
||||
|
||||
### Semantic Caching
|
||||
|
||||
The proxy caches responses for repeated queries:
|
||||
|
|
|
|||
|
|
@ -162,6 +162,76 @@ config = RollingWindowConfig(
|
|||
|
||||
---
|
||||
|
||||
## LLMLinguaCompressor (Optional)
|
||||
|
||||
ML-based compression using Microsoft's LLMLingua-2 model.
|
||||
|
||||
### When to Use
|
||||
|
||||
| Transform | Best For | Speed | Compression |
|
||||
|-----------|----------|-------|-------------|
|
||||
| SmartCrusher | JSON arrays | ~1ms | 70-90% |
|
||||
| Text Utilities | Search/logs | ~1ms | 50-90% |
|
||||
| **LLMLinguaCompressor** | Any text, max compression | 50-200ms | 80-95% |
|
||||
|
||||
### Installation
|
||||
|
||||
```bash
|
||||
pip install "headroom-ai[llmlingua]" # Adds ~2GB
|
||||
```
|
||||
|
||||
### Configuration
|
||||
|
||||
```python
|
||||
from headroom.transforms import LLMLinguaCompressor, LLMLinguaConfig
|
||||
|
||||
config = LLMLinguaConfig(
|
||||
device="auto", # auto, cuda, cpu, mps
|
||||
target_compression_rate=0.3, # Keep 30% of tokens
|
||||
min_tokens_for_compression=100, # Skip small content
|
||||
code_compression_rate=0.4, # Conservative for code
|
||||
json_compression_rate=0.35, # Moderate for JSON
|
||||
text_compression_rate=0.25, # Aggressive for text
|
||||
enable_ccr=True, # Store original for retrieval
|
||||
)
|
||||
|
||||
compressor = LLMLinguaCompressor(config)
|
||||
```
|
||||
|
||||
### Content-Aware Rates
|
||||
|
||||
LLMLinguaCompressor auto-detects content type:
|
||||
|
||||
| Content Type | Default Rate | Behavior |
|
||||
|--------------|--------------|----------|
|
||||
| Code | 0.4 | Conservative - preserves syntax |
|
||||
| JSON | 0.35 | Moderate - keeps structure |
|
||||
| Text | 0.3 | Aggressive - maximum compression |
|
||||
|
||||
### Memory Management
|
||||
|
||||
```python
|
||||
from headroom.transforms import (
|
||||
is_llmlingua_model_loaded,
|
||||
unload_llmlingua_model,
|
||||
)
|
||||
|
||||
# Check if model is loaded
|
||||
print(is_llmlingua_model_loaded()) # True/False
|
||||
|
||||
# Free ~1GB RAM when done
|
||||
unload_llmlingua_model()
|
||||
```
|
||||
|
||||
### Proxy Integration
|
||||
|
||||
```bash
|
||||
# Enable in proxy
|
||||
headroom proxy --llmlingua --llmlingua-device cuda --llmlingua-rate 0.3
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## TransformPipeline
|
||||
|
||||
Combine transforms for optimal results.
|
||||
|
|
@ -179,11 +249,36 @@ result = pipeline.transform(messages)
|
|||
print(f"Saved {result.tokens_saved} tokens")
|
||||
```
|
||||
|
||||
### With LLMLingua (Optional)
|
||||
|
||||
```python
|
||||
from headroom.transforms import (
|
||||
TransformPipeline, SmartCrusher, CacheAligner,
|
||||
RollingWindow, LLMLinguaCompressor
|
||||
)
|
||||
|
||||
pipeline = TransformPipeline([
|
||||
CacheAligner(), # 1. Stabilize prefix
|
||||
SmartCrusher(), # 2. Compress JSON arrays
|
||||
LLMLinguaCompressor(), # 3. ML compression on remaining text
|
||||
RollingWindow(), # 4. Final size constraint (always last)
|
||||
])
|
||||
```
|
||||
|
||||
### Recommended Order
|
||||
|
||||
1. **SmartCrusher** - Reduce individual messages
|
||||
2. **CacheAligner** - Optimize for caching
|
||||
3. **RollingWindow** - Final size constraint
|
||||
| Order | Transform | Purpose |
|
||||
|-------|-----------|---------|
|
||||
| 1 | CacheAligner | Stabilize prefix for caching |
|
||||
| 2 | SmartCrusher | Compress JSON tool outputs |
|
||||
| 3 | LLMLinguaCompressor | ML compression (optional) |
|
||||
| 4 | RollingWindow | Enforce token limits (always last) |
|
||||
|
||||
**Why this order?**
|
||||
- CacheAligner first to maximize prefix stability
|
||||
- SmartCrusher handles JSON arrays efficiently
|
||||
- LLMLingua compresses remaining long text
|
||||
- RollingWindow truncates only if still over limit
|
||||
|
||||
---
|
||||
|
||||
|
|
|
|||
|
|
@ -59,7 +59,17 @@ from headroom.config import CacheAlignerConfig, RollingWindowConfig, SmartCrushe
|
|||
from headroom.providers import AnthropicProvider, OpenAIProvider
|
||||
from headroom.telemetry import get_telemetry_collector
|
||||
from headroom.tokenizers import get_tokenizer
|
||||
from headroom.transforms import CacheAligner, RollingWindow, SmartCrusher, TransformPipeline
|
||||
from headroom.transforms import (
|
||||
_LLMLINGUA_AVAILABLE,
|
||||
CacheAligner,
|
||||
RollingWindow,
|
||||
SmartCrusher,
|
||||
TransformPipeline,
|
||||
)
|
||||
|
||||
# Conditionally import LLMLingua if available
|
||||
if _LLMLINGUA_AVAILABLE:
|
||||
from headroom.transforms import LLMLinguaCompressor, LLMLinguaConfig
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
|
||||
|
|
@ -145,6 +155,11 @@ class ProxyConfig:
|
|||
ccr_inject_tool: bool = True # Inject headroom_retrieve tool when compression occurs
|
||||
ccr_inject_system_instructions: bool = False # Add instructions to system message
|
||||
|
||||
# LLMLingua ML-based compression (opt-in)
|
||||
llmlingua_enabled: bool = False # 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%)
|
||||
|
||||
# Caching
|
||||
cache_enabled: bool = True
|
||||
cache_ttl_seconds: int = 3600 # 1 hour
|
||||
|
|
@ -656,6 +671,9 @@ class HeadroomProxy:
|
|||
),
|
||||
]
|
||||
|
||||
# Add LLMLingua if enabled and available
|
||||
self._llmlingua_status = self._setup_llmlingua(config, transforms)
|
||||
|
||||
self.anthropic_pipeline = TransformPipeline(
|
||||
transforms=transforms,
|
||||
provider=self.anthropic_provider,
|
||||
|
|
@ -722,6 +740,38 @@ class HeadroomProxy:
|
|||
inject_system_instructions=config.ccr_inject_system_instructions,
|
||||
)
|
||||
|
||||
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"
|
||||
|
||||
async def startup(self):
|
||||
"""Initialize async resources."""
|
||||
self.http_client = httpx.AsyncClient(
|
||||
|
|
@ -737,6 +787,18 @@ class HeadroomProxy:
|
|||
logger.info(f"Caching: {'ENABLED' if self.config.cache_enabled else 'DISABLED'}")
|
||||
logger.info(f"Rate Limiting: {'ENABLED' if self.config.rate_limit_enabled else 'DISABLED'}")
|
||||
|
||||
# LLMLingua status with helpful hint
|
||||
if self._llmlingua_status == "enabled":
|
||||
logger.info(
|
||||
f"LLMLingua: ENABLED (device={self.config.llmlingua_device}, "
|
||||
f"rate={self.config.llmlingua_target_rate})"
|
||||
)
|
||||
elif self._llmlingua_status == "available":
|
||||
logger.info(
|
||||
"LLMLingua: available but not enabled. "
|
||||
"Enable with --llmlingua for ML-based compression (3-5x better on text/logs)"
|
||||
)
|
||||
|
||||
async def shutdown(self):
|
||||
"""Cleanup async resources."""
|
||||
if self.http_client:
|
||||
|
|
@ -1818,6 +1880,21 @@ 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 "REQUESTED but not installed (pip install headroom-ai[llmlingua])"
|
||||
else:
|
||||
if _LLMLINGUA_AVAILABLE:
|
||||
return "available (enable with --llmlingua for ML compression)"
|
||||
return "DISABLED"
|
||||
|
||||
|
||||
def run_server(config: ProxyConfig | None = None):
|
||||
"""Run the proxy server."""
|
||||
if not FASTAPI_AVAILABLE:
|
||||
|
|
@ -1827,6 +1904,8 @@ def run_server(config: ProxyConfig | None = None):
|
|||
config = config or ProxyConfig()
|
||||
app = create_app(config)
|
||||
|
||||
llmlingua_status = _get_llmlingua_banner_status(config)
|
||||
|
||||
print(f"""
|
||||
╔══════════════════════════════════════════════════════════════════════╗
|
||||
║ HEADROOM PROXY SERVER ║
|
||||
|
|
@ -1840,6 +1919,7 @@ def run_server(config: ProxyConfig | None = None):
|
|||
║ 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}║
|
||||
╠══════════════════════════════════════════════════════════════════════╣
|
||||
║ USAGE: ║
|
||||
║ Claude Code: ANTHROPIC_BASE_URL=http://{config.host}:{config.port} claude ║
|
||||
|
|
@ -1893,6 +1973,25 @@ if __name__ == "__main__":
|
|||
parser.add_argument("--log-file", help="Log file path")
|
||||
parser.add_argument("--log-messages", action="store_true", help="Log full messages")
|
||||
|
||||
# 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(
|
||||
"--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%%)",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
config = ProxyConfig(
|
||||
|
|
@ -1910,6 +2009,9 @@ if __name__ == "__main__":
|
|||
budget_period=args.budget_period,
|
||||
log_file=args.log_file,
|
||||
log_full_messages=args.log_messages,
|
||||
llmlingua_enabled=args.llmlingua,
|
||||
llmlingua_device=args.llmlingua_device,
|
||||
llmlingua_target_rate=args.llmlingua_rate,
|
||||
)
|
||||
|
||||
run_server(config)
|
||||
|
|
|
|||
|
|
@ -15,6 +15,21 @@ from .smart_crusher import SmartCrusher, SmartCrusherConfig
|
|||
from .text_compressor import TextCompressionResult, TextCompressor, TextCompressorConfig
|
||||
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
|
||||
|
||||
__all__ = [
|
||||
# Base
|
||||
"Transform",
|
||||
|
|
@ -39,4 +54,19 @@ __all__ = [
|
|||
# Other transforms
|
||||
"CacheAligner",
|
||||
"RollingWindow",
|
||||
# ML-based compression (optional)
|
||||
"_LLMLINGUA_AVAILABLE",
|
||||
]
|
||||
|
||||
# Conditionally add LLMLingua exports
|
||||
if _LLMLINGUA_AVAILABLE:
|
||||
__all__.extend(
|
||||
[
|
||||
"LLMLinguaCompressor",
|
||||
"LLMLinguaConfig",
|
||||
"LLMLinguaResult",
|
||||
"compress_with_llmlingua",
|
||||
"is_llmlingua_model_loaded",
|
||||
"unload_llmlingua_model",
|
||||
]
|
||||
)
|
||||
|
|
|
|||
633
headroom/transforms/llmlingua_compressor.py
Normal file
633
headroom/transforms/llmlingua_compressor.py
Normal file
|
|
@ -0,0 +1,633 @@
|
|||
"""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 ..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 = "microsoft/llmlingua-2-xlm-roberta-large-meetingbank"
|
||||
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.4 # More conservative for code
|
||||
json_compression_rate: float = 0.35 # Slightly conservative for JSON
|
||||
text_compression_rate: float = 0.25 # More aggressive for plain text
|
||||
|
||||
|
||||
@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,
|
||||
) -> 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.
|
||||
|
||||
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:
|
||||
result = compressor.compress_prompt(
|
||||
original_prompt=content,
|
||||
rate=compression_rate,
|
||||
force_tokens=force_tokens if force_tokens else None,
|
||||
drop_consecutive=self.config.drop_consecutive,
|
||||
)
|
||||
|
||||
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:
|
||||
compressed += (
|
||||
f"\n[LLMLingua: {original_tokens}→{compressed_tokens} tokens. 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", "")
|
||||
|
||||
# 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
|
||||
|
|
@ -64,6 +64,12 @@ proxy = [
|
|||
reports = [
|
||||
"jinja2>=3.0.0",
|
||||
]
|
||||
# ML-based compression (LLMLingua-2)
|
||||
llmlingua = [
|
||||
"llmlingua>=0.2.0",
|
||||
"torch>=2.0.0",
|
||||
"transformers>=4.30.0",
|
||||
]
|
||||
# Development dependencies
|
||||
dev = [
|
||||
"pytest>=7.0.0",
|
||||
|
|
@ -76,7 +82,7 @@ dev = [
|
|||
]
|
||||
# All optional dependencies
|
||||
all = [
|
||||
"headroom-ai[relevance,proxy,reports]",
|
||||
"headroom-ai[relevance,proxy,reports,llmlingua]",
|
||||
]
|
||||
|
||||
[project.scripts]
|
||||
|
|
|
|||
458
tests/test_proxy_llmlingua.py
Normal file
458
tests/test_proxy_llmlingua.py
Normal file
|
|
@ -0,0 +1,458 @@
|
|||
"""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,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def llmlingua_config():
|
||||
"""Config with LLMLingua enabled."""
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
@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_disabled(self):
|
||||
"""LLMLingua is disabled by default."""
|
||||
config = ProxyConfig()
|
||||
|
||||
assert config.llmlingua_enabled is False
|
||||
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,
|
||||
)
|
||||
|
||||
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)
|
||||
|
||||
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)
|
||||
|
||||
with patch("headroom.proxy.server._LLMLINGUA_AVAILABLE", True):
|
||||
status = _get_llmlingua_banner_status(config)
|
||||
assert "available" in status
|
||||
assert "--llmlingua" in status
|
||||
|
||||
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)
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
# 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,
|
||||
)
|
||||
config_disabled = ProxyConfig(
|
||||
llmlingua_enabled=False,
|
||||
optimize=False,
|
||||
cache_enabled=False,
|
||||
rate_limit_enabled=False,
|
||||
)
|
||||
|
||||
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
|
||||
941
tests/test_transforms/test_llmlingua_compressor.py
Normal file
941
tests/test_transforms/test_llmlingua_compressor.py
Normal file
|
|
@ -0,0 +1,941 @@
|
|||
"""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-xlm-roberta-large-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 should be more conservative
|
||||
assert config.code_compression_rate > config.text_compression_rate
|
||||
# JSON should be between code and text
|
||||
assert config.json_compression_rate > config.text_compression_rate
|
||||
assert config.json_compression_rate < config.code_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