akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

v0.2.2: Add CCR Response Handler, Context Tracker, and restructure docs

Browse source 
Features:
- CCR Response Handler: Automatically intercepts and handles headroom_retrieve tool calls
- CCR Context Tracker: Multi-turn awareness with proactive expansion of relevant compressed content
- New CCR demo script showing before/after flow

Documentation:
- Restructured README from 885 lines to 190 lines for better DevEx
- Split detailed docs into focused guides: ccr.md, sdk.md, configuration.md,
  text-compression.md, llmlingua.md, metrics.md, errors.md
- Updated docs/README.md index with all new documentation

Tests:
- Added comprehensive tests for Response Handler (32 tests)
- Added comprehensive tests for Context Tracker (32 tests)
- All 977 tests passing
This commit is contained in:
chopratejas 2026-01-14 13:03:41 -08:00
parent 45633b69ab
commit d724f14022
19 changed files with 5227 additions and 721 deletions
Download patch file Download diff file Expand all files Collapse all files

762
README.md
View file

@ -25,45 +25,30 @@
---
## Why Headroom?
## What It Does
AI coding agents and tool-using applications generate **massive contexts**:
Headroom is a **smart compression proxy** for LLM applications:
- Tool outputs with 1000s of search results, log entries, API responses
- Long conversation histories that hit token limits
- System prompts with dynamic dates that break provider caching
- **Compresses tool outputs** — 1000 search results → 15 items (keeps errors, anomalies, relevant items)
- **Enables provider caching** — Stabilizes prefixes so cache hits actually happen
- **Manages context windows** — Prevents token limit failures without breaking tool calls
- **Reversible compression** — LLM can retrieve original data if needed ([CCR architecture](docs/ccr.md))
**Result**: You pay for tokens you don't need, and cache hits are rare.
Headroom is a **smart compression layer** that sits between your app and LLM providers:
| Transform | What It Does | Savings |
|-----------|--------------|---------|
| **SmartCrusher** | Compresses JSON tool outputs statistically (keeps errors, anomalies, relevant items) | 70-90% |
| **CacheAligner** | Stabilizes prefixes so provider caching works | Up to 10x |
| **RollingWindow** | Manages context within limits without breaking tool calls | Prevents failures |
| **Text Utilities** | Opt-in compression for search results, build logs, plain text | 50-90% |
**Zero accuracy loss** - we keep what matters: errors, anomalies, relevant items.
**Zero code changes required** — point your existing tools at the proxy.
---
## 5-Minute Quickstart
### Option 1: Proxy Server (Recommended)
Works with **any** OpenAI-compatible client without code changes:
## 30-Second Quickstart
```bash
# Install
pip install "headroom-ai[proxy]"
# Start the proxy
# Start proxy
headroom proxy --port 8787
# Verify it's running
# Verify
curl http://localhost:8787/health
# Expected: {"status": "healthy", ...}
```
**Use with your tools:**
@ -73,736 +58,121 @@ curl http://localhost:8787/health
ANTHROPIC_BASE_URL=http://localhost:8787 claude
# Cursor / Continue / any OpenAI client
OPENAI_BASE_URL=http://localhost:8787/v1 your-app
OPENAI_BASE_URL=http://localhost:8787/v1 cursor
# Python OpenAI SDK
# Python scripts
export OPENAI_BASE_URL=http://localhost:8787/v1
python your_script.py
```
### Option 2: Python SDK
Wrap your existing client for fine-grained control:
```bash
pip install headroom-ai openai
```
```python
from headroom import HeadroomClient, OpenAIProvider
from openai import OpenAI
# Create wrapped client
client = HeadroomClient(
original_client=OpenAI(),
provider=OpenAIProvider(),
default_mode="optimize", # or "audit" to observe only
)
# Use exactly like the original client
response = client.chat.completions.create(
model="gpt-4o-mini",
messages=[
{"role": "user", "content": "Hello!"},
],
)
print(response.choices[0].message.content)
# Check what happened
stats = client.get_stats()
print(f"Tokens saved this session: {stats['session']['tokens_saved_total']}")
```
**With tool outputs (where real savings happen):**
```python
import json
# Conversation with large tool output
messages = [
{"role": "user", "content": "Search for Python tutorials"},
{
"role": "assistant",
"content": None,
"tool_calls": [{
"id": "call_123",
"type": "function",
"function": {"name": "search", "arguments": '{"q": "python"}'},
}],
},
{
"role": "tool",
"tool_call_id": "call_123",
"content": json.dumps({
"results": [{"title": f"Tutorial {i}", "score": 100-i} for i in range(500)]
}),
},
{"role": "user", "content": "What are the top 3?"},
]
# Headroom compresses 500 results to ~15, keeping highest-scoring items
response = client.chat.completions.create(model="gpt-4o-mini", messages=messages)
print(f"Tokens saved: {client.get_stats()['session']['tokens_saved_total']}")
# Typical output: "Tokens saved: 3500"
```
### Option 3: LangChain Integration (Coming Soon)
```python
# Coming soon - use proxy server for now
# OPENAI_BASE_URL=http://localhost:8787/v1 python your_langchain_app.py
```
That's it. You're saving tokens.
---
## Verify It's Working
### Check Proxy Stats
```bash
curl http://localhost:8787/stats
```
```json
{
"requests": {"total": 42, "cached": 5, "rate_limited": 0, "failed": 0},
"tokens": {"input": 50000, "output": 8000, "saved": 12500, "savings_percent": 25.0},
"cost": {"total_cost_usd": 0.15, "total_savings_usd": 0.04},
"cache": {"entries": 10, "total_hits": 5}
"tokens": {"saved": 12500, "savings_percent": 25.0},
"cost": {"total_savings_usd": 0.04}
}
```
### Check SDK Stats
```python
# Quick session stats (no database query)
stats = client.get_stats()
print(stats)
# {
# "session": {"requests_total": 10, "tokens_saved_total": 5000, ...},
# "config": {"mode": "optimize", "provider": "openai", ...},
# "transforms": {"smart_crusher_enabled": True, ...}
# }
# Validate setup is correct
result = client.validate_setup()
if not result["valid"]:
print("Setup issues:", result)
```
### Enable Logging
```python
import logging
logging.basicConfig(level=logging.INFO)
# Now you'll see:
# INFO:headroom.transforms.pipeline:Pipeline complete: 45000 -> 4500 tokens (saved 40500, 90.0% reduction)
# INFO:headroom.transforms.smart_crusher:SmartCrusher applied top_n strategy: kept 15 of 1000 items
```
---
## Installation
```bash
# Core only (minimal dependencies: tiktoken, pydantic)
pip install headroom-ai
# With semantic relevance scoring (adds sentence-transformers)
pip install "headroom-ai[relevance]"
# With proxy server (adds fastapi, uvicorn)
pip install "headroom-ai[proxy]"
# With HTML reports (adds jinja2)
pip install "headroom-ai[reports]"
# Everything
pip install "headroom-ai[all]"
pip install "headroom-ai[proxy]" # Proxy server (recommended)
pip install headroom-ai # SDK only
pip install "headroom-ai[all]" # Everything
```
**Requirements**: Python 3.10+
---
## Configuration
## Features
### SDK Configuration
```python
from headroom import HeadroomClient, OpenAIProvider
from openai import OpenAI
# Full configuration example
client = HeadroomClient(
original_client=OpenAI(),
provider=OpenAIProvider(),
default_mode="optimize", # "audit" (observe only) or "optimize" (apply transforms)
enable_cache_optimizer=True, # Enable provider-specific cache optimization
enable_semantic_cache=False, # Enable query-level semantic caching
model_context_limits={ # Override default context limits
"gpt-4o": 128000,
"gpt-4o-mini": 128000,
},
# store_url defaults to temp directory; override with absolute path if needed:
# store_url="sqlite:////absolute/path/to/headroom.db",
)
```
### Proxy Configuration
```bash
# Via command line
headroom proxy \
--port 8787 \
--budget 10.00 \
--log-file headroom.jsonl
# Disable optimization (passthrough mode)
headroom proxy --no-optimize
# Disable semantic caching
headroom proxy --no-cache
# See all options
headroom proxy --help
```
### Per-Request Overrides
```python
# Override mode for specific requests
response = client.chat.completions.create(
model="gpt-4o",
messages=[...],
headroom_mode="audit", # Just observe, don't optimize
headroom_output_buffer_tokens=8000, # Reserve more for output
headroom_keep_turns=5, # Keep last 5 turns
)
```
| Feature | Description | Docs |
|---------|-------------|------|
| **SmartCrusher** | Compresses JSON tool outputs statistically | [Transforms](docs/transforms.md) |
| **CacheAligner** | Stabilizes prefixes for provider caching | [Transforms](docs/transforms.md) |
| **RollingWindow** | Manages context limits without breaking tools | [Transforms](docs/transforms.md) |
| **CCR** | Reversible compression with automatic retrieval | [CCR Guide](docs/ccr.md) |
| **Text Utilities** | Opt-in compression for search/logs | [Text Compression](docs/text-compression.md) |
| **LLMLingua-2** | ML-based 20x compression (opt-in) | [LLMLingua](docs/llmlingua.md) |
---
## Modes
| Mode | Behavior | Use Case |
|------|----------|----------|
| `audit` | Observes and logs, no modifications | Production monitoring, baseline measurement |
| `optimize` | Applies safe, deterministic transforms | Production optimization |
| `simulate` | Returns plan without API call | Testing, cost estimation |
```python
# Simulate to see what would happen
plan = client.chat.completions.simulate(
model="gpt-4o",
messages=large_conversation,
)
print(f"Would save {plan.tokens_saved} tokens")
print(f"Transforms: {plan.transforms}")
print(f"Estimated savings: {plan.estimated_savings}")
```
---
## Error Handling
Headroom provides explicit exceptions for debugging:
```python
from headroom import (
HeadroomClient,
HeadroomError, # Base class - catch all Headroom errors
ConfigurationError, # Invalid configuration
ProviderError, # Provider issues (unknown model, etc.)
StorageError, # Database/storage failures
CompressionError, # Compression failures (rare - we fail safe)
ValidationError, # Setup validation failures
)
try:
client = HeadroomClient(...)
response = client.chat.completions.create(...)
except ConfigurationError as e:
print(f"Config issue: {e}")
print(f"Details: {e.details}") # Additional context
except StorageError as e:
print(f"Storage issue: {e}")
# Headroom continues to work, just without metrics persistence
except HeadroomError as e:
print(f"Headroom error: {e}")
```
**Safety guarantee**: If compression fails, the original content passes through unchanged. Your LLM calls never fail due to Headroom.
---
## How It Works
### SmartCrusher: Statistical Compression
```python
# Before: 50KB tool response with 1000 items
{"results": [{"id": 1, "status": "ok", ...}, ... 1000 items ...]}
# After: ~2KB with important items preserved
# Headroom keeps:
# - First 3 items (context)
# - Last 2 items (recency)
# - All error items (status != "ok")
# - Statistical anomalies (values > 2 std dev from mean)
# - Items matching user's query (BM25/embedding similarity)
```
### CacheAligner: Prefix Stabilization
```python
# Before: Cache miss every day due to changing date
"You are helpful. Today is January 7, 2025."
# After: Stable prefix (cache hit!) + dynamic context moved to end
"You are helpful."
# Dynamic content: "Current date: January 7, 2025"
```
### RollingWindow: Context Management
```python
# When context exceeds limit:
# 1. Drop oldest tool outputs first (as atomic units with their calls)
# 2. Drop oldest conversation turns
# 3. NEVER drop: system prompt, last N turns, orphaned tool responses
```
---
## Text Compression Utilities (Opt-In)
For coding tasks, Headroom provides **standalone text compression utilities** that applications can use explicitly. These are **opt-in** - they're not applied automatically, giving you full control over when and how to compress text content.
> **Design Philosophy**: SmartCrusher compresses JSON automatically because it's structure-preserving and safe. Text compression is lossy and context-dependent, so applications should decide when to use it.
### Available Utilities
| Utility | Input Type | Use Case |
|---------|------------|----------|
| `SearchCompressor` | grep/ripgrep output | Search results with `file:line:content` format |
| `LogCompressor` | Build/test logs | pytest, npm, cargo, make output |
| `TextCompressor` | Generic text | Any plain text with anchor preservation |
| `detect_content_type` | Any content | Detect content type for routing decisions |
### Example: Compressing Search Results
```python
from headroom.transforms import SearchCompressor
# Your grep/ripgrep output (could be 1000s of lines)
search_results = """
src/utils.py:42:def process_data(items):
src/utils.py:43: \"\"\"Process items.\"\"\"
src/models.py:15:class DataProcessor:
src/models.py:89: def process(self, items):
... hundreds more matches ...
"""
# Explicitly compress when you decide it's appropriate
compressor = SearchCompressor()
result = compressor.compress(search_results, context="find process")
print(f"Compressed {result.original_match_count} matches to {result.compressed_match_count}")
print(result.compressed)
```
### Example: Compressing Build Logs
```python
from headroom.transforms import LogCompressor
# pytest output with 1000s of lines
build_output = """
===== test session starts =====
collected 500 items
tests/test_foo.py::test_1 PASSED
... hundreds of passed tests ...
tests/test_bar.py::test_fail FAILED
AssertionError: expected 5, got 3
===== 1 failed, 499 passed =====
"""
# Compress logs, preserving errors and stack traces
compressor = LogCompressor()
result = compressor.compress(build_output)
# Errors, stack traces, and summary are preserved
print(result.compressed)
print(f"Compression ratio: {result.compression_ratio:.1%}")
```
### Example: Content Type Detection
```python
from headroom.transforms import detect_content_type, ContentType
content = "src/main.py:42:def process():"
detection = detect_content_type(content)
if detection.content_type == ContentType.SEARCH_RESULTS:
# Route to SearchCompressor
pass
elif detection.content_type == ContentType.BUILD_OUTPUT:
# Route to LogCompressor
pass
```
### Integration Pattern
```python
from headroom.transforms import (
detect_content_type, ContentType,
SearchCompressor, LogCompressor, TextCompressor
)
def compress_tool_output(content: str, context: str = "") -> str:
"""Application-level compression with explicit control."""
detection = detect_content_type(content)
if detection.content_type == ContentType.SEARCH_RESULTS:
result = SearchCompressor().compress(content, context)
return result.compressed
elif detection.content_type == ContentType.BUILD_OUTPUT:
result = LogCompressor().compress(content)
return result.compressed
elif detection.content_type == ContentType.PLAIN_TEXT:
result = TextCompressor().compress(content, context)
return result.compressed
else:
# JSON or other - let SmartCrusher handle it automatically
return content
```
---
## 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)
```bash
curl http://localhost:8787/metrics
```
```
# HELP headroom_requests_total Total requests processed
headroom_requests_total{mode="optimize"} 1234
# HELP headroom_tokens_saved_total Total tokens saved
headroom_tokens_saved_total 5678900
# HELP headroom_compression_ratio Compression ratio histogram
headroom_compression_ratio_bucket{le="0.5"} 890
```
### Query Stored Metrics (SDK)
```python
from datetime import datetime, timedelta
# Get recent metrics
metrics = client.get_metrics(
start_time=datetime.utcnow() - timedelta(hours=1),
limit=100,
)
for m in metrics:
print(f"{m.timestamp}: {m.tokens_input_before} -> {m.tokens_input_after}")
# Get summary statistics
summary = client.get_summary()
print(f"Total requests: {summary['total_requests']}")
print(f"Total tokens saved: {summary['total_tokens_saved']}")
```
---
## Troubleshooting
### "Proxy won't start"
```bash
# Check if port is in use
lsof -i :8787
# Try a different port
headroom proxy --port 8788
# Check logs
headroom proxy --log-level debug
```
### "No token savings"
```python
# 1. Verify mode is "optimize"
stats = client.get_stats()
print(stats["config"]["mode"]) # Should be "optimize"
# 2. Check if transforms are enabled
print(stats["transforms"]) # smart_crusher_enabled should be True
# 3. Enable logging to see what's happening
import logging
logging.basicConfig(level=logging.DEBUG)
# 4. Use simulate to see what WOULD happen
plan = client.chat.completions.simulate(model="gpt-4o", messages=msgs)
print(f"Transforms that would apply: {plan.transforms}")
```
### "High latency"
```python
# Headroom adds ~1-5ms overhead. If you see more:
# 1. Check if embedding scorer is enabled (slower but better relevance)
# Switch to BM25 for faster scoring:
config.smart_crusher.relevance.tier = "bm25"
# 2. Disable transforms you don't need
config.cache_aligner.enabled = False # If you don't need cache alignment
# 3. Increase min_tokens_to_crush to skip small payloads
config.smart_crusher.min_tokens_to_crush = 500
```
### "Compression too aggressive"
```python
# Keep more items
config.smart_crusher.max_items_after_crush = 50 # Default is 15
# Or disable compression for specific tools
response = client.chat.completions.create(
model="gpt-4o",
messages=[...],
headroom_tool_profiles={
"important_tool": {"skip_compression": True}
}
)
```
---
## Supported Providers
| Provider | Token Counting | Cache Optimization | Status |
|----------|----------------|-------------------|--------|
| OpenAI | tiktoken (exact) | Automatic prefix caching | Full |
| Anthropic | Official API | cache_control blocks | Full |
| Google | Official API | Context caching | Full |
| Cohere | Official API | - | Full |
| Mistral | Official tokenizer | - | Full |
| LiteLLM | Via underlying provider | - | Full |
---
## Safety Guarantees
Headroom follows strict safety rules:
1. **Never removes human content** - User/assistant messages are never compressed
2. **Never breaks tool ordering** - Tool calls and responses stay paired as atomic units
3. **Parse failures are no-ops** - Malformed content passes through unchanged
4. **Preserves recency** - Last N turns are always kept
5. **Errors surface, don't hide** - Explicit exceptions with context
## Providers
| Provider | Token Counting | Cache Optimization |
|----------|----------------|-------------------|
| OpenAI | tiktoken (exact) | Automatic prefix caching |
| Anthropic | Official API | cache_control blocks |
| Google | Official API | Context caching |
| Cohere | Official API | - |
| Mistral | Official tokenizer | - |
---
## Performance
| Scenario | Before | After | Savings | Overhead |
|----------|--------|-------|---------|----------|
| Search results (1000 items) | 45,000 tokens | 4,500 tokens | 90% | ~2ms |
| Log analysis (500 entries) | 22,000 tokens | 3,300 tokens | 85% | ~1ms |
| API response (nested JSON) | 15,000 tokens | 2,250 tokens | 85% | ~1ms |
| Long conversation (50 turns) | 80,000 tokens | 32,000 tokens | 60% | ~3ms |
| Scenario | Before | After | Savings |
|----------|--------|-------|---------|
| Search results (1000 items) | 45,000 tokens | 4,500 tokens | 90% |
| Log analysis (500 entries) | 22,000 tokens | 3,300 tokens | 85% |
| Long conversation (50 turns) | 80,000 tokens | 32,000 tokens | 60% |
Overhead: ~1-5ms per request.
---
## Safety
- **Never removes human content** — User/assistant messages are never compressed
- **Never breaks tool ordering** — Tool calls and responses stay paired
- **Parse failures are no-ops** — Malformed content passes through unchanged
- **Compression is reversible** — LLM can retrieve original data via CCR
---
## Documentation
- **[Quickstart Guide](docs/quickstart.md)** - Complete working examples
- **[Proxy Documentation](docs/proxy.md)** - Production deployment
- **[Transform Reference](docs/transforms.md)** - How each transform works
- **[API Reference](docs/api.md)** - Complete API documentation
- **[Troubleshooting](docs/troubleshooting.md)** - Common issues and solutions
- **[Architecture](docs/ARCHITECTURE.md)** - How Headroom works internally
| Guide | Description |
|-------|-------------|
| [SDK Guide](docs/sdk.md) | Wrap your client for fine-grained control |
| [Proxy Guide](docs/proxy.md) | Production deployment |
| [Configuration](docs/configuration.md) | All configuration options |
| [CCR Guide](docs/ccr.md) | Reversible compression architecture |
| [Metrics](docs/metrics.md) | Monitoring and observability |
| [Troubleshooting](docs/troubleshooting.md) | Common issues |
| [Architecture](docs/ARCHITECTURE.md) | How it works internally |
---
## Examples
See the [`examples/`](examples/) directory for complete, runnable examples:
See [`examples/`](examples/) for runnable code:
- `basic_usage.py` - Simple SDK usage
- `proxy_integration.py` - Using the proxy with different clients
- `custom_compression.py` - Advanced compression configuration
- `metrics_dashboard.py` - Building a metrics dashboard
- `basic_usage.py` — Simple SDK usage
- `proxy_integration.py` — Using with different clients
- `ccr_demo.py` — CCR architecture demonstration
---
## Contributing
We welcome contributions!
```bash
# Development setup
git clone https://github.com/chopratejas/headroom.git
cd headroom
pip install -e ".[dev]"
# Run tests
pytest
# Run linting
ruff check .
mypy headroom
```
See [CONTRIBUTING.md](CONTRIBUTING.md) for details.
@ -811,7 +181,7 @@ See [CONTRIBUTING.md](CONTRIBUTING.md) for details.
## License
Apache License 2.0 - see [LICENSE](LICENSE) for details.
Apache License 2.0 — see [LICENSE](LICENSE).
---

182
docs/ARCHITECTURE.md
View file

@ -677,13 +677,180 @@ if self.config.use_feedback_hints and tool_name:
---
### CCR Phase 5: Response Handler (Automatic Tool Call Handling)
**Location:** `headroom/ccr/response_handler.py`
**The Problem:** When the proxy injects the `headroom_retrieve` tool, the LLM might call it. But who handles that tool call? Without response handling, the tool call would go back to the client unhandled.
**The Solution:** The Response Handler intercepts LLM responses, detects CCR tool calls, executes retrievals automatically, and continues the conversation until the LLM produces a final response.
```
┌──────────────────────────────────────────────────────────────────┐
│ RESPONSE HANDLER FLOW │
│ │
│ 1. LLM Response arrives │
│ └─ Contains: tool_use(headroom_retrieve, hash=abc123) │
│ │
│ 2. Handler detects CCR tool call │
│ └─ Extracts hash and optional query │
│ │
│ 3. Handler executes retrieval │
│ └─ Full retrieval: store.retrieve(hash) │
│ └─ Search: store.search(hash, query) │
│ │
│ 4. Handler continues conversation │
│ └─ Adds tool result to messages │
│ └─ Makes another API call │
│ │
│ 5. Repeat until no CCR tool calls │
│ └─ Max 3 rounds (configurable) │
│ │
│ 6. Return final response to client │
└──────────────────────────────────────────────────────────────────┘
```
**Key Classes:**
```python
@dataclass
class CCRToolCall:
tool_call_id: str # For matching response
hash_key: str # CCR hash to retrieve
query: str | None # Optional search query
@dataclass
class CCRToolResult:
tool_call_id: str
content: str # Retrieved data as JSON
success: bool
items_retrieved: int
was_search: bool # True if search, False if full retrieval
class CCRResponseHandler:
async def handle_response(
self,
response: dict, # Initial LLM response
messages: list, # Conversation history
tools: list, # Tool definitions
api_call_fn: Callable, # Function to make API calls
provider: str, # "anthropic" or "openai"
) -> dict:
"""Handle CCR tool calls until final response."""
```
**Streaming Support:**
The handler also supports streaming responses via `StreamingCCRHandler`:
```python
class StreamingCCRBuffer:
"""Buffers streaming chunks to detect CCR tool calls."""
chunks: list[bytes]
detected_ccr: bool
class StreamingCCRHandler:
"""Handles CCR in streaming responses."""
async def process_stream(self, stream, messages, tools, api_call_fn):
"""Yields chunks, switching to buffered mode if CCR detected."""
```
---
### CCR Phase 6: Context Tracker (Multi-Turn Awareness)
**Location:** `headroom/ccr/context_tracker.py`
**The Problem:** In multi-turn conversations, earlier compressed data might become relevant later. Without tracking, the LLM has "context amnesia" - it can't reference data that was compressed in turn 1 when answering a question in turn 5.
**The Solution:** The Context Tracker maintains awareness of all compressed content across the conversation and can proactively expand relevant data when a new query might need it.
```
┌──────────────────────────────────────────────────────────────────┐
│ CONTEXT TRACKER FLOW │
│ │
│ Turn 1: Search returns 100 files → compressed to 10 │
│ Tracker stores: hash=abc123, sample="auth.py, db.py" │
│ │
│ Turn 5: User asks "What about the authentication middleware?" │
│ Tracker analyzes query: │
│ - "authentication" matches "auth.py" in sample │
│ - Relevance score: 0.7 (above threshold) │
│ │
│ Proactive Expansion: │
│ - Retrieves abc123 before LLM responds │
│ - Adds expanded context to request │
│ │
│ Result: LLM sees full file list, can answer accurately │
└──────────────────────────────────────────────────────────────────┘
```
**Key Classes:**
```python
@dataclass
class CompressedContext:
hash_key: str # CCR hash
turn_number: int # When compression happened
timestamp: float # For age-based filtering
tool_name: str | None # Which tool was compressed
original_item_count: int
compressed_item_count: int
query_context: str # User query at compression time
sample_content: str # Preview for relevance matching
@dataclass
class ExpansionRecommendation:
hash_key: str
reason: str # Human-readable explanation
relevance_score: float # 0-1, higher = more relevant
expand_full: bool # True = full retrieval
search_query: str | None # If expand_full=False
class ContextTracker:
def track_compression(self, hash_key, turn_number, ...):
"""Track a compression event."""
def analyze_query(self, query: str) -> list[ExpansionRecommendation]:
"""Find relevant compressed contexts for a query."""
def execute_expansions(self, recommendations) -> list[dict]:
"""Execute recommended expansions."""
```
**Relevance Calculation:**
The tracker uses simple but effective heuristics:
1. **Keyword overlap with sample content** - Extract keywords from query, match against compressed content preview
2. **Keyword overlap with original query** - Match against the query that triggered compression
3. **Tool name relevance** - File operations more likely to need expansion for "file", "where", "find" queries
4. **Age discount** - Older contexts get lower scores
**Configuration:**
```python
@dataclass
class ContextTrackerConfig:
enabled: bool = True
max_tracked_contexts: int = 100 # LRU eviction
relevance_threshold: float = 0.3 # Min score to recommend
max_context_age_seconds: float = 300 # 5 minutes
proactive_expansion: bool = True
max_proactive_expansions: int = 2 # Per query
```
---
### Why CCR is a Moat
1. **Reversible**: No permanent information loss. Worst case = retrieve everything.
2. **Transparent**: LLM knows it can ask for more data.
3. **Feedback Loop**: Learn from actual needs, not guesses.
4. **Network Effect**: Retrieval patterns across users improve compression for everyone.
5. **Zero-Risk**: If compression fails, instant fallback to original data.
3. **Automatic**: Response Handler executes retrievals without client intervention.
4. **Context-Aware**: Context Tracker prevents multi-turn amnesia.
5. **Feedback Loop**: Learn from actual needs, not guesses.
6. **Network Effect**: Retrieval patterns across users improve compression for everyone.
7. **Zero-Risk**: If compression fails, instant fallback to original data.
---
@ -712,7 +879,7 @@ headroom/
│ ├── rolling_window.py # Token limit enforcement
│ └── llmlingua_compressor.py # ML-based compression (opt-in)
│
├── cache/ # CCR Architecture
├── cache/ # CCR Architecture - Caching & Storage
│ ├── compression_store.py # Phase 1: Store original content
│ ├── compression_feedback.py # Phase 4: Learn from retrievals
│ ├── anthropic.py # Anthropic cache optimizer
@ -720,6 +887,13 @@ headroom/
│ ├── google.py # Google cache optimizer
│ └── dynamic_detector.py # Dynamic content detection
│
├── ccr/ # CCR Architecture - Tool Injection & Response Handling
│ ├── __init__.py # CCR module exports
│ ├── tool_injection.py # Phase 3: Inject retrieval tool
│ ├── response_handler.py # Phase 5: Handle CCR tool calls
│ ├── context_tracker.py # Phase 6: Multi-turn context tracking
│ └── mcp_server.py # MCP server for retrieval tool
│
├── relevance/ # Relevance scoring for compression
│ ├── bm25.py # BM25 keyword scorer
│ ├── embedding.py # Semantic embedding scorer

59
docs/README.md
View file

@ -2,27 +2,58 @@
Welcome to the Headroom documentation.
## Quick Links
## Getting Started
- [Getting Started](getting-started.md)
- [Proxy Server](proxy.md)
- [Transforms](transforms.md)
- [API Reference](api.md)
- [Architecture](ARCHITECTURE.md)
| Guide | Description |
|-------|-------------|
| [Quickstart](quickstart.md) | 5-minute setup |
| [SDK Guide](sdk.md) | Python SDK usage |
| [Proxy Guide](proxy.md) | Proxy server deployment |
## Core Concepts
| Topic | Description |
|-------|-------------|
| [Transforms](transforms.md) | How compression works |
| [CCR](ccr.md) | Reversible compression architecture |
| [Configuration](configuration.md) | All configuration options |
## Advanced
| Topic | Description |
|-------|-------------|
| [Text Compression](text-compression.md) | Opt-in utilities for search/logs |
| [LLMLingua](llmlingua.md) | ML-based compression |
| [Metrics](metrics.md) | Monitoring and observability |
| [Errors](errors.md) | Error handling |
## Reference
| Topic | Description |
|-------|-------------|
| [API Reference](api.md) | Complete API docs |
| [Architecture](ARCHITECTURE.md) | Internal design |
| [Troubleshooting](troubleshooting.md) | Common issues |
## Overview
Headroom is the Context Optimization Layer for LLM applications. It reduces your LLM costs by 50-90% through intelligent context compression.
### Core Concepts
### How It Works
1. **Transforms**: Stateless functions that modify message arrays to reduce tokens
2. **Providers**: Adapters for different LLM providers (OpenAI, Anthropic, etc.)
3. **Pipeline**: Chains multiple transforms together
4. **Proxy**: HTTP server that applies transforms transparently
1. **SmartCrusher** — Compresses JSON tool outputs, keeping errors, anomalies, and relevant items
2. **CacheAligner** — Stabilizes message prefixes so provider caching works
3. **RollingWindow** — Manages context limits without breaking tool call pairs
4. **CCR** — Caches original data so compression is reversible
### Safety Guarantees
- Never removes human content
- Never breaks tool call ordering
- Parse failures pass through unchanged
- LLM can always retrieve original data
### Getting Help
- [GitHub Issues](https://github.com/headroom-sdk/headroom/issues) - Bug reports
- [GitHub Discussions](https://github.com/headroom-sdk/headroom/discussions) - Questions
- [Discord](https://discord.gg/headroom) - Community chat
- [GitHub Issues](https://github.com/chopratejas/headroom/issues) — Bug reports
- [GitHub Discussions](https://github.com/chopratejas/headroom/discussions) — Questions

146
docs/ccr.md Normal file
View file

@ -0,0 +1,146 @@
# CCR: Compress-Cache-Retrieve
Headroom's CCR architecture makes compression **reversible**. When tool outputs are compressed, the original data is cached. If the LLM needs more data, it can retrieve it instantly.
## The Problem with Traditional Compression
Traditional compression is lossy — if you guess wrong about what's important, data is lost forever. This creates a difficult tradeoff:
- **Aggressive compression**: Risk losing data the LLM needs
- **Conservative compression**: Miss out on token savings
CCR eliminates this tradeoff.
## How CCR Works
```
┌─────────────────────────────────────────────────────────────────┐
│ TOOL OUTPUT (1000 items) │
│ └─ SmartCrusher compresses to 20 items │
│ └─ Original cached with hash=abc123 │
│ └─ Retrieval tool injected into context │
└─────────────────────────────────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────────┐
│ LLM PROCESSING │
│ Option A: LLM solves task with 20 items → Done (90% savings) │
│ Option B: LLM calls headroom_retrieve(hash=abc123) │
│ → Response Handler executes retrieval automatically │
│ → LLM receives full data, responds accurately │
└─────────────────────────────────────────────────────────────────┘
```
### Phase 1: Compression Store
When SmartCrusher compresses tool output:
1. Original content is stored in an LRU cache
2. A hash key is generated for retrieval
3. A marker is added to the compressed output: `[1000 items compressed to 20. Retrieve more: hash=abc123]`
### Phase 2: Tool Injection
Headroom injects a `headroom_retrieve` tool into the LLM's available tools:
```json
{
"name": "headroom_retrieve",
"description": "Retrieve original uncompressed data from Headroom cache",
"parameters": {
"hash": "The hash key from the compression marker",
"query": "Optional: search within the cached data"
}
}
```
### Phase 3: Response Handler
When the LLM calls `headroom_retrieve`:
1. Response Handler intercepts the tool call
2. Retrieves data from the local cache (~1ms)
3. Adds the result to the conversation
4. Continues the API call automatically
**The client never sees CCR tool calls** — they're handled transparently.
### Phase 4: Context Tracker
Across multiple turns, the Context Tracker:
1. Remembers what was compressed in earlier turns
2. Analyzes new queries for relevance to compressed content
3. Proactively expands relevant data before the LLM asks
**Example:**
```
Turn 1: User searches for files
→ Tool returns 500 files
→ SmartCrusher compresses to 15, caches original (hash=abc123)
→ LLM sees 15 files, answers question
Turn 5: User asks "What about the auth middleware?"
→ Context Tracker detects "auth" might be in abc123
→ Proactively expands compressed content
→ LLM sees full file list, finds auth_middleware.py
```
## Features
| Feature | Description |
|---------|-------------|
| **Automatic Response Handling** | When LLM calls `headroom_retrieve`, the proxy handles it automatically |
| **Multi-Turn Context Tracking** | Tracks compressed content across turns, proactively expands when relevant |
| **BM25 Search** | LLM can search within compressed data: `headroom_retrieve(hash, query="errors")` |
| **Feedback Learning** | Learns from retrieval patterns to improve future compression |
## Configuration
```bash
# Proxy with CCR enabled (default)
headroom proxy --port 8787
# Disable CCR response handling
headroom proxy --no-ccr-responses
# Disable proactive expansion
headroom proxy --no-ccr-expansion
```
## Why This Matters
| Approach | Risk | Savings |
|----------|------|---------|
| No compression | None | 0% |
| Traditional compression | Data loss | 70-90% |
| CCR compression | None (reversible) | 70-90% |
CCR gives you the savings of aggressive compression with zero risk — the LLM can always retrieve the original data if needed.
## Demo
Run the CCR demonstration to see it in action:
```bash
python examples/ccr_demo.py
```
Output:
```
1. COMPRESSION STORE
Original: 100 items (7,059 chars)
Compressed: 8 items (633 chars)
Reduction: 91.0%
3. RESPONSE HANDLER
Detected CCR tool call: True
Retrieved 100 items automatically
4. CONTEXT TRACKER
Turn 5: User asks "show authentication middleware"
Tracker found 1 relevant context
→ relevance=0.73
Proactively expanded: 100 items
```
## Architecture
For implementation details, see [ARCHITECTURE.md](ARCHITECTURE.md#ccr-compress-cache-retrieve).

246
docs/configuration.md Normal file
View file

@ -0,0 +1,246 @@
# Configuration
Headroom can be configured via the SDK, proxy command line, or per-request overrides.
## SDK Configuration
```python
from headroom import HeadroomClient, OpenAIProvider
from openai import OpenAI
client = HeadroomClient(
original_client=OpenAI(),
provider=OpenAIProvider(),
# Mode: "audit" (observe only) or "optimize" (apply transforms)
default_mode="optimize",
# Enable provider-specific cache optimization
enable_cache_optimizer=True,
# Enable query-level semantic caching
enable_semantic_cache=False,
# Override default context limits per model
model_context_limits={
"gpt-4o": 128000,
"gpt-4o-mini": 128000,
},
# Database location (defaults to temp directory)
# store_url="sqlite:////absolute/path/to/headroom.db",
)
```
## Proxy Configuration
### Command Line Options
```bash
headroom proxy \
--port 8787 \ # Port to listen on
--host 0.0.0.0 \ # Host to bind to
--budget 10.00 \ # Daily budget limit in USD
--log-file headroom.jsonl # Log file path
```
### Feature Flags
```bash
# Disable optimization (passthrough mode)
headroom proxy --no-optimize
# Disable semantic caching
headroom proxy --no-cache
# Disable CCR response handling
headroom proxy --no-ccr-responses
# Disable proactive expansion
headroom proxy --no-ccr-expansion
# Enable LLMLingua ML compression
headroom proxy --llmlingua
headroom proxy --llmlingua --llmlingua-device cuda --llmlingua-rate 0.4
```
### All Options
```bash
headroom proxy --help
```
## Per-Request Overrides
Override configuration for specific requests:
```python
response = client.chat.completions.create(
model="gpt-4o",
messages=[...],
# Override mode for this request
headroom_mode="audit",
# Reserve more tokens for output
headroom_output_buffer_tokens=8000,
# Keep last N turns (don't compress)
headroom_keep_turns=5,
# Skip compression for specific tools
headroom_tool_profiles={
"important_tool": {"skip_compression": True}
}
)
```
## Modes
| Mode | Behavior | Use Case |
|------|----------|----------|
| `audit` | Observes and logs, no modifications | Production monitoring, baseline measurement |
| `optimize` | Applies safe, deterministic transforms | Production optimization |
| `simulate` | Returns plan without API call | Testing, cost estimation |
### Simulate Mode
Preview what would happen without making an API call:
```python
plan = client.chat.completions.simulate(
model="gpt-4o",
messages=large_conversation,
)
print(f"Would save {plan.tokens_saved} tokens")
print(f"Transforms: {plan.transforms}")
print(f"Estimated savings: {plan.estimated_savings}")
```
## SmartCrusher Configuration
Fine-tune JSON compression behavior:
```python
from headroom.transforms import SmartCrusherConfig
config = SmartCrusherConfig(
# Maximum items to keep after compression
max_items_after_crush=15,
# Minimum tokens before applying compression
min_tokens_to_crush=200,
# Relevance scoring tier: "bm25" (fast) or "embedding" (accurate)
relevance_tier="bm25",
# Always keep items with these field values
preserve_fields=["error", "warning", "failure"],
)
```
## Cache Aligner Configuration
Control prefix stabilization:
```python
from headroom.transforms import CacheAlignerConfig
config = CacheAlignerConfig(
# Enable/disable cache alignment
enabled=True,
# Patterns to extract from system prompt
dynamic_patterns=[
r"Today is \w+ \d+, \d{4}",
r"Current time: .*",
],
)
```
## Rolling Window Configuration
Control context window management:
```python
from headroom.transforms import RollingWindowConfig
config = RollingWindowConfig(
# Minimum turns to always keep
min_keep_turns=3,
# Reserve tokens for output
output_buffer_tokens=4000,
# Drop oldest tool outputs first
prefer_drop_tool_outputs=True,
)
```
## Environment Variables
Some settings can be configured via environment variables:
| Variable | Description | Default |
|----------|-------------|---------|
| `HEADROOM_LOG_LEVEL` | Logging level | `INFO` |
| `HEADROOM_STORE_URL` | Database URL | temp directory |
| `HEADROOM_DEFAULT_MODE` | Default mode | `optimize` |
## Provider-Specific Settings
### OpenAI
```python
from headroom import OpenAIProvider
provider = OpenAIProvider(
# Enable automatic prefix caching
enable_prefix_caching=True,
)
```
### Anthropic
```python
from headroom import AnthropicProvider
provider = AnthropicProvider(
# Enable cache_control blocks
enable_cache_control=True,
)
```
### Google
```python
from headroom import GoogleProvider
provider = GoogleProvider(
# Enable context caching
enable_context_caching=True,
)
```
## Configuration Precedence
Settings are applied in this order (later overrides earlier):
1. Default values
2. Environment variables
3. SDK constructor arguments
4. Per-request overrides
## Validation
Validate your configuration:
```python
result = client.validate_setup()
if not result["valid"]:
print("Configuration issues:")
for issue in result["issues"]:
print(f" - {issue}")
```

251
docs/errors.md Normal file
View file

@ -0,0 +1,251 @@
# Error Handling
Headroom provides explicit exceptions for debugging, with a safety guarantee that compression failures never break your LLM calls.
## Exception Hierarchy
```python
from headroom import (
HeadroomError, # Base class - catch all Headroom errors
ConfigurationError, # Invalid configuration
ProviderError, # Provider issues (unknown model, etc.)
StorageError, # Database/storage failures
CompressionError, # Compression failures (rare)
ValidationError, # Setup validation failures
)
```
## Usage
```python
from headroom import (
HeadroomClient,
HeadroomError,
ConfigurationError,
StorageError,
)
try:
client = HeadroomClient(...)
response = client.chat.completions.create(...)
except ConfigurationError as e:
print(f"Config issue: {e}")
print(f"Details: {e.details}") # Additional context
except StorageError as e:
print(f"Storage issue: {e}")
# Headroom continues to work, just without metrics persistence
except HeadroomError as e:
print(f"Headroom error: {e}")
```
## Exception Types
### ConfigurationError
Raised when configuration is invalid.
```python
# Examples:
# - Invalid mode value
# - Missing required provider
# - Invalid model context limit
try:
client = HeadroomClient(
original_client=OpenAI(),
provider=OpenAIProvider(),
default_mode="invalid_mode", # Will raise ConfigurationError
)
except ConfigurationError as e:
print(f"Config error: {e}")
print(f"Field: {e.details.get('field')}")
```
### ProviderError
Raised for provider-specific issues.
```python
# Examples:
# - Unknown model name
# - Provider API error
# - Token counting failure
try:
response = client.chat.completions.create(
model="unknown-model-xyz",
messages=[...]
)
except ProviderError as e:
print(f"Provider error: {e}")
print(f"Provider: {e.details.get('provider')}")
```
### StorageError
Raised when database operations fail.
```python
# Examples:
# - Database connection failure
# - Write permission denied
# - Disk full
try:
metrics = client.get_metrics()
except StorageError as e:
print(f"Storage error: {e}")
# Application can continue - just won't have metrics
```
### CompressionError
Raised when compression fails (rare).
```python
# Examples:
# - Malformed JSON in tool output
# - Unexpected data structure
# Note: In practice, compression errors are caught internally
# and the original content passes through unchanged.
# This exception is only raised if you explicitly enable strict mode.
```
### ValidationError
Raised when setup validation fails.
```python
result = client.validate_setup()
if not result["valid"]:
raise ValidationError(
"Setup validation failed",
details={"issues": result["issues"]}
)
```
## Safety Guarantee
**If compression fails, the original content passes through unchanged.**
This is a core design principle. Your LLM calls never fail due to Headroom:
```python
# Even if SmartCrusher encounters unexpected data:
messages = [
{"role": "tool", "content": "malformed json {{{"}
]
# This will NOT raise an exception
# Instead, the malformed content passes through unchanged
response = client.chat.completions.create(
model="gpt-4o",
messages=messages
)
```
## Logging Errors
Enable logging to see error details:
```python
import logging
logging.basicConfig(level=logging.WARNING)
# Now you'll see warnings when compression is skipped:
# WARNING:headroom.transforms.smart_crusher:Skipping compression: invalid JSON
```
## Error Details
All Headroom exceptions include a `details` dict with context:
```python
try:
client = HeadroomClient(...)
except HeadroomError as e:
print(f"Error: {e}")
print(f"Type: {type(e).__name__}")
print(f"Details: {e.details}")
# Details might include:
# - field: which config field caused the error
# - provider: which provider was involved
# - model: which model was requested
# - original_error: underlying exception
```
## Best Practices
### 1. Catch Specific Exceptions
```python
# Good: catch specific exceptions
try:
response = client.chat.completions.create(...)
except ConfigurationError:
# Handle config issues
pass
except ProviderError:
# Handle provider issues
pass
# Avoid: catching all exceptions
try:
response = client.chat.completions.create(...)
except Exception:
# Too broad - might hide real bugs
pass
```
### 2. Let StorageError Pass
```python
# Storage errors don't affect core functionality
try:
metrics = client.get_metrics()
except StorageError:
metrics = [] # Continue without historical metrics
```
### 3. Validate on Startup
```python
client = HeadroomClient(...)
# Validate once at startup
result = client.validate_setup()
if not result["valid"]:
raise SystemExit(f"Headroom setup invalid: {result['issues']}")
# Then use client normally
response = client.chat.completions.create(...)
```
## Debugging
### Enable Debug Logging
```python
import logging
logging.basicConfig(level=logging.DEBUG)
# Shows detailed transform decisions
# DEBUG:headroom.transforms.smart_crusher:Analyzing 1000 items...
# DEBUG:headroom.transforms.smart_crusher:Kept 15 items (errors: 2, anomalies: 3)
```
### Check Stats After Error
```python
try:
response = client.chat.completions.create(...)
except HeadroomError:
# Check what happened
stats = client.get_stats()
print(f"Last request stats: {stats}")
```

188
docs/llmlingua.md Normal file
View file

@ -0,0 +1,188 @@
# LLMLingua-2 Integration
For maximum 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
```
## 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)
```
## 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)
```
## Proxy Integration
Enable LLMLingua in the proxy server for automatic ML compression:
```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)
```
## Configuration Reference
| Option | Default | Description |
|--------|---------|-------------|
| `device` | `"auto"` | Device to run model on: auto, cpu, cuda, mps |
| `code_compression_rate` | `0.4` | Keep 40% of tokens for code |
| `json_compression_rate` | `0.35` | Keep 35% of tokens for JSON |
| `text_compression_rate` | `0.25` | Keep 25% of tokens for text |
| `force_tokens` | `[]` | Tokens to always preserve |
| `drop_consecutive` | `True` | Drop consecutive whitespace |
## Performance Characteristics
| Metric | Value |
|--------|-------|
| Model size | ~500MB |
| Memory usage | ~1GB RAM |
| Cold start | 10-30s (first load) |
| Inference | 50-200ms per request |
| Compression | 80-95% |
## Why Opt-In?
LLMLingua adds significant dependencies and overhead:
| Aspect | Default Proxy | With LLMLingua |
|--------|--------------|----------------|
| Dependencies | ~50MB | ~2GB |
| Cold start | <1s | 10-30s |
| Per-request | ~1-5ms | ~50-200ms |
| Compression | 70-90% | 80-95% |
The default proxy is lightweight and fast. Enable LLMLingua when you need maximum compression and can accept the tradeoffs.
## Troubleshooting
### "Model not found"
```bash
# Ensure llmlingua extra is installed
pip install "headroom-ai[llmlingua]"
```
### "CUDA out of memory"
```python
# Force CPU mode
config = LLMLinguaConfig(device="cpu")
```
### "Slow compression"
- Use GPU if available: `device="cuda"`
- Batch multiple compressions
- Consider using SmartCrusher for JSON (faster, similar results)

