benchmark.py

"""
Model and Mode Benchmark Script for SQL Query Writer Agent

This script benchmarks different LLM models AND query modes to help you
select the best configuration for the competition.

Usage:
    python benchmark.py                       # Benchmark current model with all modes
    python benchmark.py --models llama3.3     # Benchmark specific model
    python benchmark.py --modes fast,adaptive # Benchmark specific modes
    python benchmark.py --quick               # Quick benchmark (fewer tests)
    python benchmark.py --detailed            # Detailed output
    python benchmark.py --markdown            # Output markdown table for README

The script tests:
1. Query modes: fast, adaptive, accurate
2. Accuracy by difficulty: easy, medium, hard
3. Average response time
"""

import os
import sys
import time
import json
import urllib.request
import urllib.error
import duckdb
from datetime import datetime
from typing import List, Dict, Optional
from dataclasses import dataclass, field
from collections import defaultdict


def load_env():
    """Load .env file if it exists (same logic as agent.py)."""
    from pathlib import Path
    env_path = Path(__file__).parent / '.env'
    if env_path.exists():
        with open(env_path, encoding='utf-8') as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith('#') and '=' in line:
                    key, value = line.split('=', 1)
                    key = key.strip()
                    value = value.strip()
                    if key not in os.environ:
                        os.environ[key] = value


load_env()


# =============================================================================
# BENCHMARK CONFIGURATION
# =============================================================================

# Default models to benchmark
DEFAULT_MODELS = [
    'llama3.2',      # Smaller, faster
    'llama3.3',      # Larger, more accurate (recommended)
    'qwen2.5-coder', # Specialized for code
]

# Query modes to benchmark
QUERY_MODES = ['fast', 'standard', 'adaptive', 'accurate']

# Benchmark test cases - curated for evaluation coverage
BENCHMARK_CASES = [
    # Easy (should all pass)
    {"question": "How many customers are there?", "keywords": ["count", "customers"], "difficulty": "easy"},
    {"question": "List all stores", "keywords": ["stores"], "difficulty": "easy"},
    {"question": "What are the top 5 most expensive products?", "keywords": ["products", "order by", "limit 5"], "difficulty": "easy"},
    
    # Medium (most should pass)
    {"question": "Show products with their brand names", "keywords": ["products", "brands", "join"], "difficulty": "medium"},
    {"question": "Total revenue by store", "keywords": ["sum", "stores", "group by"], "difficulty": "medium"},
    {"question": "How many orders were placed in 2018?", "keywords": ["count", "orders", "2018"], "difficulty": "medium"},
    {"question": "Top 3 best selling products", "keywords": ["products", "sum", "order by", "limit 3"], "difficulty": "medium"},
    {"question": "Find orders that have not been shipped", "keywords": ["orders", "null"], "difficulty": "medium"},
    
    # Hard (differentiators)
    {"question": "Which products have never been ordered?", "keywords": ["products", "order_items"], "difficulty": "hard"},
    {"question": "Total revenue by brand", "keywords": ["brands", "sum", "group by"], "difficulty": "hard"},
    {"question": "Monthly revenue trend for 2018", "keywords": ["orders", "2018", "group by"], "difficulty": "hard"},
    {"question": "Show staff members and their managers", "keywords": ["staffs"], "difficulty": "hard"},
    {"question": "Products with above-average price", "keywords": ["products", "avg"], "difficulty": "hard"},
]

QUICK_BENCHMARK_CASES = BENCHMARK_CASES[:6]  # Subset for quick testing


def _database_ready(db_path: str = "bike_store.db") -> tuple[bool, str]:
    """Check whether the local DuckDB file exists and has tables."""
    if not os.path.exists(db_path):
        return False, f"{db_path} not found"
    try:
        con = duckdb.connect(database=db_path, read_only=True)
        tables = con.execute("SHOW TABLES").fetchall()
        con.close()
        if not tables:
            return False, f"{db_path} exists but has no tables"
        return True, f"{len(tables)} tables"
    except Exception as exc:
        return False, str(exc)


