tatlock/scripts/benchmark_steward.py

"""
Benchmark script for the Steward agent.

Tests Steward's request analysis performance with various scenarios
to ensure it meets latency targets:
- Target max: 5 seconds
- Target average: ~1.67 seconds

Usage:
    python scripts/benchmark_steward.py [--iterations N] [--verbose]
"""
import argparse
import asyncio
import statistics
from datetime import datetime
from typing import List

from src.agents.steward import analyze_request
from src.core.startup import initialize_application


class BenchmarkResult:
    """Results from a single benchmark run."""

    def __init__(self, scenario: str, duration: float, success: bool, error: str = None):
        self.scenario = scenario
        self.duration = duration
        self.success = success
        self.error = error


async def benchmark_scenario(
    name: str,
    request: str,
    history: list[dict],
    iterations: int = 10
) -> List[BenchmarkResult]:
    """
    Benchmark a specific scenario.

    Args:
        name: Scenario name
        request: User request to analyze
        history: Conversation history
        iterations: Number of times to run

    Returns:
        List of benchmark results
    """
    results = []

    print(f"\n📊 Benchmarking: {name}")
    print(f"   Request: {request[:50]}{'...' if len(request) > 50 else ''}")
    print(f"   History length: {len(history)} turns")
    print(f"   Iterations: {iterations}")

    for i in range(iterations):
        try:
            start = datetime.now()
            await analyze_request(request, history)
            duration = (datetime.now() - start).total_seconds()

            results.append(BenchmarkResult(name, duration, True))

            # Progress indicator
            print(".", end="", flush=True)

        except Exception as e:
            duration = (datetime.now() - start).total_seconds()
            results.append(BenchmarkResult(name, duration, False, str(e)))
            print("E", end="", flush=True)

    print()  # New line after progress
    return results


def analyze_results(results: List[BenchmarkResult], scenario_name: str):
    """
    Analyze and display benchmark results.

    Args:
        results: List of benchmark results
        scenario_name: Name of the scenario
    """
    successful = [r for r in results if r.success]
    failed = [r for r in results if not r.success]

    if not successful:
        print(f"\n❌ {scenario_name}: All runs failed!")
        for r in failed[:3]:  # Show first 3 errors
            print(f"   Error: {r.error}")
        return

    durations = [r.duration for r in successful]

    min_duration = min(durations)
    max_duration = max(durations)
    avg_duration = statistics.mean(durations)
    median_duration = statistics.median(durations)

    # Calculate percentiles
    sorted_durations = sorted(durations)
    p95_idx = int(len(sorted_durations) * 0.95)
    p99_idx = int(len(sorted_durations) * 0.99)
    p95 = sorted_durations[p95_idx] if p95_idx < len(sorted_durations) else max_duration
    p99 = sorted_durations[p99_idx] if p99_idx < len(sorted_durations) else max_duration

    # Targets
    target_max = 5.0
    target_avg = 1.67

    # Status emojis
    max_status = "✅" if max_duration <= target_max else "⚠️"
    avg_status = "✅" if avg_duration <= target_avg else "⚠️"

    print(f"\n   Results ({len(successful)}/{len(results)} successful):")
    print(f"   Min:     {min_duration:6.3f}s")
    print(f"   Avg:     {avg_duration:6.3f}s {avg_status} (target: ≤{target_avg}s)")
    print(f"   Median:  {median_duration:6.3f}s")
    print(f"   P95:     {p95:6.3f}s")
    print(f"   P99:     {p99:6.3f}s")
    print(f"   Max:     {max_duration:6.3f}s {max_status} (target: ≤{target_max}s)")

    if failed:
        print(f"   Failed:  {len(failed)} runs")

    return {
        "min": min_duration,
        "avg": avg_duration,
        "median": median_duration,
        "p95": p95,
        "p99": p99,
        "max": max_duration,
        "success_rate": len(successful) / len(results) * 100,
    }


