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

agents/research_agent.py

@@ -6,8 +6,10 @@ from collections.abc import Generator
 
 import anthropic
 import chromadb
+import networkx as nx
 
 from config import SYNTHESIS_MODEL
+from graph import query as kg_query
 from llm import cached_system, cached_tools
 from logging_config import get_logger
 from prompts import SYNTHESIS_SYSTEM
@@ -22,6 +24,7 @@ def stream_research_agent(
     query: str,
     collection: chromadb.Collection,
     trials: list[dict],
+    graph: nx.DiGraph | None = None,
 ) -> Generator[tuple[str, str], None, None]:
     """
     Stream the ALS research synthesis agent.
@@ -34,6 +37,8 @@ def stream_research_agent(
     messages: list[anthropic.types.MessageParam] = [
         {"role": "user", "content": query}
     ]
+    # Attach graph reference so _handle_search can use KG expansion
+    _graph = graph
 
     while True:
         stream_text = ""
@@ -91,7 +96,7 @@ def stream_research_agent(
 
             for tool_call in tool_calls:
                 if tool_call["name"] == "search_research_landscape":
-                    result = _handle_search(tool_call["input"], collection, trials)
+                    result = _handle_search(tool_call["input"], collection, trials, _graph)
                     is_error = False
                 else:
                     result = {"error": f"Unknown tool: {tool_call['name']}"}
@@ -114,16 +119,24 @@ def _handle_search(
     tool_input: dict,
     collection: chromadb.Collection,
     trials: list[dict],
+    graph: nx.DiGraph | None = None,
 ) -> dict:
-    """Execute RAG search + trial lookup and return structured context for Claude."""
+    """Execute KG expansion → RAG search → trial lookup and return structured context."""
     query_text = tool_input.get("query_text", "")
     query_entities = tool_input.get("query_entities", [])
 
-    # Semantic search
+    # Step 1: KG expansion — surface related entities Claude didn't name explicitly
+    # e.g. "tofersen" → expands to ["SOD1", "antisense oligonucleotide", "RNA splicing"]
+    if graph and query_entities:
+        expanded_entities = kg_query.expand_query_entities(graph, query_entities)
+    else:
+        expanded_entities = query_entities
+
+    # Step 2: RAG — semantic search on query text
     semantic_results = rag_retriever.search(collection, query_text, n_results=10)
 
-    # Entity-targeted search (finds papers even if query terms don't match directly)
-    entity_results = rag_retriever.search_by_entities(collection, query_entities, n_results=15)
+    # Step 3: Entity-targeted RAG using expanded entity set
+    entity_results = rag_retriever.search_by_entities(collection, expanded_entities, n_results=15)
 
     # Merge by PMID — keep best score per paper
     seen: dict[str, dict] = {}
@@ -135,19 +148,24 @@ def _handle_search(
     top_papers = sorted(seen.values(), key=lambda x: x["score"], reverse=True)[:15]
 
     _logger.info(
-        "RAG search",
+        "KG+RAG search",
         extra={"data": {
             "query_entities": query_entities,
+            "expanded_entities": len(expanded_entities),
             "semantic_hits": len(semantic_results),
             "entity_hits": len(entity_results),
             "merged": len(top_papers),
+            "kg_active": graph is not None,
         }},
     )
 
-    # Match trials by entity name in title or intervention text
-    related_trials = []
-    if query_entities:
-        entities_lower = [e.lower() for e in query_entities]
+    # 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)
+
+    if not related_trials and query_entities:
+        entities_lower = [e.lower() for e in expanded_entities]
         for trial in trials:
             iv_names = " ".join(iv.get("name", "") for iv in trial.get("interventions", []))
             trial_text = f"{trial.get('title', '')} {iv_names}".lower()
@@ -177,6 +195,8 @@ def _handle_search(
             for r in top_papers
         ],
         "query_entities": query_entities,
+        "expanded_entities": expanded_entities,
         "trials": related_trials,
         "evidence_count": len(top_papers),
+        "kg_expansion_active": graph is not None,
     }

extraction/extractor.py

@@ -1,2 +1,261 @@
-"""Claude Sonnet entity extractor — batch 10 papers/call, resumable."""
-# Stage 4 implementation
+"""
+Claude Sonnet entity extractor.
+Batches 10 papers per API call; resumable via .progress.json.
+Uses full_text when available, otherwise abstract.
+"""
+from __future__ import annotations
+
+import json
+import time
+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,
+    PAPERS_PATH,
+)
+from extraction.normalizer import CanonicalRegistry, 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."""
+    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)
+
+    results: list[PaperExtractionResult] = []
+
+    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))
+
+        for i in range(0, len(pending), EXTRACTION_BATCH_SIZE):
+            batch = pending[i : i + EXTRACTION_BATCH_SIZE]
+            batch_results = _extract_batch(client, batch, registry)
+
+            for result in batch_results:
+                out_f.write(json.dumps(result.to_dict()) + "\n")
+                done_pmids.add(result.pmid)
+                results.append(result)
+
+            _save_progress(progress_path, done_pmids)
+            registry.save()
+            progress.advance(task, len(batch))
+
+            # Respect rate limits between batches
+            if i + EXTRACTION_BATCH_SIZE < len(pending):
+                time.sleep(1.0)
+
+    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."""
+    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}
+
+    results: list[PaperExtractionResult] = []
+    for block in response.content:
+        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]
+
+    # For any paper with no Claude response, add an empty result so it's not re-processed
+    found_pmids = {r.pmid for r in results}
+    for p in batch:
+        if p.pmid not in found_pmids:
+            _logger.warning(f"No extraction result for PMID {p.pmid} — recording empty")
+            results.append(PaperExtractionResult(pmid=p.pmid, entities=[], relationships=[]))
+
+    return results
+
+
+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 3000 chars to stay within token budget for a 10-paper batch
+        excerpt = text[:3000] 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:
+        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:
+        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
+
+
+def _guess_type(name: str) -> str:
+    """Best-effort entity type guess from name for relationship source/target."""
+    from extraction.normalizer import _GENE_ALIASES, _COMPOUND_ALIASES
+    if name.strip().upper() in _GENE_ALIASES or name.strip() in _GENE_ALIASES:
+        return "Gene"
+    if name.strip() in _COMPOUND_ALIASES:
+        return "Compound"
+    return "Protein"
+
+
+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)))

extraction/normalizer.py

@@ -1,2 +1,169 @@
-"""Entity name → canonical ID mapping (HGNC aliases + PubChem + REST fallback)."""
-# Stage 4 implementation
+"""
+Entity name normalization to canonical IDs.
+Priority: static alias table → known compound map → slugify fallback.
+"""
+from __future__ import annotations
+
+import json
+import re
+from pathlib import Path
+
+from config import CANONICAL_IDS_PATH
+
+# ── Gene / protein alias table ────────────────────────────────────────────────
+_GENE_ALIASES: dict[str, str] = {
+    # TDP-43 / TARDBP
+    "TDP-43": "TARDBP", "TDP43": "TARDBP", "tdp-43": "TARDBP", "tdp43": "TARDBP",
+    "TAR DNA-binding protein 43": "TARDBP",
+    "TAR DNA binding protein 43": "TARDBP",
+    "TAR DNA-binding protein": "TARDBP",
+    # FUS
+    "FUS/TLS": "FUS", "TLS": "FUS", "TLS/FUS": "FUS", "fus": "FUS",
+    "fused in sarcoma": "FUS",
+    # C9orf72
+    "C9ORF72": "C9orf72", "c9orf72": "C9orf72", "C9": "C9orf72",
+    "chromosome 9 open reading frame 72": "C9orf72",
+    # SOD1
+    "SOD-1": "SOD1", "sod1": "SOD1",
+    "superoxide dismutase 1": "SOD1",
+    "Cu/Zn-superoxide dismutase": "SOD1",
+    "copper-zinc superoxide dismutase": "SOD1",
+    # ATXN2
+    "ataxin-2": "ATXN2", "ataxin 2": "ATXN2", "SCA2": "ATXN2",
+    # Optineurin
+    "optineurin": "OPTN",
+    # Ubiquilin
+    "ubiquilin-2": "UBQLN2", "ubiquilin2": "UBQLN2", "ubiquilin 2": "UBQLN2",
+    # SQSTM1 / p62
+    "p62": "SQSTM1", "sequestosome-1": "SQSTM1", "sequestosome 1": "SQSTM1",
+    # Profilin
+    "profilin 1": "PFN1", "profilin-1": "PFN1",
+    # Dynactin
+    "dynactin": "DCTN1", "dynactin 1": "DCTN1",
+    # Matrin
+    "matrin 3": "MATR3", "matrin-3": "MATR3",
+    # Other
+    "angiogenin": "ANG",
+    "senataxin": "SETX",
+}
+
+# ── Compound alias table ───────────────────────────────────────────────────────
+_COMPOUND_ALIASES: dict[str, str] = {
+    "riluzole": "riluzole", "Riluzole": "riluzole",
+    "edaravone": "edaravone", "Edaravone": "edaravone", "MCI-186": "edaravone",
+    "tofersen": "tofersen", "Tofersen": "tofersen",
+    "BIIB067": "tofersen", "biib067": "tofersen",
+    "AMX0035": "AMX0035", "amx0035": "AMX0035",
+    "sodium phenylbutyrate": "AMX0035",
+    "tauroursodeoxycholic acid": "AMX0035",
+    "TUDCA": "AMX0035",
+    "masitinib": "masitinib", "Masitinib": "masitinib", "AB1010": "masitinib",
+    "bosutinib": "bosutinib", "Bosutinib": "bosutinib", "SKI-606": "bosutinib",
+    "mexiletine": "mexiletine", "Mexiletine": "mexiletine",
+    "memantine": "memantine", "Memantine": "memantine",
+    "rasagiline": "rasagiline", "Rasagiline": "rasagiline",
+    "NurOwn": "NurOwn", "MSC-NTF": "NurOwn",
+    "ozanezumab": "ozanezumab",
+}
+
+# ── Mechanism normalization ────────────────────────────────────────────────────
+_MECHANISM_ALIASES: dict[str, str] = {
+    "glutamate excitotoxicity": "glutamate_excitotoxicity",
+    "excitotoxicity": "glutamate_excitotoxicity",
+    "glutamatergic excitotoxicity": "glutamate_excitotoxicity",
+    "oxidative stress": "oxidative_stress",
+    "reactive oxygen species": "oxidative_stress",
+    "ROS": "oxidative_stress",
+    "neuroinflammation": "neuroinflammation",
+    "microglial activation": "neuroinflammation",
+    "astrocyte activation": "neuroinflammation",
+    "protein aggregation": "protein_aggregation",
+    "protein misfolding": "protein_aggregation",
+    "protein inclusions": "protein_aggregation",
+    "RNA metabolism": "RNA_metabolism_dysfunction",
+    "RNA processing": "RNA_metabolism_dysfunction",
+    "RNA-binding protein dysfunction": "RNA_metabolism_dysfunction",
+    "stress granules": "RNA_metabolism_dysfunction",
+    "mitochondrial dysfunction": "mitochondrial_dysfunction",
+    "mitochondrial impairment": "mitochondrial_dysfunction",
+    "axonal transport": "axonal_transport_defect",
+    "axonal transport defect": "axonal_transport_defect",
+    "autophagy": "autophagy_impairment",
+    "mitophagy": "autophagy_impairment",
+    "ubiquitin proteasome": "autophagy_impairment",
+    "TDP-43 pathology": "TDP43_pathology",
+    "TDP-43 aggregation": "TDP43_pathology",
+    "TDP-43 mislocalization": "TDP43_pathology",
+    "antisense oligonucleotide": "antisense_oligonucleotide",
+    "ASO": "antisense_oligonucleotide",
+    "gene therapy": "gene_therapy",
+    "stem cell": "stem_cell_therapy",
+    "neurodegeneration": "neurodegeneration",
+    "apoptosis": "apoptosis",
+    "DNA damage": "DNA_damage_repair",
+}
+
+
+def normalize_entity(name: str, entity_type: str) -> str:
+    """Return a canonical_id string in the form '<prefix>:<canonical_name>'."""
+    canonical = _resolve_name(name.strip(), entity_type)
+    return f"{_prefix(entity_type)}:{canonical}"
+
+
+def _resolve_name(name: str, entity_type: str) -> str:
+    t = entity_type.lower()
+
+    if t == "gene":
+        return _GENE_ALIASES.get(name) or _GENE_ALIASES.get(name.upper()) or name.upper()
+
+    if t == "protein":
+        gene_hit = _GENE_ALIASES.get(name) or _GENE_ALIASES.get(name.upper())
+        if gene_hit:
+            return gene_hit
+        return name[0].upper() + name[1:] if name else name
+
+    if t == "compound":
+        return _COMPOUND_ALIASES.get(name) or _slugify(name)
+
+    if t == "mechanism":
+        return _MECHANISM_ALIASES.get(name) or _slugify(name)
+
+    return _slugify(name)
+
+
+def _prefix(entity_type: str) -> str:
+    return {
+        "gene": "gene",
+        "protein": "protein",
+        "compound": "compound",
+        "pathway": "pathway",
+        "phenotype": "phenotype",
+        "mechanism": "mechanism",
+    }.get(entity_type.lower(), "entity")
+
+
+def _slugify(text: str) -> str:
+    return re.sub(r"[^a-z0-9]+", "_", text.lower()).strip("_")
+
+
+class CanonicalRegistry:
+    """Persistent name→canonical_id mapping written to canonical_ids.json."""
+
+    def __init__(self, path: Path = CANONICAL_IDS_PATH) -> None:
+        self.path = path
+        self._data: dict[str, str] = {}
+        if path.exists():
+            self._data = json.loads(path.read_text())
+
+    def resolve(self, name: str, entity_type: str) -> str:
+        key = f"{entity_type.lower()}:{name}"
+        if key not in self._data:
+            self._data[key] = normalize_entity(name, entity_type)
+        return self._data[key]
+
+    def save(self) -> None:
+        self.path.parent.mkdir(parents=True, exist_ok=True)
+        self.path.write_text(json.dumps(self._data, indent=2, sort_keys=True))
+
+    def __len__(self) -> int:
+        return len(self._data)

graph/builder.py

@@ -1,2 +1,199 @@
-"""NetworkX DiGraph construction from extracted entities and trials."""
-# Stage 4 implementation
+"""
+NetworkX DiGraph construction from extracted entities + trials.
+Upserts nodes and edges (merges PMIDs, recomputes confidence average).
+Pre-populates with ALS seed entities from config.
+"""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import networkx as nx
+
+from config import (
+    ALS_SEED_ENTITIES,
+    ENTITIES_PATH,
+    KG_MIN_EDGE_CONFIDENCE,
+    TRIALS_PATH,
+)
+from extraction.normalizer import normalize_entity
+from logging_config import get_logger
+from models import PaperExtractionResult
+
+_logger = get_logger("graph.builder")
+
+
+def build_graph(
+    entities_path: Path = ENTITIES_PATH,
+    trials_path: Path = TRIALS_PATH,
+) -> nx.DiGraph:
+    """Build the ALS knowledge graph. Returns a populated DiGraph."""
+    G: nx.DiGraph = nx.DiGraph()
+
+    _add_seed_entities(G)
+    _logger.info(f"Seeded graph: {G.number_of_nodes()} seed nodes")
+
+    if entities_path.exists():
+        n_papers = _add_extracted_entities(G, entities_path)
+        _logger.info(
+            f"After extraction: {G.number_of_nodes()} nodes, "
+            f"{G.number_of_edges()} edges from {n_papers} papers"
+        )
+    else:
+        _logger.warning(f"Entities file not found: {entities_path} — skipping NER enrichment")
+
+    if trials_path.exists():
+        n_trials = _add_trials(G, trials_path)
+        _logger.info(f"Added {n_trials} trial nodes")
+    else:
+        _logger.warning(f"Trials file not found: {trials_path}")
+
+    return G
+
+
+def _add_seed_entities(G: nx.DiGraph) -> None:
+    type_map = {
+        "genes": "Gene",
+        "proteins": "Protein",
+        "compounds": "Compound",
+        "mechanisms": "Mechanism",
+        "phenotypes": "Phenotype",
+    }
+    for category, entity_type in type_map.items():
+        for name in ALS_SEED_ENTITIES.get(category, []):
+            canonical_id = normalize_entity(name, entity_type)
+            _upsert_node(G, canonical_id, {
+                "type": entity_type,
+                "display_name": name,
+                "paper_count": 0,
+                "evidence_pmids": [],
+                "is_seed": True,
+            })
+
+
+def _add_extracted_entities(G: nx.DiGraph, entities_path: Path) -> int:
+    n_papers = 0
+    with open(entities_path, encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            raw = json.loads(line)
+            pmid = raw["pmid"]
+            n_papers += 1
+
+            # Add / update entity nodes
+            for ent in raw.get("entities", []):
+                canonical_id = ent.get("canonical_id", "")
+                if not canonical_id:
+                    continue
+                if G.has_node(canonical_id):
+                    G.nodes[canonical_id]["paper_count"] += 1
+                    G.nodes[canonical_id]["evidence_pmids"].append(pmid)
+                    # Update confidence as running average
+                    cur = G.nodes[canonical_id].get("confidence", 0.7)
+                    G.nodes[canonical_id]["confidence"] = (cur + ent.get("confidence", 0.7)) / 2
+                else:
+                    _upsert_node(G, canonical_id, {
+                        "type": ent.get("type", "Unknown"),
+                        "display_name": ent.get("name", canonical_id),
+                        "paper_count": 1,
+                        "evidence_pmids": [pmid],
+                        "confidence": ent.get("confidence", 0.7),
+                        "is_seed": False,
+                    })
+
+            # Add / update relationship edges
+            for rel in raw.get("relationships", []):
+                source = rel.get("source", "")
+                target = rel.get("target", "")
+                rel_type = rel.get("relation_type", "")
+                if not source or not target or not rel_type:
+                    continue
+
+                # Ensure both endpoints exist as nodes
+                for node_id in (source, target):
+                    if not G.has_node(node_id):
+                        _upsert_node(G, node_id, {
+                            "type": "Unknown",
+                            "display_name": node_id.split(":", 1)[-1],
+                            "paper_count": 0,
+                            "evidence_pmids": [],
+                            "is_seed": False,
+                        })
+
+                conf = rel.get("confidence", 0.7)
+                if G.has_edge(source, target):
+                    edge = G[source][target]
+                    if rel_type not in edge.get("relation_types", []):
+                        edge.setdefault("relation_types", [edge.get("relation_type", rel_type)])
+                        edge["relation_types"].append(rel_type)
+                    if pmid not in edge["evidence_pmids"]:
+                        edge["evidence_pmids"].append(pmid)
+                    # Running average confidence
+                    edge["confidence"] = (edge["confidence"] + conf) / 2
+                else:
+                    G.add_edge(source, target, **{
+                        "relation_type": rel_type,
+                        "relation_types": [rel_type],
+                        "evidence_pmids": [pmid],
+                        "confidence": conf,
+                        "evidence_text": rel.get("evidence_text", ""),
+                    })
+
+    return n_papers
+
+
+def _add_trials(G: nx.DiGraph, trials_path: Path) -> int:
+    n_trials = 0
+    with open(trials_path, encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            trial = json.loads(line)
+            nct_id = trial.get("nct_id", "")
+            if not nct_id:
+                continue
+
+            node_id = f"trial:{nct_id}"
+            _upsert_node(G, node_id, {
+                "type": "ClinicalTrial",
+                "display_name": trial.get("title", nct_id)[:120],
+                "nct_id": nct_id,
+                "phase": trial.get("phase", ""),
+                "status": trial.get("status", ""),
+                "url": trial.get("url", f"https://clinicaltrials.gov/study/{nct_id}"),
+                "paper_count": 0,
+                "evidence_pmids": [],
+            })
+
+            # Link trial to known target entities
+            for target_name in trial.get("target_entities", []):
+                # Try gene, compound, protein
+                matched = False
+                for etype in ("Gene", "Compound", "Protein", "Mechanism"):
+                    candidate_id = normalize_entity(target_name, etype)
+                    if G.has_node(candidate_id):
+                        if not G.has_edge(node_id, candidate_id):
+                            G.add_edge(node_id, candidate_id, **{
+                                "relation_type": "TESTED_IN",
+                                "relation_types": ["TESTED_IN"],
+                                "evidence_pmids": [],
+                                "confidence": 1.0,
+                                "evidence_text": "",
+                            })
+                        matched = True
+                        break
+                if not matched:
+                    _logger.debug(f"Trial {nct_id}: no graph node for target {target_name!r}")
+
+            n_trials += 1
+    return n_trials
+
+
+def _upsert_node(G: nx.DiGraph, node_id: str, attrs: dict) -> None:
+    if G.has_node(node_id):
+        G.nodes[node_id].update(attrs)
+    else:
+        G.add_node(node_id, **attrs)

graph/query.py

@@ -1,2 +1,153 @@
-"""Graph traversal utilities: entity expansion, trial lookup, evidence retrieval."""
-# Stage 4 implementation
+"""Graph traversal: entity expansion, trial lookup, evidence retrieval."""
+from __future__ import annotations
+
+import networkx as nx
+
+from config import KG_EXPANSION_HOPS, KG_MIN_EDGE_CONFIDENCE
+from extraction.normalizer import normalize_entity, _GENE_ALIASES, _COMPOUND_ALIASES
+from logging_config import get_logger
+
+_logger = get_logger("graph.query")
+
+
+def expand_query_entities(
+    G: nx.DiGraph,
+    entity_names: list[str],
+    max_hops: int = KG_EXPANSION_HOPS,
+) -> list[str]:
+    """
+    Map entity names to graph node IDs, then BFS-expand up to max_hops.
+    Returns a deduplicated list of entity display_names for RAG search.
+    Only traverses edges above KG_MIN_EDGE_CONFIDENCE.
+    """
+    if not G or not entity_names:
+        return entity_names
+
+    # Phase 1: map names to canonical node IDs
+    seed_nodes: set[str] = set()
+    for name in entity_names:
+        matched = _find_node(G, name)
+        if matched:
+            seed_nodes.update(matched)
+        else:
+            _logger.debug(f"No graph node for query entity: {name!r}")
+
+    if not seed_nodes:
+        return entity_names  # fall through to lexical RAG search
+
+    # Phase 2: BFS expansion
+    frontier = set(seed_nodes)
+    expanded = set(seed_nodes)
+
+    for _ in range(max_hops):
+        next_frontier: set[str] = set()
+        for node in frontier:
+            for neighbor in list(G.successors(node)) + list(G.predecessors(node)):
+                if neighbor in expanded:
+                    continue
+                # Only follow confident edges
+                edge_data = G.get_edge_data(node, neighbor) or G.get_edge_data(neighbor, node) or {}
+                if edge_data.get("confidence", 1.0) >= KG_MIN_EDGE_CONFIDENCE:
+                    # Skip trial nodes — they inflate retrieval noise
+                    if G.nodes[neighbor].get("type") != "ClinicalTrial":
+                        next_frontier.add(neighbor)
+        expanded.update(next_frontier)
+        frontier = next_frontier
+
+    # Phase 3: convert canonical IDs back to display names for RAG text queries
+    display_names: list[str] = []
+    seen: set[str] = set()
+    for node_id in expanded:
+        if node_id.startswith("trial:"):
+            continue
+        name = G.nodes[node_id].get("display_name", node_id.split(":", 1)[-1])
+        if name not in seen:
+            display_names.append(name)
+            seen.add(name)
+
+    _logger.info(
+        f"KG expansion: {len(entity_names)} query entities → {len(display_names)} expanded",
+        extra={"data": {"seeds": list(seed_nodes), "expanded_count": len(expanded)}},
+    )
+    return display_names
+
+
+def find_trials_for_entities(
+    G: nx.DiGraph,
+    entity_names: list[str],
+    max_trials: int = 5,
+) -> list[dict]:
+    """Return clinical trials linked to the given entity names."""
+    if not G or not entity_names:
+        return []
+
+    target_nodes: set[str] = set()
+    for name in entity_names:
+        matched = _find_node(G, name)
+        target_nodes.update(matched)
+
+    trials: list[dict] = []
+    seen: set[str] = set()
+
+    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:
+                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
+
+
+def get_entity_evidence(G: nx.DiGraph, canonical_id: str) -> dict:
+    """Return node attributes + connected entity names for a canonical ID."""
+    if not G.has_node(canonical_id):
+        return {}
+    attrs = dict(G.nodes[canonical_id])
+    attrs["neighbors"] = [
+        {
+            "id": n,
+            "display_name": G.nodes[n].get("display_name", n),
+            "relation": G.get_edge_data(canonical_id, n, {}).get("relation_type", ""),
+        }
+        for n in G.successors(canonical_id)
+        if G.nodes[n].get("type") != "ClinicalTrial"
+    ]
+    return attrs
+
+
+def _find_node(G: nx.DiGraph, name: str) -> list[str]:
+    """Map a raw entity name to zero or more graph node IDs."""
+    hits: list[str] = []
+
+    # 1. Try each entity type prefix
+    for etype in ("Gene", "Protein", "Compound", "Mechanism", "Pathway", "Phenotype"):
+        candidate = normalize_entity(name, etype)
+        if G.has_node(candidate):
+            hits.append(candidate)
+
+    if hits:
+        return hits
+
+    # 2. Case-insensitive display_name match
+    name_lower = name.lower()
+    for node_id, data in G.nodes(data=True):
+        display = data.get("display_name", "").lower()
+        if display == name_lower or name_lower in display:
+            hits.append(node_id)
+
+    return hits

graph/serializer.py

@@ -1,2 +1,62 @@
-"""Save and load the ALS knowledge graph (pickle + JSON)."""
-# Stage 4 implementation
+"""Save and load the ALS knowledge graph (pickle for speed, JSON for inspection)."""
+from __future__ import annotations
+
+import json
+import pickle
+from pathlib import Path
+
+import networkx as nx
+
+from config import GRAPH_JSON_PATH, GRAPH_PICKLE_PATH
+from logging_config import get_logger
+
+_logger = get_logger("graph.serializer")
+
+
+def save_graph(
+    G: nx.DiGraph,
+    pickle_path: Path = GRAPH_PICKLE_PATH,
+    json_path: Path = GRAPH_JSON_PATH,
+) -> None:
+    for path in (pickle_path, json_path):
+        path.parent.mkdir(parents=True, exist_ok=True)
+
+    with open(pickle_path, "wb") as f:
+        pickle.dump(G, f, protocol=pickle.HIGHEST_PROTOCOL)
+
+    data = {
+        "nodes": [
+            {"id": n, **{k: _json_safe(v) for k, v in G.nodes[n].items()}}
+            for n in G.nodes
+        ],
+        "edges": [
+            {"source": u, "target": v, **{k: _json_safe(dv) for k, dv in d.items()}}
+            for u, v, d in G.edges(data=True)
+        ],
+        "stats": {
+            "nodes": G.number_of_nodes(),
+            "edges": G.number_of_edges(),
+        },
+    }
+    json_path.write_text(json.dumps(data, indent=2))
+    _logger.info(f"Saved graph: {G.number_of_nodes()} nodes, {G.number_of_edges()} edges")
+
+
+def load_graph(pickle_path: Path = GRAPH_PICKLE_PATH) -> nx.DiGraph:
+    if not pickle_path.exists():
+        raise FileNotFoundError(
+            f"Graph not found at {pickle_path}. "
+            "Run: uv run python scripts/build_graph.py"
+        )
+    with open(pickle_path, "rb") as f:
+        G = pickle.load(f)
+    _logger.info(f"Loaded graph: {G.number_of_nodes()} nodes, {G.number_of_edges()} edges")
+    return G
+
+
+def _json_safe(value: object) -> object:
+    if isinstance(value, set):
+        return list(value)
+    if isinstance(value, (list, dict, str, int, float, bool)) or value is None:
+        return value
+    return str(value)

main.py

@@ -14,7 +14,7 @@ import anthropic
 from rich.console import Console
 
 from agents.research_agent import stream_research_agent
-from config import CHROMA_COLLECTION, CHROMA_DIR, TRIALS_PATH
+from config import CHROMA_COLLECTION, CHROMA_DIR, GRAPH_PICKLE_PATH, TRIALS_PATH
 
 console = Console()
 
@@ -33,6 +33,14 @@ def _load_trials() -> list[dict]:
         return [json.loads(line) for line in f if line.strip()]
 
 
+def _load_graph():
+    try:
+        from graph.serializer import load_graph
+        return load_graph(GRAPH_PICKLE_PATH)
+    except FileNotFoundError:
+        return None
+
+
 def main() -> None:
     from rag.indexer import load_collection
 
@@ -47,12 +55,14 @@ def main() -> None:
         console.print("[yellow]Run: uv run python scripts/build_index.py[/yellow]")
         sys.exit(1)
 
+    graph = _load_graph()
     trials = _load_trials()
     client = anthropic.Anthropic()
 
     n_chunks = collection.count()
+    kg_status = f"{graph.number_of_nodes()} KG nodes" if graph else "no KG (run build_graph.py)"
     console.print(
-        f"[green]Ready.[/green] {n_chunks} chunks indexed · {len(trials)} trials loaded.{' ' * 20}"
+        f"[green]Ready.[/green] {n_chunks} chunks · {len(trials)} trials · {kg_status}.{' ' * 10}"
     )
 
     if not query:
@@ -70,7 +80,7 @@ def main() -> None:
 
     console.print()
 
-    for event_type, content in stream_research_agent(client, query, collection, trials):
+    for event_type, content in stream_research_agent(client, query, collection, trials, graph=graph):
         if event_type == "status":
             console.print(f"[dim italic]{content}[/dim italic]")
         elif event_type == "token":

rag/indexer.py

@@ -13,7 +13,9 @@ from models import ALSPaper
 
 _logger = get_logger("rag.indexer")
 
-_EMBED_FN = SentenceTransformerEmbeddingFunction(model_name="all-MiniLM-L6-v2")
+# BioLORD-2023-C: anchored to UMLS/SNOMED CT/MeSH ontologies — natively understands
+# biomedical synonyms (TARDBP = TDP-43, SOD1 = superoxide dismutase) and clinical phrasing.
+_EMBED_FN = SentenceTransformerEmbeddingFunction(model_name="FremyCompany/BioLORD-2023-C")
 
 
 def _chunk_paper(paper: ALSPaper) -> list[dict]:

scripts/build_graph.py

@@ -1,7 +1,59 @@
+#!/usr/bin/env python3
 """