def _normalize_model_tag(name: str) -> str:
    """Normalize model tags to a lowercase base name."""
    return (name or "").split(":", 1)[0].strip().lower()


def _model_available(model_name: str, available_tags: List[str]) -> bool:
    """Return True if model_name exists in available_tags (base-name tolerant)."""
    target = _normalize_model_tag(model_name)
    available = {_normalize_model_tag(tag) for tag in available_tags}
    return target in available


def _preferred_local_model(available_tags: List[str]) -> Optional[str]:
    """Choose a preferred local model tag from installed local models."""
    preferred_order = [
        "llama3.3",
        "llama3.2",
        "qwen3",
        "qwen2.5-coder",
        "mixtral",
        "mistral",
        "gemma3",
    ]
    tag_map = {_normalize_model_tag(tag): tag for tag in available_tags}
    for preferred in preferred_order:
        if preferred in tag_map:
            return tag_map[preferred]
    return available_tags[0] if available_tags else None


def _probe_ollama_tags(host: str, api_key: str = "", timeout_seconds: float = 6.0) -> tuple[bool, list[str], str]:
    """
    Probe an Ollama-compatible host for /api/tags.

    Returns:
        (ok, model_tags, error_message)
    """
    host = (host or "").strip().rstrip("/")
    if not host:
        return False, [], "host is empty"

    request = urllib.request.Request(f"{host}/api/tags", method="GET")
    if api_key:
        request.add_header("x-api-key", api_key)

    try:
        opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))
        with opener.open(request, timeout=timeout_seconds) as response:
            payload = json.loads(response.read().decode("utf-8"))
        models = payload.get("models") or []
        tags = [m.get("name", "").strip() for m in models if isinstance(m, dict) and m.get("name")]
        return True, [t for t in tags if t], ""
    except urllib.error.HTTPError as exc:
        return False, [], f"HTTP {exc.code}: {exc.reason}"
    except Exception as exc:
        return False, [], str(exc)


def resolve_runtime_for_benchmark(requested_model: str) -> tuple[str, str, list[str], bool]:
    """
    Resolve a working runtime host/model for benchmarking.

    Strategy:
    1) Try configured OLLAMA_HOST directly.
    2) If host is unreachable/unusable, fall back to local Ollama.
    3) If requested model is missing, choose a preferred installed model.

    Returns:
        (resolved_host, resolved_model, notes, runtime_ready)
    """
    notes: list[str] = []
    configured_host = os.getenv('OLLAMA_HOST', 'http://localhost:11434').strip()
    api_key = (os.getenv('OLLAMA_API_KEY') or os.getenv('RCS_API_KEY') or "").strip()

    host_ok, host_models, host_error = _probe_ollama_tags(configured_host, api_key=api_key)

    if host_ok:
        resolved_model = requested_model
        if host_models and not _model_available(requested_model, host_models):
            preferred = _preferred_local_model(host_models)
            if preferred:
                warn_msg = (
                    f"WARNING: Requested model '{requested_model}' not found on server!\n"
                    f"  Available models: {', '.join(sorted(host_models))}\n"
                    f"  Falling back to: '{preferred}'"
                )
                print(f"\n{'!'*60}")
                print(warn_msg)
                print(f"{'!'*60}\n")
                notes.append(warn_msg)
                resolved_model = preferred
        return configured_host, resolved_model, notes, True

    allow_fallback = os.getenv("BENCHMARK_ALLOW_FALLBACK", "false").strip().lower() in (
        "1",
        "true",
        "yes",
        "on",
    )

    if not allow_fallback:
        notes.append(
            f"Configured host '{configured_host}' unavailable ({host_error}). "
            "Set BENCHMARK_ALLOW_FALLBACK=true to fall back to local Ollama."
        )
        return configured_host, requested_model, notes, False

    # Fallback to local Ollama if configured host is unavailable.
    local_host = 'http://localhost:11434'
    local_ok, local_models, local_error = _probe_ollama_tags(local_host, api_key="")
    if local_ok and local_models:
        resolved_model = requested_model
        if not _model_available(requested_model, local_models):
            preferred = _preferred_local_model(local_models)
            if preferred:
                notes.append(
                    f"Configured host '{configured_host}' unavailable ({host_error}); "
                    f"falling back to local host '{local_host}' and model '{preferred}'."
                )
                resolved_model = preferred
        else:
            notes.append(
                f"Configured host '{configured_host}' unavailable ({host_error}); "
                f"falling back to local host '{local_host}'."
            )
        return local_host, resolved_model, notes, True

    notes.append(
        f"Configured host '{configured_host}' unavailable ({host_error}) and local host '{local_host}' "
        f"is unavailable ({local_error})."
    )
    return configured_host, requested_model, notes, False


