Merge pull request #7 from KevinIsInCoding/feat/rrf-cross-encoder-retrieval

agents/research_agent.py

@@ -7,8 +7,17 @@ from collections.abc import Generator
 import anthropic
 import chromadb
 import networkx as nx
-
-from config import SYNTHESIS_MODEL
+from sentence_transformers import CrossEncoder
+
+from config import (
+    CROSS_ENCODER_MODEL,
+    CROSS_ENCODER_TOP_N,
+    RETRIEVAL_ENTITY_N,
+    RETRIEVAL_SEMANTIC_N,
+    RRF_K,
+    RRF_TOP_N,
+    SYNTHESIS_MODEL,
+)
 from graph import query as kg_query
 from llm import cached_system, cached_tools
 from logging_config import get_logger
@@ -18,6 +27,9 @@ from tools import RESEARCH_TOOLS
 
 _logger = get_logger("agents.research_agent")
 
+# Loaded once at startup — ~80MB model, ~80ms/pair on CPU
+_cross_encoder = CrossEncoder(CROSS_ENCODER_MODEL)
+
 
 def stream_research_agent(
     client: anthropic.Anthropic,
@@ -63,7 +75,7 @@ def stream_research_agent(
                             "name": event.content_block.name,
                         }
                         current_input_json = ""
-                        yield ("status", "Searching ALS research knowledge base...")
+                        yield ("status", "Searching knowledge base and re-ranking results for precision...")
 
                 elif event.type == "content_block_delta":
                     if event.delta.type == "text_delta":
@@ -132,38 +144,62 @@ def _handle_search(
     else:
         expanded_entities = query_entities
 
-    # Step 2: RAG — semantic search on query text
-    semantic_results = rag_retriever.search(collection, query_text, n_results=10)
+    # Step 2: Semantic search → up to 30 papers (pure similarity, no citation weight yet)
+    semantic_results = rag_retriever.search(collection, query_text, n_results=RETRIEVAL_SEMANTIC_N)
 
-    # Step 3: Entity-targeted RAG using expanded entity set
-    entity_results = rag_retriever.search_by_entities(collection, expanded_entities, n_results=15)
+    # Step 3: Entity-targeted search → up to 30 papers (one query per expanded entity)
+    entity_results = rag_retriever.search_by_entities(
+        collection, expanded_entities, n_results=RETRIEVAL_ENTITY_N
+    )
 
-    # 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
+    # Step 4: RRF merge → top 20 papers
+    merged = rag_retriever.rrf_merge(
+        [semantic_results, entity_results], k=RRF_K, top_n=RRF_TOP_N
+    )
 
-    top_papers = sorted(seen.values(), key=lambda x: x["score"], reverse=True)[:15]
+    # Step 5: Cross-encoder rerank → top 15 papers
+    reranked = rag_retriever.cross_encoder_rerank(
+        _cross_encoder, query_text, merged, top_n=CROSS_ENCODER_TOP_N
+    )
+
+    # Step 6: Citation boost — final score = ce_score × log(citation_count + 2)
+    top_papers = rag_retriever.apply_citation_boost(reranked)
 
     _logger.info(
-        "KG+RAG search",
+        "KG+RAG+CE 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),
+            "rrf_merged": len(merged),
+            "after_cross_encoder": len(top_papers),
             "kg_active": graph is not None,
         }},
     )
 
-    # Step 4: Trial matching — prefer KG-linked trials, fall back to text match
+    # Step 7: 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=10)
 
+    if not related_trials:
+        # Exact NCT ID match first — handles "NCT06351592" style queries
+        nct_ids_in_query = {
+            w.upper() for w in query_text.split() if w.upper().startswith("NCT")
+        }
+        trial_by_nct = {t.get("nct_id", "").upper(): t for t in trials}
+        for nct_id in nct_ids_in_query:
+            if nct_id in trial_by_nct:
+                t = trial_by_nct[nct_id]
+                related_trials.append({
+                    "nct_id": t.get("nct_id", ""),
+                    "title": t.get("title", ""),
+                    "phase": t.get("phase", ""),
+                    "status": t.get("status", ""),
+                    "url": t.get("url", ""),
+                })
+
     if not related_trials and query_entities:
         entities_lower = [e.lower() for e in expanded_entities]
         for trial in trials:

app.py

@@ -37,12 +37,19 @@ def _load_trials() -> list[dict]:
         return [json.loads(line) for line in f if line.strip()]
 
 
