Update mcp/orchestrator.py

mcp/orchestrator.py

@@ -1,117 +1,122 @@
 #!/usr/bin/env python3
-"""MedGenesis – asynchronous orchestrator (v3)
-
-• Pulls literature (PubMed + arXiv)
-• Extracts keywords via spaCy → fans‑out to:
-    – MyGene.info        (gene annotations)
-    – ClinicalTrials.gov v2 (trials) with v1 fallback
-    – UMLS / openFDA / DisGeNET / MeSH
-• Summarises with OpenAI or Gemini (user‑selectable)
-• Returns a single dict ready for Streamlit UI & pydantic schemas.
+"""MedGenesis – orchestrator (v4, resilient).
+
+* Pulls PubMed + arXiv (async)  
+* spaCy keyword extraction → UMLS / openFDA / DisGeNET / MeSH fan‑out  
+* Adds MyGene.info, ClinicalTrials.gov (v2 → v1 fallback)  
+* Filters out failed enrichment calls (exceptions) so UI never crashes  
+* Summarises with OpenAI *or* Gemini; router returns `llm_used`
 """
 from __future__ import annotations
 
 import asyncio
 from typing import Any, Dict, List
 
-# ── Internal modules ────────────────────────────────────────────────
-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.mygene           import fetch_gene_info
-from mcp.ctgov            import search_trials          # v2 helper (v1 fallback inside)
-from mcp.openai_utils     import ai_summarize, ai_qa
-from mcp.gemini           import gemini_summarize, gemini_qa
-
-# -------------------------------------------------------------------
+# ── async fetchers ──────────────────────────────────────────────────
+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.mygene  import fetch_gene_info
+from mcp.ctgov   import search_trials
+
+# ── LLM helpers ────────────────────────────────────────────────────
+from mcp.openai_utils import ai_summarize, ai_qa
+from mcp.gemini       import gemini_summarize, gemini_qa
+
+# ------------------------------------------------------------------
 # LLM router
-# -------------------------------------------------------------------
-_DEF_MODEL = "openai"
+# ------------------------------------------------------------------
+_DEF = "openai"
 
-def _get_llm(llm: str | None):
-    if llm and llm.lower() == "gemini":
+def _llm_router(name: str | None):
+    if name and name.lower() == "gemini":
         return gemini_summarize, gemini_qa, "gemini"
     return ai_summarize, ai_qa, "openai"
 
-# -------------------------------------------------------------------
-# Helper: NCBI / MeSH / DisGeNET enrichment
-# -------------------------------------------------------------------
+# ------------------------------------------------------------------
+# Keyword enrichment bundle
+# ------------------------------------------------------------------
 async def _enrich_keywords(keys: List[str]) -> Dict[str, Any]:
-    tasks = []
+    tasks: List[asyncio.Future] = []
     for k in keys:
         tasks += [search_gene(k), get_mesh_definition(k), disease_to_genes(k)]
 
-    results = await asyncio.gather(*tasks, return_exceptions=True)
+    res = await asyncio.gather(*tasks, return_exceptions=True)
 
     genes, mesh_defs, disg = [], [], []
-    for idx, res in enumerate(results):
-        if isinstance(res, Exception):
+    for idx, r in enumerate(res):
+        if isinstance(r, Exception):
             continue
         bucket = idx % 3
         if bucket == 0:
-            genes.extend(res)
+            genes.extend(r)
         elif bucket == 1:
-            mesh_defs.append(res)
+            mesh_defs.append(r)
         else:
-            disg.extend(res)
+            disg.extend(r)
     return {"genes": genes, "meshes": mesh_defs, "disgenet": disg}
 
-# -------------------------------------------------------------------
-# Public API
-# -------------------------------------------------------------------
-async def orchestrate_search(query: str, *, llm: str = _DEF_MODEL) -> Dict[str, Any]:
-    """Run full async pipeline and return merged result dict."""
+# ------------------------------------------------------------------
+# Main orchestrator
+# ------------------------------------------------------------------
+async def orchestrate_search(query: str, *, llm: str = _DEF) -> Dict[str, Any]:
+    """Run entire async pipeline; never raises."""
     # 1) Literature --------------------------------------------------
-    arxiv_task  = asyncio.create_task(fetch_arxiv(query, max_results=10))
-    pubmed_task = asyncio.create_task(fetch_pubmed(query, max_results=10))
-    papers = sum(await asyncio.gather(arxiv_task, pubmed_task), [])
-
-    # 2) Keyword extraction -----------------------------------------
-    corpus = " ".join(p["summary"] for p in papers)
+    arxiv_f  = asyncio.create_task(fetch_arxiv(query, max_results=10))
+    pubmed_f = asyncio.create_task(fetch_pubmed(query, max_results=10))
+    papers   = []
+    for fut in await asyncio.gather(arxiv_f, pubmed_f, return_exceptions=True):
+        if not isinstance(fut, Exception):
+            papers.extend(fut)
+
+    # 2) Keywords ----------------------------------------------------
+    corpus   = " ".join(p.get("summary", "") for p in papers)
     keywords = extract_keywords(corpus)[:8]
 
-    # 3) Enrichment fan‑out -----------------------------------------
-    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_keywords(keywords))
-    gene_task   = asyncio.create_task(fetch_gene_info(query))
-    trials_task = asyncio.create_task(search_trials(query, max_studies=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,
-        gene_task,
-        trials_task,
+    # 3) Fan‑out enrichment -----------------------------------------
+    umls_f   = [lookup_umls(k)       for k in keywords]
+    fda_f    = [fetch_drug_safety(k) for k in keywords]
+    ncbi_f   = asyncio.create_task(_enrich_keywords(keywords))
+    mygene_f = asyncio.create_task(fetch_gene_info(query))
+    trials_f = asyncio.create_task(search_trials(query, max_studies=20))
+
+    umls, fda, ncbi, mygene, trials = await asyncio.gather(
+        asyncio.gather(*umls_f, return_exceptions=True),
+        asyncio.gather(*fda_f,  return_exceptions=True),
+        ncbi_f,
+        mygene_f,
+        trials_f,
     )
 
+    # ── filter exception objects -----------------------------------
+    umls = [u for u in umls if isinstance(u, dict)]
+    fda  = [d for d in fda  if isinstance(d, (dict, list))]
+
     # 4) LLM summary -------------------------------------------------
-    summarize_fn, _, engine = _get_llm(llm)
-    ai_summary = await summarize_fn(corpus)
+    summarize, _, engine = _llm_router(llm)
+    ai_summary = await summarize(corpus) if corpus else ""
 
-    # 5) Assemble result --------------------------------------------
+    # 5) Assemble payload -------------------------------------------
     return {
         "papers"         : papers,
         "umls"           : umls,
         "drug_safety"    : fda,
         "ai_summary"     : ai_summary,
         "llm_used"       : engine,
-
-        # genes & ontologies
-        "genes"          : (ncbi_data["genes"] or []) + ([mygene] if mygene else []),
-        "mesh_defs"      : ncbi_data["meshes"],
-        "gene_disease"   : ncbi_data["disgenet"],
-
+        # gene context
+        "genes"          : (ncbi["genes"] or []) + ([mygene] if mygene else []),
+        "mesh_defs"      : ncbi["meshes"],
+        "gene_disease"   : ncbi["disgenet"],
         # trials
         "clinical_trials": trials,
     }
 
-
-async def answer_ai_question(question: str, *, context: str, llm: str = _DEF_MODEL) -> Dict[str, str]:
-    summarize, qa_fn, engine = _get_llm(llm)
+# ------------------------------------------------------------------
+async def answer_ai_question(question: str, *, context: str, llm: str = _DEF) -> Dict[str, str]:
+    """Follow‑up QA using chosen LLM."""
+    _, qa_fn, _ = _llm_router(llm)
     return {"answer": await qa_fn(question, context)}