File size: 4,337 Bytes
6cb442b | 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 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | #!/usr/bin/env python3
"""Cheap connectivity probe for providers.json entries.
Never prints the API key. Prefer --once / --max-probes / scheduled --minutes/--hours.
python probe.py --mock
python probe.py --provider flint --once
python probe.py --provider flint --minutes 60 --hours 8
"""
from __future__ import annotations
import argparse, csv, json, os, sys, time
CSV_DEFAULT = "results/llm_probe.csv"
def _append_csv(path, row, fieldnames):
os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
new = not os.path.exists(path) or os.path.getsize(path) == 0
with open(path, "a", newline="") as f:
w = csv.DictWriter(f, fieldnames=fieldnames)
if new:
w.writeheader()
w.writerow(row)
def _scrub(err, key):
s = str(err)
if key:
s = s.replace(key, "[redacted]")
if len(key) > 8:
s = s.replace(key[:8], "[redacted]")
return s
def main():
ap = argparse.ArgumentParser(description="Probe an OpenAI-compatible provider without leaking secrets.")
ap.add_argument("--config", default="providers.json")
ap.add_argument("--provider", default="flint")
ap.add_argument("--mock", action="store_true", help="offline plumbing check only")
ap.add_argument("--once", action="store_true", help="single live call then exit")
ap.add_argument("--max-probes", type=int, default=1, help="max live calls (default 1)")
ap.add_argument("--minutes", type=float, default=0, help="spacing between probes in a long run")
ap.add_argument("--hours", type=float, default=0, help="total wall time for a long run")
ap.add_argument("--csv", default=CSV_DEFAULT, help="append probe outcomes here (never writes the key)")
ap.add_argument("--prompt", default='Return JSON {"ok":true} and nothing else.')
a = ap.parse_args()
fields = ["ts_rel", "provider", "model", "status", "seconds", "error", "usage_calls"]
if a.mock:
from interfaces import MockLLMClient
c = MockLLMClient([{"ok": True}])
out = c.complete_json(a.prompt)
print("mock_ok", bool(out.get("ok")), "calls", c.usage.get("calls"))
return 0
cfg = json.load(open(a.config))
local = "fixtures/private_unredacted/providers.local.json"
if os.path.exists(local):
loc = json.load(open(local))
if a.provider in loc:
cfg.setdefault(a.provider, {}).update({k: v for k, v in loc[a.provider].items() if not str(k).startswith("_")})
p = cfg.get(a.provider)
if not p:
print(f"provider_missing {a.provider}", file=sys.stderr); return 2
env = p.get("env", "SQ_LLM_API_KEY")
key = os.environ.get(env, "")
if not key:
print(f"key_missing env={env} provider={a.provider}"); return 3
print(f"key_present env={env} len={len(key)} provider={a.provider}")
from interfaces import OpenAICompatClient
model = p.get("model") or os.environ.get(p.get("model_env", "SQ_LLM_MODEL"), "")
base = p.get("base_url") or os.environ.get(p.get("base_url_env", "SQ_LLM_BASE_URL"), "")
client = OpenAICompatClient(model=model, base_url=base, api_key=key,
timeout=p.get("timeout", 120), sleep_s=p.get("sleep_s", 0),
json_mode=p.get("json_mode", True))
n = 1 if a.once else max(1, a.max_probes)
end = time.time() + (a.hours * 3600 if a.hours else 0)
i = 0
status = "fail"
while i < n or (a.hours and time.time() < end):
i += 1
t0 = time.time()
try:
out = client.complete_json(a.prompt)
status = "ok" if out is not None else "empty"
err = ""
except Exception as e:
status = "fail"; err = _scrub(e, key); out = None
sec = time.time() - t0
_append_csv(a.csv, {"ts_rel": f"probe-{i}", "provider": a.provider, "model": model,
"status": status, "seconds": f"{sec:.3f}", "error": err,
"usage_calls": client.usage.get("calls", 0)}, fields)
print(f"probe {i} status={status} seconds={sec:.2f} error={err}")
if a.once or (not a.hours and i >= n):
break
if a.minutes:
time.sleep(a.minutes * 60)
return 0 if status == "ok" else 1
if __name__ == "__main__":
sys.exit(main() or 0)
|