@dataclass
class BenchmarkResult:
    """Results from benchmarking a single model/mode combination."""
    model: str
    mode: str = 'adaptive'
    total_tests: int = 0
    passed: int = 0
    failed: int = 0
    total_time: float = 0.0
    avg_time: float = 0.0
    by_difficulty: Dict = field(default_factory=dict)
    errors: List[str] = field(default_factory=list)
    consistency_score: float = 0.0  # How consistent across runs
    
    @property
    def accuracy(self) -> float:
        return self.passed / self.total_tests * 100 if self.total_tests > 0 else 0
    
    @property
    def score(self) -> float:
        """Calculate overall score (accuracy weighted + speed bonus)."""
        # Accuracy is most important (70%)
        # Speed bonus (20%) - faster is better
        # Consistency (10%)
        speed_score = max(0, 100 - self.avg_time * 10)  # Lower time = higher score
        return (self.accuracy * 0.7 + 
                speed_score * 0.2 + 
                self.consistency_score * 0.1)
    
    @property
    def name(self) -> str:
        """Get display name for this benchmark."""
        return f"{self.model} ({self.mode})"


MAX_TEST_TIMEOUT = 120  # seconds per test before we skip it


def _run_with_timeout(func, timeout_sec):
    """Run func in a thread with a timeout. Returns (result, timed_out)."""
    import threading
    result_holder = [None]
    error_holder = [None]

    def target():
        try:
            result_holder[0] = func()
        except Exception as e:
            error_holder[0] = e

    t = threading.Thread(target=target, daemon=True)
    t.start()
    t.join(timeout=timeout_sec)
    if t.is_alive():
        return None, True, None
    if error_holder[0]:
        return None, False, error_holder[0]
    return result_holder[0], False, None


def run_single_test(agent, question: str, keywords: List[str], test_index: int = 0, total_tests: int = 0) -> Dict:
    """Run a single test and return results."""
    from src.utils import extract_sql
    
    prefix = f"  [{test_index}/{total_tests}]" if total_tests else "  "
    print(f"{prefix} Testing: {question[:60]}...", end=" ", flush=True)
    
    start_time = time.time()
    result = {
        'passed': False,
        'time': 0,
        'sql': None,
        'error': None
    }
    
    try:
        sql, timed_out, gen_err = _run_with_timeout(
            lambda: agent.generate_query(question), MAX_TEST_TIMEOUT
        )
        if timed_out:
            result['time'] = time.time() - start_time
            result['error'] = f"TIMEOUT (>{MAX_TEST_TIMEOUT}s)"
            print(f"TIMEOUT ({result['time']:.0f}s)")
            return result
        if gen_err:
            raise gen_err
        result['sql'] = sql
        result['time'] = time.time() - start_time
        
        # Validate syntax
        is_valid, error = agent.validator.validate(sql)
        if not is_valid:
            result['error'] = f"Syntax: {error}"
            print(f"FAIL ({result['time']:.1f}s) - {result['error']}")
            return result
        
        # Check keywords
        sql_lower = sql.lower()
        missing = [kw for kw in keywords if kw.lower() not in sql_lower]
        if missing:
            result['error'] = f"Missing keywords: {missing}"
            print(f"FAIL ({result['time']:.1f}s) - {result['error']}")
            return result
        
        # Execute the query against the actual database to verify it runs
        try:
            con = duckdb.connect(database='bike_store.db', read_only=True)
            rows = con.execute(sql).fetchall()
            con.close()
            if rows is None or (len(rows) == 0):
                result['error'] = "Execution: query returned 0 rows"
                print(f"WARN ({result['time']:.1f}s) - {result['error']}")
                # Still pass if syntax and keywords are correct but flag it
                result['passed'] = True
            else:
                result['passed'] = True
                print(f"PASS ({result['time']:.1f}s, {len(rows)} rows)")
        except Exception as exec_err:
            result['error'] = f"Execution error: {exec_err}"
            print(f"FAIL ({result['time']:.1f}s) - {result['error']}")
            return result
        
    except Exception as e:
        result['time'] = time.time() - start_time
        result['error'] = str(e)
        print(f"ERROR ({result['time']:.1f}s) - {result['error'][:60]}")
    
    return result


