"""Generate Beam reasoning traces for a prompt pool, verify answers, append to JSONL. Resumable: (id, effort) pairs already in --out are skipped. Stops launching new requests when the daily output-token or request budget (from Reflection's rate limit headers) drops below the reserve. Reads the key from REFLECTION_API_KEY. """ import argparse import asyncio import json import os import random import sys import time from pathlib import Path import httpx sys.path.insert(0, str(Path(__file__).resolve().parent)) from verify import check # noqa: E402 BASE = "https://api.reflection.ai/openai/v1" MODEL = "Beam-1.0-501B-A23B" class Budget: def __init__(self, reserve_out, reserve_req): self.reserve_out = reserve_out self.reserve_req = reserve_req self.out_left = None self.req_left = None self.reset_s = None def update(self, headers): def num(name): v = headers.get(name) return int(float(v)) if v is not None else None out_left = num("x-ratelimit-remaining-output-tokens-day") req_left = num("x-ratelimit-remaining-requests-day") if out_left is not None: self.out_left = out_left if req_left is not None: self.req_left = req_left r = headers.get("x-ratelimit-reset-output-tokens-day") if r: self.reset_s = float(r.rstrip("s")) def exhausted(self): return (self.out_left is not None and self.out_left < self.reserve_out) or ( self.req_left is not None and self.req_left < self.reserve_req ) class Pacer: """Keep under the per-minute request limit.""" def __init__(self, per_minute): self.interval = 60.0 / per_minute self.next = time.monotonic() self.lock = asyncio.Lock() async def wait(self): async with self.lock: now = time.monotonic() if self.next > now: await asyncio.sleep(self.next - now) self.next = max(now, self.next) + self.interval async def call(client, item, args, budget, pacer): body = { "model": MODEL, "messages": [{"role": "user", "content": item["prompt"]}], "max_tokens": args.max_tokens, "reasoning_effort": args.effort, "temperature": args.temperature, "top_p": args.top_p, } for attempt in range(6): await pacer.wait() t0 = time.time() try: r = await client.post(f"{BASE}/chat/completions", json=body) except httpx.HTTPError as e: err = f"transport: {e!r}" await asyncio.sleep(min(60, 2**attempt + random.random())) continue budget.update(r.headers) if r.status_code == 200: return r.json(), time.time() - t0, None err = f"HTTP {r.status_code}: {r.text[:300]}" if r.status_code in (429, 500, 502, 503, 504): if budget.exhausted(): return None, 0, err wait = r.headers.get("retry-after") await asyncio.sleep(float(wait) if wait else min(90, 2 ** (attempt + 1))) continue return None, 0, err return None, 0, err async def worker(name, queue, client, args, budget, pacer, out_f, stats): while True: item = await queue.get() try: if budget.exhausted() or stats["stop"]: continue resp, latency, err = await call(client, item, args, budget, pacer) rec = { "id": item["id"], "effort": args.effort, "domain": item["domain"], "source": item["source"], "prompt": item["prompt"], "gold": item["answer"], "answer_type": item["answer_type"], "difficulty": item.get("difficulty"), "topic": item.get("topic"), "time": time.strftime("%Y-%m-%dT%H:%M:%S"), } if resp is None: rec.update({"error": err}) stats["errors"] += 1 if "HTTP 4" in (err or "") and "429" not in err: print(f"[{name}] non-retryable error: {err}", flush=True) else: choice = resp["choices"][0] msg = choice["message"] usage = resp.get("usage") or {} rec.update( { "reasoning": msg.get("reasoning_content") or "", "content": msg.get("content") or "", "finish_reason": choice.get("finish_reason"), "usage": usage, "latency_s": round(latency, 2), "response_id": resp.get("id"), } ) ok, pred = check(rec["content"], item["answer"], item["answer_type"]) rec["pred"] = pred rec["correct"] = bool(ok) and rec["finish_reason"] == "stop" stats["done"] += 1 stats["correct"] += rec["correct"] stats["out_tokens"] += usage.get("completion_tokens", 0) out_f.write(json.dumps(rec, ensure_ascii=False) + "\n") out_f.flush() n = stats["done"] + stats["errors"] if n % args.log_every == 0: print(progress(stats, budget), flush=True) if args.target_correct and stats["correct"] >= args.target_correct: stats["stop"] = True finally: queue.task_done() def progress(stats, budget): done = max(1, stats["done"]) return ( f"{time.strftime('%H:%M:%S')} done={stats['done']} err={stats['errors']} " f"correct={stats['correct']} ({stats['correct'] / done:.0%}) " f"mean_out={stats['out_tokens'] / done:.0f} " f"out_left_day={budget.out_left} req_left_day={budget.req_left}" ) async def main(): ap = argparse.ArgumentParser() ap.add_argument("--pool", required=True) ap.add_argument("--out", required=True) ap.add_argument("--effort", default="medium") ap.add_argument("--max-tokens", type=int, default=16384) ap.add_argument("--temperature", type=float, default=0.7) ap.add_argument("--top-p", type=float, default=0.99) ap.add_argument("--concurrency", type=int, default=12) ap.add_argument("--per-minute", type=float, default=50) ap.add_argument("--limit", type=int, default=0, help="max new prompts this run") ap.add_argument("--ids", help="file with ids to run (one per line)") ap.add_argument("--target-correct", type=int, default=0) ap.add_argument("--reserve-out", type=int, default=40000) ap.add_argument("--reserve-req", type=int, default=200) ap.add_argument("--log-every", type=int, default=20) args = ap.parse_args() key = os.environ.get("REFLECTION_API_KEY") if not key: sys.exit("REFLECTION_API_KEY not set") pool = [json.loads(line) for line in open(args.pool)] if args.ids: wanted = [line.strip() for line in open(args.ids) if line.strip()] by_id = {p["id"]: p for p in pool} pool = [by_id[i] for i in wanted] done = set() if Path(args.out).exists(): for line in open(args.out): r = json.loads(line) if "error" not in r: done.add((r["id"], r["effort"])) todo = [p for p in pool if (p["id"], args.effort) not in done] if args.limit: todo = todo[: args.limit] print(f"{len(todo)} prompts to run at effort={args.effort} ({len(done)} already done)", flush=True) budget = Budget(args.reserve_out, args.reserve_req) pacer = Pacer(args.per_minute) stats = {"done": 0, "errors": 0, "correct": 0, "out_tokens": 0, "stop": False} queue = asyncio.Queue() for p in todo: queue.put_nowait(p) headers = {"Authorization": f"Bearer {key}"} timeout = httpx.Timeout(connect=30, read=900, write=60, pool=60) limits = httpx.Limits(max_connections=args.concurrency) with open(args.out, "a") as out_f: async with httpx.AsyncClient(headers=headers, timeout=timeout, limits=limits) as client: workers = [ asyncio.create_task(worker(f"w{i}", queue, client, args, budget, pacer, out_f, stats)) for i in range(args.concurrency) ] await queue.join() for w in workers: w.cancel() print("FINISHED " + progress(stats, budget), flush=True) if budget.exhausted(): print(f"STOPPED: daily budget reserve reached; resets in ~{(budget.reset_s or 0) / 3600:.1f}h", flush=True) if __name__ == "__main__": asyncio.run(main())