mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
e4a41faa33
- Fix E402: Move module-level imports to top of file - Fix F401: Add noqa for availability check imports - Fix F402: Rename loop variables shadowing imports - Fix E722: Replace bare except with except Exception - Fix B904: Add exception chaining (from e) - Fix F811: Remove duplicate imports - Fix B027: Add noqa for empty close() method - Fix E741: Rename ambiguous variable l -> label - Fix I001: Import sorting issues - Apply ruff format to all 106 files All 902 tests pass.
489 lines
17 KiB
Python
489 lines
17 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Comparison evaluation: SmartCrusher vs ToolCrusher (Naive Crushing)
|
|
|
|
This script compares:
|
|
1. ToolCrusher - Fixed rules (keep first N items)
|
|
2. SmartCrusher - Statistical analysis (preserve change points, factor constants)
|
|
|
|
We'll use the same SRE incident response scenario and see which produces
|
|
better token reduction while maintaining response quality.
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import tempfile
|
|
import time
|
|
from dataclasses import dataclass
|
|
from datetime import datetime, timedelta
|
|
|
|
from dotenv import load_dotenv
|
|
from openai import OpenAI
|
|
|
|
from headroom import HeadroomClient, OpenAIProvider, SmartCrusherConfig, ToolCrusherConfig
|
|
from headroom.config import HeadroomConfig
|
|
from headroom.transforms import TransformPipeline
|
|
|
|
load_dotenv(".env.local")
|
|
|
|
# Initialize base client
|
|
base_client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
|
|
provider = OpenAIProvider()
|
|
|
|
|
|
# =============================================================================
|
|
# CREATE THREE CLIENT CONFIGURATIONS
|
|
# =============================================================================
|
|
|
|
# 1. NAIVE CRUSHER (fixed rules - keep first 10 items)
|
|
naive_config = HeadroomConfig()
|
|
naive_config.tool_crusher = ToolCrusherConfig(
|
|
enabled=True,
|
|
min_tokens_to_crush=200,
|
|
max_array_items=10,
|
|
max_string_length=1000,
|
|
max_depth=5,
|
|
)
|
|
naive_config.smart_crusher.enabled = False
|
|
|
|
naive_client = HeadroomClient(
|
|
original_client=base_client,
|
|
provider=provider,
|
|
store_url=f"sqlite:///{tempfile.gettempdir()}/headroom_naive.db",
|
|
default_mode="audit",
|
|
)
|
|
naive_client._config = naive_config
|
|
naive_client._pipeline = TransformPipeline(naive_config, provider=provider)
|
|
|
|
# 2. SMART CRUSHER (statistical analysis)
|
|
smart_config = HeadroomConfig()
|
|
smart_config.tool_crusher.enabled = False # Disable naive
|
|
smart_config.smart_crusher = SmartCrusherConfig(
|
|
enabled=True,
|
|
min_items_to_analyze=5,
|
|
min_tokens_to_crush=200,
|
|
variance_threshold=2.0,
|
|
max_items_after_crush=15,
|
|
preserve_change_points=True,
|
|
factor_out_constants=True,
|
|
include_summaries=True,
|
|
)
|
|
|
|
smart_client = HeadroomClient(
|
|
original_client=base_client,
|
|
provider=provider,
|
|
store_url=f"sqlite:///{tempfile.gettempdir()}/headroom_smart.db",
|
|
default_mode="audit",
|
|
)
|
|
smart_client._config = smart_config
|
|
smart_client._pipeline = TransformPipeline(smart_config, provider=provider)
|
|
|
|
# 3. BASELINE (no optimization)
|
|
baseline_client = HeadroomClient(
|
|
original_client=base_client,
|
|
provider=provider,
|
|
store_url=f"sqlite:///{tempfile.gettempdir()}/headroom_baseline.db",
|
|
default_mode="audit",
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# GENERATE TEST DATA WITH CLEAR PATTERNS
|
|
# =============================================================================
|
|
|
|
|
|
def generate_metrics_with_spike() -> str:
|
|
"""
|
|
Generate metrics data with a CLEAR spike pattern.
|
|
SmartCrusher should detect and preserve the change point.
|
|
NaiveCrusher will just keep first 10 items (missing the spike).
|
|
"""
|
|
base_time = datetime.now() - timedelta(hours=1)
|
|
|
|
data_points = []
|
|
for i in range(60):
|
|
ts = base_time + timedelta(minutes=i)
|
|
|
|
# PATTERN: Stable at ~45 for first 45 minutes, then spike to 95
|
|
if i < 45:
|
|
cpu = 45 + (i % 3) # Small variance: 45-47
|
|
error_rate = 0.1
|
|
else:
|
|
cpu = 85 + (i - 45) * 2 # Spike: 85 -> 115
|
|
error_rate = 5 + (i - 45) # Error spike too
|
|
|
|
data_points.append(
|
|
{
|
|
"timestamp": ts.isoformat(),
|
|
"host": "prod-api-1", # CONSTANT - should be factored out
|
|
"region": "us-east-1", # CONSTANT - should be factored out
|
|
"datacenter": "dc-01", # CONSTANT - should be factored out
|
|
"cpu_percent": min(cpu, 99),
|
|
"memory_percent": 62, # CONSTANT
|
|
"error_rate": round(error_rate, 2),
|
|
"request_count": 1500 + (i * 10),
|
|
}
|
|
)
|
|
|
|
return json.dumps({"status": "success", "metrics": data_points, "query_time_ms": 127})
|
|
|
|
|
|
def generate_clusterable_logs() -> str:
|
|
"""
|
|
Generate logs with REPEATING patterns.
|
|
SmartCrusher should cluster and dedupe.
|
|
NaiveCrusher will just keep first 10.
|
|
"""
|
|
base_time = datetime.now() - timedelta(minutes=30)
|
|
|
|
# 4 distinct message types, repeated many times
|
|
message_templates = [
|
|
("ERROR", "Connection timeout to database-primary after 30000ms"),
|
|
("ERROR", "Connection timeout to database-primary after 30000ms"),
|
|
("ERROR", "Connection timeout to database-primary after 30000ms"),
|
|
("WARN", "Connection pool exhausted, waiting for available connection"),
|
|
("WARN", "Connection pool exhausted, waiting for available connection"),
|
|
("ERROR", "Connection timeout to database-primary after 30000ms"),
|
|
("ERROR", "Max retries exceeded for database operation"),
|
|
("INFO", "Retry attempt 1/3 for database connection"),
|
|
("INFO", "Retry attempt 2/3 for database connection"),
|
|
("ERROR", "Connection timeout to database-primary after 30000ms"),
|
|
("WARN", "Circuit breaker OPEN for database-primary"),
|
|
("ERROR", "Connection timeout to database-primary after 30000ms"),
|
|
("ERROR", "OOM killed: api-server process exceeded memory limit"), # UNIQUE
|
|
("ERROR", "Connection timeout to database-primary after 30000ms"),
|
|
("WARN", "Connection pool exhausted, waiting for available connection"),
|
|
]
|
|
|
|
logs = []
|
|
for i in range(50):
|
|
ts = base_time + timedelta(seconds=i * 36)
|
|
level, msg = message_templates[i % len(message_templates)]
|
|
|
|
logs.append(
|
|
{
|
|
"@timestamp": ts.isoformat(),
|
|
"level": level,
|
|
"message": msg,
|
|
"service": "api-server", # CONSTANT
|
|
"environment": "production", # CONSTANT
|
|
"version": "2.4.1", # CONSTANT
|
|
"host": f"prod-api-{i % 3 + 1}",
|
|
"trace_id": f"trace-{1000 + i:04d}",
|
|
}
|
|
)
|
|
|
|
return json.dumps({"took": 234, "hits": {"total": len(logs), "hits": logs}})
|
|
|
|
|
|
def generate_search_results() -> str:
|
|
"""
|
|
Generate search results with scores.
|
|
SmartCrusher should use TOP_N strategy.
|
|
"""
|
|
results = []
|
|
for i in range(30):
|
|
results.append(
|
|
{
|
|
"id": f"doc-{i + 1}",
|
|
"title": f"Result document {i + 1}",
|
|
"snippet": f"This is the snippet for document {i + 1} with relevant content...",
|
|
"score": 0.95 - (i * 0.02), # Decreasing relevance
|
|
"source": "knowledge_base", # CONSTANT
|
|
"category": "technical", # CONSTANT
|
|
}
|
|
)
|
|
|
|
return json.dumps({"results": results, "total": 30})
|
|
|
|
|
|
# =============================================================================
|
|
# BUILD TEST CONVERSATION
|
|
# =============================================================================
|
|
|
|
|
|
def build_test_conversation() -> list[dict]:
|
|
"""Build a conversation that exercises all SmartCrusher strategies."""
|
|
|
|
messages = [
|
|
{
|
|
"role": "system",
|
|
"content": """You are an SRE assistant. Analyze the data and provide insights.
|
|
Current Date: 2024-12-15T14:30:00Z""",
|
|
},
|
|
{"role": "user", "content": "Check the metrics for the last hour."},
|
|
{
|
|
"role": "assistant",
|
|
"content": None,
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_1",
|
|
"type": "function",
|
|
"function": {"name": "query_metrics", "arguments": "{}"},
|
|
}
|
|
],
|
|
},
|
|
{"role": "tool", "tool_call_id": "call_1", "content": generate_metrics_with_spike()},
|
|
{"role": "assistant", "content": "I see CPU metrics. Let me check the logs."},
|
|
{
|
|
"role": "assistant",
|
|
"content": None,
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_2",
|
|
"type": "function",
|
|
"function": {"name": "search_logs", "arguments": "{}"},
|
|
}
|
|
],
|
|
},
|
|
{"role": "tool", "tool_call_id": "call_2", "content": generate_clusterable_logs()},
|
|
{"role": "assistant", "content": "Found error patterns. Let me search docs."},
|
|
{
|
|
"role": "assistant",
|
|
"content": None,
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_3",
|
|
"type": "function",
|
|
"function": {"name": "search_docs", "arguments": "{}"},
|
|
}
|
|
],
|
|
},
|
|
{"role": "tool", "tool_call_id": "call_3", "content": generate_search_results()},
|
|
{"role": "user", "content": "What's the root cause and what should we do?"},
|
|
]
|
|
|
|
return messages
|
|
|
|
|
|
# =============================================================================
|
|
# EVALUATION
|
|
# =============================================================================
|
|
|
|
|
|
@dataclass
|
|
class EvalResult:
|
|
name: str
|
|
tokens_before: int
|
|
tokens_after: int
|
|
tokens_saved: int
|
|
reduction_pct: float
|
|
transforms: list[str]
|
|
response: str
|
|
latency_ms: float
|
|
|
|
|
|
def evaluate(client, messages: list[dict], name: str, mode: str) -> EvalResult:
|
|
"""Run evaluation with a client."""
|
|
# Simulate first to get transform info
|
|
sim = client.chat.completions.simulate(model="gpt-4o-mini", messages=messages)
|
|
|
|
# Actual call
|
|
start = time.time()
|
|
response = client.chat.completions.create(
|
|
model="gpt-4o-mini",
|
|
messages=messages,
|
|
max_tokens=1000,
|
|
headroom_mode=mode,
|
|
)
|
|
latency = (time.time() - start) * 1000
|
|
|
|
tokens_in = response.usage.prompt_tokens if response.usage else 0
|
|
|
|
return EvalResult(
|
|
name=name,
|
|
tokens_before=sim.tokens_before,
|
|
tokens_after=tokens_in,
|
|
tokens_saved=sim.tokens_before - tokens_in,
|
|
reduction_pct=(sim.tokens_before - tokens_in) / sim.tokens_before * 100
|
|
if sim.tokens_before
|
|
else 0,
|
|
transforms=sim.transforms,
|
|
response=response.choices[0].message.content or "",
|
|
latency_ms=latency,
|
|
)
|
|
|
|
|
|
def evaluate_quality(baseline_response: str, test_response: str, test_name: str) -> dict:
|
|
"""Use GPT-4o to judge response quality."""
|
|
judge_prompt = f"""Compare these two AI responses to an SRE incident investigation.
|
|
|
|
BASELINE (unoptimized):
|
|
{baseline_response}
|
|
|
|
{test_name.upper()}:
|
|
{test_response}
|
|
|
|
Score the {test_name} response on:
|
|
1. ROOT_CAUSE: Does it correctly identify the database connection issue? (1-5)
|
|
2. DATA_AWARENESS: Does it reference specific metrics (CPU spike, error rate)? (1-5)
|
|
3. ACTIONABILITY: Does it provide concrete next steps? (1-5)
|
|
|
|
Respond in JSON:
|
|
{{
|
|
"root_cause": {{"score": N, "reason": "..."}},
|
|
"data_awareness": {{"score": N, "reason": "..."}},
|
|
"actionability": {{"score": N, "reason": "..."}},
|
|
"overall": N,
|
|
"verdict": "PASS" or "FAIL"
|
|
}}
|
|
|
|
PASS = overall >= 4.0"""
|
|
|
|
response = base_client.chat.completions.create(
|
|
model="gpt-4o",
|
|
messages=[{"role": "user", "content": judge_prompt}],
|
|
max_tokens=500,
|
|
response_format={"type": "json_object"},
|
|
)
|
|
|
|
try:
|
|
return json.loads(response.choices[0].message.content)
|
|
except Exception:
|
|
return {"error": "Parse failed"}
|
|
|
|
|
|
def main():
|
|
print("=" * 70)
|
|
print("SMART CRUSHER vs NAIVE CRUSHER COMPARISON")
|
|
print("=" * 70)
|
|
print()
|
|
|
|
messages = build_test_conversation()
|
|
|
|
# Count tool outputs
|
|
tool_count = sum(1 for m in messages if m.get("role") == "tool")
|
|
print(f"Test conversation: {len(messages)} messages, {tool_count} tool outputs")
|
|
print()
|
|
|
|
# Tool output breakdown
|
|
print("Tool output patterns for SmartCrusher:")
|
|
print(" 1. Metrics: 60 data points with SPIKE at point 45")
|
|
print(" - 3 CONSTANT fields (host, region, datacenter)")
|
|
print(" - SmartCrusher should: detect spike, factor constants")
|
|
print(" 2. Logs: 50 entries with 4 REPEATING message types")
|
|
print(" - SmartCrusher should: cluster and dedupe")
|
|
print(" 3. Search: 30 results with relevance scores")
|
|
print(" - SmartCrusher should: keep top N by score")
|
|
print()
|
|
|
|
# Run evaluations
|
|
print("-" * 70)
|
|
print("RUNNING EVALUATIONS")
|
|
print("-" * 70)
|
|
|
|
print("\n1. BASELINE (no optimization)...")
|
|
baseline = evaluate(baseline_client, messages, "Baseline", "audit")
|
|
print(f" Tokens: {baseline.tokens_after:,}")
|
|
|
|
print("\n2. NAIVE CRUSHER (fixed rules: keep first 10)...")
|
|
naive = evaluate(naive_client, messages, "Naive", "optimize")
|
|
print(
|
|
f" Tokens: {naive.tokens_after:,} (saved {naive.tokens_saved:,}, {naive.reduction_pct:.1f}%)"
|
|
)
|
|
print(f" Transforms: {naive.transforms}")
|
|
|
|
print("\n3. SMART CRUSHER (statistical analysis)...")
|
|
smart = evaluate(smart_client, messages, "Smart", "optimize")
|
|
print(
|
|
f" Tokens: {smart.tokens_after:,} (saved {smart.tokens_saved:,}, {smart.reduction_pct:.1f}%)"
|
|
)
|
|
print(f" Transforms: {smart.transforms}")
|
|
|
|
# Results comparison
|
|
print()
|
|
print("=" * 70)
|
|
print("TOKEN COMPARISON")
|
|
print("=" * 70)
|
|
|
|
print(f"\n{'Method':<20} {'Tokens':>10} {'Saved':>10} {'Reduction':>10}")
|
|
print("-" * 52)
|
|
print(f"{'Baseline':<20} {baseline.tokens_after:>10,} {'-':>10} {'-':>10}")
|
|
print(
|
|
f"{'Naive Crusher':<20} {naive.tokens_after:>10,} {naive.tokens_saved:>10,} {naive.reduction_pct:>9.1f}%"
|
|
)
|
|
print(
|
|
f"{'Smart Crusher':<20} {smart.tokens_after:>10,} {smart.tokens_saved:>10,} {smart.reduction_pct:>9.1f}%"
|
|
)
|
|
|
|
# Show the difference
|
|
diff = naive.tokens_after - smart.tokens_after
|
|
if diff > 0:
|
|
print(
|
|
f"\n→ Smart Crusher saves {diff:,} MORE tokens than Naive ({diff / naive.tokens_after * 100:.1f}% better)"
|
|
)
|
|
elif diff < 0:
|
|
print(
|
|
f"\n→ Naive Crusher saves {-diff:,} MORE tokens than Smart ({-diff / smart.tokens_after * 100:.1f}% better)"
|
|
)
|
|
else:
|
|
print("\n→ Both methods produce same token count")
|
|
|
|
# Quality evaluation
|
|
print()
|
|
print("-" * 70)
|
|
print("QUALITY EVALUATION (GPT-4o Judge)")
|
|
print("-" * 70)
|
|
|
|
print("\nEvaluating Naive Crusher response...")
|
|
naive_quality = evaluate_quality(baseline.response, naive.response, "naive crusher")
|
|
|
|
print("Evaluating Smart Crusher response...")
|
|
smart_quality = evaluate_quality(baseline.response, smart.response, "smart crusher")
|
|
|
|
if "error" not in naive_quality and "error" not in smart_quality:
|
|
print(f"\n{'Criterion':<20} {'Naive':>10} {'Smart':>10}")
|
|
print("-" * 42)
|
|
for criterion in ["root_cause", "data_awareness", "actionability"]:
|
|
n_score = naive_quality.get(criterion, {}).get("score", "?")
|
|
s_score = smart_quality.get(criterion, {}).get("score", "?")
|
|
print(f"{criterion.replace('_', ' ').title():<20} {n_score:>10}/5 {s_score:>10}/5")
|
|
print("-" * 42)
|
|
print(
|
|
f"{'OVERALL':<20} {naive_quality.get('overall', '?'):>10}/5 {smart_quality.get('overall', '?'):>10}/5"
|
|
)
|
|
print(
|
|
f"{'VERDICT':<20} {naive_quality.get('verdict', '?'):>10} {smart_quality.get('verdict', '?'):>10}"
|
|
)
|
|
|
|
print("\n[Quality Analysis]")
|
|
print(f" Naive: {naive_quality.get('data_awareness', {}).get('reason', 'N/A')}")
|
|
print(f" Smart: {smart_quality.get('data_awareness', {}).get('reason', 'N/A')}")
|
|
|
|
# Summary
|
|
print()
|
|
print("=" * 70)
|
|
print("SUMMARY")
|
|
print("=" * 70)
|
|
|
|
n_verdict = naive_quality.get("verdict", "?") if "error" not in naive_quality else "ERROR"
|
|
s_verdict = smart_quality.get("verdict", "?") if "error" not in smart_quality else "ERROR"
|
|
n_overall = naive_quality.get("overall", 0) if "error" not in naive_quality else 0
|
|
s_overall = smart_quality.get("overall", 0) if "error" not in smart_quality else 0
|
|
|
|
print(f"""
|
|
SmartCrusher vs NaiveCrusher on SRE incident data:
|
|
|
|
Token Efficiency:
|
|
- Naive: {naive.reduction_pct:.1f}% reduction
|
|
- Smart: {smart.reduction_pct:.1f}% reduction
|
|
- Winner: {"SMART" if smart.reduction_pct > naive.reduction_pct else "NAIVE" if naive.reduction_pct > smart.reduction_pct else "TIE"} (+{abs(smart.reduction_pct - naive.reduction_pct):.1f}% {"more" if smart.reduction_pct > naive.reduction_pct else "less"} reduction)
|
|
|
|
Response Quality:
|
|
- Naive: {n_overall}/5 ({n_verdict})
|
|
- Smart: {s_overall}/5 ({s_verdict})
|
|
- Winner: {"SMART" if s_overall > n_overall else "NAIVE" if n_overall > s_overall else "TIE"}
|
|
|
|
Key Insight:
|
|
SmartCrusher uses statistical analysis to preserve important data:
|
|
- Change points (CPU spike at minute 45) are PRESERVED
|
|
- Constants (host, region, datacenter) are FACTORED OUT
|
|
- Logs are CLUSTERED by message similarity
|
|
- Search results keep TOP items by score
|
|
|
|
NaiveCrusher blindly keeps first N items, potentially missing the spike!
|
|
""")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|