260
docs/metrics.md Normal file
View file

@ -0,0 +1,260 @@
# Metrics & Monitoring
Headroom provides comprehensive metrics for monitoring compression performance, cost savings, and system health.
## Proxy Metrics
### Stats Endpoint
```bash
curl http://localhost:8787/stats
```
```json
{
"requests": {
"total": 42,
"cached": 5,
"rate_limited": 0,
"failed": 0
},
"tokens": {
"input": 50000,
"output": 8000,
"saved": 12500,
"savings_percent": 25.0
},
"cost": {
"total_cost_usd": 0.15,
"total_savings_usd": 0.04
},
"cache": {
"entries": 10,
"total_hits": 5
}
}
```
### Prometheus Metrics
```bash
curl http://localhost:8787/metrics
```
```prometheus
# HELP headroom_requests_total Total requests processed
headroom_requests_total{mode="optimize"} 1234
# HELP headroom_tokens_saved_total Total tokens saved
headroom_tokens_saved_total 5678900
# HELP headroom_compression_ratio Compression ratio histogram
headroom_compression_ratio_bucket{le="0.5"} 890
headroom_compression_ratio_bucket{le="0.7"} 1100
headroom_compression_ratio_bucket{le="0.9"} 1200
# HELP headroom_latency_seconds Request latency histogram
headroom_latency_seconds_bucket{le="0.01"} 800
headroom_latency_seconds_bucket{le="0.1"} 1150
# HELP headroom_cache_hits_total Cache hit counter
headroom_cache_hits_total 456
# HELP headroom_cache_misses_total Cache miss counter
headroom_cache_misses_total 778
```
### Health Check
```bash
curl http://localhost:8787/health
```
```json
{
"status": "healthy",
"version": "0.1.0",
"uptime_seconds": 3600,
"llmlingua_enabled": false
}
```
## SDK Metrics
### Session Stats
Quick stats for the current session (no database query):
```python
stats = client.get_stats()
print(stats)
```
```python
{
"session": {
"requests_total": 10,
"tokens_input_before": 50000,
"tokens_input_after": 35000,
"tokens_saved_total": 15000,
"tokens_output_total": 8000,
"cache_hits": 3,
"compression_ratio_avg": 0.70
},
"config": {
"mode": "optimize",
"provider": "openai",
"cache_optimizer_enabled": True,
"semantic_cache_enabled": False
},
"transforms": {
"smart_crusher_enabled": True,
"cache_aligner_enabled": True,
"rolling_window_enabled": True
}
}
```
### Historical Metrics
Query stored metrics from the database:
```python
from datetime import datetime, timedelta
# Get recent metrics
metrics = client.get_metrics(
start_time=datetime.utcnow() - timedelta(hours=1),
limit=100,
)
for m in metrics:
print(f"{m.timestamp}: {m.tokens_input_before} -> {m.tokens_input_after}")
```
### Summary Statistics
Aggregate statistics across all stored metrics:
```python
summary = client.get_summary()
print(f"Total requests: {summary['total_requests']}")
print(f"Total tokens saved: {summary['total_tokens_saved']}")
print(f"Average compression: {summary['avg_compression_ratio']:.1%}")
print(f"Total cost savings: ${summary['total_cost_saved_usd']:.2f}")
```
## Logging
### Enable Logging
```python
import logging
# INFO level shows compression summaries
logging.basicConfig(level=logging.INFO)
# DEBUG level shows detailed transform decisions
logging.basicConfig(level=logging.DEBUG)
```
### Log Output Examples
```
INFO:headroom.transforms.pipeline:Pipeline complete: 45000 -> 4500 tokens (saved 40500, 90.0% reduction)
INFO:headroom.transforms.smart_crusher:SmartCrusher applied top_n strategy: kept 15 of 1000 items
INFO:headroom.cache.compression_store:CCR cache hit: hash=abc123, retrieved 1000 items
DEBUG:headroom.transforms.smart_crusher:Kept items: [0,1,2,42,77,97,98,99] (errors at 42, warnings at 77)
```
### Proxy Logging
```bash
# Log to file
headroom proxy --log-file headroom.jsonl
# Increase verbosity
headroom proxy --log-level debug
```
## Grafana Dashboard
Example Grafana dashboard configuration for Prometheus metrics:
```json
{
"panels": [
{
"title": "Tokens Saved",
"type": "stat",
"targets": [{"expr": "headroom_tokens_saved_total"}]
},
{
"title": "Compression Ratio",
"type": "gauge",
"targets": [{"expr": "histogram_quantile(0.5, headroom_compression_ratio_bucket)"}]
},
{
"title": "Request Latency (p99)",
"type": "graph",
"targets": [{"expr": "histogram_quantile(0.99, headroom_latency_seconds_bucket)"}]
},
{
"title": "Cache Hit Rate",
"type": "gauge",
"targets": [{"expr": "headroom_cache_hits_total / (headroom_cache_hits_total + headroom_cache_misses_total)"}]
}
]
}
```
## Cost Tracking
### Per-Request Cost
Each request includes cost metadata in the response:
```python
response = client.chat.completions.create(...)
# Access via response metadata (if available)
# Cost is calculated based on model pricing and token counts
```
### Budget Alerts
Set a budget limit in the proxy:
```bash
headroom proxy --budget 10.00
```
When the budget is exceeded:
- Requests return a budget exceeded error
- The `/stats` endpoint shows budget status
- Logs indicate budget state
## Validation
Validate your setup is correct:
```python
result = client.validate_setup()
if result["valid"]:
print("Setup is correct!")
else:
print("Issues found:")
for issue in result["issues"]:
print(f" - {issue}")
```
## Key Metrics to Monitor
| Metric | What It Tells You | Target |
|--------|------------------|--------|
| `tokens_saved_total` | Total cost savings | Higher is better |
| `compression_ratio_avg` | Efficiency | 0.7-0.9 typical |
| `cache_hit_rate` | Cache effectiveness | >20% is good |
| `latency_p99` | Performance impact | <10ms |
| `failed_requests` | Reliability | 0 |

