Merge pull request #3 from KevinIsInCoding/fix/pmc-elink-id-mapping

agents/research_agent.py

@@ -1,2 +1,182 @@
-"""Multi-step ALS research synthesis agent (streaming)."""
-# Stage 3 implementation (RAG-only), upgraded in Stage 4 (KG+RAG)
+"""Multi-step ALS research synthesis agent (streaming). Mirrors beacon/agents/research.py pattern."""
+from __future__ import annotations
+
+import json
+from collections.abc import Generator
+
+import anthropic
+import chromadb
+
+from config import SYNTHESIS_MODEL
+from llm import cached_system, cached_tools
+from logging_config import get_logger
+from prompts import SYNTHESIS_SYSTEM
+from rag import retriever as rag_retriever
+from tools import RESEARCH_TOOLS
+
+_logger = get_logger("agents.research_agent")
+
+
+def stream_research_agent(
+    client: anthropic.Anthropic,
+    query: str,
+    collection: chromadb.Collection,
+    trials: list[dict],
+) -> Generator[tuple[str, str], None, None]:
+    """
+    Stream the ALS research synthesis agent.
+
+    Yields:
+        ("token", str)   — partial text chunk for streaming display
+        ("status", str)  — status message during tool execution
+        ("done", str)    — final complete response text
+    """
+    messages: list[anthropic.types.MessageParam] = [
+        {"role": "user", "content": query}
+    ]
+
+    while True:
+        stream_text = ""
+
+        with client.messages.stream(
+            model=SYNTHESIS_MODEL,
+            max_tokens=4096,
+            system=cached_system(SYNTHESIS_SYSTEM),
+            tools=cached_tools(RESEARCH_TOOLS),
+            messages=messages,
+        ) as stream:
+            # Accumulate tool-use input JSON alongside streaming text
+            tool_calls: list[dict] = []
+            current_tool: dict | None = None
+            current_input_json = ""
+
+            for event in stream:
+                if event.type == "content_block_start":
+                    if event.content_block.type == "tool_use":
+                        current_tool = {
+                            "id": event.content_block.id,
+                            "name": event.content_block.name,
+                        }
+                        current_input_json = ""
+                        yield ("status", "Searching ALS research knowledge base...")
+
+                elif event.type == "content_block_delta":
+                    if event.delta.type == "text_delta":
+                        chunk = event.delta.text
+                        stream_text += chunk
+                        yield ("token", chunk)
+                    elif event.delta.type == "input_json_delta" and current_tool:
+                        current_input_json += event.delta.partial_json
+
+                elif event.type == "content_block_stop":
+                    if current_tool is not None:
+                        try:
+                            current_tool["input"] = json.loads(current_input_json)
+                        except json.JSONDecodeError:
+                            current_tool["input"] = {}
+                        tool_calls.append(current_tool)
+                        current_tool = None
+                        current_input_json = ""
+
+            final_msg = stream.get_final_message()
+
+        messages.append({"role": "assistant", "content": final_msg.content})
+
+        if final_msg.stop_reason == "end_turn":
+            yield ("done", stream_text)
+            return
+
+        if final_msg.stop_reason == "tool_use" and tool_calls:
+            tool_results: list[anthropic.types.ToolResultBlockParam] = []
+
+            for tool_call in tool_calls:
+                if tool_call["name"] == "search_research_landscape":
+                    result = _handle_search(tool_call["input"], collection, trials)
+                    is_error = False
+                else:
+                    result = {"error": f"Unknown tool: {tool_call['name']}"}
+                    is_error = True
+
+                tool_results.append({
+                    "type": "tool_result",
+                    "tool_use_id": tool_call["id"],
+                    "content": json.dumps(result),
+                    "is_error": is_error,
+                })
+
+            messages.append({"role": "user", "content": tool_results})
+        else:
+            yield ("done", stream_text)
+            return
+
+
+def _handle_search(
+    tool_input: dict,
+    collection: chromadb.Collection,
+    trials: list[dict],
+) -> dict:
+    """Execute RAG search + trial lookup and return structured context for Claude."""
+    query_text = tool_input.get("query_text", "")
+    query_entities = tool_input.get("query_entities", [])
+
+    # Semantic search
+    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)
+
+    # Merge by PMID — keep best score per paper
+    seen: dict[str, dict] = {}
+    for r in semantic_results + entity_results:
+        pmid = r["pmid"]
+        if pmid not in seen or r["score"] > seen[pmid]["score"]:
+            seen[pmid] = r
+
+    top_papers = sorted(seen.values(), key=lambda x: x["score"], reverse=True)[:15]
+
+    _logger.info(
+        "RAG search",
+        extra={"data": {
+            "query_entities": query_entities,
+            "semantic_hits": len(semantic_results),
+            "entity_hits": len(entity_results),
+            "merged": len(top_papers),
+        }},
+    )
+
+    # Match trials by entity name in title or intervention text
+    related_trials = []
+    if query_entities:
+        entities_lower = [e.lower() for e in query_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()
+            if any(e in trial_text for e in entities_lower):
+                related_trials.append({
+                    "nct_id": trial.get("nct_id", ""),
+                    "title": trial.get("title", ""),
+                    "phase": trial.get("phase", ""),
+                    "status": trial.get("status", ""),
+                    "url": trial.get("url", ""),
+                })
+            if len(related_trials) >= 5:
+                break
+
+    return {
+        "papers": [
+            {
+                "pmid": r["pmid"],
+                "title": r["title"],
+                "year": r["year"],
+                "doi": r["doi"],
+                "citation_count": r["citation_count"],
+                "section": r["section"],
+                "excerpt": r["document"][:600],
+                "score": round(r["score"], 3),
+            }
+            for r in top_papers
+        ],
+        "query_entities": query_entities,
+        "trials": related_trials,
+        "evidence_count": len(top_papers),
+    }

main.py

@@ -1,2 +1,83 @@
-"""CLI entry point for candle-fire."""
-# Stage 3 implementation
+#!/usr/bin/env python3
+"""CLI interface for candle-fire ALS research synthesis."""
+from __future__ import annotations
+
+import json
+import sys
+from pathlib import Path
+
+from dotenv import load_dotenv
+
+load_dotenv()
+
+import anthropic
+from rich.console import Console
+
+from agents.research_agent import stream_research_agent
+from config import CHROMA_COLLECTION, CHROMA_DIR, TRIALS_PATH
+
+console = Console()
+
+_EXAMPLES = [
+    "What's the evidence for tofersen targeting SOD1 in ALS?",
+    "What mechanisms link TDP-43 to ALS pathology?",
+    "What compounds target glutamate excitotoxicity in ALS?",
+    "What is the role of C9orf72 repeat expansion in ALS?",
+]
+
+
+def _load_trials() -> list[dict]:
+    if not TRIALS_PATH.exists():
+        return []
+    with open(TRIALS_PATH, encoding="utf-8") as f:
+        return [json.loads(line) for line in f if line.strip()]
+
+
+def main() -> None:
+    from rag.indexer import load_collection
+
+    query = " ".join(sys.argv[1:]).strip()
+
+    # Load resources
+    console.print("[cyan]Loading knowledge base...[/cyan]", end="\r")
+    try:
+        collection = load_collection(CHROMA_DIR, CHROMA_COLLECTION)
+    except Exception:
+        console.print("[red]ChromaDB collection not found.[/red]")
+        console.print("[yellow]Run: uv run python scripts/build_index.py[/yellow]")
+        sys.exit(1)
+
+    trials = _load_trials()
+    client = anthropic.Anthropic()
+
+    n_chunks = collection.count()
+    console.print(
+        f"[green]Ready.[/green] {n_chunks} chunks indexed · {len(trials)} trials loaded.{' ' * 20}"
+    )
+
+    if not query:
+        console.print("\n[dim]Example questions:[/dim]")
+        for ex in _EXAMPLES:
+            console.print(f"[dim]  • {ex}[/dim]")
+        console.print()
+        try:
+            query = console.input("[bold]Ask about ALS research:[/bold] ").strip()
+        except (KeyboardInterrupt, EOFError):
+            sys.exit(0)
+
+    if not query:
+        sys.exit(0)
+
+    console.print()
+
+    for event_type, content in stream_research_agent(client, query, collection, trials):
+        if event_type == "status":
+            console.print(f"[dim italic]{content}[/dim italic]")
+        elif event_type == "token":
+            console.print(content, end="", highlight=False)
+        elif event_type == "done":
+            console.print()
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -14,6 +14,7 @@ dependencies = [
     "python-dotenv>=1.2.2",
     "rich>=13.0.0",
     "sentence-transformers>=3.0.0",
+    "openai>=2.44.0",
 ]
 
 [project.optional-dependencies]

rag/indexer.py

@@ -1,2 +1,146 @@
-"""ChromaDB collection builder with section-aware chunking and citation metadata."""
-# Stage 3 implementation
+"""ChromaDB collection builder with section-aware chunking."""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import chromadb
+from chromadb.utils.embedding_functions import SentenceTransformerEmbeddingFunction
+
+from config import CHROMA_COLLECTION, CHROMA_DIR, PAPERS_PATH
+from logging_config import get_logger
+from models import ALSPaper
+
+_logger = get_logger("rag.indexer")
+
+_EMBED_FN = SentenceTransformerEmbeddingFunction(model_name="all-MiniLM-L6-v2")
+
+
+def _chunk_paper(paper: ALSPaper) -> list[dict]:
+    """
+    Split a paper into indexable chunks.
+    - Full text available: one chunk per section (split on [Section Title] markers)
+    - Abstract only: single chunk = title + mesh terms + abstract
+    """
+    base_meta = {
+        "pmid": paper.pmid,
+        "title": paper.title,
+        "year": paper.year,
+        "doi": paper.doi,
+        "citation_count": paper.citation_count,
+        # ChromaDB metadata must be scalar — serialize lists as comma-separated strings
+        "entity_names": ",".join(paper.entity_names),
+        "mesh_terms": ",".join(paper.mesh_terms[:10]),  # cap to avoid huge metadata
+        "has_full_text": int(bool(paper.full_text)),  # bool not supported → int
+    }
+
+    if paper.full_text:
+        sections: list[tuple[str, str]] = []
+        current_title = "Abstract"
+        current_lines: list[str] = [paper.abstract]
+
+        for line in paper.full_text.split("\n"):
+            stripped = line.strip()
+            if stripped.startswith("[") and stripped.endswith("]") and len(stripped) < 80:
+                if current_lines:
+                    body = "\n".join(current_lines).strip()
+                    if body:
+                        sections.append((current_title, body))
+                current_title = stripped[1:-1]
+                current_lines = []
+            else:
+                current_lines.append(line)
+
+        if current_lines:
+            body = "\n".join(current_lines).strip()
+            if body:
+                sections.append((current_title, body))
+
+        chunks = []
+        for i, (section_title, section_text) in enumerate(sections):
+            doc = f"{paper.title}\n[{section_title}]\n{section_text}"
+            chunks.append({
+                "id": f"{paper.pmid}_s{i}",
+                "document": doc,
+                "metadata": {**base_meta, "section": section_title, "chunk_index": i},
+            })
+        return chunks if chunks else [_abstract_chunk(paper, base_meta)]
+
+    return [_abstract_chunk(paper, base_meta)]
+
+
+def _abstract_chunk(paper: ALSPaper, base_meta: dict) -> dict:
+    doc = f"{paper.title}\n{' '.join(paper.mesh_terms)}\n{paper.abstract}"
+    return {
+        "id": paper.pmid,
+        "document": doc,
+        "metadata": {**base_meta, "section": "abstract", "chunk_index": 0},
+    }
+
+
+def build_collection(
+    papers_path: Path = PAPERS_PATH,
+    chroma_dir: Path = CHROMA_DIR,
+    collection_name: str = CHROMA_COLLECTION,
+    reset: bool = False,
+) -> chromadb.Collection:
+    """Build ChromaDB collection from papers.jsonl. Idempotent — skips already-indexed chunks."""
+    chroma_dir.mkdir(parents=True, exist_ok=True)
+    client = chromadb.PersistentClient(path=str(chroma_dir))
+
+    if reset:
+        try:
+            client.delete_collection(collection_name)
+            _logger.info(f"Deleted collection: {collection_name}")
+        except Exception:
+            pass
+
+    collection = client.get_or_create_collection(
+        name=collection_name,
+        embedding_function=_EMBED_FN,
+        metadata={"hnsw:space": "cosine"},
+    )
+
+    papers: list[ALSPaper] = []
+    with open(papers_path, encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if line:
+                papers.append(ALSPaper.from_dict(json.loads(line)))
+
+    _logger.info(f"Loaded {len(papers)} papers")
+
+    all_chunks = []
+    for paper in papers:
+        all_chunks.extend(_chunk_paper(paper))
+
+    # Skip already-indexed chunks (safe to re-run)
+    existing_ids = set(collection.get(include=[])["ids"])
+    new_chunks = [c for c in all_chunks if c["id"] not in existing_ids]
+
+    if not new_chunks:
+        _logger.info("All chunks already indexed")
+        return collection
+
+    _logger.info(f"Indexing {len(new_chunks)} new chunks from {len(papers)} papers")
+
+    batch_size = 100
+    for i in range(0, len(new_chunks), batch_size):
+        batch = new_chunks[i : i + batch_size]
+        collection.add(
+            ids=[c["id"] for c in batch],
+            documents=[c["document"] for c in batch],
+            metadatas=[c["metadata"] for c in batch],
+        )
+
+    _logger.info(f"Collection '{collection_name}': {collection.count()} total chunks")
+    return collection
+
+
+def load_collection(
+    chroma_dir: Path = CHROMA_DIR,
+    collection_name: str = CHROMA_COLLECTION,
+) -> chromadb.Collection:
+    """Load an existing collection at query time (fast, no re-embedding)."""
+    client = chromadb.PersistentClient(path=str(chroma_dir))
+    return client.get_collection(name=collection_name, embedding_function=_EMBED_FN)

rag/retriever.py

@@ -1,2 +1,132 @@
 """ChromaDB query interface with citation-weighted re-ranking."""
-# Stage 3 implementation
+from __future__ import annotations
+
+import math
+
+import chromadb
+
+from config import CHROMA_ENTITY_N_RESULTS, CHROMA_N_RESULTS
+from logging_config import get_logger
+
+_logger = get_logger("rag.retriever")
+
+
+def search(
+    collection: chromadb.Collection,
+    query_text: str,
+    n_results: int = CHROMA_N_RESULTS,
+) -> list[dict]:
+    """
+    Semantic search with citation-weighted re-ranking.
+    Over-fetches 2× then re-ranks by: similarity * log(citation_count + 2).
+    Deduplicates to one chunk per paper (best-scoring chunk wins).
+    """
+    n_fetch = min(n_results * 2, collection.count())
+    if n_fetch == 0:
+        return []
+
+    raw = collection.query(
+        query_texts=[query_text],
+        n_results=n_fetch,
+        include=["documents", "metadatas", "distances"],
+    )
+    results = _parse_raw(raw)
+    results = _rerank(results)
+    return results[:n_results]
+
+
+def search_by_entities(
+    collection: chromadb.Collection,
+    entity_names: list[str],
+    n_results: int = CHROMA_ENTITY_N_RESULTS,
+) -> list[dict]:
+    """
+    Run one query per entity, merge and deduplicate by PMID.
+    Caps at 8 entity queries to avoid excessive API calls.
+    """
+    if not entity_names or collection.count() == 0:
+        return []
+
+    seen: dict[str, dict] = {}
+    for entity in entity_names[:8]:
+        raw = collection.query(
+            query_texts=[entity],
+            n_results=min(10, collection.count()),
+            include=["documents", "metadatas", "distances"],
+        )
+        for r in _parse_raw(raw):
+            pmid = r["pmid"]
+            if pmid not in seen or r["score"] > seen[pmid]["score"]:
+                seen[pmid] = r
+
+    merged = _rerank(list(seen.values()))
+    return merged[:n_results]
+
+
+def get_paper(collection: chromadb.Collection, pmid: str) -> dict | None:
+    """Retrieve a specific paper's abstract chunk by PMID."""
+    result = collection.get(
+        where={"$and": [{"pmid": {"$eq": pmid}}, {"chunk_index": {"$eq": 0}}]},
+        include=["documents", "metadatas"],
+    )
+    ids = result.get("ids", [])
+    if not ids:
+        return None
+    meta = result["metadatas"][0]
+    return {
+        "pmid": pmid,
+        "title": meta.get("title", ""),
+        "year": meta.get("year", 0),
+        "doi": meta.get("doi", ""),
+        "citation_count": meta.get("citation_count", 0),
+        "document": result["documents"][0],
+    }
+
+
+def _parse_raw(raw: dict) -> list[dict]:
+    """Flatten a ChromaDB query response into a list of result dicts."""
+    ids = raw.get("ids", [[]])[0]
+    docs = raw.get("documents", [[]])[0]
+    metas = raw.get("metadatas", [[]])[0]
+    distances = raw.get("distances", [[]])[0]
+
+    results = []
+    for chunk_id, doc, meta, dist in zip(ids, docs, metas, distances):
+        similarity = max(0.0, 1.0 - dist)
+        citation_count = int(meta.get("citation_count", 0))
+        results.append({
+            "chunk_id": chunk_id,
+            "pmid": meta.get("pmid", ""),
+            "title": meta.get("title", ""),
+            "year": int(meta.get("year", 0)),
+            "doi": meta.get("doi", ""),
+            "section": meta.get("section", "abstract"),
+            "citation_count": citation_count,
+            "entity_names": [e for e in meta.get("entity_names", "").split(",") if e],
+            "has_full_text": bool(meta.get("has_full_text", 0)),
+            "document": doc,
+            "similarity": similarity,
+            "score": similarity,
+        })
+    return results
+
+
+def _rerank(results: list[dict]) -> list[dict]:
+    """
+    Apply citation weighting and deduplicate to one chunk per PMID.
+    score = similarity * log(citation_count + 2)
+    log(2) ≈ 0.69 is the floor for uncited papers, so they're still ranked
+    but deprioritized relative to highly-cited work.
+    """
+    for r in results:
+        r["score"] = r["similarity"] * math.log(r["citation_count"] + 2)
+
+    results.sort(key=lambda x: x["score"], reverse=True)
+
+    # Keep best chunk per paper
+    seen: dict[str, dict] = {}
+    for r in results:
+        pmid = r["pmid"]
+        if pmid not in seen:
+            seen[pmid] = r
+    return list(seen.values())

scripts/build_index.py

@@ -1,8 +1,64 @@
+#!/usr/bin/env python3
 """
-Build the ChromaDB vector index from papers and extracted entities.
+Build the ChromaDB vector index from papers.jsonl.
+Run after scripts/ingest_papers.py.
 
 Usage:
     uv run python scripts/build_index.py
-    uv run python scripts/build_index.py --reset  # drop and rebuild collection
+    uv run python scripts/build_index.py --reset   # drop and rebuild from scratch
 """
-# Stage 3 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
+
+from rich.console import Console
+from rich.progress import track
+
+from config import CHROMA_COLLECTION, CHROMA_DIR, PAPERS_PATH
+from rag.indexer import build_collection
+
+console = Console()
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Build ChromaDB vector index from ALS papers")
+    parser.add_argument("--reset", action="store_true", help="Drop existing collection and rebuild")
+    parser.add_argument("--papers", default=str(PAPERS_PATH), help="Path to papers.jsonl")
+    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:
+        console.print("[yellow]Resetting collection...[/yellow]")
+
+    console.print(f"[cyan]Building ChromaDB index from {papers_path}...[/cyan]")
+    console.print("[dim]First run downloads the embedding model (~80MB). Subsequent runs are fast.[/dim]\n")
+
+    collection = build_collection(
+        papers_path=papers_path,
+        chroma_dir=CHROMA_DIR,
+        collection_name=CHROMA_COLLECTION,
+        reset=args.reset,
+    )
+
+    console.print(f"\n[bold green]Done![/bold green] Collection '{CHROMA_COLLECTION}' at {CHROMA_DIR}")
+    console.print(f"  Total chunks indexed: [bold]{collection.count()}[/bold]")
+    console.print("\nTest a query:")
+    console.print("  [dim]uv run python main.py \"What is the evidence for tofersen targeting SOD1?\"[/dim]")
+
+
+if __name__ == "__main__":
+    main()

uv.lock

@@ -451,6 +451,7 @@ dependencies = [
     { name = "gradio" },
     { name = "httpx" },
     { name = "networkx" },
+    { name = "openai" },
     { name = "python-dotenv" },
     { name = "rich" },
     { name = "sentence-transformers" },
@@ -472,6 +473,7 @@ requires-dist = [
     { name = "gradio", specifier = ">=6.14.0,<7.0.0" },
     { name = "httpx", specifier = ">=0.27.0" },
     { name = "networkx", specifier = ">=3.3" },
+    { name = "openai", specifier = ">=2.44.0" },
     { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0" },
     { name = "pytest-cov", marker = "extra == 'dev'" },
     { name = "pytest-httpx", marker = "extra == 'dev'", specifier = ">=0.35" },
@@ -2097,6 +2099,25 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/b7/f6/2bac21f722aa45d876d4a51f26bd0ef30e704068a3cd5021a5a7cd784271/onnxruntime-1.27.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:370d211e1ceeac4cd5f45301655463ac59e27cdc74d9f7aeb2d19ff4b7a76715", size = 18670781, upload-time = "2026-06-15T22:43:17.151Z" },
 ]
 
+[[package]]
+name = "openai"
+version = "2.44.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "anyio" },
+    { name = "distro" },
+    { name = "httpx" },
+    { name = "jiter" },
+    { name = "pydantic" },
+    { name = "sniffio" },
+    { name = "tqdm" },
+    { name = "typing-extensions" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/49/f5/7c7cb955305cb41f7f3c5fd7e0e38bf6bbf2658468863d4b7b868a5cb8df/openai-2.44.0.tar.gz", hash = "sha256:68a5a5ffad82b8ff7d451c437529fb64f7c3b8123aaf0c021966a882d9e3947d", size = 988753, upload-time = "2026-06-24T20:56:02.293Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/ae/f4/561ed79fd94876160018a5e75254cfcb9b0e62d4dded9dcb20072e86d623/openai-2.44.0-py3-none-any.whl", hash = "sha256:0a2a3ab2e29aeda368700f662ff9ba0f9df17ba4c54577a64e08b8115a3cc0ad", size = 1366216, upload-time = "2026-06-24T20:55:58.882Z" },
+]
+
 [[package]]
 name = "opentelemetry-api"
 version = "1.43.0"