chopratejas--headroom/benchmarks/real_world_agent_benchmark.py

"""
Real-World Agent Benchmark: MCP Tools + Headroom

This benchmark simulates real multi-agent workflows using actual MCP tool output formats:
1. Filesystem MCP Server - directory trees, file searches, file contents
2. GitHub MCP Server - code search, issues, PRs, commits
3. Database MCP Server - query results, schema info

We measure:
- Token usage with vs without Headroom
- Cost savings
- Answer quality (does compression hurt agent performance?)

This is NOT synthetic data - these are actual output formats from production MCP servers.
"""

import hashlib
import json
import os
import random
import time
from dataclasses import dataclass
from typing import Any

# OpenAI for agent
try:
    from openai import OpenAI  # noqa: F401

    OPENAI_AVAILABLE = True
except ImportError:
    OPENAI_AVAILABLE = False

# Headroom
try:
    from headroom import HeadroomClient, OpenAIProvider

    HEADROOM_AVAILABLE = True
except ImportError:
    HEADROOM_AVAILABLE = False


# =============================================================================
# REALISTIC MCP TOOL OUTPUT GENERATORS
# Based on actual MCP server output formats
# =============================================================================


def generate_filesystem_tree(
    path: str = "/project", depth: int = 3, files_per_dir: int = 15
) -> dict:
    """
    Generate realistic filesystem tree output (MCP filesystem server format).
    This mimics `tree` command output from @modelcontextprotocol/server-filesystem.
    """

    def generate_dir(current_path: str, current_depth: int) -> list:
        if current_depth <= 0:
            return []

        entries = []

        # Common project structure
        dir_names = [
            "src",
            "lib",
            "utils",
            "components",
            "services",
            "models",
            "controllers",
            "middleware",
            "tests",
            "config",
            "scripts",
            "api",
            "core",
            "helpers",
            "types",
            "interfaces",
        ]

        file_extensions = [".py", ".ts", ".js", ".json", ".yaml", ".md"]

        # Add some directories
        num_dirs = random.randint(2, 5) if current_depth > 1 else 0
        for i in range(num_dirs):
            dir_name = random.choice(dir_names) + (f"_{i}" if i > 0 else "")
            dir_path = f"{current_path}/{dir_name}"
            entries.append(
                {
                    "name": dir_name,
                    "type": "directory",
                    "path": dir_path,
                    "children": generate_dir(dir_path, current_depth - 1),
                }
            )

        # Add files
        for i in range(files_per_dir):
            ext = random.choice(file_extensions)
            file_name = f"module_{i}{ext}"
            entries.append(
                {
                    "name": file_name,
                    "type": "file",
                    "path": f"{current_path}/{file_name}",
                    "size": random.randint(100, 10000),
                    "modified": f"2024-01-{random.randint(1, 28):02d}T{random.randint(0, 23):02d}:{random.randint(0, 59):02d}:00Z",
                }
            )

        return entries

    return {
        "tool": "filesystem_tree",
        "path": path,
        "result": {
            "name": path.split("/")[-1] or "project",
            "type": "directory",
            "path": path,
            "children": generate_dir(path, depth),
        },
    }


def generate_filesystem_search(query: str, num_results: int = 200) -> dict:
    """
    Generate realistic file search results (MCP filesystem server format).
    Mimics search_files output with path matches and content snippets.
    """
    results = []

    # Common file paths in a real project
    paths = [
        "src/auth/jwt_handler.py",
        "src/auth/oauth_provider.py",
        "src/api/routes/users.py",
        "src/api/routes/products.py",
        "src/services/payment_processor.py",
        "src/services/email_sender.py",
        "src/middleware/rate_limiter.py",
        "src/middleware/auth_middleware.py",
        "src/models/user.py",
        "src/models/order.py",
        "tests/test_auth.py",
        "tests/test_api.py",
        "config/database.py",
        "config/settings.py",
    ]

    for i in range(num_results):
        if i < len(paths):
            path = paths[i]
        else:
            path = f"src/modules/module_{i}.py"

        # Generate realistic match context
        match_line = random.randint(10, 500)
        results.append(
            {
                "path": path,
                "type": "file",
                "size": random.randint(500, 15000),
                "modified": f"2024-01-{random.randint(1, 28):02d}",
                "matches": [
                    {
                        "line": match_line,
                        "content": f"    def process_{query.lower().replace(' ', '_')}(self, data):",
                        "context_before": "    # Process incoming request",
                        "context_after": f"        return self.handler.{query.lower()}(data)",
                    }
                ],
                "score": round(random.uniform(0.5, 1.0), 3),
            }
        )

    return {
        "tool": "search_files",
        "query": query,
        "result": {
            "total_matches": num_results,
            "files_searched": num_results * 10,
            "matches": results,
        },
    }


