Repository Public Research preview
1"""Rewrite job engine: request planning, deduplication, on-disk cache, ledger (resume), budget ceiling,2live or batch execution. Originals and rewrites are never overwritten; every rewrite row carries model,3prompt version, request id, usage and cost.45Layout under ``<job_dir>/``: plan.jsonl (requests), rewrites.jsonl (append-only results), ledger.json6(spend, counts, batch ids), batches/<id>.json (submitted custom_ids), manifest.json (provenance).7"""8from __future__ import annotations910import json11import time12from concurrent.futures import ThreadPoolExecutor, as_completed13from pathlib import Path1415from htr.data.schema import Rewrite, RewriteRequest16from htr.parse.protect import PROTECT_VERSION, protect17from htr.rewrite.prompts import PROMPT_VERSION, build_messages, prompt_fingerprint18from htr.rewrite.provider import Provider19from htr.utils.hashing import sha256_text, short_id20from htr.utils.io import append_jsonl, ensure_dir, read_json, read_jsonl, write_json, write_jsonl21from htr.utils.provenance import run_header, utc_now222324class BudgetExceeded(RuntimeError):25 pass262728def plan_requests(passages: list[dict], variants: list[str], samples: dict[str, int], provider_name: str, model: str, max_tokens: int, effort: str) -> list[RewriteRequest]:29 """One request per (passage, variant, sample). Request ids are deterministic → natural dedup across runs."""30 reqs: list[RewriteRequest] = []31 for p in passages:32 for v in variants:33 for s in range(samples.get(v, 1)):34 rid = short_id("req", p["passage_id"], v, str(s), provider_name, model, PROMPT_VERSION)35 reqs.append(RewriteRequest(request_id=rid, passage_id=p["passage_id"], variant=v, sample_index=s, provider=provider_name, model=model, prompt_version=PROMPT_VERSION, params={"max_tokens": max_tokens, "effort": effort}))36 return reqs373839def _protected_text(p: dict) -> str:40 """Passages are sent with citations/math/urls replaced by placeholders so the LLM cannot alter them; the41 engine stores the placeholder form and the validator checks placeholder multisets."""42 return protect(p["text"]).text434445def merge_plan(job_dir: Path, reqs: list[RewriteRequest]) -> dict[str, RewriteRequest]:46 """plan.jsonl accumulates every request ever planned for the job (needed to interpret batch results)."""47 existing = {r["request_id"]: RewriteRequest(**r) for r in read_jsonl(job_dir / "plan.jsonl")}48 for r in reqs:49 existing[r.request_id] = r50 write_jsonl(job_dir / "plan.jsonl", existing.values())51 return existing525354class RewriteEngine:55 def __init__(self, job_dir: Path, provider: Provider, budget_usd: float, batch: bool = False, concurrency: int = 4):56 self.job_dir = ensure_dir(job_dir)57 self.provider = provider58 self.budget = budget_usd59 self.batch = batch60 self.concurrency = concurrency61 self.ledger_path = self.job_dir / "ledger.json"62 self.ledger = read_json(self.ledger_path) if self.ledger_path.exists() else {"spent_usd": 0.0, "n_ok": 0, "n_error": 0, "n_refused": 0, "n_truncated": 0, "n_empty": 0, "batches": [], "created_at": utc_now()}63 self.done: set[str] = {r["request_id"] for r in read_jsonl(self.job_dir / "rewrites.jsonl")}6465 def _save_ledger(self) -> None:66 self.ledger["updated_at"] = utc_now()67 write_json(self.ledger_path, self.ledger)6869 def _record(self, req: RewriteRequest, passage: dict, res, batch_id: str | None = None) -> Rewrite:70 rw = Rewrite(71 rewrite_id=short_id("rw", req.request_id),72 request_id=req.request_id,73 passage_id=req.passage_id,74 doc_id=passage["doc_id"],75 variant=req.variant,76 sample_index=req.sample_index,77 provider=self.provider.name,78 model=res.model or req.model,79 prompt_version=req.prompt_version,80 text=res.text,81 status=res.status, # type: ignore[arg-type]82 error=res.error,83 provider_request_id=res.provider_request_id,84 usage={k: v for k, v in (res.usage or {}).items() if v is not None},85 cost_usd=res.cost_usd,86 stop_reason=res.stop_reason,87 created_at=utc_now(),88 latency_s=res.latency_s,89 batch_id=batch_id,90 protect_version=PROTECT_VERSION,91 prompt_input_sha256=sha256_text(_protected_text(passage)),92 )93 append_jsonl(self.job_dir / "rewrites.jsonl", rw)94 self.done.add(req.request_id)95 self.ledger["spent_usd"] = round(self.ledger["spent_usd"] + (res.cost_usd or 0.0), 6)96 self.ledger[f"n_{res.status}"] = self.ledger.get(f"n_{res.status}", 0) + 197 return rw9899 # Calibration from the pilot job (5 801 Opus 5.5 requests, 2026-11-09): measured 691 input / 323 output tokens per100 # request vs the naive chars/4 estimate (492 / 259) → multiply by 1.40 / 1.25.101 EST_IN_FACTOR = 1.40102 EST_OUT_FACTOR = 1.25103104 def estimate(self, reqs: list[RewriteRequest], passages: dict[str, dict], pricing: dict | None) -> dict:105 """Pre-flight estimate: chars/4 ≈ tokens, prompt overhead ≈ 260 tokens, calibrated on the pilot (see constants)."""106 pend = [r for r in reqs if r.request_id not in self.done]107 toks_in = int(self.EST_IN_FACTOR * sum(260 + len(_protected_text(passages[r.passage_id])) // 4 for r in pend))108 toks_out = int(self.EST_OUT_FACTOR * sum(int(1.1 * len(passages[r.passage_id]["text"]) // 4) for r in pend))109 est = None110 if pricing and self.provider.model in pricing["models"]:111 m = pricing["models"][self.provider.model]112 est = (toks_in * m["input"] + toks_out * m["output"]) / 1e6 * (pricing.get("batch_discount", 0.5) if self.batch else 1.0)113 return {"pending_requests": len(pend), "done": len(self.done), "est_input_tokens": toks_in, "est_output_tokens": toks_out, "est_cost_usd": None if est is None else round(est, 2), "budget_usd": self.budget, "spent_usd": self.ledger["spent_usd"]}114115 # ---- live ----------------------------------------------------------------------------------------------116 def run_live(self, reqs: list[RewriteRequest], passages: dict[str, dict], max_requests: int | None = None) -> dict:117 pend = [r for r in reqs if r.request_id not in self.done][: max_requests or None]118 merge_plan(self.job_dir, reqs)119 n = 0120121 def work(req: RewriteRequest):122 p = passages[req.passage_id]123 system, messages = build_messages(_protected_text(p), req.variant)124 return req, p, self.provider.rewrite(system, messages, req.params["max_tokens"], {"variant": req.variant, "sample_index": req.sample_index})125126 with ThreadPoolExecutor(max_workers=self.concurrency) as ex:127 futs = [ex.submit(work, r) for r in pend]128 for f in as_completed(futs):129 req, p, res = f.result()130 self._record(req, p, res)131 n += 1132 if n % 25 == 0:133 self._save_ledger()134 if self.ledger["spent_usd"] > self.budget:135 self._save_ledger()136 for g in futs:137 g.cancel()138 raise BudgetExceeded(f"spent {self.ledger['spent_usd']:.2f} USD > ceiling {self.budget:.2f} USD")139 self._save_ledger()140 return {"processed": n, **self.ledger}141142 # ---- batch ---------------------------------------------------------------------------------------------143 def submit_batches(self, reqs: list[RewriteRequest], passages: dict[str, dict], chunk: int = 2000, max_requests: int | None = None, pricing: dict | None = None) -> list[str]:144 if not self.provider.supports_batch():145 raise RuntimeError("provider has no batch support")146 pend = [r for r in reqs if r.request_id not in self.done]147 submitted = {cid for b in self.ledger["batches"] for cid in read_json(self.job_dir / "batches" / f"{b}.json")["custom_ids"]} if self.ledger["batches"] else set()148 pend = [r for r in pend if r.request_id not in submitted][: max_requests or None]149 merge_plan(self.job_dir, reqs)150 # pre-flight budget check: estimated cost of what is about to be submitted + spent + already-submitted estimate151 if pricing is not None:152 est_new = self.estimate(pend, passages, pricing)["est_cost_usd"] or 0.0153 est_sub = self.ledger.get("submitted_est_usd", 0.0)154 if self.ledger["spent_usd"] + est_sub + est_new > self.budget:155 raise BudgetExceeded(f"estimated {self.ledger['spent_usd'] + est_sub + est_new:.2f} USD (spent + submitted + new) > ceiling {self.budget:.2f} USD; nothing submitted")156 self.ledger["submitted_est_usd"] = round(est_sub + est_new, 4)157 ids: list[str] = []158 for i in range(0, len(pend), chunk):159 part = pend[i : i + chunk]160 items = []161 for r in part:162 p = passages[r.passage_id]163 system, messages = build_messages(_protected_text(p), r.variant)164 items.append((r.request_id, system, messages, r.params["max_tokens"]))165 bid = self.provider.submit_batch(items) # type: ignore[attr-defined]166 write_json(ensure_dir(self.job_dir / "batches") / f"{bid}.json", {"batch_id": bid, "submitted_at": utc_now(), "custom_ids": [r.request_id for r in part], "model": self.provider.model, "prompt_version": PROMPT_VERSION})167 self.ledger["batches"].append(bid)168 ids.append(bid)169 self._save_ledger()170 return ids171172 def collect_batches(self, reqs: list[RewriteRequest], passages: dict[str, dict]) -> dict:173 by_id = merge_plan(self.job_dir, reqs) # all requests ever planned, not only the current selection174 collected = 0175 unresolved = 0176 pending_batches = []177 for bid in list(self.ledger["batches"]):178 meta = read_json(self.job_dir / "batches" / f"{bid}.json")179 if meta.get("collected"):180 continue181 st = self.provider.batch_status(bid) # type: ignore[attr-defined]182 if st["status"] != "ended":183 pending_batches.append(st)184 continue185 results = self.provider.batch_results(bid) # type: ignore[attr-defined]186 for cid, res in results.items():187 if cid in self.done:188 continue189 if cid not in by_id or by_id[cid].passage_id not in passages:190 unresolved += 1191 continue192 self._record(by_id[cid], passages[by_id[cid].passage_id], res, batch_id=bid)193 collected += 1194 if unresolved:195 meta["unresolved"] = unresolved196 write_json(self.job_dir / "batches" / f"{bid}.json", meta)197 self._save_ledger()198 continue # leave the batch uncollected so a later call with the full plan can finish it199 meta["collected"] = True200 meta["collected_at"] = utc_now()201 meta["status"] = st202 write_json(self.job_dir / "batches" / f"{bid}.json", meta)203 self._save_ledger()204 if self.ledger["spent_usd"] > self.budget:205 self._save_ledger()206 raise BudgetExceeded(f"spent {self.ledger['spent_usd']:.2f} USD > ceiling {self.budget:.2f} USD (collected results are kept)")207 return {"collected": collected, "unresolved": unresolved, "pending_batches": pending_batches, **self.ledger}208209 def write_manifest(self, extra: dict) -> None:210 write_json(self.job_dir / "manifest.json", {**run_header({"job": "rewrite"}), "provider": self.provider.name, "model": self.provider.model, "prompt_version": PROMPT_VERSION, "prompt_fingerprint": prompt_fingerprint(), "budget_usd": self.budget, "batch": self.batch, **extra})211