292
docs/sdk.md Normal file
View file

@ -0,0 +1,292 @@
# SDK Guide
The Headroom SDK wraps your existing LLM client to add compression and optimization transparently.
## Installation
```bash
pip install headroom-ai openai
```
## Quick Start
```python
from headroom import HeadroomClient, OpenAIProvider
from openai import OpenAI
# Create wrapped client
client = HeadroomClient(
original_client=OpenAI(),
provider=OpenAIProvider(),
default_mode="optimize",
)
# Use exactly like the original client
response = client.chat.completions.create(
model="gpt-4o-mini",
messages=[
{"role": "user", "content": "Hello!"},
],
)
print(response.choices[0].message.content)
```
## Tool Output Compression
Real savings happen with tool outputs. Here's where Headroom shines:
```python
import json
# Conversation with large tool output
messages = [
{"role": "user", "content": "Search for Python tutorials"},
{
"role": "assistant",
"content": None,
"tool_calls": [{
"id": "call_123",
"type": "function",
"function": {"name": "search", "arguments": '{"q": "python"}'},
}],
},
{
"role": "tool",
"tool_call_id": "call_123",
"content": json.dumps({
"results": [
{"title": f"Tutorial {i}", "score": 100-i}
for i in range(500)
]
}),
},
{"role": "user", "content": "What are the top 3?"},
]
# Headroom compresses 500 results to ~15, keeping highest-scoring items
response = client.chat.completions.create(
model="gpt-4o-mini",
messages=messages
)
# Check savings
stats = client.get_stats()
print(f"Tokens saved: {stats['session']['tokens_saved_total']}")
# Typical output: "Tokens saved: 3500"
```
## Supported Providers
### OpenAI
```python
from headroom import HeadroomClient, OpenAIProvider
from openai import OpenAI
client = HeadroomClient(
original_client=OpenAI(),
provider=OpenAIProvider(),
)
```
### Anthropic
```python
from headroom import HeadroomClient, AnthropicProvider
from anthropic import Anthropic
client = HeadroomClient(
original_client=Anthropic(),
provider=AnthropicProvider(),
)
response = client.messages.create(
model="claude-3-5-sonnet-20241022",
max_tokens=1024,
messages=[{"role": "user", "content": "Hello!"}],
)
```
### Google
```python
from headroom import HeadroomClient, GoogleProvider
import google.generativeai as genai
client = HeadroomClient(
original_client=genai,
provider=GoogleProvider(),
)
```
## Check Stats
```python
# Session stats (no database query)
stats = client.get_stats()
print(stats)
# {
# "session": {"requests_total": 10, "tokens_saved_total": 5000, ...},
# "config": {"mode": "optimize", "provider": "openai", ...},
# "transforms": {"smart_crusher_enabled": True, ...}
# }
```
## Validate Setup
```python
result = client.validate_setup()
if not result["valid"]:
print("Setup issues:", result["issues"])
```
## Modes
### Optimize (Default)
Applies all safe transforms:
```python
client = HeadroomClient(
original_client=OpenAI(),
provider=OpenAIProvider(),
default_mode="optimize",
)
```
### Audit
Observes and logs without modifying:
```python
client = HeadroomClient(
original_client=OpenAI(),
provider=OpenAIProvider(),
default_mode="audit",
)
```
### Simulate
Returns a plan without making the API call:
```python
plan = client.chat.completions.simulate(
model="gpt-4o",
messages=large_conversation,
)
print(f"Would save {plan.tokens_saved} tokens")
print(f"Transforms: {plan.transforms}")
```
## Per-Request Overrides
```python
response = client.chat.completions.create(
model="gpt-4o",
messages=[...],
# Override mode for this request
headroom_mode="audit",
# Reserve more tokens for output
headroom_output_buffer_tokens=8000,
# Keep last N turns
headroom_keep_turns=5,
)
```
## Enable Logging
```python
import logging
logging.basicConfig(level=logging.INFO)
# Now you'll see:
# INFO:headroom.transforms.pipeline:Pipeline complete: 45000 -> 4500 tokens
# INFO:headroom.transforms.smart_crusher:SmartCrusher: kept 15 of 1000 items
```
## Streaming
Streaming works transparently:
```python
stream = client.chat.completions.create(
model="gpt-4o-mini",
messages=[{"role": "user", "content": "Hello!"}],
stream=True,
)
for chunk in stream:
if chunk.choices[0].delta.content:
print(chunk.choices[0].delta.content, end="")
```
## Error Handling
```python
from headroom import (
HeadroomClient,
HeadroomError,
ConfigurationError,
ProviderError,
)
try:
response = client.chat.completions.create(...)
except ConfigurationError as e:
print(f"Config issue: {e}")
except ProviderError as e:
print(f"Provider issue: {e}")
except HeadroomError as e:
print(f"Headroom error: {e}")
```
## Historical Metrics
Query stored metrics:
```python
from datetime import datetime, timedelta
metrics = client.get_metrics(
start_time=datetime.utcnow() - timedelta(hours=1),
limit=100,
)
for m in metrics:
print(f"{m.timestamp}: {m.tokens_input_before} -> {m.tokens_input_after}")
```
## Advanced Configuration
See [Configuration](configuration.md) for full options:
```python
client = HeadroomClient(
original_client=OpenAI(),
provider=OpenAIProvider(),
default_mode="optimize",
enable_cache_optimizer=True,
enable_semantic_cache=False,
model_context_limits={
"gpt-4o": 128000,
"gpt-4o-mini": 128000,
},
)
```
## Comparison with Proxy
| Aspect | SDK | Proxy |
|--------|-----|-------|
| Setup | Wrap client | Point URL |
| Control | Fine-grained | Global |
| Metrics | In-process | Centralized |
| Best for | Custom apps | Existing tools |
Use the SDK when you need fine-grained control. Use the proxy for existing tools like Claude Code, Cursor, etc.