def generate_github_code_search(query: str, num_results: int = 100) -> dict:
    """
    Generate realistic GitHub code search results (GitHub MCP server format).
    Based on actual github-mcp-server output.
    """
    repos = [
        "facebook/react",
        "microsoft/vscode",
        "tensorflow/tensorflow",
        "kubernetes/kubernetes",
        "golang/go",
        "rust-lang/rust",
        "apache/spark",
        "elastic/elasticsearch",
        "grafana/grafana",
        "prometheus/prometheus",
        "docker/docker-ce",
        "nginx/nginx",
    ]

    results = []
    for i in range(num_results):
        repo = random.choice(repos)
        results.append(
            {
                "repository": {
                    "full_name": repo,
                    "description": f"The {repo.split('/')[1]} project",
                    "stars": random.randint(1000, 100000),
                    "language": random.choice(["Python", "Go", "TypeScript", "Java", "Rust"]),
                    "updated_at": f"2024-01-{random.randint(1, 28):02d}T00:00:00Z",
                },
                "path": f"src/{query.lower().replace(' ', '_')}/handler.py",
                "sha": hashlib.sha1(f"{repo}{i}".encode()).hexdigest(),
                "url": f"https://github.com/{repo}/blob/main/src/handler.py",
                "score": round(random.uniform(10, 100), 2),
                "text_matches": [
                    {
                        "fragment": f"def {query.lower().replace(' ', '_')}(request):\n    # Implementation\n    return response",
                        "matches": [{"text": query, "indices": [4, 4 + len(query)]}],
                    }
                ],
            }
        )

    return {
        "tool": "github_search_code",
        "query": query,
        "result": {"total_count": num_results * 50, "incomplete_results": False, "items": results},
    }


def generate_github_issues(repo: str, num_issues: int = 50) -> dict:
    """
    Generate realistic GitHub issues list (GitHub MCP server format).
    """
    labels = ["bug", "enhancement", "documentation", "help wanted", "good first issue"]
    states = ["open", "open", "open", "closed"]  # Weighted toward open

    issues = []
    for i in range(num_issues):
        issues.append(
            {
                "number": 1000 + i,
                "title": f"Issue #{1000 + i}: "
                + random.choice(
                    [
                        "Fix authentication flow",
                        "Add support for OAuth2",
                        "Performance regression in v2.0",
                        "Documentation needs update",
                        "Memory leak in worker process",
                        "Add dark mode support",
                        "API rate limiting not working",
                    ]
                ),
                "state": random.choice(states),
                "user": {
                    "login": f"user{random.randint(1, 1000)}",
                    "avatar_url": f"https://avatars.githubusercontent.com/u/{random.randint(1, 100000)}",
                },
                "labels": random.sample(labels, k=random.randint(0, 3)),
                "created_at": f"2024-01-{random.randint(1, 28):02d}T{random.randint(0, 23):02d}:00:00Z",
                "updated_at": f"2024-01-{random.randint(1, 28):02d}T{random.randint(0, 23):02d}:00:00Z",
                "comments": random.randint(0, 50),
                "body": f"## Description\n\nThis issue tracks {random.choice(['a bug', 'a feature request', 'documentation update'])}.\n\n## Steps to Reproduce\n\n1. Step one\n2. Step two\n3. Step three\n\n## Expected Behavior\n\nIt should work.\n\n## Actual Behavior\n\nIt doesn't work.",
            }
        )

    return {
        "tool": "github_list_issues",
        "repository": repo,
        "result": {"total_count": num_issues, "items": issues},
    }


