#!/usr/bin/env python3 # -*- coding: utf-8 -*- """Durable Anthropic Message Batch history and spend-safe management ops. Active run files (queue / state / results) still live in util.translation. This module keeps an append-only index so batch ids and custom_id maps survive fetch and clear, enabling cancel / usage / redownload / resume without re-submitting (and re-billing) work. """ from __future__ import annotations import copy import time from datetime import datetime, timezone from pathlib import Path from typing import Any, Iterable, Optional from util.translation import ( BATCH_LOCK, BATCH_QUEUE_FILE, BATCH_RESULTS_FILE, BATCH_STATE_FILE, _batch_file_lock, _get_anthropic_client, _read_batch_file, _write_batch_file, getPricingConfig, ) BATCH_HISTORY_FILE = Path("log/batch_history.json") # Local lifecycle statuses (not 1:1 with Anthropic processing_status). STATUS_QUEUED = "queued" STATUS_SUBMITTED = "submitted" STATUS_CANCELING = "canceling" STATUS_ENDED = "ended" STATUS_FETCHED = "fetched" STATUS_CONSUMED = "consumed" STATUS_CANCELED = "canceled" STATUS_ERROR = "error" TERMINAL_STATUSES = frozenset({ STATUS_CONSUMED, STATUS_CANCELED, STATUS_ERROR, }) CANCELABLE_API = frozenset({"in_progress"}) def _utc_now() -> str: return datetime.now(timezone.utc).replace(microsecond=0).isoformat() def _empty_history() -> dict: return {"batches": []} def _read_history_unlocked() -> dict: """Read history under an existing batch file lock.""" data = _read_batch_file(BATCH_HISTORY_FILE) if not data: return _empty_history() if not isinstance(data.get("batches"), list): print(f"[BATCH] Corrupt history (missing batches list): {BATCH_HISTORY_FILE}", flush=True) return _empty_history() return data def _write_history_unlocked(data: dict) -> None: _write_batch_file(BATCH_HISTORY_FILE, data) def read_history() -> dict: """Return the full history document (copy-safe for callers).""" with _batch_file_lock(): return copy.deepcopy(_read_history_unlocked()) def _find_entry(history: dict, batch_id: str) -> Optional[dict]: for entry in history.get("batches", []): if entry.get("id") == batch_id: return entry return None def upsert_history_entry(batch_id: str, **fields: Any) -> dict: """Create or update one history row. Returns the updated entry (copy).""" with BATCH_LOCK: with _batch_file_lock(): history = _read_history_unlocked() entry = _find_entry(history, batch_id) now = _utc_now() if entry is None: entry = { "id": batch_id, "created_at": now, "updated_at": now, "status": STATUS_SUBMITTED, "model": "", "request_count": 0, "file_set": [], "cost_estimate": None, "custom_ids": {}, "usage": None, "actual_cost": None, "notes": "", "request_counts": None, "api_status": None, } history.setdefault("batches", []).append(entry) for key, value in fields.items(): if value is not None or key in ("notes", "usage", "actual_cost", "cost_estimate", "request_counts"): entry[key] = value entry["updated_at"] = now _write_history_unlocked(history) return copy.deepcopy(entry) def record_submit( batches: list[dict], *, model: str = "", file_set: Optional[list] = None, cost_estimate: Optional[dict] = None, ) -> None: """Record newly submitted Anthropic batches into durable history.""" file_set = list(file_set or []) for info in batches: bid = info.get("id") if not bid: continue custom_ids = dict(info.get("custom_ids") or {}) upsert_history_entry( bid, status=STATUS_SUBMITTED, model=model or (cost_estimate or {}).get("model") or "", request_count=len(custom_ids), file_set=file_set, cost_estimate=copy.deepcopy(cost_estimate) if cost_estimate else None, custom_ids=custom_ids, api_status="in_progress", notes="", ) def record_fetch( batch_ids: Iterable[str], *, succeeded: int = 0, errored: int = 0, usage: Optional[dict] = None, actual_cost: Optional[float] = None, ) -> None: """Mark batches as fetched after results land locally.""" note = f"fetched ok={succeeded} err={errored}" for bid in batch_ids: fields = { "status": STATUS_FETCHED, "api_status": "ended", "notes": note, } if usage is not None: fields["usage"] = copy.deepcopy(usage) if actual_cost is not None: fields["actual_cost"] = actual_cost upsert_history_entry(bid, **fields) def mark_batches_consumed(batch_ids: Iterable[str]) -> None: """Mark history rows consumed after a successful consume + clear.""" for bid in batch_ids: upsert_history_entry(bid, status=STATUS_CONSUMED, notes="consumed") def mark_batches_canceled(batch_ids: Iterable[str], *, api_status: str = "canceling") -> None: local = STATUS_CANCELED if api_status == "ended" else STATUS_CANCELING for bid in batch_ids: upsert_history_entry( bid, status=local, api_status=api_status, notes=f"cancel requested ({api_status})", ) def list_local_batches(*, refresh_live: bool = False) -> list[dict]: """Return history entries newest-first. Optionally refresh non-terminal via API.""" history = read_history() entries = list(history.get("batches") or []) if refresh_live: for entry in entries: status = entry.get("status") if status in TERMINAL_STATUSES or status == STATUS_CONSUMED: continue if status in (STATUS_FETCHED, STATUS_QUEUED): continue try: refresh_batch_status(entry["id"]) except Exception as exc: upsert_history_entry(entry["id"], notes=f"refresh failed: {exc}") history = read_history() entries = list(history.get("batches") or []) entries.sort(key=lambda e: e.get("created_at") or "", reverse=True) return entries def refresh_batch_status(batch_id: str) -> dict: """Retrieve live Anthropic status and update the history row.""" client = _get_anthropic_client() b = client.messages.batches.retrieve(batch_id) api_status = getattr(b, "processing_status", None) or "" counts = getattr(b, "request_counts", None) counts_dict = None if counts is not None: counts_dict = { k: getattr(counts, k, 0) or 0 for k in ("processing", "succeeded", "errored", "canceled", "expired") } entry = None with _batch_file_lock(): history = _read_history_unlocked() entry = _find_entry(history, batch_id) prev_status = (entry or {}).get("status") # Map API status onto local lifecycle without clobbering fetched/consumed. local_status = None if prev_status in (STATUS_FETCHED, STATUS_CONSUMED): local_status = prev_status elif api_status == "in_progress": local_status = STATUS_SUBMITTED elif api_status == "canceling": local_status = STATUS_CANCELING elif api_status == "ended": # Ended after cancel vs normal end - inspect counts when available. if counts_dict and (counts_dict.get("canceled") or 0) > 0 and (counts_dict.get("succeeded") or 0) == 0: local_status = STATUS_CANCELED else: local_status = STATUS_ENDED fields: dict[str, Any] = { "api_status": api_status, "request_counts": counts_dict, } if local_status: fields["status"] = local_status return upsert_history_entry(batch_id, **fields) def cancel_batches(batch_ids: Iterable[str]) -> list[dict]: """Cancel in-progress batches. Never submits new work. Also updates active batch_state so resume does not keep polling canceled ids. """ client = _get_anthropic_client() results = [] canceled_active = [] for bid in batch_ids: try: b = client.messages.batches.retrieve(bid) api_status = getattr(b, "processing_status", "") or "" if api_status not in CANCELABLE_API: results.append({"id": bid, "ok": False, "error": f"not cancelable ({api_status})"}) upsert_history_entry(bid, notes=f"cancel skipped: {api_status}", api_status=api_status) continue b = client.messages.batches.cancel(bid) new_status = getattr(b, "processing_status", "canceling") or "canceling" mark_batches_canceled([bid], api_status=new_status) canceled_active.append(bid) results.append({"id": bid, "ok": True, "api_status": new_status}) except Exception as exc: results.append({"id": bid, "ok": False, "error": str(exc)}) upsert_history_entry(bid, notes=f"cancel failed: {exc}", status=STATUS_ERROR) if canceled_active: _remove_active_batch_ids(canceled_active) return results def _remove_active_batch_ids(batch_ids: list[str]) -> None: """Drop canceled ids from active state so poll/resume does not wait on them.""" id_set = set(batch_ids) with BATCH_LOCK: with _batch_file_lock(): state = _read_batch_file(BATCH_STATE_FILE) batches = [b for b in (state.get("batches") or []) if b.get("id") not in id_set] if not batches: # No active submitted batches left. if state.get("status") == "fetched" or state.get("batch_ids"): state["batches"] = [] _write_batch_file(BATCH_STATE_FILE, state) else: try: if BATCH_STATE_FILE.exists(): BATCH_STATE_FILE.unlink() except Exception: pass else: state["batches"] = batches _write_batch_file(BATCH_STATE_FILE, state) def _usage_from_message(u) -> dict: cr = getattr(u, "cache_read_input_tokens", 0) or 0 cw = getattr(u, "cache_creation_input_tokens", 0) or 0 inp = getattr(u, "input_tokens", 0) or 0 out = getattr(u, "output_tokens", 0) or 0 # Adaptive thinking may surface under several names depending on SDK version. thinking = ( getattr(u, "thinking_tokens", None) or getattr(u, "output_thinking_tokens", None) or 0 ) or 0 return { "input_tokens": inp, "output_tokens": out, "cache_read_input_tokens": cr, "cache_creation_input_tokens": cw, "thinking_tokens": thinking, } def _price_usage(usage: dict, model: str) -> float: """Price real batch usage with cache multipliers and 50% batch discount.""" pricing = getPricingConfig(model) br = pricing["inputAPICost"] / 1_000_000 orr = pricing["outputAPICost"] / 1_000_000 cr = usage.get("cache_read_input_tokens", 0) or 0 cw = usage.get("cache_creation_input_tokens", 0) or 0 inp = usage.get("input_tokens", 0) or 0 out = usage.get("output_tokens", 0) or 0 thinking = usage.get("thinking_tokens", 0) or 0 # Thinking tokens are billed as output on Anthropic. raw = cr * br * 0.10 + cw * br * 2.00 + inp * br + (out + thinking) * orr return raw * 0.50 def _sum_usage_from_results(client, batch_id: str, custom_ids: dict) -> tuple[dict, int, int]: """Stream batch results and sum usage. Returns (usage, succeeded, errored).""" totals = { "input_tokens": 0, "output_tokens": 0, "cache_read_input_tokens": 0, "cache_creation_input_tokens": 0, "thinking_tokens": 0, } succeeded = 0 errored = 0 for r in client.messages.batches.results(batch_id): res = r.result if res.type != "succeeded": errored += 1 continue succeeded += 1 part = _usage_from_message(res.message.usage) for k, v in part.items(): totals[k] = totals.get(k, 0) + (v or 0) # custom_ids unused for summing but kept for API symmetry / future filters _ = custom_ids return totals, succeeded, errored def usage_for_batch(batch_id: str, model: Optional[str] = None) -> dict: """Sum real billed tokens from Anthropic results and persist onto history.""" history = read_history() entry = _find_entry(history, batch_id) if entry is None: raise ValueError(f"Unknown batch id (not in local history): {batch_id}") client = _get_anthropic_client() # Ensure batch has ended before streaming results. b = client.messages.batches.retrieve(batch_id) api_status = getattr(b, "processing_status", "") or "" if api_status != "ended": raise ValueError(f"Batch {batch_id} is not ended (status={api_status})") custom_ids = dict(entry.get("custom_ids") or {}) usage, succeeded, errored = _sum_usage_from_results(client, batch_id, custom_ids) use_model = model or entry.get("model") or "" cost = _price_usage(usage, use_model) if use_model else None updated = upsert_history_entry( batch_id, usage=usage, actual_cost=cost, api_status="ended", notes=f"usage ok={succeeded} err={errored}", ) return { "id": batch_id, "usage": usage, "actual_cost": cost, "succeeded": succeeded, "errored": errored, "model": use_model, "entry": updated, } def _result_entry_from_message(msg) -> dict: text = "".join(getattr(b, "text", "") or "" for b in msg.content) u = msg.usage part = _usage_from_message(u) cr = part["cache_read_input_tokens"] cw = part["cache_creation_input_tokens"] inp = part["input_tokens"] out = part["output_tokens"] entry = { "text": text, "prompt_tokens": inp + cr + cw, "completion_tokens": out, "cache_read_input_tokens": cr, "cache_creation_input_tokens": cw, } if part.get("thinking_tokens"): entry["thinking_tokens"] = part["thinking_tokens"] return entry def download_batch_results( batch_id: str, custom_ids: dict, *, client=None, ) -> tuple[dict, list, dict]: """Download one ended batch into a cache-key -> result map. Returns (results, errored_list, usage_totals). """ client = client or _get_anthropic_client() results, errored = {}, [] usage_totals = { "input_tokens": 0, "output_tokens": 0, "cache_read_input_tokens": 0, "cache_creation_input_tokens": 0, "thinking_tokens": 0, } for r in client.messages.batches.results(batch_id): key = custom_ids.get(r.custom_id) if key is None: continue res = r.result if res.type != "succeeded": detail = res.type err = getattr(res, "error", None) if err is not None: inner = getattr(err, "error", err) detail = ( f"{res.type} | {getattr(inner, 'type', '')}: " f"{str(getattr(inner, 'message', '') or err)[:200]}" ) errored.append((r.custom_id, detail)) continue msg = res.message results[key] = _result_entry_from_message(msg) part = _usage_from_message(msg.usage) for k, v in part.items(): usage_totals[k] = usage_totals.get(k, 0) + (v or 0) return results, errored, usage_totals def redownload_batch(batch_id: str) -> dict: """Re-fetch results into batch_results.json using stored custom_ids. No re-submit.""" history = read_history() entry = _find_entry(history, batch_id) if entry is None: raise ValueError(f"Unknown batch id (not in local history): {batch_id}") custom_ids = dict(entry.get("custom_ids") or {}) if not custom_ids: raise ValueError(f"Batch {batch_id} has no stored custom_ids - cannot redownload") client = _get_anthropic_client() b = client.messages.batches.retrieve(batch_id) api_status = getattr(b, "processing_status", "") or "" if api_status != "ended": raise ValueError(f"Batch {batch_id} is not ended (status={api_status})") results, errored, usage = download_batch_results(batch_id, custom_ids, client=client) model = entry.get("model") or "" cost = _price_usage(usage, model) if model else None import util.translation as T with T.BATCH_LOCK: with _batch_file_lock(): merged = _read_batch_file(BATCH_RESULTS_FILE) merged.update(results) _write_batch_file(BATCH_RESULTS_FILE, merged) # Activate fetched state without re-submit; keep custom_ids in history. _write_batch_file( BATCH_STATE_FILE, {"status": "fetched", "batch_ids": [batch_id], "batches": []}, ) # Queue is no longer needed for consume. try: if BATCH_QUEUE_FILE.exists(): BATCH_QUEUE_FILE.unlink() except Exception: pass T._batch_results = None record_fetch([batch_id], succeeded=len(results), errored=len(errored), usage=usage, actual_cost=cost) return { "id": batch_id, "succeeded": len(results), "errored": len(errored), "errors": errored[:20], "usage": usage, "actual_cost": cost, } def activate_for_resume(batch_id: str) -> str: """Ensure active files match a history entry for Translation-tab resume. Returns the batchRunState string to pass as batch_resume_state: 'submitted' | 'fetched'. """ history = read_history() entry = _find_entry(history, batch_id) if entry is None: raise ValueError(f"Unknown batch id: {batch_id}") status = entry.get("status") custom_ids = dict(entry.get("custom_ids") or {}) if status in (STATUS_FETCHED, STATUS_ENDED, STATUS_CONSUMED): # Prefer local results; redownload if missing. with _batch_file_lock(): results = _read_batch_file(BATCH_RESULTS_FILE) if not results: if status == STATUS_CONSUMED: raise ValueError( f"Batch {batch_id} was already consumed and local results are gone. " "Use Redownload first." ) redownload_batch(batch_id) else: with BATCH_LOCK: with _batch_file_lock(): _write_batch_file( BATCH_STATE_FILE, {"status": "fetched", "batch_ids": [batch_id], "batches": []}, ) return "fetched" if status in (STATUS_SUBMITTED, STATUS_CANCELING, STATUS_ENDED): # Restore active state so poll/fetch can continue. with BATCH_LOCK: with _batch_file_lock(): state = _read_batch_file(BATCH_STATE_FILE) batches = list(state.get("batches") or []) existing = {b.get("id") for b in batches} if batch_id not in existing: batches.append({"id": batch_id, "custom_ids": custom_ids}) state = { "batches": batches, "submitted_at": state.get("submitted_at") or entry.get("created_at"), "model": entry.get("model") or state.get("model") or "", "file_set": entry.get("file_set") or state.get("file_set") or [], "cost_estimate": entry.get("cost_estimate") or state.get("cost_estimate"), } _write_batch_file(BATCH_STATE_FILE, state) if status == STATUS_ENDED: # Already ended on Anthropic - fetch path via resume submitted → poll sees ended. return "submitted" return "submitted" if status == STATUS_CANCELED: raise ValueError(f"Batch {batch_id} was canceled - nothing to resume") raise ValueError(f"Batch {batch_id} cannot be resumed from status={status}") def missing_result_count() -> tuple[int, int]: """Return (present, expected) result counts for the active fetched run. expected is summed from history request_count for active fetched batch ids. """ with _batch_file_lock(): results = _read_batch_file(BATCH_RESULTS_FILE) state = _read_batch_file(BATCH_STATE_FILE) history = _read_history_unlocked() present = len(results) ids = list(state.get("batch_ids") or []) if not ids: ids = [b.get("id") for b in (state.get("batches") or []) if b.get("id")] expected = 0 for bid in ids: entry = _find_entry(history, bid) if entry: expected += int(entry.get("request_count") or len(entry.get("custom_ids") or {})) if not expected: expected = present return present, expected def active_fetched_batch_ids() -> list[str]: """Batch ids associated with the current fetched active run.""" with _batch_file_lock(): state = _read_batch_file(BATCH_STATE_FILE) ids = list(state.get("batch_ids") or []) if ids: return ids return [b.get("id") for b in (state.get("batches") or []) if b.get("id")] def on_clear_active_files(*, had_results: bool, fetched_ids: Optional[list] = None) -> None: """Called from clearBatchFiles - mark fetched ids consumed, never wipe history.""" if not had_results and not fetched_ids: return ids = list(fetched_ids or []) if not ids and had_results: # Fall back: mark any currently-fetched history rows as consumed. history = read_history() ids = [e["id"] for e in history.get("batches", []) if e.get("status") == STATUS_FETCHED] if ids: mark_batches_consumed(ids)