-Build the ALS knowledge graph from extracted entities and clinical trials.
+Build the ALS knowledge graph from extracted entities + trials.
+Run after scripts/extract_entities.py.
 
 Usage:
     uv run python scripts/build_graph.py
 """
-# Stage 4 implementation
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from dotenv import load_dotenv
+
+load_dotenv()
+
+from rich.console import Console
+
+from config import ENTITIES_PATH, GRAPH_JSON_PATH, GRAPH_PICKLE_PATH, TRIALS_PATH
+from graph.builder import build_graph
+from graph.serializer import save_graph
+
+console = Console()
+
+
+def main() -> None:
+    console.print("[cyan]Building ALS knowledge graph...[/cyan]")
+
+    if not ENTITIES_PATH.exists():
+        console.print(
+            "[yellow]No entities.jsonl found — graph will be seeded only (no NER enrichment).[/yellow]"
+        )
+        console.print("[dim]Run first: uv run python scripts/extract_entities.py[/dim]\n")
+
+    G = build_graph(entities_path=ENTITIES_PATH, trials_path=TRIALS_PATH)
+    save_graph(G, pickle_path=GRAPH_PICKLE_PATH, json_path=GRAPH_JSON_PATH)
+
+    console.print(f"\n[bold green]Done![/bold green]")
+    console.print(f"  Nodes: [bold]{G.number_of_nodes()}[/bold]")
+    console.print(f"  Edges: [bold]{G.number_of_edges()}[/bold]")
+    console.print(f"  Pickle: {GRAPH_PICKLE_PATH}")
+    console.print(f"  JSON:   {GRAPH_JSON_PATH}")
+
+    top = sorted(G.degree(), key=lambda x: x[1], reverse=True)[:5]
+    if top:
+        console.print("\n[dim]Most connected nodes:[/dim]")
+        for node_id, deg in top:
+            name = G.nodes[node_id].get("display_name", node_id)
+            console.print(f"  [dim]{name} ({G.nodes[node_id].get('type', '?')}) — {deg} edges[/dim]")
+
+    console.print("\nTest a query with KG expansion:")
+    console.print("  [dim]uv run python main.py \"What is the role of C9orf72 in ALS?\"[/dim]")
+
+
+if __name__ == "__main__":
+    main()

scripts/extract_entities.py

@@ -1,10 +1,89 @@
+#!/usr/bin/env python3
 """
