benchmark.py

#!/usr/bin/env python3
"""
AgentFS benchmark — measures token savings of TOC-first vs full-read strategy.

Strategies compared:
  BASELINE   : agent reads every candidate file in full
  GREP       : agent greps content, reads only matching files
  TOC-FIRST  : agent reads all TOCs, then reads only files whose TOC matches
  ADAPTIVE   : reads all TOCs, checks hit rate — if term is too generic
               (hits >THRESHOLD% of files) falls back to grep for the second step

Usage:
  python3 benchmark.py <source_dir> [--tasks tasks.txt] [--threshold 0.5]

tasks.txt: one search term per line (e.g. "parse", "authentication", "database")
           defaults to a built-in task list
"""
import os
import sys
import re
import argparse
from dataclasses import dataclass

# Add project root to path
sys.path.insert(0, os.path.dirname(__file__))

from agentfs.indexer.parser import parse_file, SUPPORTED_EXTS, TEXT_EXTS
from agentfs.virtual.toc import generate_toc

ALL_EXTS = SUPPORTED_EXTS | TEXT_EXTS

DEFAULT_TASKS = [
    "parse",
    "authentication",
    "database",
    "error",
    "import",
    "class",
    "test",
    "config",
    "read",
    "write",
]


def count_tokens(text: str) -> int:
    """Approximate token count (GPT-style: ~1 token per 0.75 words)."""
    return max(1, int(len(text.split()) / 0.75))


@dataclass
class FileStats:
    path: str
    rel_path: str
    full_tokens: int
    toc_tokens: int
    toc_text: str
    full_text: str


def collect_files(source_dir: str) -> list[FileStats]:
    stats = []
    for root, dirs, files in os.walk(source_dir):
        dirs[:] = [d for d in dirs if not d.startswith('.') and d not in
                   ('node_modules', '__pycache__', '.git', 'venv', '.venv', 'dist', 'build')]
        for fname in files:
            ext = os.path.splitext(fname)[1].lower()
            if ext not in ALL_EXTS:
                continue
            full = os.path.join(root, fname)
            rel = os.path.relpath(full, source_dir)
            try:
                with open(full, 'r', encoding='utf-8', errors='replace') as f:
                    content = f.read()
            except OSError:
                continue
            if not content.strip():
                continue
            toc = generate_toc(full)
            stats.append(FileStats(
                path=full,
                rel_path=rel,
                full_tokens=count_tokens(content),
                toc_tokens=count_tokens(toc),
                toc_text=toc,
                full_text=content,
            ))
    return stats


SELECTIVITY_THRESHOLD = 0.5  # if TOC hits >50% of files, term is too generic


def run_task(task: str, files: list[FileStats], threshold: float = SELECTIVITY_THRESHOLD) -> dict:
    """
    Simulate finding files relevant to `task`.

    BASELINE  : reads ALL files in full
    GREP      : greps content, reads only matching files (best-case without TOC)
    TOC-FIRST : reads all TOCs, then full-reads only TOC-matched files
    ADAPTIVE  : reads all TOCs, checks hit rate:
                  - if hit_rate <= threshold → use TOC hits (term is specific)
                  - if hit_rate >  threshold → grep among TOC hits (term is generic)
    """
    term = task.lower()

    # --- BASELINE ---
    baseline_tokens = sum(f.full_tokens for f in files)

    # --- GREP ---
    grep_hits = [f for f in files if term in f.full_text.lower()]
    grep_tokens = sum(f.full_tokens for f in grep_hits)

    # --- TOC-FIRST ---
    toc_scan_tokens = sum(f.toc_tokens for f in files)
    toc_hits = [f for f in files if term in f.toc_text.lower()]
    toc_full_tokens = sum(f.full_tokens for f in toc_hits)
    toc_total_tokens = toc_scan_tokens + toc_full_tokens

    # --- ADAPTIVE ---
    toc_hit_rate = len(toc_hits) / max(len(files), 1)
    if toc_hit_rate <= threshold:
        # Term is specific: TOC filtering is effective
        adaptive_full_tokens = toc_full_tokens
        adaptive_strategy = 'toc'
    else:
        # Term is generic: grep among TOC hits to narrow down
        grep_among_toc = [f for f in toc_hits if term in f.full_text.lower()]
        adaptive_full_tokens = sum(f.full_tokens for f in grep_among_toc)
        adaptive_strategy = 'grep-fallback'
    adaptive_total = toc_scan_tokens + adaptive_full_tokens

    return {
        'task': task,
        'files_total': len(files),
        'baseline_tokens': baseline_tokens,
        'grep_hits': len(grep_hits),
        'grep_tokens': grep_tokens,
        'toc_hits': len(toc_hits),
        'toc_hit_rate': toc_hit_rate,
        'toc_scan_tokens': toc_scan_tokens,
        'toc_full_tokens': toc_full_tokens,
        'toc_total_tokens': toc_total_tokens,
        'adaptive_total': adaptive_total,
        'adaptive_strategy': adaptive_strategy,
        'savings_pct_vs_baseline': round((1 - toc_total_tokens / max(baseline_tokens, 1)) * 100, 1),
        'savings_pct_vs_grep': round((1 - toc_total_tokens / max(grep_tokens, 1)) * 100, 1),
        'adaptive_pct_vs_baseline': round((1 - adaptive_total / max(baseline_tokens, 1)) * 100, 1),
        'adaptive_pct_vs_grep': round((1 - adaptive_total / max(grep_tokens, 1)) * 100, 1),
    }


