Merge pull request #6 from KevinIsInCoding/feat/stage4-kg-extraction

CLAUDE.md

@@ -59,6 +59,9 @@ uv run python scripts/build_graph.py
 
 # 5. Build ChromaDB vector index
 uv run python scripts/build_index.py
+
+# 6. run application
+uv run gradio app.py
 ```
 
 ## Key Invariants

agents/research_agent.py

@@ -162,7 +162,7 @@ def _handle_search(
     # Step 4: Trial matching — prefer KG-linked trials, fall back to text match
     related_trials: list[dict] = []
     if graph and query_entities:
-        related_trials = kg_query.find_trials_for_entities(graph, expanded_entities, max_trials=5)
+        related_trials = kg_query.find_trials_for_entities(graph, expanded_entities, max_trials=10)
 
     if not related_trials and query_entities:
         entities_lower = [e.lower() for e in expanded_entities]

app.py

@@ -129,7 +129,7 @@ with gr.Blocks(title="Candle-Fire — ALS Research Intelligence") as demo:
             height=520,
             show_label=False,
             sanitize_html=False,
-            avatar_images=(None, "https://api.dicebear.com/7.x/icons/svg?seed=candle&icon=flame"),
+            avatar_images=(None, "assets/flame.svg"),
             placeholder="Ask a question about ALS research to get started.",
         )
 

config.py

@@ -2,7 +2,7 @@ from pathlib import Path
 
 # Models
 SYNTHESIS_MODEL = "claude-sonnet-4-6"
-EXTRACTION_MODEL = "claude-sonnet-4-6"
+EXTRACTION_MODEL = "claude-haiku-4-5-20251001"
 
 # External API endpoints
 CTGOV_BASE = "https://clinicaltrials.gov/api/v2/studies"

data/tools/extract_trial_targets.json

@@ -0,0 +1,35 @@
+{
+  "type": "object",
+  "properties": {
+    "nct_id": {
+      "type": "string",
+      "description": "The ClinicalTrials.gov identifier (e.g. NCT12345678)."
+    },
+    "targets": {
+      "type": "array",
+      "description": "Primary biological targets modulated or studied in this trial.",
+      "items": {
+        "type": "object",
+        "properties": {
+          "name": {
+            "type": "string",
+            "description": "Standard name of the target (e.g. TARDBP, SOD1, riluzole, neuroinflammation). Use canonical gene symbols where applicable."
+          },
+          "type": {
+            "type": "string",
+            "enum": ["Gene", "Protein", "Compound", "Mechanism"],
+            "description": "Category of the target."
+          },
+          "confidence": {
+            "type": "number",
+            "minimum": 0.0,
+            "maximum": 1.0,
+            "description": "Confidence that this is the intended biological target (1.0 = explicit in text)."
+          }
+        },
+        "required": ["name", "type", "confidence"]
+      }
+    }
+  },
+  "required": ["nct_id", "targets"]
+}

extraction/extractor.py

@@ -99,34 +99,11 @@ def _extract_batch(
     registry: CanonicalRegistry,
 ) -> list[PaperExtractionResult]:
     """Send a batch of papers to Claude and collect one extract_entities call per paper."""
-    user_content = _format_batch(batch)
-
-    try:
-        response = client.messages.create(
-            model=EXTRACTION_MODEL,
-            max_tokens=4096,
-            system=_EXTRACTION_SYSTEM,
-            tools=EXTRACTION_TOOLS,
-            tool_choice={"type": "any"},
-            messages=[{"role": "user", "content": user_content}],
-        )
-    except anthropic.RateLimitError:
-        _logger.warning("Rate limited — sleeping 30s")
-        time.sleep(30)
-        response = client.messages.create(
-            model=EXTRACTION_MODEL,
-            max_tokens=4096,
-            system=_EXTRACTION_SYSTEM,
-            tools=EXTRACTION_TOOLS,
-            tool_choice={"type": "any"},
-            messages=[{"role": "user", "content": user_content}],
-        )
-
-    # Build a PMID→paper lookup so we can match extracted results back
     paper_by_pmid = {p.pmid: p for p in batch}
+    content_blocks = _call_claude(client, batch)
 
     results: list[PaperExtractionResult] = []
-    for block in response.content:
+    for block in content_blocks:
         if block.type != "tool_use" or block.name != "extract_entities":
             continue
 
@@ -151,16 +128,63 @@ def _extract_batch(
         # Mark paper entity_names (used downstream by RAG indexer on re-index)
         paper.entity_names = [e.canonical_id for e in entities]
 
-    # For any paper with no Claude response, add an empty result so it's not re-processed
+    # 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} — recording empty")
+            _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."""
+    try:
+        response = client.messages.create(
+            model=EXTRACTION_MODEL,
+            max_tokens=4096,
+            system=_EXTRACTION_SYSTEM,
+            tools=EXTRACTION_TOOLS,
+            tool_choice={"type": "any"},
+            messages=[{"role": "user", "content": _format_batch(batch)}],
+        )
+        return response.content
+    except anthropic.RateLimitError:
+        _logger.warning("Rate limited — sleeping 30s")
+        time.sleep(30)
+        response = client.messages.create(
+            model=EXTRACTION_MODEL,
+            max_tokens=4096,
+            system=_EXTRACTION_SYSTEM,
+            tools=EXTRACTION_TOOLS,
+            tool_choice={"type": "any"},
+            messages=[{"role": "user", "content": _format_batch(batch)}],
+        )
+        return response.content
+
+
 def _format_batch(batch: list[ALSPaper]) -> str:
     parts = [
         f"Extract entities from each of the following {len(batch)} ALS papers. "

graph/query.py

@@ -72,12 +72,19 @@ def expand_query_entities(
     return display_names
 
 
+_STATUS_RANK = {"RECRUITING": 0, "ACTIVE_NOT_RECRUITING": 1, "NOT_YET_RECRUITING": 2, "COMPLETED": 3}
+
+
 def find_trials_for_entities(
     G: nx.DiGraph,
     entity_names: list[str],
-    max_trials: int = 5,
+    max_trials: int = 10,
 ) -> list[dict]:
-    """Return clinical trials linked to the given entity names."""
+    """Return clinical trials linked to the given entity names.
+
+    Collects all matches, scores by number of linked entities, sorts by
+    status (RECRUITING first) then score, and returns the top max_trials.
+    """
     if not G or not entity_names:
         return []
 
@@ -86,31 +93,37 @@ def find_trials_for_entities(
         matched = _find_node(G, name)
         target_nodes.update(matched)
 
-    trials: list[dict] = []
-    seen: set[str] = set()
+    # score[nct_id] = number of query entities this trial links to
+    scores: dict[str, int] = {}
+    meta: dict[str, dict] = {}
 
     for node_id in target_nodes:
-        # Trials point TO their targets, so look at predecessors
         for pred in G.predecessors(node_id):
             if not pred.startswith("trial:"):
                 continue
             nct_id = G.nodes[pred].get("nct_id", "")
-            if nct_id in seen:
+            if not nct_id:
                 continue
-            seen.add(nct_id)
-            trials.append({
-                "nct_id": nct_id,
-                "title": G.nodes[pred].get("display_name", ""),
-                "phase": G.nodes[pred].get("phase", ""),
-                "status": G.nodes[pred].get("status", ""),
-                "url": G.nodes[pred].get("url", ""),
-            })
-            if len(trials) >= max_trials:
-                break
-        if len(trials) >= max_trials:
-            break
-
-    return trials
+            scores[nct_id] = scores.get(nct_id, 0) + 1
+            if nct_id not in meta:
+                status = G.nodes[pred].get("status", "")
+                meta[nct_id] = {
+                    "nct_id": nct_id,
+                    "title": G.nodes[pred].get("display_name", ""),
+                    "phase": G.nodes[pred].get("phase", ""),
+                    "status": status,
+                    "url": G.nodes[pred].get("url", ""),
+                    "_status_rank": _STATUS_RANK.get(status, 9),
+                }
+
+    ranked = sorted(
+        meta.values(),
+        key=lambda t: (t["_status_rank"], -scores[t["nct_id"]]),
+    )
+    for t in ranked:
+        del t["_status_rank"]
+
+    return ranked[:max_trials]
 
 
 def get_entity_evidence(G: nx.DiGraph, canonical_id: str) -> dict:

ingestion/clinicaltrials.py

@@ -1,47 +1,42 @@
 """ClinicalTrials.gov v2 client for ALS trials (adapted from beacon/trials_api.py)."""
 from __future__ import annotations
 
-import re
 import time
+from typing import TYPE_CHECKING
 
 import httpx
 
-from config import CTGOV_BASE
+from config import CTGOV_BASE, EXTRACTION_MODEL
 from logging_config import get_logger
 
+if TYPE_CHECKING:
+    import anthropic
+
 _logger = get_logger("ingestion.clinicaltrials")
 
-# Known ALS-relevant targets for heuristic entity linking
-_KNOWN_TARGETS: list[tuple[re.Pattern, str]] = [
-    (re.compile(r"\bSOD1\b", re.IGNORECASE), "SOD1"),
-    (re.compile(r"\bTARDBP\b", re.IGNORECASE), "TARDBP"),
-    (re.compile(r"\bTDP-?43\b", re.IGNORECASE), "TARDBP"),
-    (re.compile(r"\bFUS\b", re.IGNORECASE), "FUS"),
-    (re.compile(r"\bC9orf72\b", re.IGNORECASE), "C9orf72"),
-    (re.compile(r"\bATXN2\b", re.IGNORECASE), "ATXN2"),
-    (re.compile(r"\bTBK1\b", re.IGNORECASE), "TBK1"),
-    (re.compile(r"\bNEK1\b", re.IGNORECASE), "NEK1"),
-    (re.compile(r"\bVCP\b", re.IGNORECASE), "VCP"),
-    (re.compile(r"\briluzole\b", re.IGNORECASE), "riluzole"),
-    (re.compile(r"\bedaravone\b", re.IGNORECASE), "edaravone"),
-    (re.compile(r"\btofersen\b", re.IGNORECASE), "tofersen"),
-    (re.compile(r"\bAMX0035\b", re.IGNORECASE), "AMX0035"),
-    (re.compile(r"\bmasitinib\b", re.IGNORECASE), "masitinib"),
-    (re.compile(r"\bbosutinib\b", re.IGNORECASE), "bosutinib"),
-    (re.compile(r"\bmexiletine\b", re.IGNORECASE), "mexiletine"),
-    (re.compile(r"\bantisense oligonucleotide\b", re.IGNORECASE), "antisense oligonucleotide"),
-    (re.compile(r"\bASO\b"), "antisense oligonucleotide"),
-    (re.compile(r"\bsiRNA\b", re.IGNORECASE), "siRNA"),
-    (re.compile(r"\bstem cell\b", re.IGNORECASE), "stem cell"),
-    (re.compile(r"\bgene therapy\b", re.IGNORECASE), "gene therapy"),
-]
-
-
-def fetch_als_trials(status: str = "RECRUITING") -> list[dict]:
+_TRIAL_BATCH_SIZE = 10
+
+_TRIAL_EXTRACTION_SYSTEM = """You are a biomedical NLP expert specializing in ALS (amyotrophic lateral sclerosis) clinical trials.
+
+For each trial provided, identify the primary biological target(s) being tested or modulated:
+- Genes silenced or corrected (e.g., SOD1, TARDBP, FUS, C9orf72, NEK1, VCP, TBK1)
+- Proteins targeted (use canonical gene symbol, e.g. TARDBP for TDP-43 protein)
+- Compounds/drugs — report the molecular or pathway target, not the drug name (e.g., a trial of AMX0114 targets TARDBP)
+- Mechanisms (e.g., neuroinflammation, oxidative stress, glutamate excitotoxicity)
+
+Call extract_trial_targets once per trial. Return an empty targets list only when no specific molecular or mechanistic target is identifiable."""
+
+
+def fetch_als_trials(
+    status: str | list[str] = ("RECRUITING", "NOT_YET_RECRUITING", "ACTIVE_NOT_RECRUITING"),
+    client: "anthropic.Anthropic | None" = None,
+) -> list[dict]:
     """Fetch ALS interventional trials. Returns flat dicts ready for JSONL serialization."""
+    status_filter = ",".join(status) if isinstance(status, (list, tuple)) else status
+
     params: dict[str, str | int] = {
         "query.cond": "Amyotrophic Lateral Sclerosis",
-        "filter.overallStatus": status,
+        "filter.overallStatus": status_filter,
         "aggFilters": "studyType:int",
         "pageSize": 1000,
         "format": "json",
@@ -58,9 +53,8 @@ def fetch_als_trials(status: str = "RECRUITING") -> list[dict]:
             except httpx.HTTPError as exc:
                 if attempt == 2:
                     raise
-                wait = 2 ** attempt
                 _logger.warning(f"ClinicalTrials.gov error (attempt {attempt + 1}): {exc}")
-                time.sleep(wait)
+                time.sleep(2 ** attempt)
 
         page_studies = body.get("studies", [])
         all_studies.extend(page_studies)
@@ -74,6 +68,10 @@ def fetch_als_trials(status: str = "RECRUITING") -> list[dict]:
         params["pageToken"] = next_token
 
     trials = [_flatten_trial(s) for s in all_studies]
+
+    if client is not None:
+        _enrich_targets_llm(trials, client)
+
     _logger.info("ALS trial fetch complete", extra={"data": {"total": len(trials)}})
     return trials
 
@@ -88,16 +86,14 @@ def _flatten_trial(study: dict) -> dict:
     status_mod = proto.get("statusModule", {})
 
     nct_id = id_mod.get("nctId", "")
-    title = id_mod.get("briefTitle", "")
     interventions = [
         {"type": iv.get("type", ""), "name": iv.get("name", "")}
         for iv in arms_mod.get("interventions", [])
     ]
-    intervention_names = " ".join(iv["name"] for iv in interventions)
 
     return {
         "nct_id": nct_id,
-        "title": title,
+        "title": id_mod.get("briefTitle", ""),
         "phase": ", ".join(design_mod.get("phases", [])) or "N/A",
         "status": status_mod.get("overallStatus", ""),
         "sponsor": sponsor_mod.get("leadSponsor", {}).get("name", ""),
@@ -105,14 +101,93 @@ def _flatten_trial(study: dict) -> dict:
         "interventions": interventions,
         "start_date": status_mod.get("startDateStruct", {}).get("date", ""),
         "url": f"https://clinicaltrials.gov/study/{nct_id}" if nct_id else "",
-        "target_entities": extract_target_entities(f"{title} {intervention_names}"),
+        "target_entities": [],
     }
 
 
-def extract_target_entities(text: str) -> list[str]:
-    """Scan text for known ALS target names. Returns sorted canonical entity names."""
-    found: set[str] = set()
-    for pattern, canonical in _KNOWN_TARGETS:
-        if pattern.search(text):
-            found.add(canonical)
-    return sorted(found)
+def _enrich_targets_llm(trials: list[dict], client: "anthropic.Anthropic") -> None:
+    """Call Claude in batches to extract biological targets; mutates each trial in-place."""
+    from extraction.normalizer import normalize_entity
+    from tools import TRIAL_EXTRACTION_TOOLS
+
+    from rich.progress import BarColumn, Progress, SpinnerColumn, TaskProgressColumn, TextColumn, TimeRemainingColumn
+
+    with Progress(
+        SpinnerColumn(),
+        TextColumn("[progress.description]{task.description}"),
+        BarColumn(),
+        TaskProgressColumn(),
+        TimeRemainingColumn(),
+    ) as progress:
+        task = progress.add_task("Extracting trial targets (LLM)...", total=len(trials))
+
+        for i in range(0, len(trials), _TRIAL_BATCH_SIZE):
+            batch = trials[i : i + _TRIAL_BATCH_SIZE]
+            results = _call_claude_batch(client, batch, TRIAL_EXTRACTION_TOOLS)
+
+            for nct_id, raw_targets in results.items():
+                trial = next((t for t in batch if t["nct_id"] == nct_id), None)
+                if trial is None:
+                    continue
+                canonical: list[str] = []
+                for t in raw_targets:
+                    if t.get("confidence", 0) < 0.5:
+                        continue
+                    canon_id = normalize_entity(t["name"], t["type"])
+                    # strip prefix (e.g. "protein:TARDBP" → "TARDBP")
+                    canon_name = canon_id.split(":", 1)[-1]
+                    if canon_name and canon_name not in canonical:
+                        canonical.append(canon_name)
+                trial["target_entities"] = canonical
+
+            progress.advance(task, len(batch))
+
+            if i + _TRIAL_BATCH_SIZE < len(trials):
+                time.sleep(1.0)
+
+
+def _call_claude_batch(
+    client: "anthropic.Anthropic",
+    batch: list[dict],
+    tools: list,
+) -> dict[str, list[dict]]:
+    """Send one batch of trials to Claude; return {nct_id: [target dicts]}."""
+    lines = [
+        f"Extract targets from each of the following {len(batch)} ALS clinical trials. "
+        "Call extract_trial_targets once per trial.\n"
+    ]
+    for trial in batch:
+        iv_names = ", ".join(iv["name"] for iv in trial.get("interventions", [])) or "N/A"
+        summary = (trial.get("summary") or "")[:400]
+        lines.append(
+            f"--- NCT: {trial['nct_id']} ---\n"
+            f"Title: {trial['title']}\n"
+            f"Interventions: {iv_names}\n"
+            f"Summary: {summary}\n"
+        )
+
+    for attempt in range(3):
+        try:
+            response = client.messages.create(
+                model=EXTRACTION_MODEL,
+                max_tokens=4096,
+                system=_TRIAL_EXTRACTION_SYSTEM,
+                tools=tools,
+                tool_choice={"type": "any"},
+                messages=[{"role": "user", "content": "\n".join(lines)}],
+            )
+            break
+        except Exception as exc:
+            if attempt == 2:
+                _logger.warning(f"Claude trial extraction failed: {exc}")
+                return {}
+            time.sleep(30 if "rate" in str(exc).lower() else 2 ** attempt)
+
+    results: dict[str, list[dict]] = {}
+    for block in response.content:
+        if block.type == "tool_use" and block.name == "extract_trial_targets":
+            nct_id = block.input.get("nct_id", "")
+            if nct_id:
+                results[nct_id] = block.input.get("targets", [])
+
+    return results

scripts/ingest_trials.py

@@ -5,7 +5,7 @@ Ingest ALS clinical trials from ClinicalTrials.gov v2 API.
 Usage:
     uv run python scripts/ingest_trials.py
     uv run python scripts/ingest_trials.py --status RECRUITING
-    uv run python scripts/ingest_trials.py --status COMPLETED
+    uv run python scripts/ingest_trials.py --status RECRUITING NOT_YET_RECRUITING COMPLETED
 """
 from __future__ import annotations
 
@@ -20,6 +20,7 @@ from dotenv import load_dotenv
 
 load_dotenv()
 
+import anthropic
 from rich.console import Console
 
 from config import TRIALS_PATH
@@ -27,21 +28,30 @@ from ingestion.clinicaltrials import fetch_als_trials
 
 console = Console()
 
+_ALL_STATUSES = ["RECRUITING", "COMPLETED", "ACTIVE_NOT_RECRUITING", "NOT_YET_RECRUITING"]
+
 
 def main() -> None:
     parser = argparse.ArgumentParser(description="Ingest ALS clinical trials")
     parser.add_argument(
         "--status",
-        default="RECRUITING",
-        choices=["RECRUITING", "COMPLETED", "ACTIVE_NOT_RECRUITING", "NOT_YET_RECRUITING"],
-        help="Trial status filter (default: RECRUITING)",
+        nargs="+",
+        default=["RECRUITING", "NOT_YET_RECRUITING", "ACTIVE_NOT_RECRUITING"],
+        choices=_ALL_STATUSES,
+        metavar="STATUS",
+        help=(
+            f"One or more trial statuses to fetch (default: RECRUITING NOT_YET_RECRUITING "
+            f"ACTIVE_NOT_RECRUITING). Choices: {_ALL_STATUSES}"
+        ),
     )
     args = parser.parse_args()
 
     TRIALS_PATH.parent.mkdir(parents=True, exist_ok=True)
 
+    client = anthropic.Anthropic()
+
     console.print(f"[cyan]Fetching ALS interventional trials (status={args.status})...[/cyan]")
-    trials = fetch_als_trials(status=args.status)
+    trials = fetch_als_trials(status=args.status, client=client)
     console.print(f"[green]Fetched {len(trials)} trials[/green]")
 
     with open(TRIALS_PATH, "w", encoding="utf-8") as f:

tools.py

@@ -34,5 +34,15 @@ SEARCH_LANDSCAPE_TOOL: anthropic.types.ToolParam = {
     "input_schema": _load("search_landscape"),
 }
 
+EXTRACT_TRIAL_TARGETS_TOOL: anthropic.types.ToolParam = {
+    "name": "extract_trial_targets",
+    "description": (
+        "Extract the primary biological target(s) of an ALS clinical trial — the gene, protein, "
+        "compound, or mechanism being tested or modulated. Call once per trial."
+    ),
+    "input_schema": _load("extract_trial_targets"),
+}
+
 EXTRACTION_TOOLS: list[anthropic.types.ToolParam] = [EXTRACT_ENTITIES_TOOL]
+TRIAL_EXTRACTION_TOOLS: list[anthropic.types.ToolParam] = [EXTRACT_TRIAL_TARGETS_TOOL]
 RESEARCH_TOOLS: list[anthropic.types.ToolParam] = [SEARCH_LANDSCAPE_TOOL]