193
docs/text-compression.md Normal file
View file

@ -0,0 +1,193 @@
# Text Compression Utilities
For coding tasks, Headroom provides **standalone text compression utilities** that applications can use explicitly. These are **opt-in** — they're not applied automatically, giving you full control over when and how to compress text content.
> **Design Philosophy**: SmartCrusher compresses JSON automatically because it's structure-preserving and safe. Text compression is lossy and context-dependent, so applications should decide when to use it.
## Available Utilities
| Utility | Input Type | Use Case |
|---------|------------|----------|
| `SearchCompressor` | grep/ripgrep output | Search results with `file:line:content` format |
| `LogCompressor` | Build/test logs | pytest, npm, cargo, make output |
| `TextCompressor` | Generic text | Any plain text with anchor preservation |
| `detect_content_type` | Any content | Detect content type for routing decisions |
## SearchCompressor
Compresses search results (grep, ripgrep, ag) while preserving relevant matches.
```python
from headroom.transforms import SearchCompressor
# Your grep/ripgrep output (could be 1000s of lines)
search_results = """
src/utils.py:42:def process_data(items):
src/utils.py:43: \"\"\"Process items.\"\"\"
src/models.py:15:class DataProcessor:
src/models.py:89: def process(self, items):
... hundreds more matches ...
"""
# Explicitly compress when you decide it's appropriate
compressor = SearchCompressor()
result = compressor.compress(search_results, context="find process")
print(f"Compressed {result.original_match_count} matches to {result.compressed_match_count}")
print(result.compressed)
```
### What Gets Preserved
- **Exact query matches**: Lines containing the search term
- **High-relevance matches**: Scored by BM25 similarity to context
- **File diversity**: Ensures results from different files are kept
- **First/last matches**: Context from start and end of results
## LogCompressor
Compresses build and test output while preserving errors, warnings, and summaries.
```python
from headroom.transforms import LogCompressor
# pytest output with 1000s of lines
build_output = """
===== test session starts =====
collected 500 items
tests/test_foo.py::test_1 PASSED
... hundreds of passed tests ...
tests/test_bar.py::test_fail FAILED
AssertionError: expected 5, got 3
===== 1 failed, 499 passed =====
"""
# Compress logs, preserving errors and stack traces
compressor = LogCompressor()
result = compressor.compress(build_output)
# Errors, stack traces, and summary are preserved
print(result.compressed)
print(f"Compression ratio: {result.compression_ratio:.1%}")
```
### What Gets Preserved
- **Errors and failures**: Any line with ERROR, FAILED, Exception, etc.
- **Warnings**: Warning messages that might be important
- **Stack traces**: Full tracebacks for debugging
- **Summaries**: Test/build summary lines
- **Section headers**: Structural markers like `=====`
## TextCompressor
General-purpose text compression with anchor preservation.
```python
from headroom.transforms import TextCompressor
long_text = """
... thousands of lines of documentation ...
"""
compressor = TextCompressor()
result = compressor.compress(long_text, context="authentication")
print(result.compressed)
```
### What Gets Preserved
- **Relevant paragraphs**: Scored by similarity to context
- **Anchors**: Headers, section markers, important keywords
- **Structure**: Document organization is maintained
## Content Type Detection
Automatically detect content type to route to the right compressor.
```python
from headroom.transforms import detect_content_type, ContentType
content = "src/main.py:42:def process():"
detection = detect_content_type(content)
if detection.content_type == ContentType.SEARCH_RESULTS:
# Route to SearchCompressor
pass
elif detection.content_type == ContentType.BUILD_OUTPUT:
# Route to LogCompressor
pass
elif detection.content_type == ContentType.PLAIN_TEXT:
# Route to TextCompressor
pass
```
### Content Types
| Type | Detection Pattern |
|------|-------------------|
| `SEARCH_RESULTS` | `file:line:content` format |
| `BUILD_OUTPUT` | pytest, npm, cargo markers |
| `JSON` | Valid JSON structure |
| `PLAIN_TEXT` | Default fallback |
## Integration Pattern
```python
from headroom.transforms import (
detect_content_type, ContentType,
SearchCompressor, LogCompressor, TextCompressor
)
def compress_tool_output(content: str, context: str = "") -> str:
"""Application-level compression with explicit control."""
detection = detect_content_type(content)
if detection.content_type == ContentType.SEARCH_RESULTS:
result = SearchCompressor().compress(content, context)
return result.compressed
elif detection.content_type == ContentType.BUILD_OUTPUT:
result = LogCompressor().compress(content)
return result.compressed
elif detection.content_type == ContentType.PLAIN_TEXT:
result = TextCompressor().compress(content, context)
return result.compressed
else:
# JSON or other - let SmartCrusher handle it automatically
return content
```
## Configuration
Each compressor accepts configuration options:
```python
from headroom.transforms import SearchCompressor, SearchCompressorConfig
config = SearchCompressorConfig(
max_results=50, # Keep up to 50 matches
preserve_file_diversity=True, # Ensure different files represented
relevance_threshold=0.3, # Minimum relevance score to keep
)
compressor = SearchCompressor(config)
```
## Performance
| Compressor | Typical Input | Output | Speed |
|------------|---------------|--------|-------|
| SearchCompressor | 1000 matches | 30-50 matches | ~2ms |
| LogCompressor | 5000 lines | 100-200 lines | ~3ms |
| TextCompressor | 10000 chars | 2000 chars | ~2ms |
## When to Use
| Scenario | Recommendation |
|----------|----------------|
| JSON tool output | Let SmartCrusher handle automatically |
| grep/ripgrep results | Use SearchCompressor |
| pytest/npm/cargo output | Use LogCompressor |
| Documentation/README | Use TextCompressor |
| Unknown content | Use detect_content_type to route |

315
examples/ccr_demo.py Normal file
View file

@ -0,0 +1,315 @@
#!/usr/bin/env python3
"""Demonstration of CCR (Compress-Cache-Retrieve) architecture.
This script demonstrates:
1. How compression works with CCR caching
2. How the Response Handler automatically handles retrieval tool calls
3. How the Context Tracker enables multi-turn awareness
Run with: python examples/ccr_demo.py
"""
import asyncio
import json
from headroom.cache.compression_store import (
get_compression_store,
reset_compression_store,
)
from headroom.ccr import (
CCR_TOOL_NAME,
CCRResponseHandler,
CCRToolCall,
ContextTracker,
ContextTrackerConfig,
ResponseHandlerConfig,
create_ccr_tool_definition,
)
def print_section(title: str) -> None:
"""Print a section header."""
print("\n" + "=" * 60)
print(f" {title}")
print("=" * 60)
def demo_compression_store() -> str:
"""Demonstrate the compression store."""
print_section("1. COMPRESSION STORE - Caching Original Content")
# Reset for clean demo
reset_compression_store()
store = get_compression_store()
# Simulate tool output with 100 items
original_items = [
{"id": i, "file": f"src/module_{i}.py", "lines": 100 + i, "status": "ok"}
for i in range(100)
]
# Add some errors for interest
original_items[42]["status"] = "error"
original_items[42]["error"] = "SyntaxError: unexpected indent"
original_items[77]["status"] = "warning"
original_items[77]["warning"] = "Unused import"
original_json = json.dumps(original_items)
# SmartCrusher would compress to top 15 items (keeping errors)
compressed_items = [
original_items[0], # First few for context
original_items[1],
original_items[2],
original_items[42], # Error item - always kept!
original_items[77], # Warning item - always kept!
original_items[97], # Last few for recency
original_items[98],
original_items[99],
]
compressed_json = json.dumps(compressed_items)
# Store in CCR cache
hash_key = store.store(
original=original_json,
compressed=compressed_json,
original_item_count=100,
compressed_item_count=8,
tool_name="list_files",
)
print(f"\nOriginal: {len(original_items)} items ({len(original_json):,} chars)")
print(f"Compressed: {len(compressed_items)} items ({len(compressed_json):,} chars)")
print(f"Reduction: {100 - (len(compressed_json) / len(original_json) * 100):.1f}%")
print(f"CCR Hash: {hash_key}")
# Show that we can retrieve
entry = store.retrieve(hash_key)
print(f"\nRetrieved original: {entry.original_item_count} items")
# Show search capability
results = store.search(hash_key, "error SyntaxError")
print(f"Search for 'error SyntaxError': found {len(results)} items")
if results:
print(f" Found: {results[0]}")
return hash_key
def demo_tool_injection(hash_key: str) -> dict:
"""Demonstrate tool injection."""
print_section("2. TOOL INJECTION - Adding Retrieval Capability")
# Show the tool definition that gets injected
tool_def = create_ccr_tool_definition("anthropic")
print(f"\nInjected tool: {tool_def['name']}")
print(f"Description: {tool_def['description'][:100]}...")
# Show the marker that gets added to compressed content
marker = f"\n[100 items compressed to 8. Retrieve more: hash={hash_key}]"
print(f"\nMarker added to output:{marker}")
# Simulate an LLM response that calls the retrieval tool
simulated_response = {
"content": [
{"type": "text", "text": "I see some files. Let me get the full list."},
{
"type": "tool_use",
"id": "toolu_01ABC",
"name": CCR_TOOL_NAME,
"input": {"hash": hash_key},
},
]
}
print("\nSimulated LLM response (calls headroom_retrieve):")
print(json.dumps(simulated_response, indent=2)[:500] + "...")
return simulated_response
async def demo_response_handler(hash_key: str, initial_response: dict) -> None:
"""Demonstrate the response handler."""
print_section("3. RESPONSE HANDLER - Automatic Tool Call Handling")
print("\n--- BEFORE (without Response Handler) ---")
print("Problem: LLM calls headroom_retrieve, but no one handles it!")
print("The tool call would go back to the client unhandled.")
print("Client would need custom code to handle CCR tool calls.")
print("\n--- AFTER (with Response Handler) ---")
print("Solution: Response Handler intercepts and handles automatically!")
handler = CCRResponseHandler(
ResponseHandlerConfig(
max_retrieval_rounds=3,
)
)
# Check if response has CCR tool calls
has_ccr = handler.has_ccr_tool_calls(initial_response, "anthropic")
print(f"\nDetected CCR tool call: {has_ccr}")
# Parse the tool call
call = CCRToolCall(
tool_call_id="toolu_01ABC",
hash_key=hash_key,
)
print(f"Parsed: hash={call.hash_key}, query={call.query}")
# Execute retrieval
result = handler._execute_retrieval(call)
print(f"\nRetrieved {result.items_retrieved} items")
print(f"Success: {result.success}")
# Show what would happen in full flow
print("\nFull flow simulation:")
print("1. LLM response contains tool_use(headroom_retrieve)")
print("2. Handler detects CCR tool call")
print("3. Handler retrieves from cache (instant, ~1ms)")
print("4. Handler adds tool result to messages")
print("5. Handler makes continuation API call")
print("6. LLM responds with actual answer (no more CCR calls)")
print("7. Handler returns final response to client")
# Show handler stats
stats = handler.get_stats()
print(f"\nHandler stats: {stats}")
def demo_context_tracker(hash_key: str) -> None:
"""Demonstrate the context tracker."""
print_section("4. CONTEXT TRACKER - Multi-Turn Awareness")
print("\n--- BEFORE (without Context Tracker) ---")
print("Problem: In turn 5, LLM forgets what was compressed in turn 1!")
print("User: 'What about the authentication middleware?'")
print("LLM: 'I don't see any authentication files.'")
print("(Because auth files were in the compressed 92 items, not shown)")
print("\n--- AFTER (with Context Tracker) ---")
print("Solution: Tracker proactively expands relevant compressed content!")
config = ContextTrackerConfig(
relevance_threshold=0.1, # Lower for demo
max_context_age_seconds=300,
)
tracker = ContextTracker(config)
# Track the compression from turn 1
# Use keywords in sample_content that will match the query
tracker.track_compression(
hash_key=hash_key,
turn_number=1,
tool_name="list_files",
original_count=100,
compressed_count=8,
query_context="list all python files",
sample_content="authentication middleware handler auth_middleware.py auth_handler.py login security",
)
print(f"\nTurn 1: Tracked compression {hash_key}")
print(" Sample: 'authentication middleware handler auth_middleware.py ...'")
# Turn 5: User asks about auth
query = "show authentication middleware"
print(f"\nTurn 5: User asks '{query}'")
recommendations = tracker.analyze_query(query, current_turn=5)
print(f" Tracker found {len(recommendations)} relevant contexts")
if recommendations:
rec = recommendations[0]
print(f" → hash={rec.hash_key}")
print(f" → relevance={rec.relevance_score:.2f}")
print(f" → reason: {rec.reason}")
print(
f" → action: {'full expansion' if rec.expand_full else f'search for {rec.search_query}'}"
)
# Execute expansion
results = tracker.execute_expansions(recommendations)
if results:
print(f"\nProactively expanded: {results[0]['item_count']} items")
print("LLM now sees full file list, including auth_middleware.py!")
# Show tracker stats
stats = tracker.get_stats()
print(f"\nTracker stats: {json.dumps(stats, indent=2)}")
def demo_full_flow() -> None:
"""Show the complete CCR flow."""
print_section("5. COMPLETE CCR FLOW")
print("""
┌────────────────────────────────────────────────────────────┐
│ COMPLETE CCR ARCHITECTURE │
│ │
│ Phase 1: COMPRESSION STORE │
│ └─ Cache original content with hash │
│ └─ Enable instant retrieval (~1ms) │
│ │
│ Phase 2: TOOL INJECTION │
│ └─ Add headroom_retrieve tool to LLM context │
│ └─ Add retrieval markers to compressed output │
│ │
│ Phase 3: RESPONSE HANDLER │
│ └─ Intercept LLM responses │
│ └─ Detect CCR tool calls │
│ └─ Execute retrievals automatically │
│ └─ Continue conversation until done │
│ │
│ Phase 4: CONTEXT TRACKER │
│ └─ Track compressed content across turns │
│ └─ Analyze new queries for relevance │
│ └─ Proactively expand when needed │
│ │
│ Phase 5: FEEDBACK LOOP │
│ └─ Learn from retrieval patterns │
│ └─ Adjust compression for future requests │
└────────────────────────────────────────────────────────────┘
""")
print("KEY BENEFITS:")
print("• Reversible compression - no permanent data loss")
print("• Automatic handling - no client code changes needed")
print("• Multi-turn awareness - prevents context amnesia")
print("• Feedback learning - improves over time")
print("• Zero-risk - fallback to full data always available")
async def main() -> None:
"""Run the CCR demonstration."""
print("\n" + "=" * 60)
print(" HEADROOM CCR (Compress-Cache-Retrieve) DEMONSTRATION")
print("=" * 60)
# Demo 1: Compression Store
hash_key = demo_compression_store()
# Demo 2: Tool Injection
initial_response = demo_tool_injection(hash_key)
# Demo 3: Response Handler
await demo_response_handler(hash_key, initial_response)
# Demo 4: Context Tracker
demo_context_tracker(hash_key)
# Demo 5: Full Flow
demo_full_flow()
print("\n" + "=" * 60)
print(" DEMONSTRATION COMPLETE")
print("=" * 60)
print("\nRun the proxy with CCR enabled:")
print(" headroom proxy --port 8787")
print("\nCCR is enabled by default. The proxy will:")
print("• Cache compressed content automatically")
print("• Inject retrieval tool when compression occurs")
print("• Handle CCR tool calls in LLM responses")
print("• Track context across conversation turns")
print()
if __name__ == "__main__":
asyncio.run(main())

37
headroom/cache/compression_store.py vendored
View file