def benchmark_model(model: str, test_cases: List[Dict], mode: str = 'adaptive', runs: int = 1) -> BenchmarkResult:
    """Benchmark a single model with a specific query mode."""
    from agent import QueryWriter

    resolved_host, resolved_model, runtime_notes, runtime_ready = resolve_runtime_for_benchmark(model)
    os.environ['OLLAMA_HOST'] = resolved_host
    os.environ['OLLAMA_MODEL'] = resolved_model
    os.environ['QUERY_MODE'] = mode

    result = BenchmarkResult(model=resolved_model, mode=mode)
    result.by_difficulty = {'easy': {'passed': 0, 'total': 0},
                           'medium': {'passed': 0, 'total': 0},
                           'hard': {'passed': 0, 'total': 0}}
    for note in runtime_notes:
        result.errors.append(f"runtime: {note}")

    if not runtime_ready:
        result.errors.append("runtime: unable to reach any Ollama endpoint for benchmarking.")
        return result

    db_ok, db_msg = _database_ready("bike_store.db")
    if not db_ok:
        result.errors.append(
            "runtime: local database is not ready. "
            f"Details: {db_msg}. Initialize it once with `python main.py`."
        )
        return result
    
    try:
        # Initialize
        agent = QueryWriter('bike_store.db')
        
        all_results = []
        
        for run in range(runs):
            run_results = []
            
            for idx, tc in enumerate(test_cases, 1):
                test_result = run_single_test(agent, tc['question'], tc['keywords'],
                                              test_index=idx, total_tests=len(test_cases))
                run_results.append(test_result)
                
                result.total_tests += 1
                result.total_time += test_result['time']
                
                difficulty = tc['difficulty']
                result.by_difficulty[difficulty]['total'] += 1
                
                if test_result['passed']:
                    result.passed += 1
                    result.by_difficulty[difficulty]['passed'] += 1
                else:
                    result.failed += 1
                    if test_result['error']:
                        result.errors.append(f"{tc['question'][:40]}: {test_result['error']}")
            
            all_results.append(run_results)
        
        result.avg_time = result.total_time / result.total_tests if result.total_tests > 0 else 0
        
        # Calculate consistency (if multiple runs)
        if runs > 1:
            consistent_count = 0
            for i, tc in enumerate(test_cases):
                outcomes = [r[i]['passed'] for r in all_results]
                if all(outcomes) or not any(outcomes):
                    consistent_count += 1
            result.consistency_score = consistent_count / len(test_cases) * 100
        else:
            result.consistency_score = 100  # Assume consistent if only one run
        
    except Exception as e:
        result.errors.append(f"Model initialization error: {e}")
    
    return result


