mcp/orchestrator.py

"""
MedGenesis – parallel multi-API orchestrator
--------------------------------------------
Now pulls PubMed, arXiv, UMLS, OpenFDA, DisGeNET, ClinicalTrials.gov,
PLUS: MyGene.info, Ensembl, OpenTargets, Expression Atlas, cBioPortal,
DrugCentral & PubChem.

Call           : orchestrate_search(query, llm="openai" | "gemini")
Returns        : dict ready for Streamlit UI
Follow-up QA   : answer_ai_question(...)
"""

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   # legacy
from mcp.ncbi_turbo       import pubmed_ids
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
from mcp.mygene           import fetch_gene_info
from mcp.ensembl          import fetch_ensembl
from mcp.opentargets      import fetch_ot
from mcp.atlas            import fetch_expression
from mcp.drugcentral_ext  import fetch_drugcentral
from mcp.pubchem_ext      import fetch_compound
from mcp.cbio             import fetch_cbio

_DEF = "openai"

# ---------- LLM router ----------
def _llm_router(llm: str):
    if llm.lower() == "gemini":
        return gemini_summarize, gemini_qa, "gemini"
    return ai_summarize, ai_qa, "openai"

# ---------- gene resolver ----------
async def _resolve_gene(sym: str) -> dict:
    for fn in (fetch_gene_info, fetch_ensembl, fetch_ot):
        try:
            data = await fn(sym)
            if data:
                return data
        except Exception:
            pass
    return {}

# ---------- orchestrator ----------
async def orchestrate_search(query: str,
                             llm: str = _DEF) -> Dict[str, Any]:
    # 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  Keyword extraction
    blob = " ".join(p["summary"] for p in papers)
    keywords = extract_keywords(blob)[:10]

    # 3  Enrichment fan-out
    umls_tasks  = [lookup_umls(k)       for k in keywords]
    fda_tasks   = [fetch_drug_safety(k) for k in keywords]
    gene_tasks  = [ _resolve_gene(k)    for k in keywords]
    expr_tasks  = [ fetch_expression(k) for k in keywords]
    dcz_tasks   = [ fetch_drugcentral(k) for k in keywords]
    chem_tasks  = [ fetch_compound(k)    for k in keywords]
    cbio_tasks  = [ fetch_cbio(k)        for k in keywords[:3]]  # limit API

    genes, exprs, dcz, chems, cbio = await asyncio.gather(
        asyncio.gather(*gene_tasks, return_exceptions=True),
        asyncio.gather(*expr_tasks, return_exceptions=True),
        asyncio.gather(*dcz_tasks,   return_exceptions=True),
        asyncio.gather(*chem_tasks,  return_exceptions=True),
        asyncio.gather(*cbio_tasks,  return_exceptions=True),
    )

    umls, fda = await asyncio.gather(
        asyncio.gather(*umls_tasks, return_exceptions=True),
        asyncio.gather(*fda_tasks,  return_exceptions=True),
    )

    trials = await search_trials(query, max_studies=20)

    # 4  AI summary
    summarise, _, used_llm = _llm_router(llm)
    summary = await summarise(blob)

    return {
        "papers"          : papers,
        "ai_summary"      : summary,
        "llm_used"        : used_llm,
        "umls"            : umls,
        "drug_safety"     : fda,
        "genes_rich"      : [g for g in genes if g],
        "expr_atlas"      : [e for e in exprs if e],
        "drug_meta"       : [d for d in dcz if d],
        "chem_info"       : [c for c in chems if c],
        "clinical_trials" : trials,
        "cbio_variants"   : [v for v in cbio if v],
    }

# ---------- follow-up QA ----------
async def answer_ai_question(question: str, *,
                             context: str,
                             llm: str = _DEF) -> Dict[str, str]:
    _, qa_fn, _ = _llm_router(llm)
    return {"answer": await qa_fn(question, context)}