@ -302,6 +302,43 @@ class CompressionStore:
return result_entry
def get_metadata(
self,
hash_key: str,
) -> dict[str, Any] | None:
"""Get metadata about a stored entry without retrieving full content.
Useful for context tracking to know what was compressed without
fetching the entire original content.
Args:
hash_key: Hash key returned by store().
Returns:
Dict with metadata if found and not expired, None otherwise.
"""
with self._lock:
entry = self._store.get(hash_key)
if entry is None:
return None
if entry.is_expired():
del self._store[hash_key]
self._stale_heap_entries += 1
return None
return {
"hash": entry.hash,
"tool_name": entry.tool_name,
"original_item_count": entry.original_item_count,
"compressed_item_count": entry.compressed_item_count,
"query_context": entry.query_context,
"compressed_content": entry.compressed_content,
"created_at": entry.created_at,
"ttl": entry.ttl,
}
def search(
self,
hash_key: str,

37
headroom/ccr/__init__.py
View file

@ -3,6 +3,11 @@
This module provides tool injection and retrieval handling for the CCR architecture.
When tool outputs are compressed, the LLM can retrieve more data if needed.
Three key components:
1. Tool Injection: Proxy injects headroom_retrieve tool into requests
2. Response Handler: Intercepts responses, handles CCR tool calls automatically
3. Context Tracker: Tracks compressed content across turns, enables proactive expansion
Two distribution channels for the retrieval tool:
1. Tool Injection: Proxy injects tool into request when compression occurs
2. MCP Server: Standalone server exposes tool via MCP protocol
@ -10,6 +15,22 @@ Two distribution channels for the retrieval tool:
When MCP is configured, tool injection is skipped to avoid duplicates.
"""
from .context_tracker import (
CompressedContext,
ContextTracker,
ContextTrackerConfig,
ExpansionRecommendation,
get_context_tracker,
reset_context_tracker,
)
from .response_handler import (
CCRResponseHandler,
CCRToolCall,
CCRToolResult,
ResponseHandlerConfig,
StreamingCCRBuffer,
StreamingCCRHandler,
)
from .tool_injection import (
CCR_TOOL_NAME,
CCRToolInjector,
@ -29,11 +50,27 @@ except ImportError:
MCP_SERVER_AVAILABLE = False
__all__ = [
# Tool injection
"CCR_TOOL_NAME",
"CCRToolInjector",
"create_ccr_tool_definition",
"create_system_instructions",
"parse_tool_call",
# Response handling
"CCRResponseHandler",
"CCRToolCall",
"CCRToolResult",
"ResponseHandlerConfig",
"StreamingCCRBuffer",
"StreamingCCRHandler",
# Context tracking
"CompressedContext",
"ContextTracker",
"ContextTrackerConfig",
"ExpansionRecommendation",
"get_context_tracker",
"reset_context_tracker",
# MCP server
"CCRMCPServer",
"create_ccr_mcp_server",
"MCP_SERVER_AVAILABLE",

582
headroom/ccr/context_tracker.py Normal file
View file

@ -0,0 +1,582 @@
"""Multi-turn context tracking for CCR (Compress-Cache-Retrieve).
This module tracks compressed content across conversation turns and
provides intelligent context expansion based on query relevance.
Key features:
1. Track all compression hashes across the conversation
2. Analyze new queries to detect if they need expanded context
3. Proactively expand relevant compressed content before LLM responds
4. Prevent "context amnesia" where earlier compressed data is forgotten
Example:
Turn 1: Search returns 100 files → compressed to 10 (hash=abc123)
Turn 5: User asks "What about auth middleware?"
Without tracking: LLM doesn't know auth_middleware.py exists
With tracking: Tracker detects "auth middleware" might be in abc123,
proactively expands it, LLM gets the full context
"""
from __future__ import annotations
import json
import logging
import re
import time
from dataclasses import dataclass
from typing import Any
from ..cache.compression_store import get_compression_store
logger = logging.getLogger(__name__)
@dataclass
class CompressedContext:
"""Represents a piece of compressed context from the conversation."""
hash_key: str
turn_number: int
timestamp: float
tool_name: str | None
original_item_count: int
compressed_item_count: int
query_context: str # The query/context when compression happened
sample_content: str # Preview of what was compressed (for relevance matching)
@dataclass
class ExpansionRecommendation:
"""Recommendation to expand compressed context."""
hash_key: str
reason: str
relevance_score: float
expand_full: bool = True # True = expand all, False = search only
search_query: str | None = None
@dataclass
class ContextTrackerConfig:
"""Configuration for context tracking."""
# Whether tracking is enabled
enabled: bool = True
# Maximum contexts to track (LRU eviction)
max_tracked_contexts: int = 100
# Relevance threshold for recommending expansion (0-1)
relevance_threshold: float = 0.3
# Maximum age for contexts (seconds) - older contexts less likely to expand
max_context_age_seconds: float = 300.0 # 5 minutes
# Whether to proactively expand based on query analysis
proactive_expansion: bool = True
# Maximum items to proactively expand per turn
max_proactive_expansions: int = 2
class ContextTracker:
"""Tracks compressed contexts across conversation turns.
This tracker maintains awareness of what has been compressed
and can recommend expansions when new queries might need that data.
Usage:
tracker = ContextTracker()
# Track compression events
tracker.track_compression(
hash_key="abc123",
turn_number=1,
tool_name="Bash",
original_count=100,
compressed_count=10,
query_context="find all python files",
sample_content='["src/main.py", "src/auth.py", ...]',
)
# On new user message, check for expansion needs
recommendations = tracker.analyze_query(
query="What about the authentication code?",
current_turn=5,
)
# recommendations might suggest expanding abc123 because
# "authentication" matches "auth.py" in the sample content
"""
def __init__(self, config: ContextTrackerConfig | None = None):
self.config = config or ContextTrackerConfig()
self._contexts: dict[str, CompressedContext] = {}
self._turn_order: list[str] = [] # For LRU
self._current_turn: int = 0
def track_compression(
self,
hash_key: str,
turn_number: int,
tool_name: str | None,
original_count: int,
compressed_count: int,
query_context: str = "",
sample_content: str = "",
) -> None:
"""Track a compression event.
Args:
hash_key: The CCR hash for this compression.
turn_number: The conversation turn number.
tool_name: Name of the tool whose output was compressed.
original_count: Original item count.
compressed_count: Compressed item count.
query_context: The user query when compression happened.
sample_content: Sample of the content for relevance matching.
"""
if not self.config.enabled:
return
context = CompressedContext(
hash_key=hash_key,
turn_number=turn_number,
timestamp=time.time(),
tool_name=tool_name,
original_item_count=original_count,
compressed_item_count=compressed_count,
query_context=query_context,
sample_content=sample_content[:2000], # Limit sample size
)
# Add or update context
if hash_key in self._contexts:
self._turn_order.remove(hash_key)
self._contexts[hash_key] = context
self._turn_order.append(hash_key)
# LRU eviction
while len(self._contexts) > self.config.max_tracked_contexts:
oldest = self._turn_order.pop(0)
del self._contexts[oldest]
self._current_turn = max(self._current_turn, turn_number)
logger.debug(
f"CCR Tracker: Tracked compression {hash_key} "
f"({original_count} -> {compressed_count} items)"
)
def analyze_query(
self,
query: str,
current_turn: int | None = None,
) -> list[ExpansionRecommendation]:
"""Analyze a query to find relevant compressed contexts.
Args:
query: The user's query/message.
current_turn: Current turn number (for age calculation).
Returns:
List of expansion recommendations, sorted by relevance.
"""
if not self.config.enabled or not self.config.proactive_expansion:
return []
if current_turn is not None:
self._current_turn = current_turn
recommendations: list[ExpansionRecommendation] = []
now = time.time()
for hash_key, context in self._contexts.items():
# Check age
age = now - context.timestamp
if age > self.config.max_context_age_seconds:
continue
# Calculate relevance
relevance = self._calculate_relevance(query, context)
# Age discount: older contexts get lower scores
age_factor = 1.0 - (age / self.config.max_context_age_seconds) * 0.5
relevance *= age_factor
if relevance >= self.config.relevance_threshold:
# Determine if full expansion or search
expand_full, search_query = self._determine_expansion_type(
query, context, relevance
)
recommendations.append(
ExpansionRecommendation(
hash_key=hash_key,
reason=self._generate_reason(query, context, relevance),
relevance_score=relevance,
expand_full=expand_full,
search_query=search_query,
)
)
# Sort by relevance, limit count
recommendations.sort(key=lambda r: r.relevance_score, reverse=True)
return recommendations[: self.config.max_proactive_expansions]
def _calculate_relevance(
self,
query: str,
context: CompressedContext,
) -> float:
"""Calculate relevance score between query and compressed context.
Uses simple but effective heuristics:
1. Keyword overlap with sample content
2. Keyword overlap with original query context
3. Tool name relevance
"""
query_lower = query.lower()
query_words = set(self._extract_keywords(query_lower))
if not query_words:
return 0.0
score = 0.0
# Check sample content overlap
sample_lower = context.sample_content.lower()
sample_words = set(self._extract_keywords(sample_lower))
if sample_words:
overlap = query_words & sample_words
score += len(overlap) / len(query_words) * 0.5
# Bonus for exact substring matches
for word in query_words:
if len(word) >= 4 and word in sample_lower:
score += 0.2
# Check original query context overlap
if context.query_context:
context_lower = context.query_context.lower()
context_words = set(self._extract_keywords(context_lower))
if context_words:
overlap = query_words & context_words
score += len(overlap) / len(query_words) * 0.3
# Tool name relevance
if context.tool_name:
tool_lower = context.tool_name.lower()
# File operations more likely to need expansion
if any(w in tool_lower for w in ["find", "glob", "search", "grep", "ls"]):
if any(w in query_lower for w in ["file", "where", "find", "show", "list"]):
score += 0.1
return min(score, 1.0)
def _extract_keywords(self, text: str) -> list[str]:
"""Extract meaningful keywords from text."""
# Remove common punctuation, split into words
words = re.findall(r"\b[a-z][a-z0-9_.-]*[a-z0-9]\b|\b[a-z]{2,}\b", text)
# Filter stop words and very short words
stop_words = {
"the",
"a",
"an",
"is",
"are",
"was",
"were",
"be",
"been",
"being",
"have",
"has",
"had",
"do",
"does",
"did",
"will",
"would",
"could",
"should",
"may",
"might",
"must",
"shall",
"can",
"need",
"dare",
"ought",
"used",
"to",
"of",
"in",
"for",
"on",
"with",
"at",
"by",
"from",
"as",
"into",
"through",
"during",
"before",
"after",
"above",
"below",
"between",
"under",
"again",
"further",
"then",
"once",
"here",
"there",
"when",
"where",
"why",
"how",
"all",
"each",
"few",
"more",
"most",
"other",
"some",
"such",
"no",
"nor",
"not",
"only",
"own",
"same",
"so",
"than",
"too",
"very",
"just",
"and",
"but",
"if",
"or",
"because",
"until",
"while",
"this",
"that",
"these",
"those",
"what",
"which",
"who",
"whom",
"it",
"its",
"me",
"my",
"i",
"you",
}
return [w for w in words if w not in stop_words and len(w) >= 2]
def _determine_expansion_type(
self,
query: str,
context: CompressedContext,
relevance: float,
) -> tuple[bool, str | None]:
"""Determine whether to do full expansion or search.
Returns:
Tuple of (expand_full, search_query)
"""
# High relevance + small original count = full expansion
if relevance > 0.6 or context.original_item_count <= 50:
return True, None
# Extract specific search terms from query
keywords = self._extract_keywords(query.lower())
# Filter to most specific keywords (longer, less common)
specific_keywords = [
k
for k in keywords
if len(k) >= 4 and k not in {"file", "code", "show", "find", "list", "what"}
]
if specific_keywords:
# Use top keywords as search query
search_query = " ".join(specific_keywords[:3])
return False, search_query
# Default to full expansion if we can't form a good search
return True, None
def _generate_reason(
self,
query: str,
context: CompressedContext,
relevance: float,
) -> str:
"""Generate human-readable reason for expansion recommendation."""
parts = []
if context.tool_name:
parts.append(f"from {context.tool_name}")
parts.append(
f"{context.original_item_count} items compressed in turn {context.turn_number}"
)
if relevance > 0.5:
parts.append("high relevance to current query")
else:
parts.append("possible relevance to current query")
return ", ".join(parts)
def execute_expansions(
self,
recommendations: list[ExpansionRecommendation],
) -> list[dict[str, Any]]:
"""Execute expansion recommendations and return the expanded content.
Args:
recommendations: List of expansion recommendations.
Returns:
List of expanded content dicts with hash, content, and metadata.
"""
store = get_compression_store()
results = []
for rec in recommendations:
try:
if rec.expand_full:
entry = store.retrieve(rec.hash_key)
if entry:
results.append(
{
"hash": rec.hash_key,
"type": "full",
"content": entry.original_content,
"item_count": entry.original_item_count,
"reason": rec.reason,
}
)
logger.info(
f"CCR Tracker: Proactively expanded {rec.hash_key} "
f"({entry.original_item_count} items)"
)
else:
search_results = store.search(rec.hash_key, rec.search_query or "")
if search_results:
results.append(
{
"hash": rec.hash_key,
"type": "search",
"query": rec.search_query,
"content": search_results,
"item_count": len(search_results),
"reason": rec.reason,
}
)
logger.info(
f"CCR Tracker: Proactive search in {rec.hash_key} "
f"for '{rec.search_query}' ({len(search_results)} results)"
)
except Exception as e:
logger.warning(f"CCR Tracker: Failed to expand {rec.hash_key}: {e}")
return results
def format_expansions_for_context(
self,
expansions: list[dict[str, Any]],
) -> str:
"""Format expansions as additional context for the LLM.
Args:
expansions: Results from execute_expansions.
Returns:
Formatted string to add to context.
"""
if not expansions:
return ""
parts = ["[Proactive Context Expansion - relevant to your query]"]
for exp in expansions:
if exp["type"] == "full":
parts.append(f"\n--- Expanded from earlier ({exp['reason']}) ---")
parts.append(exp["content"])
else:
parts.append(f"\n--- Search results for '{exp['query']}' ({exp['reason']}) ---")
if isinstance(exp["content"], list):
parts.append(json.dumps(exp["content"], indent=2))
else:
parts.append(str(exp["content"]))
parts.append("\n[End Proactive Expansion]")
return "\n".join(parts)
def get_tracked_hashes(self) -> list[str]:
"""Get list of currently tracked hashes."""
return list(self._contexts.keys())
def get_stats(self) -> dict[str, Any]:
"""Get tracker statistics."""
return {
"tracked_contexts": len(self._contexts),
"current_turn": self._current_turn,
"config": {
"enabled": self.config.enabled,
"max_contexts": self.config.max_tracked_contexts,
"relevance_threshold": self.config.relevance_threshold,
"proactive_expansion": self.config.proactive_expansion,
},
"contexts": [
{
"hash": ctx.hash_key,
"turn": ctx.turn_number,
"tool": ctx.tool_name,
"items": f"{ctx.compressed_item_count}/{ctx.original_item_count}",
}
for ctx in self._contexts.values()
],
}
def clear(self) -> None:
"""Clear all tracked contexts."""
self._contexts.clear()
self._turn_order.clear()
self._current_turn = 0
# Global instance (per-session)
_context_tracker: ContextTracker | None = None
def get_context_tracker() -> ContextTracker:
"""Get the global context tracker."""
global _context_tracker
if _context_tracker is None:
_context_tracker = ContextTracker()
return _context_tracker
def reset_context_tracker() -> None:
"""Reset the global context tracker."""
global _context_tracker
if _context_tracker is not None:
_context_tracker.clear()
_context_tracker = None

772
headroom/ccr/response_handler.py Normal file
View file

@ -0,0 +1,772 @@
"""Response handling for CCR (Compress-Cache-Retrieve).
This module provides response interception and CCR tool call handling.
When the LLM calls headroom_retrieve, this handler:
1. Detects the tool call in the response
2. Retrieves content from the compression store
3. Continues the conversation with the tool result
4. Returns the final response to the client
This solves the critical gap where the proxy injects the tool but
can't handle the LLM's tool calls.
"""
from __future__ import annotations
import json
import logging
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
from typing import Any
from ..cache.compression_store import get_compression_store
from .tool_injection import CCR_TOOL_NAME, parse_tool_call
logger = logging.getLogger(__name__)
@dataclass
class CCRToolCall:
"""Represents a detected CCR tool call."""
tool_call_id: str
hash_key: str
query: str | None = None
@dataclass
class CCRToolResult:
"""Result of handling a CCR tool call."""
tool_call_id: str
content: str
success: bool
items_retrieved: int = 0
was_search: bool = False
@dataclass
class ResponseHandlerConfig:
"""Configuration for CCR response handling."""
# Whether to handle CCR tool calls automatically
enabled: bool = True
# Maximum number of CCR retrieval rounds (prevent infinite loops)
max_retrieval_rounds: int = 3
# Whether to strip CCR tool calls from final response
strip_ccr_from_response: bool = True
# Timeout for continuation requests (ms)
continuation_timeout_ms: int = 120000
class CCRResponseHandler:
"""Handles CCR tool calls in LLM responses.
This handler intercepts responses, detects CCR tool calls,
retrieves content, and continues the conversation until
the LLM produces a response without CCR tool calls.
Example flow:
1. LLM response contains: tool_use(headroom_retrieve, hash=abc123)
2. Handler detects this, retrieves original content
3. Handler makes another API call with tool result
4. LLM responds with actual content (no CCR tool call)
5. Handler returns this final response
Usage:
handler = CCRResponseHandler(config)
# Check if response needs handling
if handler.has_ccr_tool_calls(response_json):
# Handle the tool calls
final_response = await handler.handle_response(
response_json,
messages,
tools,
api_call_fn,
provider="anthropic"
)
else:
final_response = response_json
"""
def __init__(self, config: ResponseHandlerConfig | None = None):
self.config = config or ResponseHandlerConfig()
self._retrieval_count = 0
def has_ccr_tool_calls(
self,
response: dict[str, Any],
provider: str = "anthropic",
) -> bool:
"""Check if response contains CCR tool calls.
Args:
response: The API response JSON.
provider: The provider type.
Returns:
True if response contains headroom_retrieve tool calls.
"""
tool_calls = self._extract_tool_calls(response, provider)
return any(
tc.get("name") == CCR_TOOL_NAME or tc.get("function", {}).get("name") == CCR_TOOL_NAME
for tc in tool_calls
)
def _extract_tool_calls(
self,
response: dict[str, Any],
provider: str,
) -> list[dict[str, Any]]:
"""Extract tool calls from response based on provider format."""
if provider == "anthropic":
# Anthropic format: content blocks with type=tool_use
content = response.get("content", [])
if isinstance(content, list):
return [block for block in content if block.get("type") == "tool_use"]
return []
elif provider == "openai":
# OpenAI format: message.tool_calls array
message = response.get("choices", [{}])[0].get("message", {})
tool_calls = message.get("tool_calls", [])
return list(tool_calls) if tool_calls else []
return []
def _parse_ccr_tool_calls(
self,
response: dict[str, Any],
provider: str,
) -> tuple[list[CCRToolCall], list[dict[str, Any]]]:
"""Parse CCR tool calls from response, separate from other tool calls.
Returns:
Tuple of (ccr_tool_calls, other_tool_calls)
"""
all_tool_calls = self._extract_tool_calls(response, provider)
ccr_calls = []
other_calls = []
for tc in all_tool_calls:
hash_key, query = parse_tool_call(tc, provider)
if hash_key is not None:
# This is a CCR tool call
tool_call_id = tc.get("id", "")
ccr_calls.append(
CCRToolCall(
tool_call_id=tool_call_id,
hash_key=hash_key,
query=query,
)
)
else:
# Not a CCR tool call
other_calls.append(tc)
return ccr_calls, other_calls
def _execute_retrieval(self, ccr_call: CCRToolCall) -> CCRToolResult:
"""Execute a CCR retrieval.
Args:
ccr_call: The CCR tool call to execute.
Returns:
CCRToolResult with the retrieved content.
"""
store = get_compression_store()
try:
if ccr_call.query:
# Search within compressed content
results = store.search(ccr_call.hash_key, ccr_call.query)
content = json.dumps(
{
"hash": ccr_call.hash_key,
"query": ccr_call.query,
"results": results,
"count": len(results),
},
indent=2,
)
return CCRToolResult(
tool_call_id=ccr_call.tool_call_id,
content=content,
success=True,
items_retrieved=len(results),
was_search=True,
)
else:
# Full retrieval
entry = store.retrieve(ccr_call.hash_key)
if entry:
content = json.dumps(
{
"hash": ccr_call.hash_key,
"original_content": entry.original_content,
"original_item_count": entry.original_item_count,
},
indent=2,
)
return CCRToolResult(
tool_call_id=ccr_call.tool_call_id,
content=content,
success=True,
items_retrieved=entry.original_item_count,
was_search=False,
)
else:
content = json.dumps(
{
"error": "Entry not found or expired (TTL: 5 minutes)",
"hash": ccr_call.hash_key,
},
indent=2,
)
return CCRToolResult(
tool_call_id=ccr_call.tool_call_id,
content=content,
success=False,
)
except Exception as e:
logger.error(f"CCR retrieval failed for {ccr_call.hash_key}: {e}")
content = json.dumps(
{
"error": f"Retrieval failed: {str(e)}",
"hash": ccr_call.hash_key,
},
indent=2,
)
return CCRToolResult(
tool_call_id=ccr_call.tool_call_id,
content=content,
success=False,
)
def _create_tool_result_message(
self,
results: list[CCRToolResult],
provider: str,
) -> dict[str, Any]:
"""Create a tool result message from CCR results.
Args:
results: List of CCR tool results.
provider: The provider type.
Returns:
Message dict in the appropriate format.
"""
if provider == "anthropic":
# Anthropic: user message with tool_result content blocks
content_blocks = []
for result in results:
content_blocks.append(
{
"type": "tool_result",
"tool_use_id": result.tool_call_id,
"content": result.content,
}
)
return {
"role": "user",
"content": content_blocks,
}
elif provider == "openai":
# OpenAI: multiple tool messages
# Actually for OpenAI we return a list of messages
return {
"_openai_tool_results": [
{
"role": "tool",
"tool_call_id": result.tool_call_id,
"content": result.content,
}
for result in results
]
}
else:
# Generic format
return {
"role": "tool",
"content": json.dumps(
[{"tool_call_id": r.tool_call_id, "result": r.content} for r in results]
),
}
def _extract_assistant_message(
self,
response: dict[str, Any],
provider: str,
) -> dict[str, Any]:
"""Extract the assistant message from an API response.
Args:
response: The API response.
provider: The provider type.
Returns:
The assistant message dict.
"""
if provider == "anthropic":
return {
"role": "assistant",
"content": response.get("content", []),
}
elif provider == "openai":
message = response.get("choices", [{}])[0].get("message", {})
return {
"role": "assistant",
"content": message.get("content"),
"tool_calls": message.get("tool_calls"),
}
else:
return {
"role": "assistant",
"content": response.get("content", ""),
}
async def handle_response(
self,
response: dict[str, Any],
messages: list[dict[str, Any]],
tools: list[dict[str, Any]] | None,
api_call_fn: Callable[
[list[dict[str, Any]], list[dict[str, Any]] | None], Awaitable[dict[str, Any]]
],
provider: str = "anthropic",
) -> dict[str, Any]:
"""Handle CCR tool calls in a response.
This method:
1. Detects CCR tool calls
2. Executes retrievals
3. Continues conversation with tool results
4. Repeats until no CCR tool calls remain
Args:
response: The initial API response.
messages: The conversation messages.
tools: The tools list (should include CCR tool).
api_call_fn: Async function to make API calls.
Signature: (messages, tools) -> response
provider: The provider type.
Returns:
The final response (with no CCR tool calls).
"""
if not self.config.enabled:
return response
current_response = response
current_messages = list(messages) # Copy to avoid mutation
rounds = 0
while rounds < self.config.max_retrieval_rounds:
# Check for CCR tool calls
ccr_calls, other_calls = self._parse_ccr_tool_calls(current_response, provider)
if not ccr_calls:
# No CCR tool calls, we're done
break
rounds += 1
self._retrieval_count += len(ccr_calls)
logger.info(f"CCR: Handling {len(ccr_calls)} retrieval(s) in round {rounds}")
# Execute all CCR retrievals
results = [self._execute_retrieval(call) for call in ccr_calls]
# Log retrieval stats
total_items = sum(r.items_retrieved for r in results)
searches = sum(1 for r in results if r.was_search)
logger.debug(
f"CCR: Retrieved {total_items} items "
f"({searches} searches, {len(results) - searches} full)"
)
# Build continuation messages
# Add assistant message (the response that had tool calls)
assistant_msg = self._extract_assistant_message(current_response, provider)
current_messages.append(assistant_msg)
# Add tool results
tool_result_msg = self._create_tool_result_message(results, provider)
if provider == "openai" and "_openai_tool_results" in tool_result_msg:
# OpenAI uses multiple messages for tool results
current_messages.extend(tool_result_msg["_openai_tool_results"])
else:
current_messages.append(tool_result_msg)
# Make continuation API call
try:
current_response = await api_call_fn(current_messages, tools)
except Exception as e:
logger.error(f"CCR: Continuation API call failed: {e}")
# Return the response we had (with unhandled CCR calls)
# The client will see the tool_use and might handle it differently
break
if rounds >= self.config.max_retrieval_rounds:
logger.warning(
f"CCR: Hit max retrieval rounds ({self.config.max_retrieval_rounds}), "
f"returning response with possible unhandled CCR calls"
)
return current_response
def get_stats(self) -> dict[str, Any]:
"""Get handler statistics."""
return {
"total_retrievals": self._retrieval_count,
"config": {
"enabled": self.config.enabled,
"max_rounds": self.config.max_retrieval_rounds,
},
}
@dataclass
class StreamingCCRBuffer:
"""Buffer for detecting CCR tool calls in streaming responses.
Since streaming responses come in chunks, we need to buffer
until we can detect whether there's a CCR tool call.
Strategy:
1. Buffer chunks until we see a complete tool_use block
2. If it's a CCR call, switch to buffered mode
3. Handle CCR and then stream the continuation
"""
chunks: list[bytes] = field(default_factory=list)
detected_ccr: bool = False
complete_response: dict[str, Any] | None = None
# Patterns to detect tool_use in stream
_tool_use_start: bytes = b'"type":"tool_use"'
_ccr_tool_pattern: bytes = f'"{CCR_TOOL_NAME}"'.encode()
def add_chunk(self, chunk: bytes) -> bool:
"""Add a chunk and check for CCR tool calls.
Returns:
True if CCR tool call detected (should switch to buffered mode).
"""
self.chunks.append(chunk)
# Quick check: does accumulated content contain CCR tool?
accumulated = b"".join(self.chunks)
if self._tool_use_start in accumulated and self._ccr_tool_pattern in accumulated:
self.detected_ccr = True
return True
return False
def get_accumulated(self) -> bytes:
"""Get all accumulated chunks."""
return b"".join(self.chunks)
def clear(self) -> None:
"""Clear the buffer."""
self.chunks.clear()
self.detected_ccr = False
self.complete_response = None
class StreamingCCRHandler:
"""Handle CCR tool calls in streaming responses.
For streaming, we have two modes:
1. Pass-through: No CCR detected, stream chunks directly
2. Buffered: CCR detected, buffer response, handle, then stream result
The challenge is we can't know if there's a CCR call until we see
enough of the response. So we buffer initially, then decide.
"""
def __init__(
self,
response_handler: CCRResponseHandler,
provider: str = "anthropic",
) -> None:
self.response_handler = response_handler
self.provider = provider
self.buffer = StreamingCCRBuffer()
async def process_stream(
self,
stream_iterator: Any, # AsyncIterator[bytes]
messages: list[dict[str, Any]],
tools: list[dict[str, Any]] | None,
api_call_fn: Callable[
[list[dict[str, Any]], list[dict[str, Any]] | None], Awaitable[dict[str, Any]]
],
) -> Any: # AsyncGenerator[bytes, None]
"""Process a streaming response, handling CCR if needed.
This is an async generator that yields chunks.
If CCR is detected, it buffers, handles, and re-streams.
Args:
stream_iterator: Async iterator of response chunks.
messages: The conversation messages.
tools: The tools list.
api_call_fn: Function to make API calls for continuation.
Yields:
Response chunks (possibly from continuation response).
"""
# Phase 1: Initial detection
# Buffer chunks until we can determine if there's a CCR call
detection_complete = False
async for chunk in stream_iterator:
self.buffer.add_chunk(chunk)
# Check if we can determine CCR presence
# For Anthropic, tool_use blocks come after text content
# We need to see the stop_reason to know if there's a tool call
accumulated = self.buffer.get_accumulated()
# Look for stream end markers
if b'"stop_reason"' in accumulated:
detection_complete = True
if self.buffer.detected_ccr:
# CCR detected - need to handle
break
else:
# No CCR - yield all buffered chunks
for buffered_chunk in self.buffer.chunks:
yield buffered_chunk
self.buffer.clear()
# If we haven't detected anything yet and buffer is large,
# start yielding (response is probably just text)
elif len(accumulated) > 10000 and not self.buffer.detected_ccr:
for buffered_chunk in self.buffer.chunks:
yield buffered_chunk
self.buffer.clear()
# Continue streaming rest of response
if not detection_complete and not self.buffer.detected_ccr:
async for chunk in stream_iterator:
if self.buffer.detected_ccr:
self.buffer.add_chunk(chunk)
else:
yield chunk
# Phase 2: Handle CCR if detected
if self.buffer.detected_ccr:
logger.info("CCR: Detected tool call in stream, switching to buffered mode")
# Collect rest of stream
async for chunk in stream_iterator:
self.buffer.add_chunk(chunk)
# Parse the complete response
try:
# For SSE streams, we need to parse the accumulated data
complete_data = self._parse_sse_stream(self.buffer.get_accumulated())
# Handle CCR
final_response = await self.response_handler.handle_response(
complete_data,
messages,
tools,
api_call_fn,
self.provider,
)
# Re-stream the final response
# Convert back to SSE format
async for chunk in self._response_to_sse(final_response):
yield chunk
except Exception as e:
logger.error(f"CCR: Failed to handle streamed CCR: {e}")
# Fall back to yielding original buffered content
yield self.buffer.get_accumulated()
def _parse_sse_stream(self, data: bytes) -> dict[str, Any]:
"""Parse SSE stream data into a response dict.
SSE format: data: {...}\n\n
"""
# Accumulate all event data
events = []
for line in data.decode("utf-8", errors="replace").split("\n"):
if line.startswith("data: "):
event_data = line[6:]
if event_data.strip() and event_data.strip() != "[DONE]":
try:
events.append(json.loads(event_data))
except json.JSONDecodeError:
pass
# Reconstruct response from events
# This is provider-specific
if self.provider == "anthropic":
return self._reconstruct_anthropic_response(events)
else:
return self._reconstruct_openai_response(events)
def _reconstruct_anthropic_response(
self,
events: list[dict[str, Any]],
) -> dict[str, Any]:
"""Reconstruct Anthropic response from stream events."""
response: dict[str, Any] = {
"content": [],
"stop_reason": None,
"usage": {},
}
current_text = ""
current_tool: dict[str, Any] | None = None
for event in events:
event_type = event.get("type", "")
if event_type == "content_block_start":
block = event.get("content_block", {})
if block.get("type") == "text":
current_text = block.get("text", "")
elif block.get("type") == "tool_use":
current_tool = {
"type": "tool_use",
"id": block.get("id", ""),
"name": block.get("name", ""),
"input": {},
}
elif event_type == "content_block_delta":
delta = event.get("delta", {})
if delta.get("type") == "text_delta":
current_text += delta.get("text", "")
elif delta.get("type") == "input_json_delta":
# Accumulate JSON for tool input
if current_tool is not None:
partial = delta.get("partial_json", "")
# This is tricky - partial JSON needs accumulation
# For simplicity, we'll try to parse when complete
current_tool["_partial_json"] = (
current_tool.get("_partial_json", "") + partial
)
elif event_type == "content_block_stop":
if current_text:
response["content"].append(
{
"type": "text",
"text": current_text,
}
)
current_text = ""
if current_tool:
# Parse accumulated JSON
partial = current_tool.pop("_partial_json", "")
if partial:
try:
current_tool["input"] = json.loads(partial)
except json.JSONDecodeError:
current_tool["input"] = {}
response["content"].append(current_tool)
current_tool = None
elif event_type == "message_delta":
delta = event.get("delta", {})
if "stop_reason" in delta:
response["stop_reason"] = delta["stop_reason"]
elif event_type == "message_stop":
pass
return response
def _reconstruct_openai_response(
self,
events: list[dict[str, Any]],
) -> dict[str, Any]:
"""Reconstruct OpenAI response from stream events."""
message: dict[str, Any] = {
"role": "assistant",
"content": "",
"tool_calls": [],
}
tool_calls_map: dict[int, dict[str, Any]] = {}
for event in events:
choices = event.get("choices", [])
if not choices:
continue
delta = choices[0].get("delta", {})
if "content" in delta and delta["content"]:
message["content"] = (message.get("content") or "") + delta["content"]
if "tool_calls" in delta:
for tc_delta in delta["tool_calls"]:
idx = tc_delta.get("index", 0)
if idx not in tool_calls_map:
tool_calls_map[idx] = {
"id": "",
"type": "function",
"function": {"name": "", "arguments": ""},
}
tc = tool_calls_map[idx]
if "id" in tc_delta:
tc["id"] = tc_delta["id"]
if "function" in tc_delta:
fn = tc_delta["function"]
if "name" in fn:
tc["function"]["name"] = fn["name"]
if "arguments" in fn:
tc["function"]["arguments"] += fn["arguments"]
message["tool_calls"] = [tool_calls_map[i] for i in sorted(tool_calls_map.keys())]
if not message["tool_calls"]:
del message["tool_calls"]
if not message["content"]:
message["content"] = None
return {
"choices": [{"message": message, "finish_reason": "stop"}],
}
async def _response_to_sse(
self,
response: dict[str, Any],
) -> Any: # AsyncGenerator[bytes, None]
"""Convert a response back to SSE format for streaming.
This is a simplified version - in practice you might want
to chunk the response more granularly.
"""
if self.provider == "anthropic":
# Anthropic SSE format
yield b"event: message_start\n"
yield f"data: {json.dumps({'type': 'message_start', 'message': response})}\n\n".encode()
yield b"event: message_stop\n"
yield b'data: {"type": "message_stop"}\n\n'
else:
# OpenAI SSE format
yield f"data: {json.dumps(response)}\n\n".encode()
yield b"data: [DONE]\n\n"

182
headroom/proxy/server.py
View file

@ -35,7 +35,7 @@ from collections import defaultdict
from dataclasses import asdict, dataclass
from datetime import datetime, timedelta
from pathlib import Path
from typing import Literal
from typing import Any, Literal
import httpx
@ -54,7 +54,15 @@ sys.path.insert(0, str(Path(__file__).parent.parent.parent))
from headroom.cache.compression_feedback import get_compression_feedback
from headroom.cache.compression_store import get_compression_store
from headroom.ccr import CCR_TOOL_NAME, CCRToolInjector, parse_tool_call
from headroom.ccr import (
CCR_TOOL_NAME,
CCRResponseHandler,
CCRToolInjector,
ContextTracker,
ContextTrackerConfig,
ResponseHandlerConfig,
parse_tool_call,
)
from headroom.config import CacheAlignerConfig, RollingWindowConfig, SmartCrusherConfig
from headroom.providers import AnthropicProvider, OpenAIProvider
from headroom.telemetry import get_telemetry_collector
@ -155,6 +163,15 @@ 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
# CCR Response Handling (intercept and handle CCR tool calls automatically)
ccr_handle_responses: bool = True # Handle headroom_retrieve calls in responses
ccr_max_retrieval_rounds: int = 3 # Max rounds of retrieval before returning
# CCR Context Tracking (track compressed content across turns)
ccr_context_tracking: bool = True # Track compressed contexts for proactive expansion
ccr_proactive_expansion: bool = True # Proactively expand based on query relevance
ccr_max_proactive_expansions: int = 2 # Max contexts to proactively expand per turn
# LLMLingua ML-based compression (opt-in)
llmlingua_enabled: bool = False # Enable LLMLingua-2 for ML-based compression
llmlingua_device: str = "auto" # Device: 'auto', 'cuda', 'cpu', 'mps'
@ -740,6 +757,34 @@ class HeadroomProxy:
inject_system_instructions=config.ccr_inject_system_instructions,
)
# CCR Response Handler (handles CCR tool calls automatically)
self.ccr_response_handler = (
CCRResponseHandler(
ResponseHandlerConfig(
enabled=True,
max_retrieval_rounds=config.ccr_max_retrieval_rounds,
)
)
if config.ccr_handle_responses
else None
)
# CCR Context Tracker (tracks compressed content across turns)
self.ccr_context_tracker = (
ContextTracker(
ContextTrackerConfig(
enabled=True,
proactive_expansion=config.ccr_proactive_expansion,
max_proactive_expansions=config.ccr_max_proactive_expansions,
)
)
if config.ccr_context_tracking
else None
)
# Turn counter for context tracking
self._turn_counter = 0
def _setup_llmlingua(self, config: ProxyConfig, transforms: list) -> str:
"""Set up LLMLingua compression if enabled.
@ -799,6 +844,21 @@ class HeadroomProxy:
"Enable with --llmlingua for ML-based compression (3-5x better on text/logs)"
)
# CCR status
ccr_features = []
if self.config.ccr_inject_tool:
ccr_features.append("tool_injection")
if self.config.ccr_handle_responses:
ccr_features.append("response_handling")
if self.config.ccr_context_tracking:
ccr_features.append("context_tracking")
if self.config.ccr_proactive_expansion:
ccr_features.append("proactive_expansion")
if ccr_features:
logger.info(f"CCR (Compress-Cache-Retrieve): ENABLED ({', '.join(ccr_features)})")
else:
logger.info("CCR: DISABLED")
async def shutdown(self):
"""Cleanup async resources."""
if self.http_client:
@ -1017,6 +1077,65 @@ class HeadroomProxy:
f"[{request_id}] CCR: Tool already present (MCP?), skipped injection for hashes: {injector.detected_hashes}"
)
# Track compression in context tracker for multi-turn awareness
if self.ccr_context_tracker:
self._turn_counter += 1
for hash_key in injector.detected_hashes:
# Get compression metadata from store
store = get_compression_store()
entry = store.get_metadata(hash_key)
if entry:
self.ccr_context_tracker.track_compression(
hash_key=hash_key,
turn_number=self._turn_counter,
tool_name=entry.get("tool_name"),
original_count=entry.get("original_item_count", 0),
compressed_count=entry.get("compressed_item_count", 0),
query_context=entry.get("query_context", ""),
sample_content=entry.get("compressed_content", "")[:500],
)
# CCR Proactive Expansion: Check if current query needs expanded context
if self.ccr_context_tracker and self.config.ccr_proactive_expansion:
# Extract user query from messages
user_query = ""
for msg in reversed(messages):
if msg.get("role") == "user":
content = msg.get("content", "")
if isinstance(content, str):
user_query = content
elif isinstance(content, list):
for block in content:
if isinstance(block, dict) and block.get("type") == "text":
user_query = block.get("text", "")
break
break
if user_query:
recommendations = self.ccr_context_tracker.analyze_query(
user_query, self._turn_counter
)
if recommendations:
expansions = self.ccr_context_tracker.execute_expansions(recommendations)
if expansions:
# Add expanded context to the system message or as additional context
expansion_text = self.ccr_context_tracker.format_expansions_for_context(
expansions
)
logger.info(
f"[{request_id}] CCR: Proactively expanded {len(expansions)} context(s) "
f"based on query relevance"
)
# Append to the last user message
if optimized_messages and optimized_messages[-1].get("role") == "user":
last_msg = optimized_messages[-1]
content = last_msg.get("content", "")
if isinstance(content, str):
optimized_messages[-1] = {
**last_msg,
"content": content + "\n\n" + expansion_text,
}
# Update body
body["messages"] = optimized_messages
if tools is not None:
@ -1043,16 +1162,67 @@ class HeadroomProxy:
)
else:
response = await self._retry_request("POST", url, headers, body)
# Parse response for CCR handling
resp_json = None
try:
resp_json = response.json()
except Exception:
pass
# CCR Response Handling: Handle headroom_retrieve tool calls automatically
if (
self.ccr_response_handler
and resp_json
and response.status_code == 200
and self.ccr_response_handler.has_ccr_tool_calls(resp_json, "anthropic")
):
logger.info(f"[{request_id}] CCR: Detected retrieval tool call, handling...")
# Create API call function for continuation
async def api_call_fn(
msgs: list[dict], tls: list[dict] | None
) -> dict[str, Any]:
continuation_body = {
**body,
"messages": msgs,
}
if tls is not None:
continuation_body["tools"] = tls
cont_response = await self._retry_request(
"POST", url, headers, continuation_body
)
result: dict[str, Any] = cont_response.json()
return result
# Handle CCR tool calls
try:
final_resp_json = await self.ccr_response_handler.handle_response(
resp_json,
optimized_messages,
tools,
api_call_fn,
provider="anthropic",
)
# Update response content with final response
resp_json = final_resp_json
response = httpx.Response(
status_code=200,
content=json.dumps(final_resp_json).encode(),
headers=dict(response.headers),
)
logger.info(f"[{request_id}] CCR: Retrieval handled successfully")
except Exception as e:
logger.warning(f"[{request_id}] CCR: Response handling failed: {e}")
# Continue with original response
total_latency = (time.time() - start_time) * 1000
# Parse response for output tokens
output_tokens = 0
try:
resp_json = response.json()
if resp_json:
usage = resp_json.get("usage", {})
output_tokens = usage.get("output_tokens", 0)
except Exception:
pass
# Calculate cost
cost_usd = None

