Update mcp/orchestrator.py

mcp/orchestrator.py

@@ -1,131 +1,117 @@
 #!/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 – 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.
-"""
-
 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
-
-# ------------------------------------------------------------------
+# ── 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
+
+# -------------------------------------------------------------------
 # LLM router
-# ------------------------------------------------------------------
+# -------------------------------------------------------------------
+_DEF_MODEL = "openai"
 
-def _get_llm(llm: str):
-    """Return (summarize_fn, qa_fn) based on requested engine."""
+def _get_llm(llm: str | None):
     if llm and llm.lower() == "gemini":
-        return gemini_summarize, gemini_qa
-    return ai_summarize, ai_qa  # default → OpenAI
+        return gemini_summarize, gemini_qa, "gemini"
+    return ai_summarize, ai_qa, "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]
+# -------------------------------------------------------------------
+# Helper: NCBI / MeSH / DisGeNET enrichment
+# -------------------------------------------------------------------
+async def _enrich_keywords(keys: List[str]) -> Dict[str, Any]:
+    tasks = []
+    for k in keys:
+        tasks += [search_gene(k), get_mesh_definition(k), disease_to_genes(k)]
 
-    results = await asyncio.gather(*jobs, return_exceptions=True)
+    results = await asyncio.gather(*tasks, return_exceptions=True)
 
-    genes, mesh_defs, disg_links = [], [], []
-    n = len(keys)
+    genes, mesh_defs, disg = [], [], []
     for idx, res in enumerate(results):
         if isinstance(res, Exception):
             continue
-        bucket = idx // n  # 0 = gene, 1 = mesh, 2 = disg
+        bucket = idx % 3
         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}
-
-# ------------------------------------------------------------------
-# 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))
+            disg.extend(res)
+    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."""
+    # 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 -----------------------------------------
+    # 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]
+    # 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_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))
+    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,
-        mygene_task,
+        gene_task,
         trials_task,
     )
 
-    # 4) LLM summary -------------------------------------------------
-    summarize_fn, _ = _get_llm(llm)
+    # 4) LLM summary -------------------------------------------------
+    summarize_fn, _, engine = _get_llm(llm)
     ai_summary = await summarize_fn(corpus)
 
-    # 5) Assemble payload -------------------------------------------
+    # 5) Assemble result --------------------------------------------
     return {
         "papers"         : papers,
         "umls"           : umls,
         "drug_safety"    : fda,
         "ai_summary"     : ai_summary,
-        "llm_used"       : llm.lower(),
+        "llm_used"       : engine,
 
-        # Gene & variant context
+        # genes & ontologies
         "genes"          : (ncbi_data["genes"] or []) + ([mygene] if mygene else []),
         "mesh_defs"      : ncbi_data["meshes"],
         "gene_disease"   : ncbi_data["disgenet"],
 
-        # Clinical trials
+        # 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)
+
+async def answer_ai_question(question: str, *, context: str, llm: str = _DEF_MODEL) -> Dict[str, str]:
+    summarize, qa_fn, engine = _get_llm(llm)
     return {"answer": await qa_fn(question, context)}