-_collection = load_collection(CHROMA_DIR, CHROMA_COLLECTION)
+def _load_collection():
+    try:
+        return load_collection(CHROMA_DIR, CHROMA_COLLECTION)
+    except Exception:
+        _logger.warning("ChromaDB collection not found — running in demo mode (no data)")
+        return None
+
+_collection = _load_collection()
 _graph = _load_graph()
 _trials = _load_trials()
 _client = anthropic.Anthropic()
 
-_n_chunks = _collection.count()
+_n_chunks = _collection.count() if _collection else 0
 _n_trials = len(_trials)
 _kg_nodes = _graph.number_of_nodes() if _graph else 0
 
@@ -64,6 +71,12 @@ def respond(message: str, history: list[dict]):
         yield history, gr.update(value="", interactive=True)
         return
 
+    if _collection is None:
+        history = history + [{"role": "user", "content": message}]
+        history = history + [{"role": "assistant", "content": "⚠️ The knowledge base has not been loaded yet. The pipeline data (ChromaDB index, knowledge graph, papers) needs to be uploaded to this Space. Please contact the Space administrator."}]
+        yield history, gr.update(value="", interactive=True)
+        return
+
     history = history + [{"role": "user", "content": message}]
     history = history + [{"role": "assistant", "content": ""}]
     yield history, gr.update(value="", interactive=False)

config.py

@@ -34,9 +34,20 @@ PUBMED_BATCH_SIZE = 200  # PMIDs per Entrez efetch call
 # Entity extraction
 EXTRACTION_BATCH_SIZE = 10  # papers per Claude call
 
-# RAG
-CHROMA_N_RESULTS = 10
-CHROMA_ENTITY_N_RESULTS = 15
+# RAG — retrieval counts per stage
+CHROMA_N_RESULTS = 10          # legacy default (kept for backward compat)
+CHROMA_ENTITY_N_RESULTS = 15   # legacy default (kept for backward compat)
+RETRIEVAL_SEMANTIC_N = 30      # semantic search candidate pool
+RETRIEVAL_ENTITY_N = 30        # entity search candidate pool
+RETRIEVAL_ENTITY_QUERY_CAP = 12  # max entity names to query individually
+
+# RRF merge
+RRF_K = 10       # lower k → stronger rank differentiation (k=60 is too flat for 30-item lists)
+RRF_TOP_N = 20   # candidates passed to cross-encoder
+
+# Cross-encoder reranking
+CROSS_ENCODER_MODEL = "cross-encoder/ms-marco-MiniLM-L-6-v2"
+CROSS_ENCODER_TOP_N = 15  # final papers sent to Claude for synthesis
 
 # Knowledge graph
 KG_EXPANSION_HOPS = 1  # hops for query entity expansion

rag/retriever.py

@@ -1,11 +1,20 @@
-"""ChromaDB query interface with citation-weighted re-ranking."""
+"""ChromaDB query interface with RRF merge and cross-encoder reranking."""
 from __future__ import annotations
 
 import math
 
 import chromadb
 
-from config import CHROMA_ENTITY_N_RESULTS, CHROMA_N_RESULTS
+from config import (
+    CHROMA_ENTITY_N_RESULTS,
+    CHROMA_N_RESULTS,
+    CROSS_ENCODER_TOP_N,
+    RETRIEVAL_ENTITY_N,
+    RETRIEVAL_ENTITY_QUERY_CAP,
+    RETRIEVAL_SEMANTIC_N,
+    RRF_K,
+    RRF_TOP_N,
+)
 from logging_config import get_logger
 
 _logger = get_logger("rag.retriever")
@@ -14,12 +23,12 @@ _logger = get_logger("rag.retriever")
 def search(
     collection: chromadb.Collection,
     query_text: str,
-    n_results: int = CHROMA_N_RESULTS,
+    n_results: int = RETRIEVAL_SEMANTIC_N,
 ) -> 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).
+    Semantic search — returns pure similarity-ranked results (no citation weighting).
+    Over-fetches 2× then deduplicates to one chunk per paper.
+    Citation boost is applied downstream after cross-encoder reranking.
     """
     n_fetch = min(n_results * 2, collection.count())
     if n_fetch == 0:
@@ -31,24 +40,24 @@ def search(
         include=["documents", "metadatas", "distances"],
     )
     results = _parse_raw(raw)
-    results = _rerank(results)
+    results = _dedup_by_pmid(results)
     return results[:n_results]
 
 
 def search_by_entities(
     collection: chromadb.Collection,
     entity_names: list[str],
-    n_results: int = CHROMA_ENTITY_N_RESULTS,
+    n_results: int = RETRIEVAL_ENTITY_N,
 ) -> list[dict]:
     """
