Update mcp/orchestrator.py

mcp/orchestrator.py

@@ -1,92 +1,131 @@
+#!/usr/bin/env python3
 """
-MedGenesis – dual-LLM orchestrator
-----------------------------------
-• Accepts `llm` arg ("openai" | "gemini")
-• Defaults to "openai" if arg omitted
+MedGenesis – dual‑LLM asynchronous orchestrator
+==============================================
+• Accepts `llm` argument ("openai" | "gemini"), defaults to "openai".
+• Harvests literature (PubMed + arXiv) → extracts keywords.
+• Fans‑out to open APIs for genes, trials, safety, ontology:
+    – **MyGene.info** for live gene annotations
+    – **ClinicalTrials.gov v2** for recruiting & completed studies
+    – UMLS / openFDA / DisGeNET / MeSH (existing helpers)
+    – Optional Open Targets & DrugCentral via `multi_enrich` if needed.
+• Returns a single JSON‑serialisable dict consumed by the Streamlit UI.
 """
 
-import asyncio, httpx
-from typing import Dict, Any, List
-
-from mcp.arxiv          import fetch_arxiv
-from mcp.pubmed         import fetch_pubmed
-from mcp.nlp            import extract_keywords
-from mcp.umls           import lookup_umls
-from mcp.openfda        import fetch_drug_safety
-from mcp.ncbi           import search_gene, get_mesh_definition
-from mcp.disgenet       import disease_to_genes
-from mcp.clinicaltrials import search_trials
-from mcp.openai_utils   import ai_summarize, ai_qa
-from mcp.gemini         import gemini_summarize, gemini_qa   # make sure gemini.py exists
-
-# ---------------- LLM router ----------------
+from __future__ import annotations
+
+import asyncio
+from typing import Any, Dict, List
+
+# ── Literature fetchers ─────────────────────────────────────────────
+from mcp.arxiv  import fetch_arxiv
+from mcp.pubmed import fetch_pubmed
+
+# ── NLP & legacy enrichers ─────────────────────────────────────────
+from mcp.nlp      import extract_keywords
+from mcp.umls     import lookup_umls
+from mcp.openfda  import fetch_drug_safety
+from mcp.ncbi     import search_gene, get_mesh_definition
+from mcp.disgenet import disease_to_genes
+
+# ── Modern high‑throughput APIs ────────────────────────────────────
+from mcp.mygene   import fetch_gene_info            # MyGene.info
+from mcp.ctgov    import search_trials_v2           # ClinicalTrials.gov v2
+# from mcp.targets  import fetch_ot_associations    # (optional future use)
+
+# ── LLM utilities ─────────────────────────────────────────────────
+from mcp.openai_utils import ai_summarize, ai_qa
+from mcp.gemini       import gemini_summarize, gemini_qa
+
+# ------------------------------------------------------------------
+# LLM router
+# ------------------------------------------------------------------
+
 def _get_llm(llm: str):
-    if llm.lower() == "gemini":
+    """Return (summarize_fn, qa_fn) based on requested engine."""
+    if llm and llm.lower() == "gemini":
         return gemini_summarize, gemini_qa
-    return ai_summarize, ai_qa                         # default → OpenAI
+    return ai_summarize, ai_qa  # default → OpenAI
 
+# ------------------------------------------------------------------
+# Helper: batch NCBI / MeSH / DisGeNET enrichment for keyword list
+# ------------------------------------------------------------------
+async def _enrich_ncbi_mesh_disg(keys: List[str]) -> Dict[str, Any]:
+    jobs = [search_gene(k) for k in keys] + \
+           [get_mesh_definition(k) for k in keys] + \
+           [disease_to_genes(k) for k in keys]
 
-async def _enrich_genes_mesh_disg(keys: List[str]) -> Dict[str, Any]:
-    jobs = []
-    for k in keys:
-        jobs += [search_gene(k), get_mesh_definition(k), disease_to_genes(k)]
-    res = await asyncio.gather(*jobs, return_exceptions=True)
+    results = await asyncio.gather(*jobs, return_exceptions=True)
 
-    genes, meshes, disg = [], [], []
-    for i, r in enumerate(res):
-        if isinstance(r, Exception):  # skip failures quietly
+    genes, mesh_defs, disg_links = [], [], []
+    n = len(keys)
+    for idx, res in enumerate(results):
+        if isinstance(res, Exception):
             continue
-        if   i % 3 == 0: genes.extend(r)
-        elif i % 3 == 1: meshes.append(r)
-        else:            disg.extend(r)
-    return {"genes": genes, "meshes": meshes, "disgenet": disg}
+        bucket = idx // n  # 0 = gene, 1 = mesh, 2 = disg
+        if bucket == 0:
+            genes.extend(res)
+        elif bucket == 1:
+            mesh_defs.append(res)
+        else:
+            disg_links.extend(res)
 