def generate_database_query_results(query: str, num_rows: int = 500) -> dict:
    """
    Generate realistic database query results (Database MCP server format).
    """
    # Simulate a user analytics query
    rows = []
    for i in range(num_rows):
        rows.append(
            {
                "user_id": f"user_{10000 + i}",
                "email": f"user{10000 + i}@example.com",
                "created_at": f"2024-01-{random.randint(1, 28):02d}",
                "last_login": f"2024-01-{random.randint(1, 28):02d}T{random.randint(0, 23):02d}:00:00Z",
                "total_orders": random.randint(0, 100),
                "total_revenue": round(random.uniform(0, 10000), 2),
                "status": random.choice(["active", "active", "active", "inactive", "suspended"]),
                "country": random.choice(["US", "UK", "DE", "FR", "JP", "AU", "CA"]),
                "subscription_tier": random.choice(
                    ["free", "free", "basic", "premium", "enterprise"]
                ),
            }
        )

    # Add some anomalies (high-value users)
    for _ in range(3):
        rows[random.randint(0, len(rows) - 1)]["total_revenue"] = round(
            random.uniform(50000, 100000), 2
        )
        rows[random.randint(0, len(rows) - 1)]["status"] = "suspended"

    return {
        "tool": "database_query",
        "query": query,
        "result": {
            "columns": [
                "user_id",
                "email",
                "created_at",
                "last_login",
                "total_orders",
                "total_revenue",
                "status",
                "country",
                "subscription_tier",
            ],
            "row_count": num_rows,
            "rows": rows,
            "execution_time_ms": random.randint(50, 500),
        },
    }


def generate_log_search(query: str, num_entries: int = 300) -> dict:
    """
    Generate realistic log search results (Logging MCP server format).
    """
    log_levels = ["INFO", "INFO", "INFO", "INFO", "WARN", "ERROR", "DEBUG"]
    services = [
        "api-gateway",
        "auth-service",
        "payment-service",
        "user-service",
        "notification-service",
    ]

    entries = []

    for i in range(num_entries):
        level = random.choice(log_levels)
        service = random.choice(services)

        if level == "ERROR":
            message = random.choice(
                [
                    f"Connection refused to {service}:8080 - ECONNREFUSED",
                    f"Timeout waiting for response from {service} after 30000ms",
                    "Failed to process request: NullPointerException",
                    "Database connection pool exhausted",
                ]
            )
        elif level == "WARN":
            message = random.choice(
                [
                    "High latency detected: 2500ms (threshold: 1000ms)",
                    "Rate limit approaching: 450/500 requests",
                    "Memory usage at 85%",
                    f"Retry attempt 3/5 for {service}",
                ]
            )
        else:
            message = random.choice(
                [
                    "Request processed successfully",
                    "Health check passed",
                    f"Cache hit for key: user_session_{random.randint(1000, 9999)}",
                    f"Authenticated user: user_{random.randint(1000, 9999)}",
                ]
            )

        entries.append(
            {
                "timestamp": f"2024-01-15T{10 + (i // 60):02d}:{i % 60:02d}:00Z",
                "level": level,
                "service": service,
                "message": message,
                "trace_id": hashlib.md5(f"{i}".encode()).hexdigest()[:16],  # nosec B324
                "metadata": {
                    "host": f"pod-{service}-{random.randint(1, 5)}",
                    "region": random.choice(["us-east-1", "us-west-2", "eu-west-1"]),
                },
            }
        )

    return {
        "tool": "search_logs",
        "query": query,
        "result": {"total_hits": num_entries * 10, "returned": num_entries, "entries": entries},
    }


# =============================================================================
# AGENT SCENARIOS
# =============================================================================