-    Run one query per entity, merge and deduplicate by PMID.
-    Caps at 8 entity queries to avoid excessive API calls.
+    Run one ChromaDB query per entity name, merge and deduplicate by PMID.
+    Caps at RETRIEVAL_ENTITY_QUERY_CAP (12) entity queries to bound latency.
     """
     if not entity_names or collection.count() == 0:
         return []
 
     seen: dict[str, dict] = {}
-    for entity in entity_names[:8]:
+    for entity in entity_names[:RETRIEVAL_ENTITY_QUERY_CAP]:
         raw = collection.query(
             query_texts=[entity],
             n_results=min(10, collection.count()),
@@ -56,13 +65,79 @@ def search_by_entities(
         )
         for r in _parse_raw(raw):
             pmid = r["pmid"]
-            if pmid not in seen or r["score"] > seen[pmid]["score"]:
+            if pmid not in seen or r["similarity"] > seen[pmid]["similarity"]:
                 seen[pmid] = r
 
-    merged = _rerank(list(seen.values()))
+    merged = _dedup_by_pmid(list(seen.values()))
     return merged[:n_results]
 
 
+def rrf_merge(
+    ranked_lists: list[list[dict]],
+    k: int = RRF_K,
+    top_n: int = RRF_TOP_N,
+) -> list[dict]:
+    """
+    Reciprocal Rank Fusion — combines N ranked lists into one.
+    score(pmid) = Σ 1 / (k + rank_in_list_i + 1)
+    Preserves the best-scoring dict per PMID from all input lists.
+    """
+    rrf_scores: dict[str, float] = {}
+    best: dict[str, dict] = {}
+
+    for ranked in ranked_lists:
+        for rank, result in enumerate(ranked):
+            pmid = result["pmid"]
+            rrf_scores[pmid] = rrf_scores.get(pmid, 0.0) + 1.0 / (k + rank + 1)
+            if pmid not in best or result["similarity"] > best[pmid]["similarity"]:
+                best[pmid] = result
+
+    sorted_pmids = sorted(rrf_scores, key=lambda p: rrf_scores[p], reverse=True)
+    merged = []
+    for pmid in sorted_pmids[:top_n]:
+        r = best[pmid].copy()
+        r["rrf_score"] = round(rrf_scores[pmid], 6)
+        merged.append(r)
+    return merged
+
+
+def cross_encoder_rerank(
+    model,
+    query: str,
+    candidates: list[dict],
+    top_n: int = CROSS_ENCODER_TOP_N,
+) -> list[dict]:
+    """
+    Cross-encoder reranking — scores (query, document) pairs jointly.
+    Truncates document text to 1800 chars (~450 tokens) so query+doc fits
+    within the ms-marco model's 512-token limit.
+    """
+    if not candidates:
+        return []
+
+    pairs = [(query, r["document"][:1800]) for r in candidates]
+    ce_scores = model.predict(pairs, show_progress_bar=False)
+
+    for r, score in zip(candidates, ce_scores):
+        r["ce_score"] = float(score)
+
+    candidates.sort(key=lambda x: x["ce_score"], reverse=True)
+    return candidates[:top_n]
+
+
+def apply_citation_boost(results: list[dict]) -> list[dict]:
+    """
+    Final score = cross_encoder_score × log(citation_count + 2).
+    Applied after cross-encoder so citation quality amplifies — not corrupts — relevance.
+    log(2) ≈ 0.69 is the floor for uncited papers.
+    """
+    for r in results:
+        base = r.get("ce_score", r.get("similarity", 0.0))
+        r["score"] = base * math.log(r["citation_count"] + 2)
+    results.sort(key=lambda x: x["score"], reverse=True)
+    return results
+
+
 def get_paper(collection: chromadb.Collection, pmid: str) -> dict | None:
     """Retrieve a specific paper's abstract chunk by PMID."""
     result = collection.get(
@@ -111,19 +186,9 @@ def _parse_raw(raw: dict) -> list[dict]:
     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
+def _dedup_by_pmid(results: list[dict]) -> list[dict]:
+    """Keep best-similarity chunk per paper, sorted by similarity descending."""
+    results.sort(key=lambda x: x["similarity"], reverse=True)
     seen: dict[str, dict] = {}
     for r in results:
         pmid = r["pmid"]