chopratejas--headroom/benchmarks/proxy_mode_benchmark.py

#!/usr/bin/env python3
"""Local benchmark for proxy run modes (no API calls).

Compares:
- baseline: no compression
- token mode: prioritize compression
- cache mode: preserve prior-turn prefix stability

Includes an optional real-test harness printout for Claude Code, but does not
invoke external APIs unless the user does so manually.
"""

from __future__ import annotations

import argparse
import copy
import json
import logging
from dataclasses import dataclass
from typing import Any

from headroom.cache.compression_cache import CompressionCache
from headroom.cache.prefix_tracker import PrefixCacheTracker
from headroom.proxy.handlers.anthropic import AnthropicHandlerMixin
from headroom.proxy.models import ProxyConfig
from headroom.proxy.modes import PROXY_MODE_CACHE, PROXY_MODE_TOKEN
from headroom.proxy.server import HeadroomProxy
from headroom.tokenizers import get_tokenizer
from headroom.utils import extract_user_query

MODEL = "claude-sonnet-4-6"


@dataclass
class ModeBenchmarkResult:
    mode: str
    total_original_tokens: int = 0
    total_sent_tokens: int = 0
    total_tokens_saved: int = 0
    total_cache_read_tokens: int = 0
    total_cache_write_tokens: int = 0
    total_uncached_tokens: int = 0

    @property
    def compression_pct(self) -> float:
        if self.total_original_tokens <= 0:
            return 0.0
        return self.total_tokens_saved / self.total_original_tokens * 100.0

    @property
    def cache_hit_pct(self) -> float:
        total = (
            self.total_cache_read_tokens
            + self.total_cache_write_tokens
            + self.total_uncached_tokens
        )
        if total <= 0:
            return 0.0
        return self.total_cache_read_tokens / total * 100.0


def _build_tool_result(turn: int, rows: int = 240) -> str:
    payload = []
    for i in range(rows):
        payload.append(
            {
                "id": f"{turn:02d}-{i:04d}",
                "status": "ok" if i % 37 else "warning",
                "service": "auth-api" if i % 2 else "gateway",
                "latency_ms": 100 + (i % 13),
                "hint": "retry with exponential backoff" if i % 89 == 0 else "none",
            }
        )
    return json.dumps(payload)


def _build_conversation(turn: int) -> list[dict[str, Any]]:
    messages: list[dict[str, Any]] = []
    for t in range(1, turn):
        messages.extend(
            [
                {
                    "role": "user",
                    "content": f"Analyze tool output turn {t} and summarize anomalies.",
                },
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "tool_result",
                            "tool_use_id": f"tool-{t}",
                            "content": _build_tool_result(t),
                        }
                    ],
                },
                {"role": "assistant", "content": f"Turn {t} acknowledged."},
            ]
        )
    # Current turn: user request + fresh tool output, no assistant response yet.
    messages.extend(
        [
            {
                "role": "user",
                "content": f"Analyze tool output turn {turn} and summarize anomalies.",
            },
            {
                "role": "user",
                "content": [
                    {
                        "type": "tool_result",
                        "tool_use_id": f"tool-{turn}",
                        "content": _build_tool_result(turn),
                    }
                ],
            },
        ]
    )
    return messages


def _common_prefix_tokens(
    prev: list[dict[str, Any]], curr: list[dict[str, Any]], tokenizer: Any
) -> tuple[int, list[int]]:
    common = 0
    counts: list[int] = []
    for msg in curr:
        counts.append(tokenizer.count_message(msg))
    for i, (a, b) in enumerate(zip(prev, curr)):
        if a != b:
            break
        common += counts[i]
    return common, counts


def _make_proxy(mode: str) -> HeadroomProxy:
    cfg = ProxyConfig(
        mode=mode,
        optimize=True,
        image_optimize=False,
        smart_routing=False,
        code_aware_enabled=False,
        read_lifecycle=False,
        cache_enabled=False,
        rate_limit_enabled=False,
        cost_tracking_enabled=False,
        log_requests=False,
        ccr_inject_tool=False,
        ccr_handle_responses=False,
        ccr_context_tracking=False,
    )
    return HeadroomProxy(cfg)