@dataclass
class AgentScenario:
    """A realistic agent workflow scenario."""

    name: str
    description: str
    system_prompt: str
    user_query: str
    tools: list[dict]  # Tool outputs in sequence
    expected_answer_contains: list[str]  # Key phrases expected in good answer


def create_sre_debugging_scenario() -> AgentScenario:
    """
    SRE agent debugging a production incident.
    Multiple tool calls with large outputs.
    """
    return AgentScenario(
        name="SRE Incident Debugging",
        description="Debug a production incident using logs, metrics, and deployment info",
        system_prompt="""You are an SRE assistant helping debug production incidents.
You have access to tools for searching logs, querying metrics, and checking deployments.
Analyze the data carefully and identify the root cause.""",
        user_query="We're seeing 500 errors on the payment service. Can you investigate and find the root cause?",
        tools=[
            generate_log_search("payment error", num_entries=300),
            generate_database_query_results(
                "SELECT * FROM service_metrics WHERE service='payment'", num_rows=200
            ),
            generate_filesystem_search("payment", num_results=150),
        ],
        expected_answer_contains=["payment", "error", "connection", "timeout"],
    )


def create_codebase_exploration_scenario() -> AgentScenario:
    """
    Developer agent exploring a new codebase.
    File tree + search + code reading.
    """
    return AgentScenario(
        name="Codebase Exploration",
        description="Explore a codebase to understand authentication implementation",
        system_prompt="""You are a developer assistant helping explore codebases.
You have access to file system tools and code search.
Help the user understand how the codebase is structured.""",
        user_query="I need to understand how authentication is implemented. Can you find the relevant files and explain the flow?",
        tools=[
            generate_filesystem_tree("/project", depth=3, files_per_dir=20),
            generate_filesystem_search("authentication", num_results=200),
            generate_github_code_search("JWT authentication middleware", num_results=100),
        ],
        expected_answer_contains=["auth", "jwt", "middleware", "handler"],
    )


def create_issue_triage_scenario() -> AgentScenario:
    """
    GitHub agent triaging issues and finding related code.
    """
    return AgentScenario(
        name="GitHub Issue Triage",
        description="Triage GitHub issues and find related code",
        system_prompt="""You are a GitHub assistant helping triage issues.
Analyze issues, find patterns, and identify related code.""",
        user_query="Can you analyze the open issues and identify any patterns or high-priority bugs we should focus on?",
        tools=[
            generate_github_issues("myorg/myrepo", num_issues=100),
            generate_github_code_search("bug fix", num_results=80),
            generate_log_search("exception", num_entries=200),
        ],
        expected_answer_contains=["bug", "issue", "priority"],
    )


# =============================================================================
# BENCHMARK RUNNER
# =============================================================================


@dataclass
class BenchmarkResult:
    """Result from running a scenario."""

    scenario_name: str
    mode: str  # "baseline" or "headroom"
    total_input_tokens: int
    total_output_tokens: int
    total_tokens: int
    cost_usd: float
    latency_ms: float
    answer_quality: float  # 0-1 based on expected keywords
    num_tool_calls: int


def count_tokens_simple(text: str) -> int:
    """Simple token estimation (4 chars per token)."""
    return len(text) // 4


