""" 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%. """ from __future__ import annotations import json import threading import time from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path import anthropic from rich.progress import BarColumn, MofNCompleteColumn, Progress, TextColumn, TimeElapsedColumn from config import ( ENTITIES_PATH, EXTRACTION_BATCH_SIZE, EXTRACTION_MODEL, EXTRACTION_PROGRESS_PATH, EXTRACTION_WORKERS, 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") _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, client: anthropic.Anthropic | None = None, ) -> list[PaperExtractionResult]: """Extract entities from all papers. Skips already-processed PMIDs. Runs EXTRACTION_WORKERS batches in parallel. A lock serializes file writes and progress saves so threads don't corrupt each other. """ if client is None: client = anthropic.Anthropic() papers = _load_papers(papers_path) 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) batches = [pending[i : i + EXTRACTION_BATCH_SIZE] for i in range(0, len(pending), EXTRACTION_BATCH_SIZE)] 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])) return results def _extract_batch( client: anthropic.Anthropic, batch: list[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) results: list[PaperExtractionResult] = [] for block in content_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") 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(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=[])) 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 _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)))