tirth8205--code-review-graph/code_review_graph/eval/token_benchmark.py

"""Measures total tokens consumed by agent workflows against benchmark repos."""

from __future__ import annotations

import json
import logging
from typing import Any, Callable

logger = logging.getLogger(__name__)


def estimate_tokens(obj: Any) -> int:
    """Estimate token count from JSON-serializable object.

    Uses character count / 4 as a rough approximation for English + code.
    """
    return len(json.dumps(obj, default=str)) // 4


def benchmark_review_workflow(repo_root: str, base: str = "HEAD~1") -> dict:
    """Simulate a review workflow and measure total tokens consumed."""
    from ..tools.context import get_minimal_context
    from ..tools.review import detect_changes_func

    total_tokens = 0
    calls = []

    # Step 1: get_minimal_context
    result = get_minimal_context(task="review changes", repo_root=repo_root, base=base)
    tokens = estimate_tokens(result)
    total_tokens += tokens
    calls.append({"tool": "get_minimal_context", "tokens": tokens})

    # Step 2: detect_changes (minimal)
    result = detect_changes_func(base=base, repo_root=repo_root, detail_level="minimal")
    tokens = estimate_tokens(result)
    total_tokens += tokens
    calls.append({"tool": "detect_changes_minimal", "tokens": tokens})

    return {
        "workflow": "review",
        "total_tokens": total_tokens,
        "tool_calls": len(calls),
        "calls": calls,
    }


def benchmark_architecture_workflow(repo_root: str) -> dict:
    """Simulate an architecture exploration workflow."""
    from ..tools.community_tools import list_communities_func
    from ..tools.context import get_minimal_context
    from ..tools.flows_tools import list_flows

    total_tokens = 0
    calls = []

    result = get_minimal_context(task="map architecture", repo_root=repo_root)
    tokens = estimate_tokens(result)
    total_tokens += tokens
    calls.append({"tool": "get_minimal_context", "tokens": tokens})

    result = list_communities_func(repo_root=repo_root, detail_level="minimal")
    tokens = estimate_tokens(result)
    total_tokens += tokens
    calls.append({"tool": "list_communities_minimal", "tokens": tokens})

    result = list_flows(repo_root=repo_root, detail_level="minimal")
    tokens = estimate_tokens(result)
    total_tokens += tokens
    calls.append({"tool": "list_flows_minimal", "tokens": tokens})

    return {
        "workflow": "architecture",
        "total_tokens": total_tokens,
        "tool_calls": len(calls),
        "calls": calls,
    }


def benchmark_debug_workflow(repo_root: str) -> dict:
    """Simulate a debug workflow."""
    from ..tools.context import get_minimal_context
    from ..tools.query import semantic_search_nodes

    total_tokens = 0
    calls = []

    result = get_minimal_context(task="debug login bug", repo_root=repo_root)
    tokens = estimate_tokens(result)
    total_tokens += tokens
    calls.append({"tool": "get_minimal_context", "tokens": tokens})

    result = semantic_search_nodes(
        query="login", repo_root=repo_root, detail_level="minimal",
    )
    tokens = estimate_tokens(result)
    total_tokens += tokens
    calls.append({"tool": "semantic_search_minimal", "tokens": tokens})

    return {
        "workflow": "debug",
        "total_tokens": total_tokens,
        "tool_calls": len(calls),
        "calls": calls,
    }


def benchmark_onboard_workflow(repo_root: str) -> dict:
    """Simulate an onboarding workflow."""
    from ..tools.context import get_minimal_context
    from ..tools.query import list_graph_stats

    total_tokens = 0
    calls = []

    result = get_minimal_context(task="onboard developer", repo_root=repo_root)
    tokens = estimate_tokens(result)
    total_tokens += tokens
    calls.append({"tool": "get_minimal_context", "tokens": tokens})

    result = list_graph_stats(repo_root=repo_root)
    tokens = estimate_tokens(result)
    total_tokens += tokens
    calls.append({"tool": "list_graph_stats", "tokens": tokens})

    return {
        "workflow": "onboard",
        "total_tokens": total_tokens,
        "tool_calls": len(calls),
        "calls": calls,
    }


def benchmark_pre_merge_workflow(repo_root: str, base: str = "HEAD~1") -> dict:
    """Simulate a pre-merge check workflow."""
    from ..tools.context import get_minimal_context
    from ..tools.review import detect_changes_func

    total_tokens = 0
    calls = []

    result = get_minimal_context(task="pre-merge check", repo_root=repo_root, base=base)
    tokens = estimate_tokens(result)
    total_tokens += tokens
    calls.append({"tool": "get_minimal_context", "tokens": tokens})

    result = detect_changes_func(base=base, repo_root=repo_root, detail_level="minimal")
    tokens = estimate_tokens(result)
    total_tokens += tokens
    calls.append({"tool": "detect_changes_minimal", "tokens": tokens})

    return {
        "workflow": "pre_merge",
        "total_tokens": total_tokens,
        "tool_calls": len(calls),
        "calls": calls,
    }


ALL_WORKFLOWS: dict[str, Callable[..., dict]] = {
    "review": benchmark_review_workflow,
    "architecture": benchmark_architecture_workflow,
    "debug": benchmark_debug_workflow,
    "onboard": benchmark_onboard_workflow,
    "pre_merge": benchmark_pre_merge_workflow,
}


def run_all_benchmarks(repo_root: str, base: str = "HEAD~1") -> list[dict]:
    """Run all workflow benchmarks and return results."""
    results = []
    for name, fn in ALL_WORKFLOWS.items():
        try:
            if "base" in fn.__code__.co_varnames:
                result = fn(repo_root=repo_root, base=base)
            else:
                result = fn(repo_root=repo_root)
            results.append(result)
        except Exception as e:
            logger.warning("Benchmark %s failed: %s", name, e)
            results.append({"workflow": name, "error": str(e)})
    return results