KevinIsInCoding Claude Sonnet 4.6 commited on
Commit
56fa10c
·
1 Parent(s): 83263a0

feat: Stage 3 — RAG pipeline + working end-to-end query

Browse files

- rag/indexer.py: ChromaDB collection builder with section-aware chunking
(one chunk per PMC XML section, falls back to abstract)
- rag/retriever.py: semantic search + entity search with citation-weighted
re-ranking (score = similarity * log(citation_count + 2))
- agents/research_agent.py: streaming synthesis agent (mirrors beacon pattern)
uses search_research_landscape tool → RAG + trial match → Claude synthesis
- scripts/build_index.py: CLI to build ChromaDB index from papers.jsonl
- main.py: Rich CLI interface with streaming output
- Add openai dependency (required by llm.py provider abstraction)

Verified: 510 chunks indexed from 499 papers, end-to-end query returns
structured research landscape with PMIDs, mechanisms, and trial links.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Files changed (7) hide show
  1. agents/research_agent.py +182 -2
  2. main.py +83 -2
  3. pyproject.toml +1 -0
  4. rag/indexer.py +146 -2
  5. rag/retriever.py +131 -1
  6. scripts/build_index.py +59 -3
  7. uv.lock +21 -0
agents/research_agent.py CHANGED
@@ -1,2 +1,182 @@
1
- """Multi-step ALS research synthesis agent (streaming)."""
2
- # Stage 3 implementation (RAG-only), upgraded in Stage 4 (KG+RAG)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Multi-step ALS research synthesis agent (streaming). Mirrors beacon/agents/research.py pattern."""
2
+ from __future__ import annotations
3
+
4
+ import json
5
+ from collections.abc import Generator
6
+
7
+ import anthropic
8
+ import chromadb
9
+
10
+ from config import SYNTHESIS_MODEL
11
+ from llm import cached_system, cached_tools
12
+ from logging_config import get_logger
13
+ from prompts import SYNTHESIS_SYSTEM
14
+ from rag import retriever as rag_retriever
15
+ from tools import RESEARCH_TOOLS
16
+
17
+ _logger = get_logger("agents.research_agent")
18
+
19
+
20
+ def stream_research_agent(
21
+ client: anthropic.Anthropic,
22
+ query: str,
23
+ collection: chromadb.Collection,
24
+ trials: list[dict],
25
+ ) -> Generator[tuple[str, str], None, None]:
26
+ """
27
+ Stream the ALS research synthesis agent.
28
+
29
+ Yields:
30
+ ("token", str) — partial text chunk for streaming display
31
+ ("status", str) — status message during tool execution
32
+ ("done", str) — final complete response text
33
+ """
34
+ messages: list[anthropic.types.MessageParam] = [
35
+ {"role": "user", "content": query}
36
+ ]
37
+
38
+ while True:
39
+ stream_text = ""
40
+
41
+ with client.messages.stream(
42
+ model=SYNTHESIS_MODEL,
43
+ max_tokens=4096,
44
+ system=cached_system(SYNTHESIS_SYSTEM),
45
+ tools=cached_tools(RESEARCH_TOOLS),
46
+ messages=messages,
47
+ ) as stream:
48
+ # Accumulate tool-use input JSON alongside streaming text
49
+ tool_calls: list[dict] = []
50
+ current_tool: dict | None = None
51
+ current_input_json = ""
52
+
53
+ for event in stream:
54
+ if event.type == "content_block_start":
55
+ if event.content_block.type == "tool_use":
56
+ current_tool = {
57
+ "id": event.content_block.id,
58
+ "name": event.content_block.name,
59
+ }
60
+ current_input_json = ""
61
+ yield ("status", "Searching ALS research knowledge base...")
62
+
63
+ elif event.type == "content_block_delta":
64
+ if event.delta.type == "text_delta":
65
+ chunk = event.delta.text
66
+ stream_text += chunk
67
+ yield ("token", chunk)
68
+ elif event.delta.type == "input_json_delta" and current_tool:
69
+ current_input_json += event.delta.partial_json
70
+
71
+ elif event.type == "content_block_stop":
72
+ if current_tool is not None:
73
+ try:
74
+ current_tool["input"] = json.loads(current_input_json)
75
+ except json.JSONDecodeError:
76
+ current_tool["input"] = {}
77
+ tool_calls.append(current_tool)
78
+ current_tool = None
79
+ current_input_json = ""
80
+
81
+ final_msg = stream.get_final_message()
82
+
83
+ messages.append({"role": "assistant", "content": final_msg.content})
84
+
85
+ if final_msg.stop_reason == "end_turn":
86
+ yield ("done", stream_text)
87
+ return
88
+
89
+ if final_msg.stop_reason == "tool_use" and tool_calls:
90
+ tool_results: list[anthropic.types.ToolResultBlockParam] = []
91
+
92
+ for tool_call in tool_calls:
93
+ if tool_call["name"] == "search_research_landscape":
94
+ result = _handle_search(tool_call["input"], collection, trials)
95
+ is_error = False
96
+ else:
97
+ result = {"error": f"Unknown tool: {tool_call['name']}"}
98
+ is_error = True
99
+
100
+ tool_results.append({
101
+ "type": "tool_result",
102
+ "tool_use_id": tool_call["id"],
103
+ "content": json.dumps(result),
104
+ "is_error": is_error,
105
+ })
106
+
107
+ messages.append({"role": "user", "content": tool_results})
108
+ else:
109
+ yield ("done", stream_text)
110
+ return
111
+
112
+
113
+ def _handle_search(
114
+ tool_input: dict,
115
+ collection: chromadb.Collection,
116
+ trials: list[dict],
117
+ ) -> dict:
118
+ """Execute RAG search + trial lookup and return structured context for Claude."""
119
+ query_text = tool_input.get("query_text", "")
120
+ query_entities = tool_input.get("query_entities", [])
121
+
122
+ # Semantic search
123
+ semantic_results = rag_retriever.search(collection, query_text, n_results=10)
124
+
125
+ # Entity-targeted search (finds papers even if query terms don't match directly)
126
+ entity_results = rag_retriever.search_by_entities(collection, query_entities, n_results=15)
127
+
128
+ # Merge by PMID — keep best score per paper
129
+ seen: dict[str, dict] = {}
130
+ for r in semantic_results + entity_results:
131
+ pmid = r["pmid"]
132
+ if pmid not in seen or r["score"] > seen[pmid]["score"]:
133
+ seen[pmid] = r
134
+
135
+ top_papers = sorted(seen.values(), key=lambda x: x["score"], reverse=True)[:15]
136
+
137
+ _logger.info(
138
+ "RAG search",
139
+ extra={"data": {
140
+ "query_entities": query_entities,
141
+ "semantic_hits": len(semantic_results),
142
+ "entity_hits": len(entity_results),
143
+ "merged": len(top_papers),
144
+ }},
145
+ )
146
+
147
+ # Match trials by entity name in title or intervention text
148
+ related_trials = []
149
+ if query_entities:
150
+ entities_lower = [e.lower() for e in query_entities]
151
+ for trial in trials:
152
+ iv_names = " ".join(iv.get("name", "") for iv in trial.get("interventions", []))
153
+ trial_text = f"{trial.get('title', '')} {iv_names}".lower()
154
+ if any(e in trial_text for e in entities_lower):
155
+ related_trials.append({
156
+ "nct_id": trial.get("nct_id", ""),
157
+ "title": trial.get("title", ""),
158
+ "phase": trial.get("phase", ""),
159
+ "status": trial.get("status", ""),
160
+ "url": trial.get("url", ""),
161
+ })
162
+ if len(related_trials) >= 5:
163
+ break
164
+
165
+ return {
166
+ "papers": [
167
+ {
168
+ "pmid": r["pmid"],
169
+ "title": r["title"],
170
+ "year": r["year"],
171
+ "doi": r["doi"],
172
+ "citation_count": r["citation_count"],
173
+ "section": r["section"],
174
+ "excerpt": r["document"][:600],
175
+ "score": round(r["score"], 3),
176
+ }
177
+ for r in top_papers
178
+ ],
179
+ "query_entities": query_entities,
180
+ "trials": related_trials,
181
+ "evidence_count": len(top_papers),
182
+ }
main.py CHANGED
@@ -1,2 +1,83 @@
1
- """CLI entry point for candle-fire."""
2
- # Stage 3 implementation
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """CLI interface for candle-fire ALS research synthesis."""
3
+ from __future__ import annotations
4
+
5
+ import json
6
+ import sys
7
+ from pathlib import Path
8
+
9
+ from dotenv import load_dotenv
10
+
11
+ load_dotenv()
12
+
13
+ import anthropic
14
+ from rich.console import Console
15
+
16
+ from agents.research_agent import stream_research_agent
17
+ from config import CHROMA_COLLECTION, CHROMA_DIR, TRIALS_PATH
18
+
19
+ console = Console()
20
+
21
+ _EXAMPLES = [
22
+ "What's the evidence for tofersen targeting SOD1 in ALS?",
23
+ "What mechanisms link TDP-43 to ALS pathology?",
24
+ "What compounds target glutamate excitotoxicity in ALS?",
25
+ "What is the role of C9orf72 repeat expansion in ALS?",
26
+ ]
27
+
28
+
29
+ def _load_trials() -> list[dict]:
30
+ if not TRIALS_PATH.exists():
31
+ return []
32
+ with open(TRIALS_PATH, encoding="utf-8") as f:
33
+ return [json.loads(line) for line in f if line.strip()]
34
+
35
+
36
+ def main() -> None:
37
+ from rag.indexer import load_collection
38
+
39
+ query = " ".join(sys.argv[1:]).strip()
40
+
41
+ # Load resources
42
+ console.print("[cyan]Loading knowledge base...[/cyan]", end="\r")
43
+ try:
44
+ collection = load_collection(CHROMA_DIR, CHROMA_COLLECTION)
45
+ except Exception:
46
+ console.print("[red]ChromaDB collection not found.[/red]")
47
+ console.print("[yellow]Run: uv run python scripts/build_index.py[/yellow]")
48
+ sys.exit(1)
49
+
50
+ trials = _load_trials()
51
+ client = anthropic.Anthropic()
52
+
53
+ n_chunks = collection.count()
54
+ console.print(
55
+ f"[green]Ready.[/green] {n_chunks} chunks indexed · {len(trials)} trials loaded.{' ' * 20}"
56
+ )
57
+
58
+ if not query:
59
+ console.print("\n[dim]Example questions:[/dim]")
60
+ for ex in _EXAMPLES:
61
+ console.print(f"[dim] • {ex}[/dim]")
62
+ console.print()
63
+ try:
64
+ query = console.input("[bold]Ask about ALS research:[/bold] ").strip()
65
+ except (KeyboardInterrupt, EOFError):
66
+ sys.exit(0)
67
+
68
+ if not query:
69
+ sys.exit(0)
70
+
71
+ console.print()
72
+
73
+ for event_type, content in stream_research_agent(client, query, collection, trials):
74
+ if event_type == "status":
75
+ console.print(f"[dim italic]{content}[/dim italic]")
76
+ elif event_type == "token":
77
+ console.print(content, end="", highlight=False)
78
+ elif event_type == "done":
79
+ console.print()
80
+
81
+
82
+ if __name__ == "__main__":
83
+ main()
pyproject.toml CHANGED
@@ -14,6 +14,7 @@ dependencies = [
14
  "python-dotenv>=1.2.2",
15
  "rich>=13.0.0",
16
  "sentence-transformers>=3.0.0",
 
17
  ]