-Extract biomedical entities from papers using Claude Sonnet.
-Resumable — safe to interrupt and re-run.
+Extract biomedical entities from ALS papers using Claude Sonnet.
+Run after scripts/ingest_papers.py.
 
 Usage:
-    uv run python scripts/extract_entities.py            # all papers
-    uv run python scripts/extract_entities.py --limit 10 # test run
-    uv run python scripts/extract_entities.py --reset    # clear progress and restart
+    uv run python scripts/extract_entities.py
+    uv run python scripts/extract_entities.py --max 50     # test with first 50 papers
+    uv run python scripts/extract_entities.py --reset      # clear progress and re-extract
 """
-# Stage 4 implementation
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from dotenv import load_dotenv
+
+load_dotenv()
+
+import argparse
+import json
+
+import anthropic
+from rich.console import Console
+
+from config import (
+    CANONICAL_IDS_PATH,
+    ENTITIES_PATH,
+    EXTRACTION_PROGRESS_PATH,
+    PAPERS_PATH,
+)
+from extraction.extractor import extract_all
+
+console = Console()
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Extract entities from ALS papers")
+    parser.add_argument("--max", type=int, default=None, help="Only process first N papers")
+    parser.add_argument("--reset", action="store_true", help="Clear progress and re-extract")
+    parser.add_argument("--papers", default=str(PAPERS_PATH))
+    args = parser.parse_args()
+
+    papers_path = Path(args.papers)
+    if not papers_path.exists():
+        console.print(f"[red]Papers file not found: {papers_path}[/red]")
+        console.print("[yellow]Run first: uv run python scripts/ingest_papers.py[/yellow]")
+        sys.exit(1)
+
+    if args.reset:
+        for p in (ENTITIES_PATH, EXTRACTION_PROGRESS_PATH, CANONICAL_IDS_PATH):
+            if p.exists():
+                p.unlink()
+                console.print(f"[yellow]Deleted {p}[/yellow]")
+
+    if args.max:
+        all_lines = papers_path.read_text().strip().splitlines()
+        sliced = Path("/tmp/papers_slice.jsonl")
+        sliced.write_text("\n".join(all_lines[: args.max]))
+        papers_path = sliced
+        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")
+
+    client = anthropic.Anthropic()
+    results = extract_all(papers_path=papers_path, client=client)
+
+    total_entities = sum(len(r.entities) for r in results)
+    total_rels = sum(len(r.relationships) for r in results)
+
+    console.print(f"\n[bold green]Done![/bold green]")
+    console.print(f"  Papers processed this run: [bold]{len(results)}[/bold]")
+    console.print(f"  Total entities extracted:  [bold]{total_entities}[/bold]")
+    console.print(f"  Total relationships:        [bold]{total_rels}[/bold]")
+    console.print(f"  Output: {ENTITIES_PATH}")
+
+    if CANONICAL_IDS_PATH.exists():
+        registry = json.loads(CANONICAL_IDS_PATH.read_text())
+        console.print(f"  Canonical IDs: {len(registry)}")
+
+    console.print("\nNext step:")
+    console.print("  [dim]uv run python scripts/build_graph.py[/dim]")
+
+
+if __name__ == "__main__":
+    main()