TeX 59%
Python 40.8%
1"""Anthropic Messages API wrapper for WP20 extraction.23Two modes share one request builder:4 - direct(): synchronous call (smoke tests, retries, ≤ 5 tries), cached by request hash.5 - batch(): Message Batches API (50 % price). Requests are written to results/raw/batches/<name>.requests.jsonl6 and results to results/raw/batches/<name>.results.jsonl; nothing is lost if the process dies.7Secrets: ANTHROPIC_API_KEY from the environment only (loaded from a .env OUTSIDE the repo if WP20_ENV_FILE is set).8Usage accounting: every response's usage block is appended to results/raw/api_usage.jsonl.9"""10from __future__ import annotations1112import hashlib13import json14import os15import time16from pathlib import Path1718import anthropic19import yaml2021ROOT = Path(__file__).resolve().parents[2]22RAW = ROOT / "results" / "raw"23RAW.mkdir(parents=True, exist_ok=True)24(RAW / "batches").mkdir(exist_ok=True)25CACHE = ROOT / "data" / "cache"26CACHE.mkdir(parents=True, exist_ok=True)2728PROMPTS = yaml.safe_load(open(ROOT / "configs" / "prompts.yaml"))29MODELS = yaml.safe_load(open(ROOT / "configs" / "models.yaml"))303132def _load_env() -> None:33 f = os.environ.get("WP20_ENV_FILE")34 if f and os.path.exists(f) and "ANTHROPIC_API_KEY" not in os.environ:35 for line in open(f, encoding="utf-8"):36 if line.startswith("ANTHROPIC_API_KEY="):37 os.environ["ANTHROPIC_API_KEY"] = line.split("=", 1)[1].strip()383940_load_env()41client = anthropic.Anthropic(max_retries=3, timeout=180)424344def model_params(tier: str) -> dict:45 """Per-model request parameters (thinking / effort) from configs/models.yaml."""46 m = MODELS["tiers"][tier]47 p = {"model": m["id"], "max_tokens": m["max_tokens"]}48 if m.get("thinking"):49 p["thinking"] = m["thinking"]50 if m.get("effort"):51 p["output_config"] = {"effort": m["effort"]}52 return p535455def image_request(tier: str, url: str, schema: dict) -> dict:56 p = model_params(tier)57 p["system"] = PROMPTS["image_system"]58 p["messages"] = [{"role": "user", "content": [59 {"type": "image", "source": {"type": "url", "url": url}},60 {"type": "text", "text": PROMPTS["image_user"]}]}]61 oc = p.get("output_config", {})62 oc["format"] = {"type": "json_schema", "schema": schema}63 p["output_config"] = oc64 return p656667def text_request(tier: str, text: str, schema: dict) -> dict:68 p = model_params(tier)69 p["system"] = PROMPTS["text_system"]70 p["messages"] = [{"role": "user", "content": PROMPTS["text_user"].format(text=text)}]71 oc = p.get("output_config", {})72 oc["format"] = {"type": "json_schema", "schema": schema}73 p["output_config"] = oc74 return p757677def _hash(params: dict) -> str:78 return hashlib.sha256(json.dumps(params, sort_keys=True, ensure_ascii=False).encode()).hexdigest()[:24]798081def _log_usage(tag: str, model: str, usage, extra: dict | None = None) -> None:82 rec = {"ts": time.time(), "tag": tag, "model": model,83 "input_tokens": getattr(usage, "input_tokens", None), "output_tokens": getattr(usage, "output_tokens", None),84 "cache_read": getattr(usage, "cache_read_input_tokens", None)}85 if extra:86 rec.update(extra)87 with open(RAW / "api_usage.jsonl", "a") as fh:88 fh.write(json.dumps(rec) + "\n")899091def parse_json_text(message) -> dict | None:92 if getattr(message, "stop_reason", None) == "refusal":93 return None94 for b in message.content:95 if b.type == "text":96 try:97 return json.loads(b.text)98 except json.JSONDecodeError:99 return None100 return None101102103def direct(params: dict, tag: str = "direct") -> dict:104 """Synchronous call with on-disk cache. Returns {'ok', 'data', 'usage', 'model', 'stop_reason', 'error'}."""105 h = _hash(params)106 cf = CACHE / f"{h}.json"107 if cf.exists():108 return json.load(open(cf))109 try:110 msg = client.messages.create(**params)111 except anthropic.APIStatusError as e:112 return {"ok": False, "error": f"{e.status_code}: {e.message}", "hash": h}113 data = parse_json_text(msg)114 out = {"ok": data is not None, "data": data, "model": msg.model, "stop_reason": msg.stop_reason,115 "usage": {"input_tokens": msg.usage.input_tokens, "output_tokens": msg.usage.output_tokens}, "hash": h}116 _log_usage(tag, msg.model, msg.usage)117 json.dump(out, open(cf, "w"), ensure_ascii=False)118 return out119120121# ---------------------------------------------------------------- batches122def batch_submit(name: str, requests: list[tuple[str, dict]]) -> list[str]:123 """Submit requests [(custom_id, params)] in chunks (≤ 10,000 per batch, ≤ 200 MB). Returns batch ids."""124 from anthropic.types.message_create_params import MessageCreateParamsNonStreaming125 from anthropic.types.messages.batch_create_params import Request126127 reqf = RAW / "batches" / f"{name}.requests.jsonl"128 with open(reqf, "w") as fh:129 for cid, p in requests:130 fh.write(json.dumps({"custom_id": cid, "params": p}, ensure_ascii=False) + "\n")131 ids = []132 chunk = 4000133 for i in range(0, len(requests), chunk):134 part = requests[i:i + chunk]135 b = client.messages.batches.create(requests=[136 Request(custom_id=cid, params=MessageCreateParamsNonStreaming(**p)) for cid, p in part])137 ids.append(b.id)138 with open(RAW / "batches" / f"{name}.batch_ids.txt", "a") as fh:139 fh.write(f"{b.id}\t{len(part)}\t{time.strftime('%Y-%m-%dT%H:%M:%S')}\n")140 return ids141142143def batch_status(batch_id: str):144 return client.messages.batches.retrieve(batch_id)145146147def batch_collect(name: str, batch_id: str) -> int:148 """Append results of an ended batch to <name>.results.jsonl (idempotent by custom_id). Returns n written."""149 resf = RAW / "batches" / f"{name}.results.jsonl"150 seen = set()151 if resf.exists():152 for line in open(resf):153 try:154 seen.add(json.loads(line)["custom_id"])155 except Exception:156 pass157 n = 0158 with open(resf, "a") as fh:159 for r in client.messages.batches.results(batch_id):160 if r.custom_id in seen:161 continue162 rec = {"custom_id": r.custom_id, "type": r.result.type, "batch_id": batch_id}163 if r.result.type == "succeeded":164 msg = r.result.message165 rec["model"] = msg.model166 rec["stop_reason"] = msg.stop_reason167 rec["data"] = parse_json_text(msg)168 rec["usage"] = {"input_tokens": msg.usage.input_tokens, "output_tokens": msg.usage.output_tokens}169 _log_usage(f"batch:{name}", msg.model, msg.usage)170 else:171 err = getattr(r.result, "error", None)172 inner = getattr(err, "error", None) if err else None173 rec["error"] = getattr(inner, "type", None) or (getattr(err, "type", None) if err else None)174 rec["error_message"] = getattr(inner, "message", None) if inner else (str(err) if err else None)175 fh.write(json.dumps(rec, ensure_ascii=False) + "\n")176 n += 1177 return n178