#!/usr/bin/env python3 """Benchmark runner for the web API. Usage: uv run python bench.py # basic run (default scenario) uv run python bench.py -n 5 # 5 iterations with stats uv run python bench.py --all -n 3 # all scenarios uv run python bench.py --group search -n 3 # all search scenarios uv run python bench.py -c 3 -n 5 # concurrent load test uv run python bench.py -o out.json --tag "before-refactor" # save with label """ from __future__ import annotations import argparse import json import math import subprocess import sys import tempfile import time from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import asdict, dataclass, field from datetime import datetime, timezone # --------------------------------------------------------------------------- # Benchmark URLs (reliable, fast, publicly available) # --------------------------------------------------------------------------- BENCH_URLS = [ "https://www.python.org/", "https://en.wikipedia.org/wiki/Python_(programming_language)", "https://docs.python.org/3/tutorial/index.html", "https://httpbin.org/html", "https://example.com", "https://www.reuters.com/", "https://en.wikipedia.org/wiki/Machine_learning", "https://en.wikipedia.org/wiki/Artificial_intelligence", "https://www.bbc.com/news", "https://docs.python.org/3/library/asyncio.html", ] # --------------------------------------------------------------------------- # Scenarios # --------------------------------------------------------------------------- SCENARIOS: dict[str, dict] = { # ------------------------------------------------------------------------- # Health baseline # ------------------------------------------------------------------------- "health": { "endpoint": "/health", "method": "GET", "body": None, }, # ------------------------------------------------------------------------- # Search endpoint variations # ------------------------------------------------------------------------- "search-single": { "endpoint": "/v1/search", "method": "POST", "body": { "queries": ["Python programming language popularity"], "max_results": 10, }, }, "search-multi-query": { "endpoint": "/v1/search", "method": "POST", "body": { "queries": [ "Python popularity 2025", "JavaScript frameworks comparison", "Rust programming adoption", "machine learning trends", "cloud computing market", ], "max_results": 10, }, }, "search-max-queries": { "endpoint": "/v1/search", "method": "POST", "body": { "queries": [ "Python web frameworks", "JavaScript runtime performance", "Rust memory safety", "Go concurrency patterns", "TypeScript adoption rate", "Kotlin multiplatform", "Swift server side", "C++ modern standards", "Java virtual threads", "Ruby on Rails 2025", ], "max_results": 5, }, }, "search-large-results": { "endpoint": "/v1/search", "method": "POST", "body": { "queries": ["artificial intelligence news", "climate change research"], "max_results": 50, }, }, "search-freshness-day": { "endpoint": "/v1/search", "method": "POST", "body": { "queries": ["breaking news today"], "max_results": 10, "freshness": "day", }, }, "search-freshness-week": { "endpoint": "/v1/search", "method": "POST", "body": { "queries": ["technology announcements"], "max_results": 10, "freshness": "week", }, }, "search-site-filter": { "endpoint": "/v1/search", "method": "POST", "body": { "queries": ["Python programming"], "max_results": 20, "site_allowlist": ["wikipedia.org", "python.org", "realpython.com"], }, }, "search-non-english": { "endpoint": "/v1/search", "method": "POST", "body": { "queries": ["intelligence artificielle actualités"], "max_results": 10, "language": "fr", "country": "FR", }, }, # ------------------------------------------------------------------------- # Fetch endpoint variations # ------------------------------------------------------------------------- "fetch-single": { "endpoint": "/v1/fetch", "method": "POST", "body": { "urls": [BENCH_URLS[0]], "timeout_seconds": 30, }, }, "fetch-multi": { "endpoint": "/v1/fetch", "method": "POST", "body": { "urls": BENCH_URLS[:5], "parallel_fetches": 5, "timeout_seconds": 30, }, }, "fetch-large": { "endpoint": "/v1/fetch", "method": "POST", "body": { "urls": BENCH_URLS, "parallel_fetches": 10, "timeout_seconds": 45, }, }, "fetch-serial": { "endpoint": "/v1/fetch", "method": "POST", "body": { "urls": BENCH_URLS[:5], "parallel_fetches": 1, "timeout_seconds": 60, }, }, "fetch-no-text": { "endpoint": "/v1/fetch", "method": "POST", "body": { "urls": BENCH_URLS[:3], "extract_text": False, "include_html": True, "timeout_seconds": 30, }, }, "fetch-no-fallback": { "endpoint": "/v1/fetch", "method": "POST", "body": { "urls": BENCH_URLS[:3], "auto_fallback": False, "method": "http", "timeout_seconds": 30, }, }, "fetch-short-timeout": { "endpoint": "/v1/fetch", "method": "POST", "body": { "urls": BENCH_URLS[:3], "timeout_seconds": 5, }, }, # ------------------------------------------------------------------------- # Image search endpoint variations # ------------------------------------------------------------------------- "image-search-small": { "endpoint": "/v1/image-search", "method": "POST", "body": { "queries": ["cute cats"], "max_results": 5, }, }, "image-search-large": { "endpoint": "/v1/image-search", "method": "POST", "body": { "queries": ["nature photography", "city skyline"], "max_results": 100, }, }, # ------------------------------------------------------------------------- # Gather endpoint variations (unified pipeline) # ------------------------------------------------------------------------- "gather-minimal": { "endpoint": "/v1/gather", "method": "POST", "body": { "claim": "Python is the most popular programming language in 2025", "search_queries": ["Python popularity 2025"], "max_search_results": 5, "max_evidence_items": 3, "extract_snippets": False, "timeout_seconds": 60, }, }, "gather-default": { "endpoint": "/v1/gather", "method": "POST", "body": { "claim": "Python is the most popular programming language in 2025", "search_queries": [ "Python popularity 2025", "TIOBE index programming languages", ], "max_search_results": 10, "max_evidence_items": 8, "extract_snippets": False, "timeout_seconds": 90, }, }, "gather-snippets": { "endpoint": "/v1/gather", "method": "POST", "body": { "claim": "Python is the most popular programming language in 2025", "search_queries": ["Python popularity 2025"], "max_search_results": 5, "max_evidence_items": 5, "extract_snippets": True, "timeout_seconds": 120, }, }, "gather-high-parallel": { "endpoint": "/v1/gather", "method": "POST", "body": { "claim": "Artificial intelligence is transforming healthcare", "search_queries": [ "AI healthcare applications", "machine learning medical diagnosis", ], "max_search_results": 20, "max_evidence_items": 10, "parallel_fetches": 10, "extract_snippets": False, "timeout_seconds": 90, }, }, "gather-serial-fetch": { "endpoint": "/v1/gather", "method": "POST", "body": { "claim": "Electric vehicles are becoming mainstream", "search_queries": ["electric vehicle adoption 2025"], "max_search_results": 10, "max_evidence_items": 5, "parallel_fetches": 1, "extract_snippets": False, "timeout_seconds": 120, }, }, "gather-no-fallback": { "endpoint": "/v1/gather", "method": "POST", "body": { "claim": "Renewable energy costs are declining", "search_queries": ["renewable energy cost trends"], "max_search_results": 10, "max_evidence_items": 5, "fetch_method": "http", "auto_fallback": False, "extract_snippets": False, "timeout_seconds": 60, }, }, "gather-site-restricted": { "endpoint": "/v1/gather", "method": "POST", "body": { "claim": "Python is widely used in data science", "search_queries": ["Python data science"], "max_search_results": 15, "max_evidence_items": 8, "site_allowlist": ["wikipedia.org", "bbc.com", "reuters.com"], "extract_snippets": False, "timeout_seconds": 90, }, }, "gather-max-evidence": { "endpoint": "/v1/gather", "method": "POST", "body": { "claim": "Climate change is accelerating", "search_queries": [ "climate change scientific evidence", "global warming data 2025", ], "max_search_results": 30, "max_evidence_items": 25, "extract_snippets": False, "timeout_seconds": 180, }, }, # Legacy aliases for backward compatibility "default": { "endpoint": "/v1/gather", "method": "POST", "body": { "claim": "Python is the most popular programming language in 2025", "search_queries": [ "Python popularity 2025", "TIOBE index programming languages", ], "max_search_results": 5, "max_evidence_items": 5, "extract_snippets": True, "timeout_seconds": 90, }, }, "gather-large": { "endpoint": "/v1/gather", "method": "POST", "body": { "claim": "Python is the most popular programming language in 2025", "search_queries": [ "Python popularity 2025", "TIOBE index programming languages", ], "max_search_results": 20, "max_evidence_items": 15, "extract_snippets": False, "timeout_seconds": 120, }, }, "search-only": { "endpoint": "/v1/search", "method": "POST", "body": { "queries": ["Python popularity 2025", "TIOBE index 2025"], "max_results": 10, }, }, "image-search": { "endpoint": "/v1/image-search", "method": "POST", "body": { "queries": [ "Python programming language", "machine learning visualization", ], "max_results": 10, }, }, "fetch-only": { "endpoint": "/v1/fetch", "method": "POST", "body": { "urls": BENCH_URLS[:3], "extract_text": True, "timeout_seconds": 30, }, }, } # --------------------------------------------------------------------------- # Scenario Groups # --------------------------------------------------------------------------- SCENARIO_GROUPS: dict[str, list[str]] = { "search": [ "search-single", "search-multi-query", "search-max-queries", "search-large-results", "search-freshness-day", "search-freshness-week", "search-site-filter", "search-non-english", ], "fetch": [ "fetch-single", "fetch-multi", "fetch-large", "fetch-serial", "fetch-no-text", "fetch-no-fallback", "fetch-short-timeout", ], "image": [ "image-search-small", "image-search-large", ], "gather": [ "gather-minimal", "gather-default", "gather-snippets", "gather-high-parallel", "gather-serial-fetch", "gather-no-fallback", "gather-site-restricted", "gather-max-evidence", ], "quick": [ "health", "search-single", "fetch-single", "gather-minimal", ], } # curl -w format string for transport-level timing CURL_WRITE_OUT = json.dumps( { "status_code": "%{http_code}", "time_total": "%{time_total}", "time_connect": "%{time_connect}", "time_starttransfer": "%{time_starttransfer}", "size_download": "%{size_download}", } ) # --------------------------------------------------------------------------- # RunResult # --------------------------------------------------------------------------- @dataclass class RunResult: scenario: str iteration: int status_code: int = 0 curl_time_total: float = 0.0 curl_time_connect: float = 0.0 curl_time_starttransfer: float = 0.0 api_time_ms: float | None = None stages: dict[str, float] = field(default_factory=dict) evidence_items: int | None = None failed_items: int | None = None queries_processed: int | None = None tokens_used: int | None = None error: str | None = None # --------------------------------------------------------------------------- # Runner # --------------------------------------------------------------------------- def run_single(base_url: str, scenario_name: str, timeout: int) -> RunResult: """Execute a single benchmark request via curl.""" scenario = SCENARIOS[scenario_name] url = f"{base_url}{scenario['endpoint']}" result = RunResult(scenario=scenario_name, iteration=0) with tempfile.NamedTemporaryFile(suffix=".json", delete=True) as body_file: cmd = [ "curl", "-s", "--max-time", str(timeout), "-w", CURL_WRITE_OUT, "-o", body_file.name, ] if scenario["method"] == "POST" and scenario["body"] is not None: cmd += [ "-X", "POST", "-H", "Content-Type: application/json", "-d", json.dumps(scenario["body"]), ] cmd.append(url) try: proc = subprocess.run( cmd, capture_output=True, text=True, timeout=timeout + 10 ) except subprocess.TimeoutExpired: result.error = "subprocess timeout" return result except Exception as exc: result.error = str(exc) return result # Parse curl write-out (appended to stdout) try: curl_stats = json.loads(proc.stdout) result.status_code = int(curl_stats["status_code"]) result.curl_time_total = float(curl_stats["time_total"]) result.curl_time_connect = float(curl_stats["time_connect"]) result.curl_time_starttransfer = float(curl_stats["time_starttransfer"]) except (json.JSONDecodeError, KeyError, ValueError): result.error = f"failed to parse curl output: {proc.stdout[:200]}" return result if result.status_code == 0: stderr_snippet = proc.stderr[:200] if proc.stderr else "no stderr" result.error = f"curl failed (code 0): {stderr_snippet}" return result # Parse response body try: with open(body_file.name) as f: body = json.load(f) except (json.JSONDecodeError, OSError): body = None if body and isinstance(body, dict): result.api_time_ms = body.get("execution_time_ms") # Gather-specific fields (stages, evidence stats) for stage in body.get("stages", []): if isinstance(stage, dict) and "stage" in stage: result.stages[stage["stage"]] = stage.get("duration_ms", 0.0) stats = body.get("evidence_stats") if isinstance(stats, dict): result.evidence_items = stats.get("output_items") result.tokens_used = stats.get("tokens_used") # Search / image-search specific if "total_results" in body: result.evidence_items = body["total_results"] if "queries_processed" in body: result.queries_processed = body["queries_processed"] # Fetch-specific if "total_fetched" in body: result.evidence_items = body["total_fetched"] if "total_failed" in body: result.failed_items = body["total_failed"] # Error in response body if result.status_code >= 400: detail = body.get("detail", body.get("error", "")) if detail: result.error = str(detail)[:200] return result # --------------------------------------------------------------------------- # Stats # --------------------------------------------------------------------------- def _percentile(sorted_vals: list[float], p: float) -> float: """Compute the p-th percentile from a sorted list.""" if not sorted_vals: return 0.0 k = (len(sorted_vals) - 1) * (p / 100.0) f = math.floor(k) c = math.ceil(k) if f == c: return sorted_vals[int(k)] return sorted_vals[f] * (c - k) + sorted_vals[c] * (k - f) def compute_stats(values: list[float]) -> dict[str, float]: """Compute min/mean/median/p95/max for a list of values.""" if not values: return {"min": 0, "mean": 0, "median": 0, "p95": 0, "max": 0} s = sorted(values) return { "min": round(s[0], 2), "mean": round(sum(s) / len(s), 2), "median": round(_percentile(s, 50), 2), "p95": round(_percentile(s, 95), 2), "max": round(s[-1], 2), } def compute_all_stats(results: list[RunResult]) -> dict[str, dict[str, float]]: """Compute stats across all metric dimensions.""" out: dict[str, dict[str, float]] = {} api_times = [r.api_time_ms for r in results if r.api_time_ms is not None] if api_times: out["total_ms"] = compute_stats(api_times) # Collect all stage names stage_names: set[str] = set() for r in results: stage_names.update(r.stages.keys()) for name in sorted(stage_names): vals = [r.stages[name] for r in results if name in r.stages] if vals: out[f"{name}_ms"] = compute_stats(vals) ttfb = [r.curl_time_starttransfer for r in results if r.curl_time_starttransfer > 0] if ttfb: out["curl_ttfb_s"] = compute_stats(ttfb) totals = [r.curl_time_total for r in results if r.curl_time_total > 0] if totals: out["curl_total_s"] = compute_stats(totals) return out # --------------------------------------------------------------------------- # Reporter # --------------------------------------------------------------------------- def _fmt_stages(r: RunResult) -> str: if not r.stages: return "" parts = [f"{k}:{v:.0f}" for k, v in r.stages.items()] return f" [{' '.join(parts)}]" def _fmt_items(r: RunResult) -> str: parts = [] if r.evidence_items is not None: parts.append(f"{r.evidence_items} items") if r.failed_items is not None and r.failed_items > 0: parts.append(f"{r.failed_items} failed") if r.queries_processed is not None: parts.append(f"{r.queries_processed} queries") return f" ({', '.join(parts)})" if parts else "" def print_run(r: RunResult, quiet: bool) -> None: """Print a single run result line.""" if quiet: return if r.error: print(f" #{r.iteration} ERROR: {r.error}") return time_str = ( f"{r.api_time_ms:.0f}ms" if r.api_time_ms is not None else f"{r.curl_time_total:.2f}s" ) print(f" #{r.iteration} {time_str}{_fmt_stages(r)}{_fmt_items(r)}") def print_stats_table(stats: dict[str, dict[str, float]]) -> None: """Print the statistics summary table.""" if not stats: return header = ( f" {'Metric':<16} {'min':>8} {'mean':>8} {'median':>8} {'p95':>8} {'max':>8}" ) print() print(header) print(f" {'─' * 56}") for metric, vals in stats.items(): fmt = ".0f" if metric.endswith("_ms") else ".2f" print( f" {metric:<16}" f" {vals['min']:>8{fmt}}" f" {vals['mean']:>8{fmt}}" f" {vals['median']:>8{fmt}}" f" {vals['p95']:>8{fmt}}" f" {vals['max']:>8{fmt}}" ) print() def print_docker_logs(container: str) -> None: """Print recent docker logs.""" print(f"=== Recent logs ({container}) ===") try: proc = subprocess.run( ["docker", "logs", container, "--tail", "10"], capture_output=True, text=True, timeout=5, ) output = proc.stdout or proc.stderr or "(no output)" print(output.rstrip()) except (subprocess.TimeoutExpired, FileNotFoundError): print(" (docker logs unavailable)") print() # --------------------------------------------------------------------------- # Main execution # --------------------------------------------------------------------------- def run_scenario( base_url: str, scenario_name: str, iterations: int, warmup: int, concurrency: int, timeout: int, quiet: bool, ) -> list[RunResult]: """Run a scenario for the given number of iterations.""" scenario = SCENARIOS[scenario_name] method = scenario["method"] endpoint = scenario["endpoint"] print(f"=== {scenario_name} ({method} {endpoint}) ===") # Warmup if warmup > 0: sys.stdout.write(" warmup...") sys.stdout.flush() t0 = time.monotonic() for _ in range(warmup): run_single(base_url, scenario_name, timeout) elapsed = time.monotonic() - t0 print(f"done ({elapsed:.1f}s)") results: list[RunResult] = [] if concurrency <= 1: # Sequential mode for i in range(1, iterations + 1): r = run_single(base_url, scenario_name, timeout) r.iteration = i results.append(r) print_run(r, quiet) else: # Concurrent mode for i in range(1, iterations + 1): batch: list[RunResult] = [] with ThreadPoolExecutor(max_workers=concurrency) as pool: futures = { pool.submit(run_single, base_url, scenario_name, timeout): c for c in range(concurrency) } for fut in as_completed(futures): r = fut.result() r.iteration = i batch.append(r) results.extend(batch) if not quiet: ok = sum(1 for r in batch if r.error is None) times = [r.curl_time_total for r in batch if r.error is None] avg = sum(times) / len(times) if times else 0 print( f" #{i} {concurrency} reqs {ok}/{concurrency} ok avg {avg:.2f}s" ) # Stats successful = [r for r in results if r.error is None] if len(successful) >= 2: stats = compute_all_stats(successful) print_stats_table(stats) elif successful: print() else: print(" No successful runs.\n") # Concurrency throughput summary if concurrency > 1 and successful: total_time = sum(r.curl_time_total for r in successful) wall_time = total_time / concurrency rps = len(successful) / wall_time if wall_time > 0 else 0 print( f" throughput: ~{rps:.1f} req/s ({len(successful)} reqs, {concurrency} concurrent)" ) print() return results def build_json_output( base_url: str, all_results: dict[str, list[RunResult]], tag: str | None = None, ) -> dict: """Build the JSON output structure.""" scenarios_out = {} for name, results in all_results.items(): successful = [r for r in results if r.error is None] stats = compute_all_stats(successful) if len(successful) >= 2 else {} scenarios_out[name] = { "config": SCENARIOS[name], "runs": [asdict(r) for r in results], "stats": stats, } output: dict = { "timestamp": datetime.now(timezone.utc).isoformat(), "base_url": base_url, "scenarios": scenarios_out, } if tag: output["tag"] = tag return output def parse_args(argv: list[str] | None = None) -> argparse.Namespace: parser = argparse.ArgumentParser( description="web API benchmark runner", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=( f"Available scenarios: {', '.join(SCENARIOS)}\n" f"Available groups: {', '.join(SCENARIO_GROUPS)}" ), ) parser.add_argument( "--base-url", default="http://localhost:51100", help="API base URL (default: http://localhost:51100)", ) parser.add_argument( "-n", "--iterations", type=int, default=3, help="runs per scenario (default: 3)", ) parser.add_argument( "--warmup", type=int, default=1, help="warmup runs excluded from stats (default: 1)", ) parser.add_argument( "-c", "--concurrency", type=int, default=1, help="parallel requests per iteration (default: 1)", ) parser.add_argument( "-s", "--scenario", default=None, choices=list(SCENARIOS.keys()), help="single scenario to run", ) parser.add_argument( "--group", default=None, choices=list(SCENARIO_GROUPS.keys()), help="run all scenarios in a group", ) parser.add_argument( "--all", action="store_true", help="run all scenarios", ) parser.add_argument( "--timeout", type=int, default=120, help="curl --max-time in seconds (default: 120)", ) parser.add_argument( "-o", "--output", default=None, help="write JSON results to file", ) parser.add_argument( "--tag", default=None, help="label for this benchmark run (saved in JSON output)", ) parser.add_argument( "--docker-logs", default="didiAI-web-api", help="container name for log tail (empty to skip, default: didiAI-web-api)", ) parser.add_argument( "-q", "--quiet", action="store_true", help="suppress per-run output", ) return parser.parse_args(argv) def main(argv: list[str] | None = None) -> None: args = parse_args(argv) # Determine which scenarios to run if args.all: # All scenarios except legacy aliases legacy = {"default", "search-only", "image-search", "fetch-only"} scenario_names = [s for s in SCENARIOS if s not in legacy] elif args.group: scenario_names = SCENARIO_GROUPS[args.group] elif args.scenario: scenario_names = [args.scenario] else: scenario_names = ["gather-default"] print("bench.py — web API benchmark") print(f"base_url: {args.base_url}") if args.tag: print(f"tag: {args.tag}") print(f"scenarios: {', '.join(scenario_names)}") warmup_note = f" (+ {args.warmup} warmup)" if args.warmup else "" conc_note = f", concurrency {args.concurrency}" if args.concurrency > 1 else "" print(f"iterations: {args.iterations}{warmup_note}{conc_note}") print() all_results: dict[str, list[RunResult]] = {} for name in scenario_names: results = run_scenario( base_url=args.base_url, scenario_name=name, iterations=args.iterations, warmup=args.warmup, concurrency=args.concurrency, timeout=args.timeout, quiet=args.quiet, ) all_results[name] = results # Docker logs if args.docker_logs: print_docker_logs(args.docker_logs) # JSON output if args.output: output = build_json_output(args.base_url, all_results, args.tag) with open(args.output, "w") as f: json.dump(output, f, indent=2) print(f"Results written to {args.output}") if __name__ == "__main__": main()