def print_results(results: List[BenchmarkResult], detailed: bool = False):
    """Print benchmark results in a formatted table."""
    print("\n" + "=" * 80)
    print("BENCHMARK RESULTS")
    print("=" * 80)
    
    # Sort by score
    results = sorted(results, key=lambda r: r.score, reverse=True)
    
    print(f"\n{'Model':<20} {'Mode':<10} {'Accuracy':>10} {'Avg Time':>10} {'Score':>10}")
    print("-" * 65)
    
    for r in results:
        print(f"{r.model:<20} {r.mode:<10} {r.accuracy:>9.1f}% {r.avg_time:>9.2f}s {r.score:>10.1f}")
    
    # Print detailed breakdown
    if detailed:
        print("\n" + "-" * 65)
        print("BREAKDOWN BY DIFFICULTY")
        print("-" * 65)
        
        for r in results:
            print(f"\n{r.name}:")
            for diff in ['easy', 'medium', 'hard']:
                if diff in r.by_difficulty:
                    stats = r.by_difficulty[diff]
                    if stats['total'] > 0:
                        pct = stats['passed'] / stats['total'] * 100
                        print(f"  {diff:>8}: {stats['passed']}/{stats['total']} ({pct:.0f}%)")
        
        # Print errors
        print("\n" + "-" * 65)
        print("ERRORS (first 5 per configuration)")
        print("-" * 65)
        
        for r in results:
            if r.errors:
                print(f"\n{r.name}:")
                for err in r.errors[:5]:
                    print(f"  - {err[:70]}")
    
    # Recommendation
    print("\n" + "=" * 80)
    print("RECOMMENDATION")
    print("=" * 80)
    
    if results:
        best = results[0]
        print(f"\nBest configuration: {best.name}")
        print(f"  - Accuracy: {best.accuracy:.1f}%")
        print(f"  - Avg response time: {best.avg_time:.2f}s")
        print(f"  - Overall score: {best.score:.1f}")
        
        if len(results) > 1:
            runner_up = results[1]
            if runner_up.accuracy > best.accuracy - 5:
                print(f"\nClose runner-up: {runner_up.name}")
                print(f"  Consider if you need faster responses.")
        
        print(f"\nTo use this configuration, set in your .env file:")
        print(f"  OLLAMA_MODEL={best.model}")
        print(f"  QUERY_MODE={best.mode}")


def print_markdown_table(results: List[BenchmarkResult]):
    """Output benchmark results as a markdown table for README."""
    print("\n## Performance Benchmarks\n")
    print(f"Tested on {results[0].total_tests if results else 0} queries across easy/medium/hard difficulties.\n")
    
    # Sort by accuracy descending
    results = sorted(results, key=lambda r: r.accuracy, reverse=True)
    
    # Header
    print("| Mode | Accuracy | Avg Time | Easy | Medium | Hard |")
    print("|------|----------|----------|------|--------|------|")
    
    for r in results:
        easy_pct = r.by_difficulty.get('easy', {})
        med_pct = r.by_difficulty.get('medium', {})
        hard_pct = r.by_difficulty.get('hard', {})
        
        easy_str = f"{easy_pct.get('passed', 0)}/{easy_pct.get('total', 0)}" if easy_pct else "-"
        med_str = f"{med_pct.get('passed', 0)}/{med_pct.get('total', 0)}" if med_pct else "-"
        hard_str = f"{hard_pct.get('passed', 0)}/{hard_pct.get('total', 0)}" if hard_pct else "-"
        
        mode_display = r.mode.capitalize()
        print(f"| {mode_display} | {r.accuracy:.0f}% | {r.avg_time:.1f}s | {easy_str} | {med_str} | {hard_str} |")
    
    print("\n*Results may vary based on model and LLM server performance.*\n")