def run_agent_scenario(
    client: Any, scenario: AgentScenario, model: str = "gpt-4o-mini"
) -> BenchmarkResult:
    """Run a scenario and measure token usage."""

    messages = [
        {"role": "system", "content": scenario.system_prompt},
        {"role": "user", "content": scenario.user_query},
    ]

    # Add tool results with proper OpenAI format
    for tool_output in scenario.tools:
        tool_call_id = f"call_{hashlib.md5(tool_output['tool'].encode()).hexdigest()[:8]}"  # nosec B324
        # Assistant message with tool_calls (required by OpenAI)
        messages.append(
            {
                "role": "assistant",
                "content": None,
                "tool_calls": [
                    {
                        "id": tool_call_id,
                        "type": "function",
                        "function": {"name": tool_output["tool"], "arguments": "{}"},
                    }
                ],
            }
        )
        messages.append(
            {
                "role": "tool",
                "tool_call_id": tool_call_id,
                "content": json.dumps(tool_output["result"], indent=2),
            }
        )

    # Add final question
    messages.append(
        {
            "role": "user",
            "content": "Based on all this information, what's your analysis and recommendation?",
        }
    )

    # Count input tokens
    input_text = json.dumps(messages)
    input_tokens = count_tokens_simple(input_text)

    # Make API call
    start = time.time()

    # Determine if using HeadroomClient
    is_headroom = isinstance(client, HeadroomClient) if HEADROOM_AVAILABLE else False
    mode = "headroom" if is_headroom else "baseline"

    try:
        response = client.chat.completions.create(
            model=model,
            messages=messages,
            max_tokens=1000,
        )
        latency = (time.time() - start) * 1000

        answer = response.choices[0].message.content
        output_tokens = (
            response.usage.completion_tokens
            if hasattr(response, "usage")
            else count_tokens_simple(answer)
        )
        actual_input_tokens = (
            response.usage.prompt_tokens if hasattr(response, "usage") else input_tokens
        )

        # Calculate answer quality
        answer_lower = answer.lower()
        matches = sum(1 for kw in scenario.expected_answer_contains if kw.lower() in answer_lower)
        quality = matches / len(scenario.expected_answer_contains)

        # Estimate cost (gpt-4o-mini pricing)
        cost = (actual_input_tokens * 0.00015 + output_tokens * 0.0006) / 1000

    except Exception as e:
        print(f"   Error: {e}")
        return BenchmarkResult(
            scenario_name=scenario.name,
            mode=mode,
            total_input_tokens=input_tokens,
            total_output_tokens=0,
            total_tokens=input_tokens,
            cost_usd=0.0,
            latency_ms=0,
            answer_quality=0.0,
            num_tool_calls=len(scenario.tools),
        )

    return BenchmarkResult(
        scenario_name=scenario.name,
        mode=mode,
        total_input_tokens=actual_input_tokens,
        total_output_tokens=output_tokens,
        total_tokens=actual_input_tokens + output_tokens,
        cost_usd=cost,
        latency_ms=latency,
        answer_quality=quality,
        num_tool_calls=len(scenario.tools),
    )