+    return {"genes": genes, "meshes": mesh_defs, "disgenet": disg_links}
 
 # ------------------------------------------------------------------
-async def orchestrate_search(query: str, llm: str = "openai") -> Dict[str, Any]:
-    """
-    Main orchestrator – returns dict for UI.
-    """
-    # 1) Literature ---------------------------------------------------
-    arxiv_f  = asyncio.create_task(fetch_arxiv(query))
-    pubmed_f = asyncio.create_task(fetch_pubmed(query))
-    papers   = sum(await asyncio.gather(arxiv_f, pubmed_f), [])
-
-    # 2) Keywords -----------------------------------------------------
-    blob = " ".join(p["summary"] for p in papers)
-    keys = extract_keywords(blob)[:8]
-
-    # 3) Enrichment ---------------------------------------------------
-    umls_f   = [lookup_umls(k)       for k in keys]
-    fda_f    = [fetch_drug_safety(k) for k in keys]
-    genes_f  = asyncio.create_task(_enrich_genes_mesh_disg(keys))
-    trials_f = asyncio.create_task(search_trials(query, max_studies=10))
-
-    umls, fda, genes, trials = await asyncio.gather(
-        asyncio.gather(*umls_f,  return_exceptions=True),
-        asyncio.gather(*fda_f,   return_exceptions=True),
-        genes_f,
-        trials_f,
+# Main orchestrator
+# ------------------------------------------------------------------
+async def orchestrate_search(query: str, *, llm: str = "openai") -> Dict[str, Any]:
+    """Master async pipeline – returns dict consumed by UI."""
+
+    # 1) Literature --------------------------------------------------
+    arxiv_task  = asyncio.create_task(fetch_arxiv(query))
+    pubmed_task = asyncio.create_task(fetch_pubmed(query))
+    papers = sum(await asyncio.gather(arxiv_task, pubmed_task), [])
+
+    # 2) Keyword extraction -----------------------------------------
+    corpus = " ".join(p["summary"] for p in papers)
+    keywords = extract_keywords(corpus)[:8]
+
+    # 3) Fan‑out enrichment -----------------------------------------
+    umls_tasks  = [lookup_umls(k)       for k in keywords]
+    fda_tasks   = [fetch_drug_safety(k) for k in keywords]
+
+    ncbi_task   = asyncio.create_task(_enrich_ncbi_mesh_disg(keywords))
+    mygene_task = asyncio.create_task(fetch_gene_info(query))           # top gene hit
+    trials_task = asyncio.create_task(search_trials_v2(query, max_n=20))
+
+    umls, fda, ncbi_data, mygene, trials = await asyncio.gather(
+        asyncio.gather(*umls_tasks, return_exceptions=True),
+        asyncio.gather(*fda_tasks,  return_exceptions=True),
+        ncbi_task,
+        mygene_task,
+        trials_task,
     )
 
-    # 4) AI summary ---------------------------------------------------
-    summarize, _ = _get_llm(llm)
-    summary = await summarize(blob)
+    # 4) LLM summary -------------------------------------------------
+    summarize_fn, _ = _get_llm(llm)
+    ai_summary = await summarize_fn(corpus)
 
+    # 5) Assemble payload -------------------------------------------
     return {
-        "papers"          : papers,
-        "umls"            : umls,
-        "drug_safety"     : fda,
-        "ai_summary"      : summary,
-        "llm_used"        : llm.lower(),
-        "genes"           : genes["genes"],
-        "mesh_defs"       : genes["meshes"],
-        "gene_disease"    : genes["disgenet"],
-        "clinical_trials" : trials,
-    }
+        "papers"         : papers,
+        "umls"           : umls,
+        "drug_safety"    : fda,
+        "ai_summary"     : ai_summary,
+        "llm_used"       : llm.lower(),
 
+        # Gene & variant context
+        "genes"          : (ncbi_data["genes"] or []) + ([mygene] if mygene else []),
+        "mesh_defs"      : ncbi_data["meshes"],
+        "gene_disease"   : ncbi_data["disgenet"],
 
-async def answer_ai_question(question: str, context: str, llm: str = "openai") -> Dict[str, str]:
-    """One-shot follow-up Q-A via chosen engine."""
-    _, qa = _get_llm(llm)
-    return {"answer": await qa(question, context)}
+        # Clinical trials
+        "clinical_trials": trials,
+    }
+
+# ------------------------------------------------------------------
+async def answer_ai_question(question: str, *, context: str, llm: str = "openai") -> Dict[str, str]:
+    """One‑shot follow‑up QA using selected engine."""
+    _, qa_fn = _get_llm(llm)
+    return {"answer": await qa_fn(question, context)}