def _simulate_mode(turns: int, mode: str) -> ModeBenchmarkResult:
    tokenizer = get_tokenizer(MODEL)
    result = ModeBenchmarkResult(mode=mode)

    if mode == "baseline":
        prev_forwarded: list[dict[str, Any]] = []
        for turn in range(1, turns + 1):
            messages = _build_conversation(turn)
            before = tokenizer.count_messages(messages)
            common, counts = _common_prefix_tokens(prev_forwarded, messages, tokenizer)
            uncached = max(0, before - common)

            result.total_original_tokens += before
            result.total_sent_tokens += before
            result.total_cache_read_tokens += common
            result.total_cache_write_tokens += 0
            result.total_uncached_tokens += uncached
            prev_forwarded = copy.deepcopy(messages)
        return result

    proxy = _make_proxy(mode)
    prefix_tracker = PrefixCacheTracker("anthropic")
    comp_cache = CompressionCache()
    prev_forwarded = []

    for turn in range(1, turns + 1):
        messages = _build_conversation(turn)
        before = tokenizer.count_messages(messages)

        frozen = prefix_tracker.get_frozen_message_count()
        if mode == PROXY_MODE_CACHE:
            frozen = AnthropicHandlerMixin._strict_previous_turn_frozen_count(messages, frozen)

        working = messages
        if mode == PROXY_MODE_TOKEN:
            working = comp_cache.apply_cached(messages)
            frozen = min(frozen, comp_cache.compute_frozen_count(messages))

        context_limit = proxy.anthropic_provider.get_context_limit(MODEL)
        pipeline_result = proxy.anthropic_pipeline.apply(
            messages=working,
            model=MODEL,
            model_limit=context_limit,
            context=extract_user_query(working),
            frozen_message_count=frozen,
        )
        forwarded = pipeline_result.messages

        if mode == PROXY_MODE_TOKEN:
            comp_cache.update_from_result(messages, forwarded)
        if mode == PROXY_MODE_CACHE:
            forwarded, _ = AnthropicHandlerMixin._restore_frozen_prefix(
                messages, forwarded, frozen_message_count=frozen
            )

        after = tokenizer.count_messages(forwarded)
        common, msg_counts = _common_prefix_tokens(prev_forwarded, forwarded, tokenizer)
        uncached = max(0, after - common)

        result.total_original_tokens += before
        result.total_sent_tokens += after
        result.total_tokens_saved += max(0, before - after)
        result.total_cache_read_tokens += common
        result.total_uncached_tokens += uncached

        prefix_tracker.update_from_response(
            cache_read_tokens=common,
            cache_write_tokens=uncached,
            messages=forwarded,
            message_token_counts=msg_counts,
        )
        result.total_cache_write_tokens += uncached
        prev_forwarded = copy.deepcopy(forwarded)

    return result


def run_local_benchmark(turns: int = 12) -> dict[str, ModeBenchmarkResult]:
    return {
        "baseline": _simulate_mode(turns, "baseline"),
        PROXY_MODE_TOKEN: _simulate_mode(turns, PROXY_MODE_TOKEN),
        PROXY_MODE_CACHE: _simulate_mode(turns, PROXY_MODE_CACHE),
    }


def _print_results(results: dict[str, ModeBenchmarkResult]) -> None:
    print(
        "\nMode benchmark (higher compression + higher cache_hit is better for total cost):\n"
        "mode      orig_tok   sent_tok   saved_tok   compression   cache_hit   uncached_tok"
    )
    for key in ("baseline", PROXY_MODE_TOKEN, PROXY_MODE_CACHE):
        r = results[key]
        print(
            f"{r.mode:<9} {r.total_original_tokens:>9,} {r.total_sent_tokens:>10,} "
            f"{r.total_tokens_saved:>10,} {r.compression_pct:>10.1f}% "
            f"{r.cache_hit_pct:>9.1f}% {r.total_uncached_tokens:>12,}"
        )

    token = results[PROXY_MODE_TOKEN]
    cache = results[PROXY_MODE_CACHE]
    print("\nDelta (cache - token):")
    print(f"  cache_hit_pct: {cache.cache_hit_pct - token.cache_hit_pct:+.1f}%")
    print(f"  compression_pct: {cache.compression_pct - token.compression_pct:+.1f}%")
    print(f"  uncached_tokens: {cache.total_uncached_tokens - token.total_uncached_tokens:+,}")


def _print_real_harness() -> None:
    print("\nReal test harness (manual; optional, not executed by this benchmark):")
    print("  1) Start proxy in cache mode:  HEADROOM_MODE=cache headroom proxy --port 8787")
    print("  2) Start proxy in token mode:  HEADROOM_MODE=token headroom proxy --port 8787")
    print("  3) Run Claude Code against each:")
    print("     ANTHROPIC_BASE_URL=http://localhost:8787 claude")
    print("  4) Compare /stats prefix_cache and compression sections per run.")


def main() -> None:
    logging.getLogger("headroom").setLevel(logging.WARNING)
    logging.getLogger("sentence_transformers").setLevel(logging.WARNING)
    logging.getLogger("huggingface_hub").setLevel(logging.WARNING)

    parser = argparse.ArgumentParser(description="Local benchmark for proxy token/cache modes")
    parser.add_argument("--turns", type=int, default=12, help="Conversation turns to simulate")
    parser.add_argument(
        "--show-real-harness",
        action="store_true",
        help="Print manual steps for optional Claude Code real testing",
    )
    args = parser.parse_args()

    results = run_local_benchmark(turns=args.turns)
    _print_results(results)
    if args.show_real_harness:
        _print_real_harness()


if __name__ == "__main__":
    main()