atlas-iac/scripts/ops/test_planning/synthetic_probe.py

87 lines
4.1 KiB
Python
Executable File

#!/usr/bin/env python3
"""Run only the bundled fictional cases against the private batch endpoint."""
import argparse
import hashlib
import json
from pathlib import Path
import sys
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from hermes_batch_client import BatchClient, MODELS, canonical
from contracts import FAMILY, PROFILE, validate_families, validate_profile
ROOT = Path(__file__).resolve().parent
def source_cases():
"""Load built-in synthetic records, never a workbook or a user-selected file."""
fixture = json.loads((ROOT / "synthetic_cases.json").read_text())
assert fixture["synthetic"] is True and fixture["campaign_id"] == "SYNTHETIC"
return [{"campaign_id": fixture["campaign_id"], "suite_id": fixture["suite_id"], **case}
for case in fixture["cases"]]
def envelope(model, mode, sources, thinking):
"""Build versioned synthetic input with explicit sampling and context budgets."""
if mode == "transport":
prompt = 'Return only JSON with the field "status" equal to "LOCAL_OK".'
schema = {"type": "object", "properties": {"status": {"type": "string", "enum": ["LOCAL_OK"]}},
"required": ["status"], "additionalProperties": False}
count, context = 128, 16384
else:
name = "profile" if mode == "profile" else "family"
prompt = (ROOT / (name + "_prompt.txt")).read_text()
prompt += "\nOutput schema:\n" + json.dumps(PROFILE if mode == "profile" else FAMILY)
prompt += "\nSynthetic source records:\n" + json.dumps(sources, ensure_ascii=False)
schema = PROFILE if mode == "profile" else FAMILY
count, context = (2048, 16384) if mode == "profile" else (8192, 32768)
return {"campaign_id": "SYNTHETIC", "suite_id": "SYN-SUITE-01",
"source_sha256": hashlib.sha256(canonical(sources)).hexdigest(),
"prompt_version": "synthetic-" + mode + "-v1",
"request": {"model": model, "prompt": prompt, "stream": False, "think": thinking,
"format": schema, "options": {
"num_ctx": context, "num_predict": count,
"temperature": 0.6 if thinking else 0, "top_p": 0.95 if thinking else 1,
"top_k": 20, "seed": 42}}}
def main():
"""Emit synthetic timing and contract results without displaying source text."""
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--token-file", required=True)
parser.add_argument("--cache-dir", required=True)
parser.add_argument("--mode", choices=("transport", "profile", "group"), default="transport")
parser.add_argument("--model", choices=list(MODELS), required=True)
parser.add_argument("--case-id", default="SYN-010")
parser.add_argument("--think", action="store_true")
args = parser.parse_args()
sources = source_cases()
if args.mode == "profile":
sources = [source for source in sources if source["case_id"] == args.case_id]
if len(sources) != 1:
parser.error("case-id must select one bundled synthetic case")
request = envelope(args.model, args.mode, sources, args.think)
client = BatchClient(args.token_file, args.cache_dir)
outcome = client.run(request)
record = outcome["record"]
result = record["result"]
content = json.loads(result["response"])
if args.mode == "profile":
validate_profile(content, sources[0])
elif args.mode == "group":
validate_families(content, sources)
elif content != {"status": "LOCAL_OK"}:
raise ValueError("synthetic transport check failed")
print(json.dumps({"synthetic": True, "mode": args.mode, "model": args.model,
"cache_hit": outcome["cache_hit"], "result_file": outcome["path"],
"wall_seconds": record["client_wall_seconds"],
"load_seconds": result.get("load_duration", 0) / 1e9,
"prompt_tokens": result.get("prompt_eval_count"),
"output_tokens": result.get("eval_count"), "contract_valid": True}), flush=True)
if __name__ == "__main__":
main()