File size: 1,110 Bytes
914e970 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 | #!/usr/bin/env python3
"""Micro-benchmark to compare initialization + generate latency for local vs cloud path.
Usage:
python3 scripts/bench_latency.py --mode local
python3 scripts/bench_latency.py --mode cloud
It runs a small DriftBrain init and one `think()` call and reports elapsed time.
"""
import os
import subprocess
import sys
import time
MODE = "local"
if len(sys.argv) > 1:
MODE = sys.argv[1]
# Determine env overrides
env = os.environ.copy()
if MODE == "local":
env["DRIFT_PREFER_LOCAL"] = "true"
else:
env["DRIFT_PREFER_LOCAL"] = "false"
snippet = r"""
import time, os
from infj_bot.core.brain import DriftBrain
b = DriftBrain()
start = time.time()
# run a quick think — use short prompt
try:
out = b.think('Hello test latency')
except Exception as e:
out = str(e)
print('elapsed', time.time()-start)
print('sdk', getattr(b,'sdk',None))
"""
cmd = [sys.executable, "-c", snippet]
print(f"Running benchmark mode={MODE}...")
proc = subprocess.run(cmd, env=env, capture_output=True, text=True)
print(proc.stdout)
if proc.stderr:
print(proc.stderr, file=sys.stderr)
|