diff --git a/README.md b/README.md index 5595b0530..c04b907f8 100644 --- a/README.md +++ b/README.md @@ -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). --- diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index a4b01f6b5..7a5f92972 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -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 diff --git a/docs/README.md b/docs/README.md index 68aed4ff7..be80126b3 100644 --- a/docs/README.md +++ b/docs/README.md @@ -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 diff --git a/docs/ccr.md b/docs/ccr.md new file mode 100644 index 000000000..2753d3c93 --- /dev/null +++ b/docs/ccr.md @@ -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). diff --git a/docs/configuration.md b/docs/configuration.md new file mode 100644 index 000000000..f6da3c09e --- /dev/null +++ b/docs/configuration.md @@ -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}") +``` diff --git a/docs/errors.md b/docs/errors.md new file mode 100644 index 000000000..345e593a5 --- /dev/null +++ b/docs/errors.md @@ -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}") +``` diff --git a/docs/llmlingua.md b/docs/llmlingua.md new file mode 100644 index 000000000..be9cee960 --- /dev/null +++ b/docs/llmlingua.md @@ -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) diff --git a/docs/metrics.md b/docs/metrics.md new file mode 100644 index 000000000..5c41571cd --- /dev/null +++ b/docs/metrics.md @@ -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 | diff --git a/docs/sdk.md b/docs/sdk.md new file mode 100644 index 000000000..81fe5af90 --- /dev/null +++ b/docs/sdk.md @@ -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. diff --git a/docs/text-compression.md b/docs/text-compression.md new file mode 100644 index 000000000..8fd4d743a --- /dev/null +++ b/docs/text-compression.md @@ -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 | diff --git a/examples/ccr_demo.py b/examples/ccr_demo.py new file mode 100644 index 000000000..9ac6ae2e7 --- /dev/null +++ b/examples/ccr_demo.py @@ -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()) diff --git a/headroom/cache/compression_store.py b/headroom/cache/compression_store.py index ab767c206..5515d0bb7 100644 --- a/headroom/cache/compression_store.py +++ b/headroom/cache/compression_store.py @@ -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, diff --git a/headroom/ccr/__init__.py b/headroom/ccr/__init__.py index c3ae014a9..fe612e1b0 100644 --- a/headroom/ccr/__init__.py +++ b/headroom/ccr/__init__.py @@ -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", diff --git a/headroom/ccr/context_tracker.py b/headroom/ccr/context_tracker.py new file mode 100644 index 000000000..596b80ae3 --- /dev/null +++ b/headroom/ccr/context_tracker.py @@ -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 diff --git a/headroom/ccr/response_handler.py b/headroom/ccr/response_handler.py new file mode 100644 index 000000000..9589e706d --- /dev/null +++ b/headroom/ccr/response_handler.py @@ -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" diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index bd73209be..5bc739516 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -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 diff --git a/pyproject.toml b/pyproject.toml index fd98ce0b0..3773340f2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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" diff --git a/tests/test_ccr_context_tracker.py b/tests/test_ccr_context_tracker.py new file mode 100644 index 000000000..c517d0acd --- /dev/null +++ b/tests/test_ccr_context_tracker.py @@ -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 diff --git a/tests/test_ccr_response_handler.py b/tests/test_ccr_response_handler.py new file mode 100644 index 000000000..de1d96f8f --- /dev/null +++ b/tests/test_ccr_response_handler.py @@ -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"}]