18
 
19
  [project.optional-dependencies]
 
14
  "python-dotenv>=1.2.2",
15
  "rich>=13.0.0",
16
  "sentence-transformers>=3.0.0",
17
+ "openai>=2.44.0",
18
  ]
19
 
20
  [project.optional-dependencies]
rag/indexer.py CHANGED
@@ -1,2 +1,146 @@
1
- """ChromaDB collection builder with section-aware chunking and citation metadata."""
2
- # Stage 3 implementation
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """ChromaDB collection builder with section-aware chunking."""
2
+ from __future__ import annotations
3
+
4
+ import json
5
+ from pathlib import Path
6
+
7
+ import chromadb
8
+ from chromadb.utils.embedding_functions import SentenceTransformerEmbeddingFunction
9
+
10
+ from config import CHROMA_COLLECTION, CHROMA_DIR, PAPERS_PATH
11
+ from logging_config import get_logger
12
+ from models import ALSPaper
13
+
14
+ _logger = get_logger("rag.indexer")
15
+
16
+ _EMBED_FN = SentenceTransformerEmbeddingFunction(model_name="all-MiniLM-L6-v2")
17
+
18
+
19
+ def _chunk_paper(paper: ALSPaper) -> list[dict]:
20
+ """
21
+ Split a paper into indexable chunks.
22
+ - Full text available: one chunk per section (split on [Section Title] markers)
23
+ - Abstract only: single chunk = title + mesh terms + abstract
24
+ """
25
+ base_meta = {
26
+ "pmid": paper.pmid,
27
+ "title": paper.title,
28
+ "year": paper.year,
29
+ "doi": paper.doi,
30
+ "citation_count": paper.citation_count,
31
+ # ChromaDB metadata must be scalar — serialize lists as comma-separated strings
32
+ "entity_names": ",".join(paper.entity_names),
33
+ "mesh_terms": ",".join(paper.mesh_terms[:10]), # cap to avoid huge metadata
34
+ "has_full_text": int(bool(paper.full_text)), # bool not supported → int
35
+ }
36
+
37
+ if paper.full_text:
38
+ sections: list[tuple[str, str]] = []
39
+ current_title = "Abstract"
40
+ current_lines: list[str] = [paper.abstract]
41
+
42
+ for line in paper.full_text.split("\n"):
43
+ stripped = line.strip()
44
+ if stripped.startswith("[") and stripped.endswith("]") and len(stripped) < 80:
45
+ if current_lines:
46
+ body = "\n".join(current_lines).strip()
47
+ if body:
48
+ sections.append((current_title, body))
49
+ current_title = stripped[1:-1]
50
+ current_lines = []
51
+ else:
52
+ current_lines.append(line)
53
+
54
+ if current_lines:
55
+ body = "\n".join(current_lines).strip()
56
+ if body:
57
+ sections.append((current_title, body))
58
+
59
+ chunks = []
60
+ for i, (section_title, section_text) in enumerate(sections):
61
+ doc = f"{paper.title}\n[{section_title}]\n{section_text}"
62
+ chunks.append({
63
+ "id": f"{paper.pmid}_s{i}",
64
+ "document": doc,
65
+ "metadata": {**base_meta, "section": section_title, "chunk_index": i},
66
+ })
67
+ return chunks if chunks else [_abstract_chunk(paper, base_meta)]
68
+
69
+ return [_abstract_chunk(paper, base_meta)]
70
+
71
+
72
+ def _abstract_chunk(paper: ALSPaper, base_meta: dict) -> dict:
73
+ doc = f"{paper.title}\n{' '.join(paper.mesh_terms)}\n{paper.abstract}"
74
+ return {
75
+ "id": paper.pmid,
76
+ "document": doc,
77
+ "metadata": {**base_meta, "section": "abstract", "chunk_index": 0},
78
+ }
79
+
80
+
81
+ def build_collection(
82
+ papers_path: Path = PAPERS_PATH,
83
+ chroma_dir: Path = CHROMA_DIR,
84
+ collection_name: str = CHROMA_COLLECTION,
85
+ reset: bool = False,
86
+ ) -> chromadb.Collection:
87
+ """Build ChromaDB collection from papers.jsonl. Idempotent — skips already-indexed chunks."""
88
+ chroma_dir.mkdir(parents=True, exist_ok=True)
89
+ client = chromadb.PersistentClient(path=str(chroma_dir))
90
+
91
+ if reset:
92
+ try:
93
+ client.delete_collection(collection_name)
94
+ _logger.info(f"Deleted collection: {collection_name}")
95
+ except Exception:
96
+ pass
97
+
98
+ collection = client.get_or_create_collection(
99
+ name=collection_name,
100
+ embedding_function=_EMBED_FN,
101
+ metadata={"hnsw:space": "cosine"},
102
+ )
103
+
104
+ papers: list[ALSPaper] = []
105
+ with open(papers_path, encoding="utf-8") as f:
106
+ for line in f:
107
+ line = line.strip()
108
+ if line:
109
+ papers.append(ALSPaper.from_dict(json.loads(line)))
110
+
111
+ _logger.info(f"Loaded {len(papers)} papers")
112
+
113
+ all_chunks = []
114
+ for paper in papers:
115
+ all_chunks.extend(_chunk_paper(paper))
116
+
117
+ # Skip already-indexed chunks (safe to re-run)
118
+ existing_ids = set(collection.get(include=[])["ids"])
119
+ new_chunks = [c for c in all_chunks if c["id"] not in existing_ids]
120
+
121
+ if not new_chunks:
122
+ _logger.info("All chunks already indexed")
123
+ return collection
124
+
125
+ _logger.info(f"Indexing {len(new_chunks)} new chunks from {len(papers)} papers")
126
+
127
+ batch_size = 100
128
+ for i in range(0, len(new_chunks), batch_size):
129
+ batch = new_chunks[i : i + batch_size]
130
+ collection.add(
131
+ ids=[c["id"] for c in batch],
132
+ documents=[c["document"] for c in batch],
133
+ metadatas=[c["metadata"] for c in batch],
134
+ )
135
+
136
+ _logger.info(f"Collection '{collection_name}': {collection.count()} total chunks")
137
+ return collection
138
+
139
+
140
+ def load_collection(
141
+ chroma_dir: Path = CHROMA_DIR,
142
+ collection_name: str = CHROMA_COLLECTION,
143
+ ) -> chromadb.Collection:
144
+ """Load an existing collection at query time (fast, no re-embedding)."""
145
+ client = chromadb.PersistentClient(path=str(chroma_dir))
146
+ return client.get_collection(name=collection_name, embedding_function=_EMBED_FN)
rag/retriever.py CHANGED
@@ -1,2 +1,132 @@
1
  """ChromaDB query interface with citation-weighted re-ranking."""
