mcp/orchestrator.py

# mcp/orchestrator.py

import asyncio
from mcp.pubmed import fetch_pubmed
from mcp.arxiv import fetch_arxiv
from mcp.umls import extract_umls_concepts
from mcp.openfda import fetch_drug_safety
from mcp.ncbi import search_gene, get_mesh_definition
from mcp.mygene import fetch_gene_info
from mcp.ensembl import fetch_ensembl
from mcp.opentargets import fetch_ot
from mcp.clinicaltrials import search_trials
from mcp.cbio import fetch_cbio
from mcp.gemini import gemini_summarize, gemini_qa
from mcp.openai_utils import ai_summarize, ai_qa
from mcp.disgenet import disease_to_genes

async def orchestrate_search(query, llm="openai"):
    # --- Literature: PubMed + arXiv
    pubmed_task = asyncio.create_task(fetch_pubmed(query, max_results=7))
    arxiv_task = asyncio.create_task(fetch_arxiv(query, max_results=7))
    # --- UMLS, OpenFDA, Gene, Mesh
    umls_task = asyncio.create_task(extract_umls_concepts(query))
    fda_task = asyncio.create_task(fetch_drug_safety(query))
    gene_ncbi_task = asyncio.create_task(search_gene(query))
    mygene_task = asyncio.create_task(fetch_gene_info(query))
    ensembl_task = asyncio.create_task(fetch_ensembl(query))
    ot_task = asyncio.create_task(fetch_ot(query))
    mesh_task = asyncio.create_task(get_mesh_definition(query))
    # --- Trials, cBio, DisGeNET
    trials_task = asyncio.create_task(search_trials(query, max_studies=10))
    cbio_task = asyncio.create_task(fetch_cbio(query))
    disgenet_task = asyncio.create_task(disease_to_genes(query))

    # Run
    pubmed, arxiv, umls, fda, ncbi, mygene, ensembl, ot, mesh, trials, cbio, disgenet = await asyncio.gather(
        pubmed_task, arxiv_task, umls_task, fda_task, gene_ncbi_task,
        mygene_task, ensembl_task, ot_task, mesh_task, trials_task, cbio_task, disgenet_task
    )
    # Genes: flatten and deduplicate
    genes = []
    for g in (ncbi, mygene, ensembl, ot):
        if isinstance(g, list):
            genes.extend(g)
        elif isinstance(g, dict) and g:
            genes.append(g)
    genes = [g for i, g in enumerate(genes) if g and genes.index(g) == i]  # dedup

    # --- AI summary (LLM engine select)
    papers = (pubmed or []) + (arxiv or [])
    if llm == "gemini":
        ai_summary = await gemini_summarize(" ".join([p.get("summary", "") for p in papers]))
        llm_used = "gemini"
    else:
        ai_summary = await ai_summarize(" ".join([p.get("summary", "") for p in papers]))
        llm_used = "openai"

    return {
        "papers": papers,
        "genes": genes,
        "umls": umls or [],
        "gene_disease": disgenet if isinstance(disgenet, list) else [],
        "mesh_defs": [mesh] if isinstance(mesh, str) and mesh else [],
        "drug_safety": fda or [],
        "clinical_trials": trials or [],
        "variants": cbio if isinstance(cbio, list) else [],
        "ai_summary": ai_summary,
        "llm_used": llm_used
    }

async def answer_ai_question(question, context="", llm="openai"):
    # Gemini fallback if OpenAI quota is exceeded
    try:
        if llm == "gemini":
            answer = await gemini_qa(question, context)
        else:
            answer = await ai_qa(question, context)
    except Exception as e:
        answer = f"LLM unavailable or quota exceeded. ({e})"
    return {"answer": answer}