#!/usr/bin/env python3 """ Quality Retention Evaluation for Intelligent Context Management This eval verifies that when we intelligently drop/compress content, we RETAIN critical information that the model needs to answer correctly. Methodology: 1. NEEDLE-IN-HAYSTACK: Embed specific facts in large tool outputs 2. COMPRESS: Apply IntelligentContextManager 3. VERIFY: Ask questions requiring those facts 4. SCORE: Compare answers before/after compression Key Metrics: - Retention Rate: % of critical facts preserved - Answer Accuracy: % of verification questions answered correctly - Quality-Adjusted Savings: compression_ratio * retention_rate Usage: ANTHROPIC_API_KEY=sk-... python examples/quality_retention_eval.py """ import json import os import sys from dataclasses import dataclass, field from typing import Any # Check for API key early API_KEY = os.environ.get("ANTHROPIC_API_KEY") if not API_KEY: print("ERROR: ANTHROPIC_API_KEY environment variable required") sys.exit(1) from anthropic import Anthropic # noqa: E402 from headroom import AnthropicProvider, HeadroomClient # noqa: E402 from headroom.config import HeadroomConfig, IntelligentContextConfig # noqa: E402 from headroom.tokenizer import Tokenizer # noqa: E402 from headroom.tokenizers import TiktokenCounter # noqa: E402 # ============================================================================= # TEST CASE DEFINITIONS # ============================================================================= @dataclass class CriticalFact: """A fact that MUST be retained after compression.""" description: str value: str verification_question: str expected_answer_contains: list[str] # Answer should contain these strings @dataclass class EvalTestCase: """A test case for quality retention evaluation.""" name: str description: str messages: list[dict[str, Any]] critical_facts: list[CriticalFact] system_prompt: str = "You are a helpful assistant. Answer questions accurately based on the information provided." @dataclass class EvalResult: """Result of evaluating a single test case.""" test_name: str tokens_before: int tokens_after: int compression_ratio: float facts_tested: int facts_retained: int retention_rate: float quality_adjusted_savings: float details: list[dict[str, Any]] = field(default_factory=list) error: str | None = None # ============================================================================= # TEST CASE GENERATORS # ============================================================================= def create_search_results_with_needles() -> EvalTestCase: """ Test case: 100 search results with 3 critical "needles" hidden inside. The critical facts are specific error reports that should be retained even when most results are compressed away. """ results = [] # Needle 1: Critical memory error at position 42 # Needle 2: Database failure at position 17 # Needle 3: Security incident at position 73 for i in range(100): if i == 42: result = { "id": f"INC-{i:04d}", "title": "CRITICAL: OutOfMemory in production worker pool", "severity": "P0", "timestamp": "2024-01-15T10:42:37Z", "server": "prod-worker-7", "details": { "error_type": "java.lang.OutOfMemoryError", "heap_used_gb": 7.8, "heap_max_gb": 8.0, "thread_count": 847, "affected_users": 12847, }, "root_cause": "Thread pool executor not releasing completed tasks", "resolution": "Restart required, fix deployed in v2.4.1", } elif i == 17: result = { "id": f"INC-{i:04d}", "title": "Database connection pool exhausted", "severity": "P1", "timestamp": "2024-01-15T09:17:22Z", "server": "db-primary-3", "details": { "error_type": "ConnectionPoolExhausted", "active_connections": 500, "max_connections": 500, "waiting_queries": 2341, "longest_wait_ms": 45000, }, "root_cause": "Slow query from analytics job holding connections", "resolution": "Killed analytics query, added connection timeout", } elif i == 73: result = { "id": f"INC-{i:04d}", "title": "SECURITY: Unauthorized API access attempt blocked", "severity": "P0", "timestamp": "2024-01-15T14:23:55Z", "source_ip": "203.0.113.42", "details": { "attack_type": "credential_stuffing", "attempts": 15847, "accounts_targeted": 892, "accounts_compromised": 0, "blocked_by": "rate_limiter_v2", }, "root_cause": "Stolen credentials from third-party breach", "resolution": "IP blocked, affected users notified to reset passwords", } else: result = { "id": f"INC-{i:04d}", "title": f"Routine alert #{i}: {'CPU spike' if i % 3 == 0 else 'Latency increase' if i % 3 == 1 else 'Disk usage warning'}", "severity": "P3", "timestamp": f"2024-01-15T{10 + (i % 8):02d}:{i % 60:02d}:00Z", "server": f"app-server-{i % 20}", "details": { "metric": "cpu_percent" if i % 3 == 0 else "latency_p99" if i % 3 == 1 else "disk_usage", "value": 75 + (i % 20), "threshold": 80, "duration_minutes": 5 + (i % 10), }, "root_cause": "Normal traffic variation", "resolution": "Auto-resolved", } results.append(result) messages = [ {"role": "user", "content": "Search for all incidents from today"}, { "role": "assistant", "content": [ { "type": "tool_use", "id": "call_search", "name": "search_incidents", "input": {"date": "2024-01-15"}, } ], }, { "role": "user", "content": [ { "type": "tool_result", "tool_use_id": "call_search", "content": json.dumps(results, indent=2), } ], }, { "role": "user", "content": "What were the critical issues today? I need specific details.", }, ] critical_facts = [ CriticalFact( description="Memory error server identification", value="prod-worker-7", verification_question="Which server experienced the OutOfMemory error?", expected_answer_contains=["prod-worker-7"], ), CriticalFact( description="Memory error affected users count", value="12847", verification_question="How many users were affected by the memory error?", expected_answer_contains=["12847", "12,847"], ), CriticalFact( description="Database error connection count", value="500 active connections", verification_question="How many active database connections were there when the pool was exhausted?", expected_answer_contains=["500"], ), CriticalFact( description="Security attack source IP", value="203.0.113.42", verification_question="What IP address was the source of the credential stuffing attack?", expected_answer_contains=["203.0.113.42"], ), CriticalFact( description="Security attack attempt count", value="15847 attempts", verification_question="How many credential stuffing attempts were made?", expected_answer_contains=["15847", "15,847"], ), ] return EvalTestCase( name="search_results_needles", description="100 search results with 3 critical incidents (needles) at positions 17, 42, 73", messages=messages, critical_facts=critical_facts, ) def create_log_analysis_with_needles() -> EvalTestCase: """ Test case: 200 log entries with critical error buried in the middle. """ logs = [] for i in range(200): if i == 127: # The critical needle - a specific error with unique identifiers log = { "timestamp": "2024-01-15T11:27:33.847Z", "level": "ERROR", "service": "payment-gateway", "trace_id": "abc123def456", "message": "Payment processing failed: Card declined", "details": { "transaction_id": "TXN-98765432", "amount_cents": 15999, "currency": "USD", "error_code": "CARD_DECLINED_INSUFFICIENT_FUNDS", "customer_id": "CUST-789012", "retry_count": 3, "final_status": "FAILED", }, } elif i == 45: # Another needle - rate limit hit log = { "timestamp": "2024-01-15T10:45:12.123Z", "level": "WARN", "service": "api-gateway", "message": "Rate limit exceeded for client", "details": { "client_id": "CLIENT-ACME-001", "endpoint": "/api/v2/bulk-upload", "requests_per_minute": 1500, "limit": 1000, "blocked_duration_seconds": 300, }, } else: log = { "timestamp": f"2024-01-15T{10 + (i % 4):02d}:{i % 60:02d}:{i % 60:02d}.{i % 1000:03d}Z", "level": "INFO", "service": ["api", "auth", "worker", "cache", "db"][i % 5], "message": f"Request processed successfully (id={i})", "details": {"latency_ms": 50 + (i % 100), "status": 200}, } logs.append(log) messages = [ {"role": "user", "content": "Get the logs from the last hour"}, { "role": "assistant", "content": [ { "type": "tool_use", "id": "call_logs", "name": "get_logs", "input": {"timerange": "1h"}, } ], }, { "role": "user", "content": [ { "type": "tool_result", "tool_use_id": "call_logs", "content": json.dumps(logs, indent=2), } ], }, {"role": "user", "content": "Are there any errors or warnings I should know about?"}, ] critical_facts = [ CriticalFact( description="Failed transaction ID", value="TXN-98765432", verification_question="What was the transaction ID of the failed payment?", expected_answer_contains=["TXN-98765432"], ), CriticalFact( description="Payment failure amount", value="$159.99 (15999 cents)", verification_question="What was the amount of the failed payment in dollars?", expected_answer_contains=["159.99", "159"], ), CriticalFact( description="Rate limited client", value="CLIENT-ACME-001", verification_question="Which client was rate limited?", expected_answer_contains=["CLIENT-ACME-001", "ACME"], ), CriticalFact( description="Rate limit endpoint", value="/api/v2/bulk-upload", verification_question="Which API endpoint triggered the rate limit?", expected_answer_contains=["bulk-upload", "/api/v2/bulk-upload"], ), ] return EvalTestCase( name="log_analysis_needles", description="200 log entries with critical error at position 127 and warning at position 45", messages=messages, critical_facts=critical_facts, ) def create_code_review_with_needles() -> EvalTestCase: """ Test case: Code review results with security vulnerabilities hidden in large output. """ files = [] for i in range(80): if i == 23: # Critical security vulnerability file_result = { "path": "src/auth/login.py", "issues": [ { "line": 47, "severity": "CRITICAL", "type": "SQL_INJECTION", "message": "User input directly concatenated into SQL query", "code_snippet": "query = f\"SELECT * FROM users WHERE username = '{username}'\"", "fix": "Use parameterized queries: cursor.execute('SELECT * FROM users WHERE username = ?', (username,))", }, { "line": 89, "severity": "HIGH", "type": "HARDCODED_SECRET", "message": "API key hardcoded in source code", "code_snippet": 'API_KEY = "sk-prod-a1b2c3d4e5f6g7h8i9j0"', "fix": "Use environment variables or secret management", }, ], "metrics": {"complexity": 34, "coverage": 0.12}, } elif i == 56: # Another critical issue file_result = { "path": "src/api/upload.py", "issues": [ { "line": 112, "severity": "CRITICAL", "type": "PATH_TRAVERSAL", "message": "File path not sanitized, allows directory traversal", "code_snippet": "file_path = os.path.join(UPLOAD_DIR, user_filename)", "fix": "Use secure_filename() and validate path stays within UPLOAD_DIR", } ], "metrics": {"complexity": 28, "coverage": 0.45}, } else: file_result = { "path": f"src/module_{i}/handler.py", "issues": [ { "line": 10 + (i % 50), "severity": "LOW", "type": "STYLE", "message": "Line too long (> 100 characters)", } ] if i % 4 == 0 else [], "metrics": {"complexity": 5 + (i % 15), "coverage": 0.7 + (i % 30) / 100}, } files.append(file_result) messages = [ {"role": "user", "content": "Run a security scan on the codebase"}, { "role": "assistant", "content": [ { "type": "tool_use", "id": "call_scan", "name": "security_scan", "input": {"path": "src/"}, } ], }, { "role": "user", "content": [ { "type": "tool_result", "tool_use_id": "call_scan", "content": json.dumps(files, indent=2), } ], }, { "role": "user", "content": "What are the critical security vulnerabilities I need to fix immediately?", }, ] critical_facts = [ CriticalFact( description="SQL injection file location", value="src/auth/login.py line 47", verification_question="In which file and line number is the SQL injection vulnerability?", expected_answer_contains=["login.py", "47"], ), CriticalFact( description="Hardcoded API key value", value="sk-prod-a1b2c3d4e5f6g7h8i9j0", verification_question="What is the hardcoded API key that was found?", expected_answer_contains=["sk-prod", "a1b2c3d4"], ), CriticalFact( description="Path traversal file", value="src/api/upload.py", verification_question="Which file has the path traversal vulnerability?", expected_answer_contains=["upload.py"], ), CriticalFact( description="Path traversal line", value="line 112", verification_question="What line number has the path traversal issue in upload.py?", expected_answer_contains=["112"], ), ] return EvalTestCase( name="code_review_needles", description="80 file scan results with critical vulnerabilities at positions 23 and 56", messages=messages, critical_facts=critical_facts, ) # ============================================================================= # EVALUATION ENGINE # ============================================================================= class QualityRetentionEvaluator: """Evaluates whether compression retains critical information.""" # Tool definitions required by Anthropic API when messages contain tool results TOOL_DEFINITIONS = [ { "name": "search_incidents", "description": "Search for incidents by date", "input_schema": { "type": "object", "properties": { "date": {"type": "string", "description": "Date to search (YYYY-MM-DD)"} }, "required": ["date"], }, }, { "name": "get_logs", "description": "Get logs for a time range", "input_schema": { "type": "object", "properties": { "timerange": {"type": "string", "description": "Time range (e.g., '1h', '24h')"} }, "required": ["timerange"], }, }, { "name": "security_scan", "description": "Run security scan on codebase", "input_schema": { "type": "object", "properties": {"path": {"type": "string", "description": "Path to scan"}}, "required": ["path"], }, }, ] def __init__(self, api_key: str): self.api_key = api_key self.base_client = Anthropic(api_key=api_key) self.provider = AnthropicProvider() self.tokenizer = Tokenizer(TiktokenCounter()) def _create_client(self, config: HeadroomConfig) -> HeadroomClient: """Create a HeadroomClient with the given config.""" return HeadroomClient( original_client=Anthropic(api_key=self.api_key), provider=self.provider, default_mode="optimize", config=config, ) def _ask_question( self, client: HeadroomClient, messages: list[dict], question: str, system_prompt: str, ) -> str: """Ask a verification question and get the response.""" # Add the verification question to the conversation eval_messages = messages + [{"role": "user", "content": question}] response = client.messages.create( model="claude-sonnet-4-20250514", messages=eval_messages, system=system_prompt, tools=self.TOOL_DEFINITIONS, # Required when messages contain tool results max_tokens=300, ) if hasattr(response, "content") and response.content: if isinstance(response.content, list): return response.content[0].text if response.content else "" return str(response.content) return str(response) def _check_answer(self, answer: str, expected_contains: list[str]) -> bool: """Check if answer contains expected content.""" answer_lower = answer.lower() for expected in expected_contains: if expected.lower() in answer_lower: return True return False def evaluate_test_case( self, test_case: EvalTestCase, config: HeadroomConfig, ) -> EvalResult: """Evaluate a single test case.""" print(f"\n{'=' * 60}") print(f"EVAL: {test_case.name}") print(f"{'=' * 60}") print(f"Description: {test_case.description}") # Count tokens before compression tokens_before = self.tokenizer.count_messages(test_case.messages) print(f"Tokens before: {tokens_before:,}") # Create client with compression enabled client = self._create_client(config) # Simulate compression to see what we get try: sim_result = client.messages.simulate( model="claude-sonnet-4-20250514", messages=test_case.messages, system=test_case.system_prompt, tools=self.TOOL_DEFINITIONS, # Required when messages contain tool results ) tokens_after = sim_result.tokens_after compression_ratio = 1 - (tokens_after / tokens_before) if tokens_before > 0 else 0 print(f"Tokens after: {tokens_after:,}") print(f"Compression: {compression_ratio * 100:.1f}%") print(f"Transforms: {sim_result.transforms[:3]}...") # First 3 except Exception as e: return EvalResult( test_name=test_case.name, tokens_before=tokens_before, tokens_after=0, compression_ratio=0, facts_tested=len(test_case.critical_facts), facts_retained=0, retention_rate=0, quality_adjusted_savings=0, error=str(e), ) # Now verify each critical fact print(f"\nVerifying {len(test_case.critical_facts)} critical facts...") facts_retained = 0 details = [] for fact in test_case.critical_facts: print(f"\n Fact: {fact.description}") print(f" Question: {fact.verification_question}") try: # Ask the verification question (with compression applied) answer = self._ask_question( client, test_case.messages, fact.verification_question, test_case.system_prompt, ) # Check if answer contains expected content retained = self._check_answer(answer, fact.expected_answer_contains) if retained: facts_retained += 1 print(" Result: ✅ RETAINED") print(f" Answer: {answer[:100]}...") else: print(" Result: ❌ LOST") print(f" Expected: {fact.expected_answer_contains}") print(f" Got: {answer[:150]}...") details.append( { "fact": fact.description, "question": fact.verification_question, "expected": fact.expected_answer_contains, "answer": answer[:200], "retained": retained, } ) except Exception as e: print(f" Result: ❌ ERROR: {e}") details.append( { "fact": fact.description, "error": str(e), "retained": False, } ) retention_rate = ( facts_retained / len(test_case.critical_facts) if test_case.critical_facts else 0 ) quality_adjusted_savings = compression_ratio * retention_rate print(f"\n{'=' * 60}") print( f"RESULT: {facts_retained}/{len(test_case.critical_facts)} facts retained ({retention_rate * 100:.0f}%)" ) print(f"Quality-Adjusted Savings: {quality_adjusted_savings * 100:.1f}%") print(f"{'=' * 60}") return EvalResult( test_name=test_case.name, tokens_before=tokens_before, tokens_after=tokens_after, compression_ratio=compression_ratio, facts_tested=len(test_case.critical_facts), facts_retained=facts_retained, retention_rate=retention_rate, quality_adjusted_savings=quality_adjusted_savings, details=details, ) def run_full_eval(self, config: HeadroomConfig) -> list[EvalResult]: """Run evaluation on all test cases.""" test_cases = [ create_search_results_with_needles(), create_log_analysis_with_needles(), create_code_review_with_needles(), ] results = [] for test_case in test_cases: result = self.evaluate_test_case(test_case, config) results.append(result) return results # ============================================================================= # MAIN # ============================================================================= def main(): print("\n" + "=" * 70) print("QUALITY RETENTION EVALUATION") print("=" * 70) print("Verifying that intelligent compression retains critical information") print("=" * 70) # Create config with intelligent context enabled config = HeadroomConfig() config.intelligent_context = IntelligentContextConfig( enabled=True, use_importance_scoring=True, compress_threshold=0.10, summarize_threshold=0.25, ) config.rolling_window.enabled = False config.smart_crusher.enabled = True # Run evaluation evaluator = QualityRetentionEvaluator(API_KEY) results = evaluator.run_full_eval(config) # Summary print("\n" + "=" * 70) print("EVALUATION SUMMARY") print("=" * 70) total_facts = sum(r.facts_tested for r in results) total_retained = sum(r.facts_retained for r in results) total_tokens_before = sum(r.tokens_before for r in results) total_tokens_after = sum(r.tokens_after for r in results) print(f"\n{'Test Case':<30} {'Compression':<15} {'Retention':<15} {'Quality-Adj':<15}") print("-" * 75) for result in results: status = ( "✅" if result.retention_rate >= 0.8 else "⚠️" if result.retention_rate >= 0.5 else "❌" ) print( f"{result.test_name:<30} " f"{result.compression_ratio * 100:>6.1f}% " f"{result.retention_rate * 100:>6.0f}% ({result.facts_retained}/{result.facts_tested}) " f"{result.quality_adjusted_savings * 100:>6.1f}% {status}" ) print("-" * 75) overall_compression = ( 1 - (total_tokens_after / total_tokens_before) if total_tokens_before > 0 else 0 ) overall_retention = total_retained / total_facts if total_facts > 0 else 0 overall_quality_adj = overall_compression * overall_retention print( f"{'OVERALL':<30} " f"{overall_compression * 100:>6.1f}% " f"{overall_retention * 100:>6.0f}% ({total_retained}/{total_facts}) " f"{overall_quality_adj * 100:>6.1f}%" ) print("\n" + "=" * 70) if overall_retention >= 0.8: print("✅ PASS: Critical information retention is good (>=80%)") elif overall_retention >= 0.5: print("⚠️ WARNING: Some critical information was lost (50-80% retention)") else: print("❌ FAIL: Significant critical information loss (<50% retention)") print( f"Tokens saved: {total_tokens_before - total_tokens_after:,} ({overall_compression * 100:.1f}% compression)" ) print("=" * 70) return 0 if overall_retention >= 0.8 else 1 if __name__ == "__main__": sys.exit(main())