2
- # Stage 3 implementation
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  """ChromaDB query interface with citation-weighted re-ranking."""
2
+ from __future__ import annotations
3
+
4
+ import math
5
+
6
+ import chromadb
7
+
8
+ from config import CHROMA_ENTITY_N_RESULTS, CHROMA_N_RESULTS
9
+ from logging_config import get_logger
10
+
11
+ _logger = get_logger("rag.retriever")
12
+
13
+
14
+ def search(
15
+ collection: chromadb.Collection,
16
+ query_text: str,
17
+ n_results: int = CHROMA_N_RESULTS,
18
+ ) -> list[dict]:
19
+ """
20
+ Semantic search with citation-weighted re-ranking.
21
+ Over-fetches 2× then re-ranks by: similarity * log(citation_count + 2).
22
+ Deduplicates to one chunk per paper (best-scoring chunk wins).
23
+ """
24
+ n_fetch = min(n_results * 2, collection.count())
25
+ if n_fetch == 0:
26
+ return []
27
+
28
+ raw = collection.query(
29
+ query_texts=[query_text],
30
+ n_results=n_fetch,
31
+ include=["documents", "metadatas", "distances"],
32
+ )
33
+ results = _parse_raw(raw)
34
+ results = _rerank(results)
35
+ return results[:n_results]
36
+
37
+
38
+ def search_by_entities(
39
+ collection: chromadb.Collection,
40
+ entity_names: list[str],
41
+ n_results: int = CHROMA_ENTITY_N_RESULTS,
42
+ ) -> list[dict]:
43
+ """
44
+ Run one query per entity, merge and deduplicate by PMID.
45
+ Caps at 8 entity queries to avoid excessive API calls.
46
+ """
47
+ if not entity_names or collection.count() == 0:
48
+ return []
49
+
50
+ seen: dict[str, dict] = {}
51
+ for entity in entity_names[:8]:
52
+ raw = collection.query(
53
+ query_texts=[entity],
54
+ n_results=min(10, collection.count()),
55
+ include=["documents", "metadatas", "distances"],
56
+ )
57
+ for r in _parse_raw(raw):
58
+ pmid = r["pmid"]
59
+ if pmid not in seen or r["score"] > seen[pmid]["score"]:
60
+ seen[pmid] = r
61
+
62
+ merged = _rerank(list(seen.values()))
63
+ return merged[:n_results]
64
+
65
+
66
+ def get_paper(collection: chromadb.Collection, pmid: str) -> dict | None:
67
+ """Retrieve a specific paper's abstract chunk by PMID."""
68
+ result = collection.get(
69
+ where={"$and": [{"pmid": {"$eq": pmid}}, {"chunk_index": {"$eq": 0}}]},
70
+ include=["documents", "metadatas"],
71
+ )
72
+ ids = result.get("ids", [])
73
+ if not ids:
74
+ return None
75
+ meta = result["metadatas"][0]
76
+ return {
77
+ "pmid": pmid,
78
+ "title": meta.get("title", ""),
79
+ "year": meta.get("year", 0),
80
+ "doi": meta.get("doi", ""),
81
+ "citation_count": meta.get("citation_count", 0),
82
+ "document": result["documents"][0],
83
+ }
84
+
85
+
86
+ def _parse_raw(raw: dict) -> list[dict]:
87
+ """Flatten a ChromaDB query response into a list of result dicts."""
88
+ ids = raw.get("ids", [[]])[0]
89
+ docs = raw.get("documents", [[]])[0]
90
+ metas = raw.get("metadatas", [[]])[0]
91
+ distances = raw.get("distances", [[]])[0]
92
+
93
+ results = []
94
+ for chunk_id, doc, meta, dist in zip(ids, docs, metas, distances):
95
+ similarity = max(0.0, 1.0 - dist)
96
+ citation_count = int(meta.get("citation_count", 0))
97
+ results.append({
98
+ "chunk_id": chunk_id,
99
+ "pmid": meta.get("pmid", ""),
100
+ "title": meta.get("title", ""),
101
+ "year": int(meta.get("year", 0)),
102
+ "doi": meta.get("doi", ""),
103
+ "section": meta.get("section", "abstract"),
104
+ "citation_count": citation_count,
105
+ "entity_names": [e for e in meta.get("entity_names", "").split(",") if e],
106
+ "has_full_text": bool(meta.get("has_full_text", 0)),
107
+ "document": doc,
108
+ "similarity": similarity,
109
+ "score": similarity,
110
+ })
111
+ return results
112
+
113
+
114
+ def _rerank(results: list[dict]) -> list[dict]:
115
+ """
116
+ Apply citation weighting and deduplicate to one chunk per PMID.
117
+ score = similarity * log(citation_count + 2)
118
+ log(2) ≈ 0.69 is the floor for uncited papers, so they're still ranked
119
+ but deprioritized relative to highly-cited work.
120
+ """
121
+ for r in results:
122
+ r["score"] = r["similarity"] * math.log(r["citation_count"] + 2)
123
+
124
+ results.sort(key=lambda x: x["score"], reverse=True)
125
+
126
+ # Keep best chunk per paper
127
+ seen: dict[str, dict] = {}
128
+ for r in results:
129
+ pmid = r["pmid"]
130
+ if pmid not in seen:
131
+ seen[pmid] = r
132
+ return list(seen.values())
scripts/build_index.py CHANGED
@@ -1,8 +1,64 @@
 