async def run_benchmarks(iterations: int = 10, verbose: bool = False):
    """
    Run comprehensive Steward benchmarks.

    Args:
        iterations: Number of iterations per scenario
        verbose: Enable verbose output
    """
    print("=" * 60)
    print("🔬 Steward Performance Benchmark")
    print("=" * 60)
    print(f"\nTargets:")
    print(f"  - Maximum response time: ≤5.0s")
    print(f"  - Average response time: ≤1.67s")
    print(f"\nIterations per scenario: {iterations}")

    # Initialize application
    print("\n🚀 Initializing application...")
    initialize_application()

    all_stats = {}

    # Scenario 1: Simple greeting (no capabilities needed)
    results = await benchmark_scenario(
        "Simple Greeting",
        "Hello!",
        [],
        iterations
    )
    all_stats["simple_greeting"] = analyze_results(results, "Simple Greeting")

    # Scenario 2: Single tool request (calculator)
    results = await benchmark_scenario(
        "Calculator Request",
        "What's sqrt(144) + 25?",
        [],
        iterations
    )
    all_stats["calculator"] = analyze_results(results, "Calculator Request")

    # Scenario 3: Web search request
    results = await benchmark_scenario(
        "Web Search Request",
        "Search for the latest Python 3.12 features",
        [],
        iterations
    )
    all_stats["web_search"] = analyze_results(results, "Web Search Request")

    # Scenario 4: Request with conversation history (short)
    short_history = [
        {"role": "user", "content": "What's 15 times 7?"},
        {"role": "assistant", "content": "105"},
    ]
    results = await benchmark_scenario(
        "With Short History",
        "And what's that divided by 3?",
        short_history,
        iterations
    )
    all_stats["short_history"] = analyze_results(results, "With Short History")

    # Scenario 5: Request with longer conversation history
    long_history = [
        {"role": "user", "content": f"Question {i}"} if i % 2 == 0
        else {"role": "assistant", "content": f"Answer {i}"}
        for i in range(20)
    ]
    results = await benchmark_scenario(
        "With Long History",
        "What was the first question I asked?",
        long_history,
        iterations
    )
    all_stats["long_history"] = analyze_results(results, "With Long History")

    # Scenario 6: Complex request
    results = await benchmark_scenario(
        "Complex Request",
        "Calculate the compound interest on $5000 at 4.5% over 10 years, "
        "then search for current savings account rates to compare",
        [],
        iterations
    )
    all_stats["complex"] = analyze_results(results, "Complex Request")

    # Scenario 7: Missing capabilities
    results = await benchmark_scenario(
        "Missing Capabilities",
        "Generate an image of a sunset over mountains",
        [],
        iterations
    )
    all_stats["missing_caps"] = analyze_results(results, "Missing Capabilities")

    # Summary
    print("\n" + "=" * 60)
    print("📈 SUMMARY")
    print("=" * 60)

    # Calculate overall stats
    all_avgs = [stats["avg"] for stats in all_stats.values() if stats]
    all_maxs = [stats["max"] for stats in all_stats.values() if stats]

    if all_avgs:
        overall_avg = statistics.mean(all_avgs)
        overall_max = max(all_maxs)

        avg_status = "✅" if overall_avg <= 1.67 else "⚠️"
        max_status = "✅" if overall_max <= 5.0 else "⚠️"

        print(f"\nOverall Performance:")
        print(f"  Average of averages: {overall_avg:.3f}s {avg_status}")
        print(f"  Maximum observed:    {overall_max:.3f}s {max_status}")

        # Performance verdict
        print(f"\n{'=' * 60}")
        if overall_avg <= 1.67 and overall_max <= 5.0:
            print("✅ PERFORMANCE TARGETS MET!")
            print(f"   The Steward is operating within target parameters.")
        elif overall_max <= 5.0:
            print("⚠️  PARTIAL SUCCESS")
            print(f"   Max response time is good, but average is above target.")
            print(f"   Average: {overall_avg:.3f}s (target: ≤1.67s)")
            print(f"\n   Recommendations:")
            print(f"   - Consider using a faster model")
            print(f"   - Optimize system prompt length")
            print(f"   - Review tool call limits")
        else:
            print("❌ PERFORMANCE TARGETS NOT MET")
            print(f"   Max: {overall_max:.3f}s (target: ≤5.0s)")
            print(f"   Avg: {overall_avg:.3f}s (target: ≤1.67s)")
            print(f"\n   Recommendations:")
            print(f"   - Switch to a faster model (current: mistral-nemo)")
            print(f"   - Reduce system prompt complexity")
            print(f"   - Limit tool calls (currently limited to 3)")
            print(f"   - Consider caching household registry responses")

    print("=" * 60)


async def main():
    """Main entry point."""
    parser = argparse.ArgumentParser(description="Benchmark Steward agent performance")
    parser.add_argument(
        "--iterations",
        type=int,
        default=10,
        help="Number of iterations per scenario (default: 10)"
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="Enable verbose output"
    )

    args = parser.parse_args()

    await run_benchmarks(iterations=args.iterations, verbose=args.verbose)


if __name__ == "__main__":
    asyncio.run(main())