extraction/extractor.py

"""
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 time
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,
    PAPERS_PATH,
)
from extraction.normalizer import CanonicalRegistry, guess_entity_type, normalize_entity
from logging_config import get_logger
from models import ALSPaper, ExtractedEntity, EntityRelationship, PaperExtractionResult
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.
Call it once per paper. Use the full text when provided — it is richer than the abstract alone.

Entity types: Gene, Protein, Compound, Pathway, Phenotype, Mechanism.
Relationship types: BINDS, INHIBITS, ASSOCIATED_WITH, TESTED_IN, EXPRESSED_IN, CO_OCCURS.

Be precise. Only extract entities explicitly mentioned. Return pmid exactly as given.
"""


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 via the Batch API. Skips done PMIDs.

    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"
    )

    if not pending:
        return []

    registry = CanonicalRegistry()
    entities_path.parent.mkdir(parents=True, exist_ok=True)

    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] = []
    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 _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]:
    """Parse extract_entities tool_use blocks from one response into results."""
    results: list[PaperExtractionResult] = []
    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 request — skipping")
            continue

        paper = paper_by_pmid[pmid]
        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.entity_names = [e.canonical_id for e in entities]
        _logger.info(f"PMID {pmid}: {len(entities)} entities, {len(relationships)} relationships")

    return results


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:
    parts = [
        f"Extract entities from each of the following {len(batch)} ALS papers. "
        "Call extract_entities once per paper.\n"
    ]
    for paper in batch:
        text = paper.full_text if paper.full_text else paper.abstract
        # Cap at 2000 chars — 20-paper batches at ~500 tokens each stay well under 8192 output limit
        excerpt = text[:2000] if text else paper.abstract[:1000]
        parts.append(
            f"--- PMID:{paper.pmid} ---\n"
            f"Title: {paper.title}\n\n"
            f"{excerpt}\n"
        )
    return "\n".join(parts)


def _parse_entities(
    raw: list[dict],
    pmid: str,
    registry: CanonicalRegistry,
) -> list[ExtractedEntity]:
    entities = []
    for item in raw:
        if not isinstance(item, dict):
            continue
        name = item.get("name", "").strip()
        entity_type = item.get("type", "").strip()
        if not name or not entity_type:
            continue
        canonical_id = registry.resolve(name, entity_type)
        entities.append(
            ExtractedEntity(
                type=entity_type,
                name=name,
                canonical_id=canonical_id,
                confidence=float(item.get("confidence", 0.7)),
                mentions=int(item.get("mentions", 1)),
            )
        )
    return entities


def _parse_relationships(
    raw: list[dict],
    pmid: str,
    registry: CanonicalRegistry,
) -> list[EntityRelationship]:
    rels = []
    for item in raw:
        if not isinstance(item, dict):
            continue
        source_name = item.get("source", "").strip()
        target_name = item.get("target", "").strip()
        rel_type = item.get("type", "").strip()
        if not source_name or not target_name or not rel_type:
            continue
        # We don't know entity types for source/target here — infer from name
        source_id = registry.resolve(source_name, _guess_type(source_name))
        target_id = registry.resolve(target_name, _guess_type(target_name))
        rels.append(
            EntityRelationship(
                source=source_id,
                target=target_id,
                relation_type=rel_type,
                evidence_pmids=[pmid],
                confidence=0.7,
                evidence_text=item.get("evidence_text", "")[:300],
            )
        )
    return rels


_guess_type = guess_entity_type


def _load_papers(path: Path) -> list[ALSPaper]:
    papers = []
    with open(path, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                papers.append(ALSPaper.from_dict(json.loads(line)))
    return papers


def _load_progress(path: Path) -> set[str]:
    if path.exists():
        return set(json.loads(path.read_text()))
    return set()


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)