2
pyproject.toml
View file

@ -4,7 +4,7 @@ build-backend = "hatchling.build"
[project]
name = "headroom-ai"
version = "0.2.1"
version = "0.2.2"
description = "The Context Optimization Layer for LLM Applications - Cut costs by 50-90%"
readme = "README.md"
license = "Apache-2.0"

741
tests/test_ccr_context_tracker.py Normal file
View file

@ -0,0 +1,741 @@
"""Tests for CCR context tracker.
These tests verify that:
1. Compression events are tracked correctly
2. Query analysis finds relevant compressed contexts
3. Proactive expansion works as expected
4. LRU eviction and TTL work correctly
5. Expansion recommendations are appropriate
"""
import json
import time
import pytest
from headroom.cache.compression_store import (
get_compression_store,
reset_compression_store,
)
from headroom.ccr.context_tracker import (
CompressedContext,
ContextTracker,
ContextTrackerConfig,
ExpansionRecommendation,
get_context_tracker,
reset_context_tracker,
)
class TestContextTrackerBasics:
"""Test basic context tracking functionality."""
@pytest.fixture(autouse=True)
def reset_trackers(self):
"""Reset trackers before each test."""
reset_context_tracker()
reset_compression_store()
yield
reset_context_tracker()
reset_compression_store()
def test_track_compression(self):
"""Track a compression event."""
tracker = ContextTracker()
tracker.track_compression(
hash_key="abc123",
turn_number=1,
tool_name="Bash",
original_count=100,
compressed_count=10,
query_context="find all python files",
sample_content='["src/main.py", "src/auth.py"]',
)
assert "abc123" in tracker.get_tracked_hashes()
stats = tracker.get_stats()
assert stats["tracked_contexts"] == 1
def test_track_multiple_compressions(self):
"""Track multiple compression events."""
tracker = ContextTracker()
for i in range(5):
tracker.track_compression(
hash_key=f"hash_{i}",
turn_number=i,
tool_name="Bash",
original_count=100,
compressed_count=10,
)
hashes = tracker.get_tracked_hashes()
assert len(hashes) == 5
assert all(f"hash_{i}" in hashes for i in range(5))
def test_tracking_disabled(self):
"""Tracking disabled doesn't store contexts."""
config = ContextTrackerConfig(enabled=False)
tracker = ContextTracker(config)
tracker.track_compression(
hash_key="abc123",
turn_number=1,
tool_name="Bash",
original_count=100,
compressed_count=10,
)
assert len(tracker.get_tracked_hashes()) == 0
def test_update_existing_hash(self):
"""Updating existing hash updates the context data."""
tracker = ContextTracker()
tracker.track_compression(
hash_key="first", turn_number=1, tool_name=None, original_count=50, compressed_count=5
)
tracker.track_compression(
hash_key="second", turn_number=2, tool_name=None, original_count=50, compressed_count=5
)
tracker.track_compression(
hash_key="first", turn_number=3, tool_name=None, original_count=60, compressed_count=6
)
# Should still have 2 unique hashes
hashes = tracker.get_tracked_hashes()
assert len(hashes) == 2
assert "first" in hashes
assert "second" in hashes
# The updated context should have the new values
stats = tracker.get_stats()
first_ctx = next(c for c in stats["contexts"] if c["hash"] == "first")
assert first_ctx["items"] == "6/60" # Updated values
assert first_ctx["turn"] == 3
class TestLRUEviction:
"""Test LRU eviction at capacity."""
def test_eviction_at_capacity(self):
"""Oldest entries evicted when at capacity."""
config = ContextTrackerConfig(max_tracked_contexts=3)
tracker = ContextTracker(config)
for i in range(5):
tracker.track_compression(
hash_key=f"hash_{i}",
turn_number=i,
tool_name=None,
original_count=100,
compressed_count=10,
)
time.sleep(0.01) # Ensure different timestamps
hashes = tracker.get_tracked_hashes()
assert len(hashes) == 3
# Should have the last 3
assert "hash_2" in hashes
assert "hash_3" in hashes
assert "hash_4" in hashes
# First 2 should be evicted
assert "hash_0" not in hashes
assert "hash_1" not in hashes
class TestQueryAnalysis:
"""Test query analysis for relevance detection."""
@pytest.fixture(autouse=True)
def reset_trackers(self):
"""Reset trackers before each test."""
reset_context_tracker()
reset_compression_store()
yield
reset_context_tracker()
reset_compression_store()
def test_analyze_query_finds_relevant_context(self):
"""Query analysis finds relevant compressed context."""
# Use lower relevance threshold to make test more reliable
config = ContextTrackerConfig(relevance_threshold=0.1)
tracker = ContextTracker(config)
tracker.track_compression(
hash_key="auth_hash",
turn_number=1,
tool_name="Bash",
original_count=100,
compressed_count=10,
query_context="find authentication files",
# Use more explicit content with keywords that will match
sample_content="authentication middleware handler login security",
)
recommendations = tracker.analyze_query(
query="show authentication middleware",
current_turn=2,
)
assert len(recommendations) >= 1
assert recommendations[0].hash_key == "auth_hash"
def test_analyze_query_no_match(self):
"""Query analysis returns empty for unrelated query."""
tracker = ContextTracker()
tracker.track_compression(
hash_key="db_hash",
turn_number=1,
tool_name="Bash",
original_count=100,
compressed_count=10,
query_context="find database files",
sample_content='["database.py", "models.py"]',
)
recommendations = tracker.analyze_query(
query="What is the weather like?",
current_turn=2,
)
# Should not match unrelated query
assert len(recommendations) == 0
def test_analyze_query_keyword_overlap(self):
"""Query matches based on keyword overlap."""
tracker = ContextTracker()
tracker.track_compression(
hash_key="python_files",
turn_number=1,
tool_name="Glob",
original_count=200,
compressed_count=20,
query_context="find python files",
sample_content='["main.py", "utils.py", "config.py", "test_main.py"]',
)
# Query with overlapping keywords
recommendations = tracker.analyze_query(
query="Show me the main python file",
current_turn=2,
)
assert len(recommendations) >= 1
def test_analyze_query_proactive_disabled(self):
"""No recommendations when proactive expansion disabled."""
config = ContextTrackerConfig(proactive_expansion=False)
tracker = ContextTracker(config)
tracker.track_compression(
hash_key="abc123",
turn_number=1,
tool_name="Bash",
original_count=100,
compressed_count=10,
sample_content='["relevant.py"]',
)
recommendations = tracker.analyze_query(
query="Show me relevant files",
current_turn=2,
)
assert len(recommendations) == 0
def test_analyze_query_respects_age(self):
"""Old contexts get lower relevance scores."""
config = ContextTrackerConfig(max_context_age_seconds=2.0)
tracker = ContextTracker(config)
tracker.track_compression(
hash_key="old_context",
turn_number=1,
tool_name="Bash",
original_count=100,
compressed_count=10,
sample_content='["auth.py"]',
)
# Wait for context to age
time.sleep(2.1)
recommendations = tracker.analyze_query(
query="Show me the authentication code",
current_turn=5,
)
# Should not recommend aged-out context
assert len(recommendations) == 0
def test_analyze_query_max_recommendations(self):
"""Respects max proactive expansions limit."""
config = ContextTrackerConfig(max_proactive_expansions=2)
tracker = ContextTracker(config)
# Track many relevant contexts
for i in range(5):
tracker.track_compression(
hash_key=f"hash_{i}",
turn_number=i,
tool_name="Bash",
original_count=100,
compressed_count=10,
sample_content=f'["python_{i}.py", "main.py"]',
)
recommendations = tracker.analyze_query(
query="Show me the python main file",
current_turn=10,
)
assert len(recommendations) <= 2
class TestRelevanceCalculation:
"""Test relevance score calculation."""
def test_extract_keywords(self):
"""Keywords are extracted correctly."""
tracker = ContextTracker()
keywords = tracker._extract_keywords("Find authentication middleware files")
assert "authentication" in keywords
assert "middleware" in keywords
assert "files" in keywords
# Stop words should be filtered
assert "the" not in keywords
def test_exact_substring_match_bonus(self):
"""Exact substring matches get bonus score."""
tracker = ContextTracker()
tracker.track_compression(
hash_key="exact_match",
turn_number=1,
tool_name=None,
original_count=100,
compressed_count=10,
sample_content="authentication_middleware.py, auth_handler.py",
)
# Query with exact substring match
recommendations = tracker.analyze_query(
query="authentication middleware",
current_turn=2,
)
assert len(recommendations) >= 1
# Should have high relevance
assert recommendations[0].relevance_score > 0.3
class TestExpansionTypeDetection:
"""Test determination of expansion type (full vs search)."""
def test_full_expansion_high_relevance(self):
"""High relevance triggers full expansion."""
tracker = ContextTracker()
context = CompressedContext(
hash_key="test",
turn_number=1,
timestamp=time.time(),
tool_name="Bash",
original_item_count=50,
compressed_item_count=5,
query_context="find files",
sample_content="auth.py, middleware.py",
)
expand_full, search_query = tracker._determine_expansion_type(
query="authentication middleware",
context=context,
relevance=0.8, # High relevance
)
assert expand_full is True
assert search_query is None
def test_full_expansion_small_count(self):
"""Small original item count triggers full expansion."""
tracker = ContextTracker()
context = CompressedContext(
hash_key="test",
turn_number=1,
timestamp=time.time(),
tool_name="Bash",
original_item_count=30, # Small
compressed_item_count=5,
query_context="find files",
sample_content="file.py",
)
expand_full, search_query = tracker._determine_expansion_type(
query="some query",
context=context,
relevance=0.4,
)
assert expand_full is True
def test_search_expansion_large_count(self):
"""Large original count with specific keywords triggers search."""
tracker = ContextTracker()
context = CompressedContext(
hash_key="test",
turn_number=1,
timestamp=time.time(),
tool_name="Bash",
original_item_count=500, # Large
compressed_item_count=20,
query_context="find all files",
sample_content="many files...",
)
expand_full, search_query = tracker._determine_expansion_type(
query="find authentication middleware handler",
context=context,
relevance=0.4, # Medium relevance
)
# Should use search for large datasets
if not expand_full:
assert search_query is not None
assert "authentication" in search_query or "middleware" in search_query
class TestExpansionExecution:
"""Test execution of expansion recommendations."""
@pytest.fixture(autouse=True)
def reset_stores(self):
"""Reset stores before each test."""
reset_context_tracker()
reset_compression_store()
yield
reset_context_tracker()
reset_compression_store()
def test_execute_full_expansion(self):
"""Execute full expansion retrieval."""
store = get_compression_store()
original = json.dumps([{"id": i} for i in range(100)])
hash_key = store.store(
original=original,
compressed="[]",
original_item_count=100,
)
tracker = ContextTracker()
recommendations = [
ExpansionRecommendation(
hash_key=hash_key,
reason="relevant to query",
relevance_score=0.8,
expand_full=True,
)
]
results = tracker.execute_expansions(recommendations)
assert len(results) == 1
assert results[0]["type"] == "full"
assert results[0]["item_count"] == 100
def test_execute_search_expansion(self):
"""Execute search expansion."""
store = get_compression_store()
items = [
{"id": 1, "content": "authentication code"},
{"id": 2, "content": "database operations"},
{"id": 3, "content": "authentication middleware"},
]
original = json.dumps(items)
hash_key = store.store(
original=original,
compressed="[]",
original_item_count=3,
)
tracker = ContextTracker()
recommendations = [
ExpansionRecommendation(
hash_key=hash_key,
reason="relevant to query",
relevance_score=0.5,
expand_full=False,
search_query="authentication",
)
]
results = tracker.execute_expansions(recommendations)
assert len(results) == 1
assert results[0]["type"] == "search"
assert results[0]["query"] == "authentication"
def test_execute_nonexistent_hash(self):
"""Handle expansion of nonexistent hash gracefully."""
tracker = ContextTracker()
recommendations = [
ExpansionRecommendation(
hash_key="nonexistent123",
reason="test",
relevance_score=0.5,
expand_full=True,
)
]
results = tracker.execute_expansions(recommendations)
# Should handle gracefully, no results
assert len(results) == 0
class TestExpansionFormatting:
"""Test formatting of expansions for context."""
def test_format_full_expansion(self):
"""Format full expansion for LLM context."""
tracker = ContextTracker()
expansions = [
{
"hash": "abc123",
"type": "full",
"content": '[{"id": 1}, {"id": 2}]',
"item_count": 2,
"reason": "relevant to query",
}
]
formatted = tracker.format_expansions_for_context(expansions)
assert "[Proactive Context Expansion" in formatted
assert "Expanded from earlier" in formatted
assert '[{"id": 1}, {"id": 2}]' in formatted
def test_format_search_expansion(self):
"""Format search expansion for LLM context."""
tracker = ContextTracker()
expansions = [
{
"hash": "def456",
"type": "search",
"query": "authentication",
"content": [{"id": 1, "content": "auth"}],
"item_count": 1,
"reason": "matched query",
}
]
formatted = tracker.format_expansions_for_context(expansions)
assert "Search results for 'authentication'" in formatted
def test_format_empty_expansions(self):
"""Empty expansions return empty string."""
tracker = ContextTracker()
formatted = tracker.format_expansions_for_context([])
assert formatted == ""
class TestGlobalTracker:
"""Test global tracker singleton."""
@pytest.fixture(autouse=True)
def reset_tracker(self):
"""Reset global tracker."""
reset_context_tracker()
yield
reset_context_tracker()
def test_singleton_pattern(self):
"""Global tracker uses singleton pattern."""
tracker1 = get_context_tracker()
tracker2 = get_context_tracker()
assert tracker1 is tracker2
def test_reset_clears_tracker(self):
"""Reset clears the global tracker."""
tracker = get_context_tracker()
tracker.track_compression(
hash_key="test",
turn_number=1,
tool_name=None,
original_count=10,
compressed_count=1,
)
assert len(tracker.get_tracked_hashes()) == 1
reset_context_tracker()
new_tracker = get_context_tracker()
assert len(new_tracker.get_tracked_hashes()) == 0
class TestContextTrackerConfig:
"""Test context tracker configuration."""
def test_default_config(self):
"""Default config values."""
config = ContextTrackerConfig()
assert config.enabled is True
assert config.max_tracked_contexts == 100
assert config.relevance_threshold == 0.3
assert config.max_context_age_seconds == 300.0
assert config.proactive_expansion is True
assert config.max_proactive_expansions == 2
def test_custom_config(self):
"""Custom config values."""
config = ContextTrackerConfig(
enabled=False,
max_tracked_contexts=50,
relevance_threshold=0.5,
max_proactive_expansions=5,
)
assert config.enabled is False
assert config.max_tracked_contexts == 50
assert config.relevance_threshold == 0.5
assert config.max_proactive_expansions == 5
class TestCompressedContextDataClass:
"""Test CompressedContext dataclass."""
def test_create_context(self):
"""Create compressed context."""
context = CompressedContext(
hash_key="abc123",
turn_number=5,
timestamp=1234567890.0,
tool_name="Bash",
original_item_count=100,
compressed_item_count=10,
query_context="find files",
sample_content='["file1.py", "file2.py"]',
)
assert context.hash_key == "abc123"
assert context.turn_number == 5
assert context.tool_name == "Bash"
assert context.original_item_count == 100
assert context.compressed_item_count == 10
class TestExpansionRecommendationDataClass:
"""Test ExpansionRecommendation dataclass."""
def test_full_expansion_recommendation(self):
"""Create full expansion recommendation."""
rec = ExpansionRecommendation(
hash_key="abc123",
reason="high relevance",
relevance_score=0.9,
expand_full=True,
)
assert rec.expand_full is True
assert rec.search_query is None
def test_search_expansion_recommendation(self):
"""Create search expansion recommendation."""
rec = ExpansionRecommendation(
hash_key="def456",
reason="partial match",
relevance_score=0.5,
expand_full=False,
search_query="authentication",
)
assert rec.expand_full is False
assert rec.search_query == "authentication"
class TestContextTrackerStats:
"""Test tracker statistics."""
def test_stats_structure(self):
"""Stats have expected structure."""
tracker = ContextTracker()
tracker.track_compression(
hash_key="test",
turn_number=1,
tool_name="Bash",
original_count=100,
compressed_count=10,
)
stats = tracker.get_stats()
assert "tracked_contexts" in stats
assert "current_turn" in stats
assert "config" in stats
assert "contexts" in stats
assert stats["tracked_contexts"] == 1
assert len(stats["contexts"]) == 1
def test_stats_context_details(self):
"""Stats include context details."""
tracker = ContextTracker()
tracker.track_compression(
hash_key="abc123",
turn_number=3,
tool_name="Glob",
original_count=50,
compressed_count=5,
)
stats = tracker.get_stats()
context_stat = stats["contexts"][0]
assert context_stat["hash"] == "abc123"
assert context_stat["turn"] == 3
assert context_stat["tool"] == "Glob"
assert context_stat["items"] == "5/50"
class TestTrackerClear:
"""Test tracker clear functionality."""
def test_clear_removes_all(self):
"""Clear removes all tracked contexts."""
tracker = ContextTracker()
for i in range(5):
tracker.track_compression(
hash_key=f"hash_{i}",
turn_number=i,
tool_name=None,
original_count=10,
compressed_count=1,
)
assert len(tracker.get_tracked_hashes()) == 5
tracker.clear()
assert len(tracker.get_tracked_hashes()) == 0
stats = tracker.get_stats()
assert stats["current_turn"] == 0

