Update mcp/orchestrator.py

mcp/orchestrator.py

@@ -1,47 +1,79 @@
 # mcp/orchestrator.py
 import asyncio
-from mcp.mygene    import mygene
-from mcp.opentargets import ot
-from mcp.cbio      import cbio
-# … import pubmed, umls, clinicaltrials, etc …
-from typing import Dict, Any
-
-async def orchestrate_search(query: str, *, llm: str="openai") -> Dict[str,Any]:
-    # 1) fetch papers + abstracts
-    papers_task = asyncio.create_task(fetch_papers(query))
-    # 2) pull UMLS concepts
-    from mcp.nlp import extract_keywords
-    kws = extract_keywords(query)[:5]
-    umls_tasks = [lookup_umls(k) for k in kws]
-    # 3) fetch gene info + associations
-    gene_task    = asyncio.create_task(mygene.fetch(query))
-    ot_task      = asyncio.create_task(ot.fetch(query))
-    # 4) fetch variants
-    cbio_task    = asyncio.create_task(cbio.fetch_variants(query))
-    # 5) clinical trials
-    trials_task  = asyncio.create_task(search_trials(query))
-
-    # wait all
-    papers = await papers_task
-    umls   = await asyncio.gather(*umls_tasks, return_exceptions=True)
-    gene, assoc, vars_, trials = await asyncio.gather(
-        gene_task, ot_task, cbio_task, trials_task, return_exceptions=True
+from typing import Any, Dict
+from mcp.arxiv          import fetch_arxiv
+from mcp.pubmed         import fetch_pubmed
+from mcp.nlp            import extract_umls_concepts
+from mcp.umls           import lookup_umls
+from mcp.umls_rel       import fetch_relations
+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.mygene         import mygene
+from mcp.opentargets    import ot
+from mcp.cbio           import cbio
+from mcp.openai_utils   import ai_summarize, ai_qa
+from mcp.gemini         import gemini_summarize, gemini_qa
+
+def _get_llm(llm: str):
+    return (gemini_summarize, gemini_qa) if llm.lower()=="gemini" else (ai_summarize, ai_qa)
+
+async def orchestrate_search(query: str, llm: str="openai") -> Dict[str,Any]:
+    # 1) literature
+    arxiv_t, pubmed_t = fetch_arxiv(query), fetch_pubmed(query)
+    papers = []
+    for res in await asyncio.gather(arxiv_t, pubmed_t, return_exceptions=True):
+        if isinstance(res, list):
+            papers.extend(res)
+
+    # 2) UMLS concept linking
+    blob = " ".join(p.get("summary","") for p in papers)
+    umls = extract_umls_concepts(blob)
+    rels = await asyncio.gather(*[fetch_relations(c["cui"]) for c in umls], return_exceptions=True)
+
+    # 3) enrichment
+    keys = [c["name"] for c in umls]
+    fda_t    = [fetch_drug_safety(k) for k in keys]
+    genes_t  = search_gene(keys[0]) if keys else asyncio.sleep(0, result=[])
+    mesh_t   = get_mesh_definition(keys[0]) if keys else asyncio.sleep(0, result="")
+    dis_t    = disease_to_genes(keys[0])    if keys else asyncio.sleep(0, result=[])
+    trials_t = search_trials(query)
+    ot_t     = ot.fetch(keys[0])            if keys else asyncio.sleep(0, result=[])
+    var_t    = cbio.fetch_variants(keys[0]) if keys else asyncio.sleep(0, result=[])
+
+    fda, genes, mesh, dis, trials, ot_assoc, variants = await asyncio.gather(
+        asyncio.gather(*fda_t, return_exceptions=True),
+        genes_t, mesh_t, dis_t, trials_t, ot_t, var_t,
+        return_exceptions=False
     )
 
-    # 6) call your chosen LLM
-    from mcp.ai import ai_summarize, gemini_summarize
-    if llm=="openai":
-        summary = await ai_summarize("\n\n".join(p["summary"] for p in papers))
-    else:
-        summary = await gemini_summarize("\n\n".join(p["summary"] for p in papers))
+    # 4) AI summary
+    summarize, _ = _get_llm(llm)
+    try:
+        summary = await summarize(blob)
+    except:
+        summary = "LLM unavailable."
 
     return {
         "papers": papers,
-        "umls": [u for u in umls if not isinstance(u, Exception)],
-        "gene": gene if not isinstance(gene, Exception) else {},
-        "associations": assoc if not isinstance(assoc, Exception) else [],
-        "variants": vars_ if not isinstance(vars_, Exception) else [],
-        "trials": trials if not isinstance(trials, Exception) else [],
+        "umls": umls,
+        "umls_relations": rels,
+        "drug_safety": fda,
+        "genes": [genes],
+        "mesh_defs": [mesh],
+        "gene_disease": dis,
+        "clinical_trials": trials,
+        "ot_associations": ot_assoc,
+        "variants": variants,
         "ai_summary": summary,
-        "llm_used": llm
+        "llm_used": llm.lower()
     }
+
+async def answer_ai_question(question: str, context: str="", llm: str="openai"):
+    _, qa = _get_llm(llm)
+    try:
+        ans = await qa(question, context)
+    except:
+        ans = "LLM unavailable."
+    return {"answer": ans}