def main():
    """Main benchmark function."""
    args = sys.argv[1:]
    
    # Parse arguments
    quick = '--quick' in args
    detailed = '--detailed' in args
    markdown = '--markdown' in args
    modes_only = '--modes-only' in args  # Only test modes, not models
    
    # Get models to benchmark
    models = None
    for i, arg in enumerate(args):
        if arg == '--models' and i + 1 < len(args):
            models = [m.strip() for m in args[i + 1].split(',')]
    
    # Get modes to benchmark
    modes = QUERY_MODES
    for i, arg in enumerate(args):
        if arg == '--modes' and i + 1 < len(args):
            modes = [m.strip() for m in args[i + 1].split(',')]
    
    # If no models specified, use current model from env
    if models is None:
        models = [os.getenv('OLLAMA_MODEL', 'llama3.2')]
    
    # Select test cases
    test_cases = QUICK_BENCHMARK_CASES if quick else BENCHMARK_CASES
    
    print("=" * 80)
    print("SQL Query Writer Agent - Benchmark")
    print("=" * 80)
    
    host = os.getenv('OLLAMA_HOST', 'http://localhost:11434')
    print(f"\nOllama Host: {host}")
    print(f"Models to test: {', '.join(models)}")
    print(f"Modes to test: {', '.join(modes)}")
    print(f"Test cases: {len(test_cases)}")
    
    # Run benchmarks
    results = []
    
    for model in models:
        for mode in modes:
            print(f"\n{'='*60}")
            print(f"Benchmarking: {model} with mode={mode}")
            print('='*60)
            
            result = benchmark_model(model, test_cases, mode=mode)
            results.append(result)
            
            print(f"  Accuracy: {result.accuracy:.1f}%")
            print(f"  Avg time: {result.avg_time:.2f}s")
            
            # Show difficulty breakdown
            for diff in ['easy', 'medium', 'hard']:
                stats = result.by_difficulty.get(diff, {})
                if stats.get('total', 0) > 0:
                    pct = stats['passed'] / stats['total'] * 100
                    print(f"    {diff}: {stats['passed']}/{stats['total']} ({pct:.0f}%)")
    
    # Print final results
    if markdown:
        print_markdown_table(results)
    else:
        print_results(results, detailed=detailed)
    
    # Save results
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    results_file = f"benchmark_{timestamp}.json"
    
    results_data = []
    for r in results:
        results_data.append({
            'model': r.model,
            'mode': r.mode,
            'accuracy': r.accuracy,
            'avg_time': r.avg_time,
            'score': r.score,
            'passed': r.passed,
            'failed': r.failed,
            'by_difficulty': r.by_difficulty,
            'errors': r.errors[:10]  # Limit errors
        })
    
    with open(results_file, 'w') as f:
        json.dump({
            'timestamp': timestamp,
            'host': host,
            'models': models,
            'modes': modes,
            'test_cases': len(test_cases),
            'results': results_data
        }, f, indent=2)
    
    print(f"\nResults saved to {results_file}")
    
    # Generate markdown file if requested
    if markdown:
        md_file = f"BENCHMARK_RESULTS.md"
        with open(md_file, 'w') as f:
            f.write("# Benchmark Results\n\n")
            f.write(f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M')}\n\n")
            f.write(f"Model: {', '.join(models)}\n\n")
            f.write("## Performance by Mode\n\n")
            f.write("| Mode | Accuracy | Avg Time | Easy | Medium | Hard |\n")
            f.write("|------|----------|----------|------|--------|------|\n")
            
            for r in sorted(results, key=lambda x: x.accuracy, reverse=True):
                easy = r.by_difficulty.get('easy', {})
                med = r.by_difficulty.get('medium', {})
                hard = r.by_difficulty.get('hard', {})
                
                easy_str = f"{easy.get('passed', 0)}/{easy.get('total', 0)}"
                med_str = f"{med.get('passed', 0)}/{med.get('total', 0)}"
                hard_str = f"{hard.get('passed', 0)}/{hard.get('total', 0)}"
                
                f.write(f"| {r.mode.capitalize()} | {r.accuracy:.0f}% | {r.avg_time:.1f}s | {easy_str} | {med_str} | {hard_str} |\n")
            
            f.write("\n*Results may vary based on model and LLM server performance.*\n")
        
        print(f"Markdown results saved to {md_file}")


if __name__ == "__main__":
    main()