feat(extraction): Batch API extractor (~50% cheaper), fix stale cost banner

Route entity extraction through the Message Batches API instead of
synchronous calls. Offline pipeline, so the 50% batch discount applies
with no quality change. Resumable via .batch_state.json (resumes an
in-flight batch instead of resubmitting). Drop the no-op prompt-cache
prefix (below Haiku 4.5's 4096-token cacheable minimum) and correct the
cost banner to ~$0.015 per 10 papers.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

config.py

@@ -17,6 +17,7 @@ TRIALS_PATH = DATA_DIR / "trials" / "trials.jsonl"
 ENTITIES_PATH = DATA_DIR / "extracted" / "entities.jsonl"
 CANONICAL_IDS_PATH = DATA_DIR / "extracted" / "canonical_ids.json"
 EXTRACTION_PROGRESS_PATH = DATA_DIR / "extracted" / ".progress.json"
+EXTRACTION_BATCH_STATE_PATH = DATA_DIR / "extracted" / ".batch_state.json"
 GRAPH_PICKLE_PATH = DATA_DIR / "graph" / "als_graph.pkl"
 GRAPH_JSON_PATH = DATA_DIR / "graph" / "als_graph.json"
 CHROMA_DIR = DATA_DIR / "chroma"

extraction/extractor.py

@@ -1,26 +1,30 @@
 """
-Claude Haiku entity extractor.
-Batches 20 papers per API call; resumable via .progress.json.
-Uses full_text when available, otherwise abstract.
-Prompt caching on system + tools reduces per-call cost ~40%.
+Claude Haiku entity extractor — Batch API.
+
+Submits extraction requests through the Message Batches API (50% cheaper than
+synchronous calls) since this is an offline, non-latency-sensitive pipeline.
+Batches 20 papers per request; uses full_text when available, else abstract.
+Resumable at two levels: completed PMIDs are tracked in .progress.json, and an
+in-flight batch id is persisted in .batch_state.json so an interrupted run
+resumes polling the same (already-paid-for) batch instead of resubmitting.
 """
 from __future__ import annotations
 
 import json
-import threading
 import time
-from concurrent.futures import ThreadPoolExecutor, as_completed
 from pathlib import Path
 
 import anthropic
+from anthropic.types.message_create_params import MessageCreateParamsNonStreaming
+from anthropic.types.messages.batch_create_params import Request
 from rich.progress import BarColumn, MofNCompleteColumn, Progress, TextColumn, TimeElapsedColumn
 
 from config import (
     ENTITIES_PATH,
     EXTRACTION_BATCH_SIZE,
+    EXTRACTION_BATCH_STATE_PATH,
     EXTRACTION_MODEL,
     EXTRACTION_PROGRESS_PATH,
-    EXTRACTION_WORKERS,
     PAPERS_PATH,
 )
 from extraction.normalizer import CanonicalRegistry, guess_entity_type, normalize_entity
@@ -30,6 +34,10 @@ from tools import EXTRACTION_TOOLS
 
 _logger = get_logger("extraction.extractor")
 
+# Seconds between batch status polls. Batches usually finish in well under an
+# hour; the ceiling is 24h.
+_POLL_INTERVAL_S = 30
+
 _EXTRACTION_SYSTEM = """\
 You are a biomedical NLP expert specializing in ALS (amyotrophic lateral sclerosis).
 Extract entities and relationships from each paper using the extract_entities tool.
@@ -46,21 +54,26 @@ def extract_all(
     papers_path: Path = PAPERS_PATH,
     entities_path: Path = ENTITIES_PATH,
     progress_path: Path = EXTRACTION_PROGRESS_PATH,
+    batch_state_path: Path = EXTRACTION_BATCH_STATE_PATH,
     client: anthropic.Anthropic | None = None,
 ) -> list[PaperExtractionResult]:
-    """Extract entities from all papers. Skips already-processed PMIDs.
+    """Extract entities from all papers via the Batch API. Skips done PMIDs.
 
-    Runs EXTRACTION_WORKERS batches in parallel. A lock serializes file writes
-    and progress saves so threads don't corrupt each other.
+    Runs one main batch round (20 papers/request), then an individual retry
+    round for any papers Claude skipped, then records empty results for papers
+    still missing so they aren't re-attempted on the next run.
     """
     if client is None:
         client = anthropic.Anthropic()
 
     papers = _load_papers(papers_path)
+    paper_by_pmid = {p.pmid: p for p in papers}
     done_pmids = _load_progress(progress_path)
 
     pending = [p for p in papers if p.pmid not in done_pmids]
-    _logger.info(f"{len(papers)} papers total; {len(done_pmids)} already processed; {len(pending)} pending")
+    _logger.info(
+        f"{len(papers)} papers total; {len(done_pmids)} already processed; {len(pending)} pending"
+    )
 
     if not pending:
         return []
@@ -68,140 +81,202 @@ def extract_all(
     registry = CanonicalRegistry()
     entities_path.parent.mkdir(parents=True, exist_ok=True)
 
-    batches = [pending[i : i + EXTRACTION_BATCH_SIZE] for i in range(0, len(pending), EXTRACTION_BATCH_SIZE)]
+    all_results: list[PaperExtractionResult] = []
+    with open(entities_path, "a", encoding="utf-8") as out_f:
+        # Round 1 — main batches of EXTRACTION_BATCH_SIZE papers each.
+        batches = [
+            pending[i : i + EXTRACTION_BATCH_SIZE]
+            for i in range(0, len(pending), EXTRACTION_BATCH_SIZE)
+        ]
+        main_map = {f"batch-{i}": batch for i, batch in enumerate(batches)}
+        round1 = _run_batch_round(client, main_map, registry, paper_by_pmid, batch_state_path)
+        _write_results(out_f, round1, done_pmids, progress_path, registry)
+        all_results.extend(round1)
+
+        found = {r.pmid for r in round1}
+        missed = [p for p in pending if p.pmid not in found]
+
+        # Round 2 — retry missed papers one per request.
+        if missed:
+            _logger.info(f"Retrying {len(missed)} missed papers individually")
+            retry_map = {f"retry-{p.pmid}": [p] for p in missed}
+            round2 = _run_batch_round(client, retry_map, registry, paper_by_pmid, batch_state_path)
+            _write_results(out_f, round2, done_pmids, progress_path, registry)
+            all_results.extend(round2)
+            found |= {r.pmid for r in round2}
+
+        # Record empty results for anything still missing after retry.
+        still_missing = [p for p in pending if p.pmid not in found]
+        if still_missing:
+            empties = []
+            for p in still_missing:
+                _logger.warning(
+                    f"No extraction result for PMID {p.pmid} after retry — recording empty"
+                )
+                empties.append(PaperExtractionResult(pmid=p.pmid, entities=[], relationships=[]))
+            _write_results(out_f, empties, done_pmids, progress_path, registry)
+            all_results.extend(empties)
+
+    return all_results
+
+
+def _run_batch_round(
+    client: anthropic.Anthropic,
+    custom_id_to_papers: dict[str, list[ALSPaper]],
+    registry: CanonicalRegistry,
+    paper_by_pmid: dict[str, ALSPaper],
+    state_path: Path,
+) -> list[PaperExtractionResult]:
+    """Submit (or resume) one batch, poll to completion, and parse its results.
+
+    Persists the batch id + custom_id→PMID mapping to state_path on submit so an
+    interrupted process resumes the same batch. Clears the state on completion.
+    """
+    batch = None
+    state = _load_batch_state(state_path)
+    if state and state.get("batch_id"):
+        try:
+            existing = client.messages.batches.retrieve(state["batch_id"])
+        except anthropic.NotFoundError:
+            _logger.warning("Persisted batch id not found — submitting a fresh batch")
+        else:
+            if existing.processing_status in {"in_progress", "validating", "finalizing", "ended"}:
+                _logger.info(f"Resuming in-flight batch {existing.id}")
+                batch = existing
+                # Rebuild the mapping from persisted PMIDs so results match.
+                custom_id_to_papers = {
+                    cid: [paper_by_pmid[pmid] for pmid in pmids if pmid in paper_by_pmid]
+                    for cid, pmids in state.get("papers", {}).items()
+                }
+
+    if batch is None:
+        requests = [
+            Request(custom_id=cid, params=_build_params(papers))
+            for cid, papers in custom_id_to_papers.items()
+        ]
+        batch = client.messages.batches.create(requests=requests)
+        _save_batch_state(
+            state_path,
+            {
+                "batch_id": batch.id,
+                "papers": {
+                    cid: [p.pmid for p in papers] for cid, papers in custom_id_to_papers.items()
+                },
+            },
+        )
+        _logger.info(f"Submitted batch {batch.id} with {len(requests)} requests")
+
+    batch = _poll_until_done(client, batch)
+
     results: list[PaperExtractionResult] = []
-    write_lock = threading.Lock()
-
-    with (
-        open(entities_path, "a", encoding="utf-8") as out_f,
-        Progress(
-            TextColumn("[cyan]{task.description}[/cyan]"),
-            BarColumn(),
-            MofNCompleteColumn(),
-            TimeElapsedColumn(),
-        ) as progress,
-    ):
-        task = progress.add_task("Extracting entities", total=len(pending))
-
-        def _process_batch(batch: list[ALSPaper]) -> list[PaperExtractionResult]:
-            return _extract_batch(client, batch, registry)
-
-        with ThreadPoolExecutor(max_workers=EXTRACTION_WORKERS) as pool:
-            futures = {pool.submit(_process_batch, b): b for b in batches}
-            for future in as_completed(futures):
-                batch_results = future.result()
-                with write_lock:
-                    for result in batch_results:
-                        out_f.write(json.dumps(result.to_dict()) + "\n")
-                        done_pmids.add(result.pmid)
-                        results.append(result)
-                    out_f.flush()
-                    _save_progress(progress_path, done_pmids)
-                    registry.save()
-                    progress.advance(task, len(futures[future]))
+    for res in client.messages.batches.results(batch.id):
+        papers = custom_id_to_papers.get(res.custom_id, [])
+        local_by_pmid = {p.pmid: p for p in papers}
+        if res.result.type == "succeeded":
+            results.extend(
+                _parse_response_blocks(res.result.message.content, local_by_pmid, registry)
+            )
+        elif res.result.type == "errored":
+            _logger.warning(f"Batch request {res.custom_id} errored: {res.result.error}")
+        else:
+            _logger.warning(f"Batch request {res.custom_id} {res.result.type}")
 
+    _clear_batch_state(state_path)
     return results
 
 
-def _extract_batch(
-    client: anthropic.Anthropic,
-    batch: list[ALSPaper],
+def _poll_until_done(client: anthropic.Anthropic, batch) -> object:
+    """Poll a batch until it reaches a terminal status, showing progress."""
+    total = (
+        batch.request_counts.processing
+        + batch.request_counts.succeeded
+        + batch.request_counts.errored
+        + batch.request_counts.canceled
+        + batch.request_counts.expired
+    )
+    with Progress(
+        TextColumn("[cyan]{task.description}[/cyan]"),
+        BarColumn(),
+        MofNCompleteColumn(),
+        TimeElapsedColumn(),
+    ) as progress:
+        task = progress.add_task("Extracting entities (batch)", total=total or None)
+        while batch.processing_status != "ended":
+            if batch.processing_status in {"canceling", "canceled", "expired"}:
+                _logger.warning(f"Batch {batch.id} ended early with status {batch.processing_status}")
+                break
+            time.sleep(_POLL_INTERVAL_S)
+            batch = client.messages.batches.retrieve(batch.id)
+            counts = batch.request_counts
+            completed = counts.succeeded + counts.errored + counts.canceled + counts.expired
+            progress.update(task, completed=completed)
+        progress.update(task, completed=total)
+    return batch
+
+
+def _build_params(batch: list[ALSPaper]) -> MessageCreateParamsNonStreaming:
+    """Build the per-request Messages params for a batch of papers.
+
+    system + tools are identical across every request, but on Haiku 4.5 the
+    combined prefix is far below the 4096-token minimum cacheable size, so
+    prompt caching would silently no-op — we don't set cache_control here.
+    """
+    return MessageCreateParamsNonStreaming(
+        model=EXTRACTION_MODEL,
+        max_tokens=8192,
+        system=_EXTRACTION_SYSTEM,
+        tools=list(EXTRACTION_TOOLS),
+        tool_choice={"type": "any"},
+        messages=[{"role": "user", "content": _format_batch(batch)}],
+    )
+
+
+def _parse_response_blocks(
+    blocks: list,
+    paper_by_pmid: dict[str, ALSPaper],
     registry: CanonicalRegistry,
 ) -> list[PaperExtractionResult]:
-    """Send a batch of papers to Claude and collect one extract_entities call per paper."""
-    paper_by_pmid = {p.pmid: p for p in batch}
-    content_blocks = _call_claude(client, batch)
-
+    """Parse extract_entities tool_use blocks from one response into results."""
     results: list[PaperExtractionResult] = []
-    for block in content_blocks:
+    for block in blocks:
         if block.type != "tool_use" or block.name != "extract_entities":
             continue
 
         inp = block.input
         pmid = str(inp.get("pmid", ""))
         if not pmid or pmid not in paper_by_pmid:
-            _logger.warning(f"Extracted PMID {pmid!r} not in batch — skipping")
+            _logger.warning(f"Extracted PMID {pmid!r} not in request — skipping")
             continue
 
         paper = paper_by_pmid[pmid]
         entities = _parse_entities(inp.get("entities", []), pmid, registry)
         relationships = _parse_relationships(inp.get("relationships", []), pmid, registry)
 
-        result = PaperExtractionResult(
-            pmid=pmid,
-            entities=entities,
-            relationships=relationships,
+        results.append(
+            PaperExtractionResult(pmid=pmid, entities=entities, relationships=relationships)
         )
-        results.append(result)
-        _logger.info(f"PMID {pmid}: {len(entities)} entities, {len(relationships)} relationships")
-
-        # Mark paper entity_names (used downstream by RAG indexer on re-index)
         paper.entity_names = [e.canonical_id for e in entities]
-
-    # Retry any papers Claude missed — send them individually
-    found_pmids = {r.pmid for r in results}
-    missed = [p for p in batch if p.pmid not in found_pmids]
-    if missed:
-        _logger.info(f"Retrying {len(missed)} missed papers individually")
-        for paper in missed:
-            retry_results = _call_claude(client, [paper])
-            for block in retry_results:
-                if block.type != "tool_use" or block.name != "extract_entities":
-                    continue
-                inp = block.input
-                pmid = str(inp.get("pmid", ""))
-                if not pmid or pmid not in paper_by_pmid:
-                    continue
-                entities = _parse_entities(inp.get("entities", []), pmid, registry)
-                relationships = _parse_relationships(inp.get("relationships", []), pmid, registry)
-                results.append(PaperExtractionResult(pmid=pmid, entities=entities, relationships=relationships))
-                paper_by_pmid[pmid].entity_names = [e.canonical_id for e in entities]
-                found_pmids.add(pmid)
-                _logger.info(f"Retry succeeded for PMID {pmid}")
-            time.sleep(0.5)
-
-    # Any still-missing after retry → record empty so they're not re-attempted
-    for p in batch:
-        if p.pmid not in found_pmids:
-            _logger.warning(f"No extraction result for PMID {p.pmid} after retry — recording empty")
-            results.append(PaperExtractionResult(pmid=p.pmid, entities=[], relationships=[]))
+        _logger.info(f"PMID {pmid}: {len(entities)} entities, {len(relationships)} relationships")
 
     return results
 
 
-def _call_claude(client: anthropic.Anthropic, batch: list[ALSPaper]) -> list:
-    """Raw Claude call — returns response.content blocks.
-
-    Prompt caching: system and tools are static across all calls; adding
-    cache_control to the last tool + system block caches the entire prefix
-    (tools render before system in the API token order). Cache reads cost
-    ~10% of normal input price, halving the effective per-call overhead.
-    """
-    # Cache the static system+tools prefix across batch calls
-    cached_system = [{"type": "text", "text": _EXTRACTION_SYSTEM, "cache_control": {"type": "ephemeral"}}]
-    cached_tools = list(EXTRACTION_TOOLS)
-    if cached_tools:
-        last = dict(cached_tools[-1])
-        last["cache_control"] = {"type": "ephemeral"}
-        cached_tools[-1] = last
-
-    def _request() -> list:
-        response = client.messages.create(
-            model=EXTRACTION_MODEL,
-            max_tokens=8192,
-            system=cached_system,
-            tools=cached_tools,
-            tool_choice={"type": "any"},
-            messages=[{"role": "user", "content": _format_batch(batch)}],
-        )
-        return response.content
-
-    try:
-        return _request()
-    except anthropic.RateLimitError:
-        _logger.warning("Rate limited — sleeping 30s")
-        time.sleep(30)
-        return _request()
+def _write_results(
+    out_f,
+    results: list[PaperExtractionResult],
+    done_pmids: set[str],
+    progress_path: Path,
+    registry: CanonicalRegistry,
+) -> None:
+    """Append results to the output file and advance the resumability trackers."""
+    if not results:
+        return
+    for result in results:
+        out_f.write(json.dumps(result.to_dict()) + "\n")
+        done_pmids.add(result.pmid)
+    out_f.flush()
+    _save_progress(progress_path, done_pmids)
+    registry.save()
 
 
 def _format_batch(batch: list[ALSPaper]) -> str:
@@ -299,3 +374,18 @@ def _load_progress(path: Path) -> set[str]:
 def _save_progress(path: Path, done: set[str]) -> None:
     path.parent.mkdir(parents=True, exist_ok=True)
     path.write_text(json.dumps(sorted(done)))
+
+
+def _load_batch_state(path: Path) -> dict | None:
+    if path.exists():
+        return json.loads(path.read_text())
+    return None
+
+
+def _save_batch_state(path: Path, state: dict) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(state))
+
+
+def _clear_batch_state(path: Path) -> None:
+    path.unlink(missing_ok=True)

scripts/extract_entities.py

@@ -63,7 +63,10 @@ def main() -> None:
         console.print(f"[dim]Testing with {args.max} papers[/dim]")
 
     console.print(f"[cyan]Starting entity extraction from {papers_path}...[/cyan]")
-    console.print("[dim]Rate-limited to ~1 batch/second. Cost: ~$0.03 per 10 papers.[/dim]\n")
+    console.print(
+        "[dim]Submitted via the Batch API (async, usually <1h). "
+        "Cost: ~$0.015 per 10 papers (~50% off vs. synchronous).[/dim]\n"
+    )
 
     client = anthropic.Anthropic()
     results = extract_all(papers_path=papers_path, client=client)