|
| 1 | +#!/usr/bin/env python3 |
| 2 | +"""ProofLayer LangGraph hot-path latency benchmark.""" |
| 3 | + |
| 4 | +import argparse |
| 5 | +import json |
| 6 | +import statistics |
| 7 | +import sys |
| 8 | +import time |
| 9 | +from pathlib import Path |
| 10 | +from typing import Dict, TypedDict |
| 11 | + |
| 12 | +PROJECT_ROOT = Path(__file__).parent.parent |
| 13 | +sys.path.insert(0, str(PROJECT_ROOT)) |
| 14 | + |
| 15 | +from langgraph.graph import END, StateGraph |
| 16 | + |
| 17 | +from prooflayer.integrations.langgraph import SecurityConfig, SecurityMiddleware |
| 18 | + |
| 19 | + |
| 20 | +class BenchState(TypedDict): |
| 21 | + """State for the benchmark graph.""" |
| 22 | + |
| 23 | + input: str |
| 24 | + answer: str |
| 25 | + |
| 26 | + |
| 27 | +def answer(state: BenchState) -> Dict[str, str]: |
| 28 | + """Return a deterministic answer for a representative node.""" |
| 29 | + return {"answer": f"ok:{state['input']}"} |
| 30 | + |
| 31 | + |
| 32 | +def percentile(samples: list[float], p: float) -> float: |
| 33 | + """Nearest-rank percentile, p in [0, 100].""" |
| 34 | + if not samples: |
| 35 | + return 0.0 |
| 36 | + sorted_samples = sorted(samples) |
| 37 | + index = max( |
| 38 | + 0, |
| 39 | + min(len(sorted_samples) - 1, int(round(p / 100 * (len(sorted_samples) - 1)))), |
| 40 | + ) |
| 41 | + return sorted_samples[index] |
| 42 | + |
| 43 | + |
| 44 | +def build_secured_graph(): |
| 45 | + """Build a representative secured LangGraph.""" |
| 46 | + graph = StateGraph(BenchState) |
| 47 | + graph.add_node("answer", answer) |
| 48 | + graph.set_entry_point("answer") |
| 49 | + graph.add_edge("answer", END) |
| 50 | + |
| 51 | + middleware = SecurityMiddleware( |
| 52 | + SecurityConfig( |
| 53 | + prompt_injection="block", |
| 54 | + jailbreak="block", |
| 55 | + tool_abuse="block", |
| 56 | + exfil="block", |
| 57 | + scope_drift="warn", |
| 58 | + state_manipulation="block", |
| 59 | + multi_turn="warn", |
| 60 | + emit_to=["stdout"], |
| 61 | + ) |
| 62 | + ) |
| 63 | + return middleware.wrap(graph.compile()) |
| 64 | + |
| 65 | + |
| 66 | +def main() -> int: |
| 67 | + """Run the LangGraph hot-path benchmark.""" |
| 68 | + parser = argparse.ArgumentParser( |
| 69 | + description="ProofLayer LangGraph latency benchmark" |
| 70 | + ) |
| 71 | + parser.add_argument("--n", type=int, default=200, help="number of invokes") |
| 72 | + parser.add_argument("--json", action="store_true", help="emit JSON, not text") |
| 73 | + args = parser.parse_args() |
| 74 | + |
| 75 | + graph = build_secured_graph() |
| 76 | + payload = {"input": "summarize the approved runbook", "answer": ""} |
| 77 | + config = {"configurable": {"thread_id": "bench-thread"}} |
| 78 | + |
| 79 | + for _ in range(min(50, args.n)): |
| 80 | + graph.invoke(payload, config=config) |
| 81 | + |
| 82 | + latencies_ms: list[float] = [] |
| 83 | + start = time.perf_counter() |
| 84 | + for _ in range(args.n): |
| 85 | + t0 = time.perf_counter() |
| 86 | + graph.invoke(payload, config=config) |
| 87 | + latencies_ms.append((time.perf_counter() - t0) * 1000.0) |
| 88 | + wall_s = time.perf_counter() - start |
| 89 | + |
| 90 | + summary = { |
| 91 | + "samples": len(latencies_ms), |
| 92 | + "wall_seconds": round(wall_s, 3), |
| 93 | + "throughput_per_second": round(len(latencies_ms) / wall_s, 1), |
| 94 | + "p50_ms": round(percentile(latencies_ms, 50), 4), |
| 95 | + "p95_ms": round(percentile(latencies_ms, 95), 4), |
| 96 | + "p99_ms": round(percentile(latencies_ms, 99), 4), |
| 97 | + "mean_ms": round(statistics.fmean(latencies_ms), 4), |
| 98 | + "max_ms": round(max(latencies_ms), 4), |
| 99 | + } |
| 100 | + |
| 101 | + if args.json: |
| 102 | + print(json.dumps(summary, indent=2)) |
| 103 | + else: |
| 104 | + print("=" * 60) |
| 105 | + print(" ProofLayer LangGraph Hot-Path Benchmark") |
| 106 | + print("=" * 60) |
| 107 | + print(f" samples : {summary['samples']}") |
| 108 | + print(f" wall time : {summary['wall_seconds']} s") |
| 109 | + print(f" throughput : {summary['throughput_per_second']} invokes/s") |
| 110 | + print(f" p50 latency : {summary['p50_ms']} ms") |
| 111 | + print(f" p95 latency : {summary['p95_ms']} ms") |
| 112 | + print(f" p99 latency : {summary['p99_ms']} ms") |
| 113 | + print(f" mean latency : {summary['mean_ms']} ms") |
| 114 | + print(f" max latency : {summary['max_ms']} ms") |
| 115 | + print("=" * 60) |
| 116 | + |
| 117 | + return 0 |
| 118 | + |
| 119 | + |
| 120 | +if __name__ == "__main__": |
| 121 | + sys.exit(main()) |
0 commit comments