def run_full_benchmark(api_key: str = None) -> dict:
    """Run complete benchmark comparing baseline vs Headroom."""

    if api_key is None:
        api_key = os.environ.get("OPENAI_API_KEY")

    if not api_key:
        raise ValueError("OPENAI_API_KEY required")

    if not HEADROOM_AVAILABLE:
        raise RuntimeError("Headroom not available")

    # Create clients
    import tempfile

    from openai import OpenAI

    baseline_client = OpenAI(api_key=api_key)

    # Headroom-wrapped client
    db_path = os.path.join(tempfile.gettempdir(), "headroom_benchmark.db")
    headroom_client = HeadroomClient(
        original_client=OpenAI(api_key=api_key),
        provider=OpenAIProvider(),
        store_url=f"sqlite:///{db_path}",
        default_mode="optimize",
    )

    scenarios = [
        create_sre_debugging_scenario(),
        create_codebase_exploration_scenario(),
        create_issue_triage_scenario(),
    ]

    results = []

    print("\n" + "=" * 70)
    print("REAL-WORLD AGENT BENCHMARK: MCP Tools + Headroom")
    print("=" * 70)

    for scenario in scenarios:
        print(f"\n{'=' * 60}")
        print(f"Scenario: {scenario.name}")
        print(f"Description: {scenario.description}")
        print(f"Tool calls: {len(scenario.tools)}")
        print(f"{'=' * 60}")

        # Estimate raw data size
        raw_size = sum(len(json.dumps(t["result"])) for t in scenario.tools)
        print(f"\nRaw tool output size: {raw_size:,} chars (~{raw_size // 4:,} tokens)")

        # Run baseline
        print("\n[1/2] Running BASELINE (no compression)...")
        baseline_result = run_agent_scenario(baseline_client, scenario)
        print(f"   Input tokens: {baseline_result.total_input_tokens:,}")
        print(f"   Output tokens: {baseline_result.total_output_tokens:,}")
        print(f"   Cost: ${baseline_result.cost_usd:.4f}")
        print(f"   Answer quality: {baseline_result.answer_quality:.1%}")
        results.append(baseline_result)

        # Run with Headroom
        print("\n[2/2] Running HEADROOM (optimized)...")
        headroom_result = run_agent_scenario(headroom_client, scenario)
        print(f"   Input tokens: {headroom_result.total_input_tokens:,}")
        print(f"   Output tokens: {headroom_result.total_output_tokens:,}")
        print(f"   Cost: ${headroom_result.cost_usd:.4f}")
        print(f"   Answer quality: {headroom_result.answer_quality:.1%}")
        results.append(headroom_result)

        # Calculate savings
        if baseline_result.total_input_tokens > 0:
            token_savings = 1 - (
                headroom_result.total_input_tokens / baseline_result.total_input_tokens
            )
            cost_savings = (
                1 - (headroom_result.cost_usd / baseline_result.cost_usd)
                if baseline_result.cost_usd > 0
                else 0
            )
            print("\n   📊 SAVINGS:")
            print(f"   Token reduction: {token_savings:.1%}")
            print(f"   Cost reduction: {cost_savings:.1%}")
            print(
                f"   Quality preserved: {'✓' if headroom_result.answer_quality >= baseline_result.answer_quality * 0.9 else '✗'}"
            )

    # Summary
    print("\n" + "=" * 70)
    print("BENCHMARK SUMMARY")
    print("=" * 70)

    baseline_results = [r for r in results if r.mode == "baseline"]
    headroom_results = [r for r in results if r.mode == "headroom"]

    total_baseline_tokens = sum(r.total_input_tokens for r in baseline_results)
    total_headroom_tokens = sum(r.total_input_tokens for r in headroom_results)
    total_baseline_cost = sum(r.cost_usd for r in baseline_results)
    total_headroom_cost = sum(r.cost_usd for r in headroom_results)

    print(f"\n{'Metric':<25} {'Baseline':>15} {'Headroom':>15} {'Savings':>15}")
    print("-" * 70)

    token_savings = (
        (1 - total_headroom_tokens / total_baseline_tokens) if total_baseline_tokens > 0 else 0
    )
    cost_savings = (1 - total_headroom_cost / total_baseline_cost) if total_baseline_cost > 0 else 0

    print(
        f"{'Total Input Tokens':<25} {total_baseline_tokens:>15,} {total_headroom_tokens:>15,} {token_savings:>14.1%}"
    )
    print(
        f"{'Total Cost':<25} ${total_baseline_cost:>14.4f} ${total_headroom_cost:>14.4f} {cost_savings:>14.1%}"
    )

    avg_baseline_quality = (
        sum(r.answer_quality for r in baseline_results) / len(baseline_results)
        if baseline_results
        else 0
    )
    avg_headroom_quality = (
        sum(r.answer_quality for r in headroom_results) / len(headroom_results)
        if headroom_results
        else 0
    )
    print(
        f"{'Avg Answer Quality':<25} {avg_baseline_quality:>14.1%} {avg_headroom_quality:>14.1%} {'preserved' if avg_headroom_quality >= avg_baseline_quality * 0.9 else 'degraded':>15}"
    )

    return {
        "baseline": [r.__dict__ for r in baseline_results],
        "headroom": [r.__dict__ for r in headroom_results],
        "summary": {
            "total_baseline_tokens": total_baseline_tokens,
            "total_headroom_tokens": total_headroom_tokens,
            "token_savings": token_savings,
            "total_baseline_cost": total_baseline_cost,
            "total_headroom_cost": total_headroom_cost,
            "cost_savings": cost_savings,
        },
    }


if __name__ == "__main__":
    results = run_full_benchmark()

    # Save results
    with open("real_world_benchmark_results.json", "w") as f:
        json.dump(results, f, indent=2)

    print("\nResults saved to real_world_benchmark_results.json")