1
  """
2
- Build the ChromaDB vector index from papers and extracted entities.
 
3
 
4
  Usage:
5
  uv run python scripts/build_index.py
6
- uv run python scripts/build_index.py --reset # drop and rebuild collection
7
  """
8
- # Stage 3 implementation
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
  """
3
+ Build the ChromaDB vector index from papers.jsonl.
4
+ Run after scripts/ingest_papers.py.
5
 
6
  Usage:
7
  uv run python scripts/build_index.py
8
+ uv run python scripts/build_index.py --reset # drop and rebuild from scratch
9
  """
10
+ from __future__ import annotations
11
+
12
+ import sys
13
+ from pathlib import Path
14
+
15
+ sys.path.insert(0, str(Path(__file__).parent.parent))
16
+
17
+ from dotenv import load_dotenv
18
+
19
+ load_dotenv()
20
+
21
+ import argparse
22
+
23
+ from rich.console import Console
24
+ from rich.progress import track
25
+
26
+ from config import CHROMA_COLLECTION, CHROMA_DIR, PAPERS_PATH
27
+ from rag.indexer import build_collection
28
+
29
+ console = Console()
30
+
31
+
32
+ def main() -> None:
33
+ parser = argparse.ArgumentParser(description="Build ChromaDB vector index from ALS papers")
34
+ parser.add_argument("--reset", action="store_true", help="Drop existing collection and rebuild")
35
+ parser.add_argument("--papers", default=str(PAPERS_PATH), help="Path to papers.jsonl")
36
+ args = parser.parse_args()
37
+
38
+ papers_path = Path(args.papers)
39
+ if not papers_path.exists():
40
+ console.print(f"[red]Papers file not found: {papers_path}[/red]")
41
+ console.print("[yellow]Run first: uv run python scripts/ingest_papers.py[/yellow]")
42
+ sys.exit(1)
43
+
44
+ if args.reset:
45
+ console.print("[yellow]Resetting collection...[/yellow]")
46
+
47
+ console.print(f"[cyan]Building ChromaDB index from {papers_path}...[/cyan]")
48
+ console.print("[dim]First run downloads the embedding model (~80MB). Subsequent runs are fast.[/dim]\n")
49
+
50
+ collection = build_collection(
51
+ papers_path=papers_path,
52
+ chroma_dir=CHROMA_DIR,
53
+ collection_name=CHROMA_COLLECTION,
54
+ reset=args.reset,
55
+ )
56
+
57
+ console.print(f"\n[bold green]Done![/bold green] Collection '{CHROMA_COLLECTION}' at {CHROMA_DIR}")
58
+ console.print(f" Total chunks indexed: [bold]{collection.count()}[/bold]")
59
+ console.print("\nTest a query:")
60
+ console.print(" [dim]uv run python main.py \"What is the evidence for tofersen targeting SOD1?\"[/dim]")
61
+
62
+
63
+ if __name__ == "__main__":
64
+ main()
uv.lock CHANGED
@@ -451,6 +451,7 @@ dependencies = [
451
  { name = "gradio" },
452
  { name = "httpx" },
453
  { name = "networkx" },
 
454
  { name = "python-dotenv" },
455
  { name = "rich" },
456
  { name = "sentence-transformers" },
@@ -472,6 +473,7 @@ requires-dist = [
472
  { name = "gradio", specifier = ">=6.14.0,<7.0.0" },
473
  { name = "httpx", specifier = ">=0.27.0" },
474
  { name = "networkx", specifier = ">=3.3" },
 
475
  { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0" },
476
  { name = "pytest-cov", marker = "extra == 'dev'" },
477
  { name = "pytest-httpx", marker = "extra == 'dev'", specifier = ">=0.35" },
@@ -2097,6 +2099,25 @@ wheels = [
2097
  { 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" },
2098
  ]
2099
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2100
  [[package]]
2101
  name = "opentelemetry-api"
2102
  version = "1.43.0"
 
451
  { name = "gradio" },
452
  { name = "httpx" },
453
  { name = "networkx" },
454
+ { name = "openai" },
455
  { name = "python-dotenv" },
456
  { name = "rich" },
457
  { name = "sentence-transformers" },
 
473
  { name = "gradio", specifier = ">=6.14.0,<7.0.0" },
474
  { name = "httpx", specifier = ">=0.27.0" },
475
  { name = "networkx", specifier = ">=3.3" },
476
+ { name = "openai", specifier = ">=2.44.0" },
477
  { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0" },
478
  { name = "pytest-cov", marker = "extra == 'dev'" },
479
  { name = "pytest-httpx", marker = "extra == 'dev'", specifier = ">=0.35" },
 
2099
  { 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" },
2100
  ]
2101
 
2102
+ [[package]]
2103
+ name = "openai"
2104
+ version = "2.44.0"
2105
+ source = { registry = "https://pypi.org/simple" }
2106
+ dependencies = [
2107
+ { name = "anyio" },
2108
+ { name = "distro" },
2109
+ { name = "httpx" },
2110
+ { name = "jiter" },
2111
+ { name = "pydantic" },
2112
+ { name = "sniffio" },
2113
+ { name = "tqdm" },
2114
+ { name = "typing-extensions" },
2115
+ ]
2116
+ 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" }
2117
+ wheels = [
2118
+ { 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" },
2119
+ ]
2120
+
2121
  [[package]]
2122
  name = "opentelemetry-api"
2123
  version = "1.43.0"