701
tests/test_ccr_response_handler.py Normal file
View file

@ -0,0 +1,701 @@
"""Tests for CCR response handler.
These tests verify that:
1. CCR tool calls are correctly detected in responses
2. Retrieval execution works for both full and search modes
3. Continuation flow handles multiple rounds
4. Provider-specific formats are handled correctly
5. Streaming buffer detection works
"""
import json
import pytest
from headroom.cache.compression_store import (
get_compression_store,
reset_compression_store,
)
from headroom.ccr.response_handler import (
CCRResponseHandler,
CCRToolCall,
CCRToolResult,
ResponseHandlerConfig,
StreamingCCRBuffer,
)
from headroom.ccr.tool_injection import CCR_TOOL_NAME
class TestCCRToolCallDetection:
"""Test detection of CCR tool calls in responses."""
@pytest.fixture(autouse=True)
def reset_store(self):
"""Reset global store before each test."""
reset_compression_store()
yield
reset_compression_store()
def test_detect_anthropic_ccr_tool_call(self):
"""Detect CCR tool call in Anthropic format."""
handler = CCRResponseHandler()
response = {
"content": [
{"type": "text", "text": "Let me retrieve that data."},
{
"type": "tool_use",
"id": "tool_123",
"name": CCR_TOOL_NAME,
"input": {"hash": "abc123"},
},
]
}
assert handler.has_ccr_tool_calls(response, "anthropic")
def test_detect_openai_ccr_tool_call(self):
"""Detect CCR tool call in OpenAI format."""
handler = CCRResponseHandler()
response = {
"choices": [
{
"message": {
"role": "assistant",
"content": "Let me retrieve that data.",
"tool_calls": [
{
"id": "call_123",
"type": "function",
"function": {
"name": CCR_TOOL_NAME,
"arguments": '{"hash": "abc123"}',
},
}
],
}
}
]
}
assert handler.has_ccr_tool_calls(response, "openai")
def test_no_ccr_tool_call_anthropic(self):
"""No false positive when no CCR tool call present."""
handler = CCRResponseHandler()
response = {
"content": [
{"type": "text", "text": "Here is the data."},
{
"type": "tool_use",
"id": "tool_123",
"name": "some_other_tool",
"input": {"param": "value"},
},
]
}
assert not handler.has_ccr_tool_calls(response, "anthropic")
def test_no_ccr_tool_call_openai(self):
"""No false positive when no CCR tool call present in OpenAI format."""
handler = CCRResponseHandler()
response = {
"choices": [
{
"message": {
"role": "assistant",
"content": "Here is the data.",
"tool_calls": [
{
"id": "call_123",
"type": "function",
"function": {
"name": "other_tool",
"arguments": '{"param": "value"}',
},
}
],
}
}
]
}
assert not handler.has_ccr_tool_calls(response, "openai")
def test_text_only_response(self):
"""No false positive for text-only responses."""
handler = CCRResponseHandler()
response = {"content": [{"type": "text", "text": "Just plain text."}]}
assert not handler.has_ccr_tool_calls(response, "anthropic")
def test_empty_response(self):
"""Handle empty response gracefully."""
handler = CCRResponseHandler()
assert not handler.has_ccr_tool_calls({}, "anthropic")
assert not handler.has_ccr_tool_calls({"content": []}, "anthropic")
class TestCCRToolCallParsing:
"""Test parsing of CCR tool calls."""
def test_parse_anthropic_full_retrieval(self):
"""Parse full retrieval call from Anthropic format."""
handler = CCRResponseHandler()
response = {
"content": [
{
"type": "tool_use",
"id": "tool_123",
"name": CCR_TOOL_NAME,
"input": {"hash": "abc123"},
}
]
}
ccr_calls, other_calls = handler._parse_ccr_tool_calls(response, "anthropic")
assert len(ccr_calls) == 1
assert ccr_calls[0].tool_call_id == "tool_123"
assert ccr_calls[0].hash_key == "abc123"
assert ccr_calls[0].query is None
assert len(other_calls) == 0
def test_parse_anthropic_search_retrieval(self):
"""Parse search retrieval call from Anthropic format."""
handler = CCRResponseHandler()
response = {
"content": [
{
"type": "tool_use",
"id": "tool_456",
"name": CCR_TOOL_NAME,
"input": {"hash": "def456", "query": "authentication error"},
}
]
}
ccr_calls, other_calls = handler._parse_ccr_tool_calls(response, "anthropic")
assert len(ccr_calls) == 1
assert ccr_calls[0].hash_key == "def456"
assert ccr_calls[0].query == "authentication error"
def test_parse_mixed_tool_calls(self):
"""Parse response with both CCR and other tool calls."""
handler = CCRResponseHandler()
response = {
"content": [
{
"type": "tool_use",
"id": "tool_1",
"name": CCR_TOOL_NAME,
"input": {"hash": "abc123"},
},
{
"type": "tool_use",
"id": "tool_2",
"name": "read_file",
"input": {"path": "/etc/config"},
},
]
}
ccr_calls, other_calls = handler._parse_ccr_tool_calls(response, "anthropic")
assert len(ccr_calls) == 1
assert len(other_calls) == 1
assert other_calls[0]["name"] == "read_file"
class TestCCRRetrievalExecution:
"""Test CCR retrieval execution."""
@pytest.fixture(autouse=True)
def reset_store(self):
"""Reset global store before each test."""
reset_compression_store()
yield
reset_compression_store()
def test_full_retrieval_success(self):
"""Successfully retrieve full content."""
store = get_compression_store()
original = json.dumps([{"id": i} for i in range(100)])
compressed = json.dumps([{"id": i} for i in range(10)])
hash_key = store.store(
original=original,
compressed=compressed,
original_item_count=100,
compressed_item_count=10,
)
handler = CCRResponseHandler()
call = CCRToolCall(tool_call_id="test_id", hash_key=hash_key)
result = handler._execute_retrieval(call)
assert result.success
assert result.items_retrieved == 100
assert not result.was_search
# Check content structure
content = json.loads(result.content)
assert content["hash"] == hash_key
assert "original_content" in content
def test_search_retrieval_success(self):
"""Successfully search within cached content."""
store = get_compression_store()
# Use items with more searchable content
items = [
{"id": 1, "text": "Python programming language tutorial"},
{"id": 2, "text": "JavaScript web development framework"},
{"id": 3, "text": "Python data science machine learning"},
{"id": 4, "text": "Ruby programming language basics"},
{"id": 5, "text": "Python web framework django flask"},
]
original = json.dumps(items)
compressed = json.dumps(items[:1])
hash_key = store.store(
original=original,
compressed=compressed,
original_item_count=5,
compressed_item_count=1,
)
handler = CCRResponseHandler()
# Use a more specific query
call = CCRToolCall(tool_call_id="test_id", hash_key=hash_key, query="Python programming")
result = handler._execute_retrieval(call)
assert result.success
assert result.was_search
content = json.loads(result.content)
assert content["query"] == "Python programming"
# The search should return results (may be 0 depending on BM25 behavior)
assert "results" in content
def test_retrieval_nonexistent_hash(self):
"""Handle retrieval of nonexistent hash."""
handler = CCRResponseHandler()
call = CCRToolCall(tool_call_id="test_id", hash_key="nonexistent123")
result = handler._execute_retrieval(call)
assert not result.success
assert result.items_retrieved == 0
content = json.loads(result.content)
assert "error" in content
class TestCCRToolResultMessage:
"""Test tool result message creation."""
def test_anthropic_tool_result_format(self):
"""Create tool result message in Anthropic format."""
handler = CCRResponseHandler()
results = [
CCRToolResult(
tool_call_id="tool_123",
content='{"data": "retrieved"}',
success=True,
items_retrieved=10,
)
]
message = handler._create_tool_result_message(results, "anthropic")
assert message["role"] == "user"
assert len(message["content"]) == 1
assert message["content"][0]["type"] == "tool_result"
assert message["content"][0]["tool_use_id"] == "tool_123"
def test_openai_tool_result_format(self):
"""Create tool result messages in OpenAI format."""
handler = CCRResponseHandler()
results = [
CCRToolResult(
tool_call_id="call_123",
content='{"data": "retrieved"}',
success=True,
),
CCRToolResult(
tool_call_id="call_456",
content='{"data": "more data"}',
success=True,
),
]
message = handler._create_tool_result_message(results, "openai")
assert "_openai_tool_results" in message
assert len(message["_openai_tool_results"]) == 2
assert message["_openai_tool_results"][0]["role"] == "tool"
class TestCCRResponseHandling:
"""Test the full response handling flow."""
@pytest.fixture(autouse=True)
def reset_store(self):
"""Reset global store before each test."""
reset_compression_store()
yield
reset_compression_store()
@pytest.mark.asyncio
async def test_handle_response_no_ccr(self):
"""Handle response with no CCR calls (pass-through)."""
handler = CCRResponseHandler()
response = {"content": [{"type": "text", "text": "Just text."}]}
async def mock_api_call(messages, tools):
return {"content": [{"type": "text", "text": "Response"}]}
result = await handler.handle_response(response, [], None, mock_api_call, "anthropic")
# Should return original response unchanged
assert result == response
@pytest.mark.asyncio
async def test_handle_response_with_ccr(self):
"""Handle response containing CCR tool call."""
store = get_compression_store()
original = json.dumps([{"id": i} for i in range(50)])
hash_key = store.store(
original=original,
compressed="[]",
original_item_count=50,
)
handler = CCRResponseHandler()
# Initial response with CCR tool call
initial_response = {
"content": [
{"type": "text", "text": "Let me get that data."},
{
"type": "tool_use",
"id": "tool_123",
"name": CCR_TOOL_NAME,
"input": {"hash": hash_key},
},
]
}
# Final response after tool result
final_response = {"content": [{"type": "text", "text": "Here is all 50 items of data."}]}
call_count = 0
async def mock_api_call(messages, tools):
nonlocal call_count
call_count += 1
return final_response
result = await handler.handle_response(
initial_response,
[{"role": "user", "content": "Get me the data"}],
None,
mock_api_call,
"anthropic",
)
# Should have made continuation call
assert call_count == 1
# Should return final response
assert result == final_response
@pytest.mark.asyncio
async def test_handle_response_max_rounds(self):
"""Respects max retrieval rounds limit."""
store = get_compression_store()
hash_key = store.store(original="[1,2,3]", compressed="[]")
config = ResponseHandlerConfig(max_retrieval_rounds=2)
handler = CCRResponseHandler(config)
# Response that always has CCR tool call (simulating infinite loop)
ccr_response = {
"content": [
{
"type": "tool_use",
"id": "tool_123",
"name": CCR_TOOL_NAME,
"input": {"hash": hash_key},
}
]
}
call_count = 0
async def mock_api_call(messages, tools):
nonlocal call_count
call_count += 1
return ccr_response
await handler.handle_response(ccr_response, [], None, mock_api_call, "anthropic")
# Should stop after max rounds
assert call_count == 2
@pytest.mark.asyncio
async def test_handle_response_disabled(self):
"""Disabled handler returns response unchanged."""
config = ResponseHandlerConfig(enabled=False)
handler = CCRResponseHandler(config)
response = {
"content": [
{
"type": "tool_use",
"id": "tool_123",
"name": CCR_TOOL_NAME,
"input": {"hash": "abc123"},
}
]
}
async def mock_api_call(messages, tools):
raise AssertionError("Should not be called")
result = await handler.handle_response(response, [], None, mock_api_call, "anthropic")
assert result == response
class TestCCRResponseHandlerStats:
"""Test handler statistics."""
@pytest.fixture(autouse=True)
def reset_store(self):
"""Reset global store before each test."""
reset_compression_store()
yield
reset_compression_store()
@pytest.mark.asyncio
async def test_retrieval_count_tracking(self):
"""Track total retrieval count."""
store = get_compression_store()
hash_key = store.store(original="[1,2,3]", compressed="[]")
handler = CCRResponseHandler()
initial_response = {
"content": [
{
"type": "tool_use",
"id": "tool_123",
"name": CCR_TOOL_NAME,
"input": {"hash": hash_key},
}
]
}
final_response = {"content": [{"type": "text", "text": "Done"}]}
async def mock_api_call(messages, tools):
return final_response
await handler.handle_response(initial_response, [], None, mock_api_call, "anthropic")
stats = handler.get_stats()
assert stats["total_retrievals"] == 1
class TestStreamingCCRBuffer:
"""Test streaming buffer for CCR detection."""
def test_buffer_accumulation(self):
"""Buffer accumulates chunks."""
buffer = StreamingCCRBuffer()
buffer.add_chunk(b"part1")
buffer.add_chunk(b"part2")
buffer.add_chunk(b"part3")
assert buffer.get_accumulated() == b"part1part2part3"
def test_detect_ccr_tool_in_stream(self):
"""Detect CCR tool call in streaming chunks."""
buffer = StreamingCCRBuffer()
# Simulate streaming response with tool_use
chunk1 = b'{"type":"content_block_start","content_block":{"type":"tool_use"'
chunk2 = f',"name":"{CCR_TOOL_NAME}"'.encode()
detected = buffer.add_chunk(chunk1)
assert not detected # Not complete yet
detected = buffer.add_chunk(chunk2)
assert detected # Now detected
assert buffer.detected_ccr
def test_no_false_positive_detection(self):
"""No false positive for non-CCR tool calls."""
buffer = StreamingCCRBuffer()
chunk = b'{"type":"content_block_start","content_block":{"type":"tool_use","name":"other_tool"}}'
detected = buffer.add_chunk(chunk)
assert not detected
assert not buffer.detected_ccr
def test_buffer_clear(self):
"""Buffer clears state correctly."""
buffer = StreamingCCRBuffer()
buffer.add_chunk(b"data")
buffer.detected_ccr = True
buffer.clear()
assert buffer.get_accumulated() == b""
assert not buffer.detected_ccr
class TestResponseHandlerConfig:
"""Test response handler configuration."""
def test_default_config(self):
"""Default config values."""
config = ResponseHandlerConfig()
assert config.enabled is True
assert config.max_retrieval_rounds == 3
assert config.strip_ccr_from_response is True
assert config.continuation_timeout_ms == 120000
def test_custom_config(self):
"""Custom config values."""
config = ResponseHandlerConfig(
enabled=False,
max_retrieval_rounds=5,
)
assert config.enabled is False
assert config.max_retrieval_rounds == 5
class TestCCRToolCallDataClass:
"""Test CCRToolCall dataclass."""
def test_full_retrieval_call(self):
"""Create full retrieval call."""
call = CCRToolCall(
tool_call_id="test_123",
hash_key="abc123",
)
assert call.tool_call_id == "test_123"
assert call.hash_key == "abc123"
assert call.query is None
def test_search_retrieval_call(self):
"""Create search retrieval call."""
call = CCRToolCall(
tool_call_id="test_456",
hash_key="def456",
query="authentication",
)
assert call.query == "authentication"
class TestCCRToolResultDataClass:
"""Test CCRToolResult dataclass."""
def test_successful_result(self):
"""Create successful result."""
result = CCRToolResult(
tool_call_id="test_123",
content='{"data": "content"}',
success=True,
items_retrieved=50,
was_search=False,
)
assert result.success
assert result.items_retrieved == 50
assert not result.was_search
def test_search_result(self):
"""Create search result."""
result = CCRToolResult(
tool_call_id="test_456",
content='{"results": []}',
success=True,
items_retrieved=5,
was_search=True,
)
assert result.was_search
def test_failed_result(self):
"""Create failed result."""
result = CCRToolResult(
tool_call_id="test_789",
content='{"error": "not found"}',
success=False,
)
assert not result.success
assert result.items_retrieved == 0
class TestExtractAssistantMessage:
"""Test extraction of assistant messages from responses."""
def test_extract_anthropic_message(self):
"""Extract assistant message from Anthropic response."""
handler = CCRResponseHandler()
response = {
"content": [
{"type": "text", "text": "Hello"},
{"type": "tool_use", "id": "123", "name": "test", "input": {}},
]
}
message = handler._extract_assistant_message(response, "anthropic")
assert message["role"] == "assistant"
assert message["content"] == response["content"]
def test_extract_openai_message(self):
"""Extract assistant message from OpenAI response."""
handler = CCRResponseHandler()
response = {
"choices": [
{
"message": {
"role": "assistant",
"content": "Hello",
"tool_calls": [{"id": "123"}],
}
}
]
}
message = handler._extract_assistant_message(response, "openai")
assert message["role"] == "assistant"
assert message["content"] == "Hello"
assert message["tool_calls"] == [{"id": "123"}]