def print_file_stats(files: list[FileStats]):
    total_full = sum(f.full_tokens for f in files)
    total_toc = sum(f.toc_tokens for f in files)
    avg_savings = round((1 - total_toc / max(total_full, 1)) * 100, 1)

    print("\n=== FILE-LEVEL TOKEN COMPRESSION ===")
    print(f"{'File':<45} {'Full':>7} {'TOC':>7} {'Saved':>7}")
    print("-" * 70)
    for f in sorted(files, key=lambda x: x.full_tokens, reverse=True):
        saved = round((1 - f.toc_tokens / max(f.full_tokens, 1)) * 100)
        print(f"{f.rel_path:<45} {f.full_tokens:>7,} {f.toc_tokens:>7,} {saved:>6}%")
    print("-" * 70)
    print(f"{'TOTAL':<45} {total_full:>7,} {total_toc:>7,} {avg_savings:>6}%")


def print_task_results(results: list[dict], threshold: float):
    print("\n\n=== TASK-LEVEL TOKEN USAGE ===")
    print(f"{'Task':<18} {'HitRate':>8} {'Strat':>13} {'Baseline':>10} {'Grep':>10} {'TOC':>10} {'Adaptive':>10} {'Adp/Base':>10} {'Adp/Grep':>10}")
    print("-" * 103)
    for r in results:
        hit_rate = f"{r['toc_hit_rate']:.0%}"
        strat = f"[{r['adaptive_strategy']}]"
        print(
            f"{r['task']:<18} {hit_rate:>8} {strat:>13}"
            f" {r['baseline_tokens']:>10,} {r['grep_tokens']:>10,}"
            f" {r['toc_total_tokens']:>10,} {r['adaptive_total']:>10,}"
            f" {r['adaptive_pct_vs_baseline']:>+9.0f}% {r['adaptive_pct_vs_grep']:>+9.0f}%"
        )
    print("-" * 103)

    avg_adp_base = sum(r['adaptive_pct_vs_baseline'] for r in results) / len(results)
    avg_adp_grep = sum(r['adaptive_pct_vs_grep'] for r in results) / len(results)
    avg_toc_base = sum(r['savings_pct_vs_baseline'] for r in results) / len(results)
    avg_toc_grep = sum(r['savings_pct_vs_grep'] for r in results) / len(results)
    total_b = sum(r['baseline_tokens'] for r in results)
    total_g = sum(r['grep_tokens'] for r in results)
    total_t = sum(r['toc_total_tokens'] for r in results)
    total_a = sum(r['adaptive_total'] for r in results)
    print(
        f"{'AVERAGE':<18} {'':>8} {'':>13}"
        f" {total_b:>10,} {total_g:>10,} {total_t:>10,} {total_a:>10,}"
        f" {avg_adp_base:>+9.1f}% {avg_adp_grep:>+9.1f}%"
    )

    print(f"""
=== VERDICT ===
TOC-first  vs baseline : {avg_toc_base:+.1f}%
TOC-first  vs grep     : {avg_toc_grep:+.1f}%
Adaptive   vs baseline : {avg_adp_base:+.1f}%
Adaptive   vs grep     : {avg_adp_grep:+.1f}%

Adaptive strategy (threshold={threshold:.0%}):
  - [toc]           term hit rate <= {threshold:.0%} → use TOC hits directly (term is specific)
  - [grep-fallback] term hit rate >  {threshold:.0%} → grep among TOC hits (term is generic)

Notes:
- Token counts are approximate (~0.75 words/token)
- Grep cost assumes grep is free (only counts full-file reads of matches)
""")


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument('source_dir', help='Directory to benchmark')
    parser.add_argument('--tasks', help='File with one task/term per line')
    parser.add_argument('--threshold', type=float, default=SELECTIVITY_THRESHOLD,
                        help=f'TOC hit rate above which adaptive falls back to grep (default: {SELECTIVITY_THRESHOLD})')
    args = parser.parse_args()

    source_dir = os.path.realpath(args.source_dir)
    if not os.path.isdir(source_dir):
        print(f"Not a directory: {source_dir}", file=sys.stderr)
        sys.exit(1)

    if args.tasks:
        with open(args.tasks) as f:
            tasks = [l.strip() for l in f if l.strip() and not l.startswith('#')]
    else:
        tasks = DEFAULT_TASKS

    print(f"Benchmarking: {source_dir}")
    print(f"Tasks: {tasks}")

    files = collect_files(source_dir)
    if not files:
        print("No supported files found.")
        sys.exit(1)

    print_file_stats(files)

    results = [run_task(task, files, threshold=args.threshold) for task in tasks]
    print_task_results(results, threshold=args.threshold)


if __name__ == '__main__':
    main()