Update mcp/orchestrator.py

mcp/orchestrator.py

@@ -1,6 +1,6 @@
-# mcp/orchestrator.py
-
 import asyncio
+from typing import Any, Dict, List, Literal, Union
+
 from mcp.pubmed import fetch_pubmed
 from mcp.arxiv import fetch_arxiv
 from mcp.umls import extract_umls_concepts
@@ -11,70 +11,123 @@ 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.disgenet import disease_to_genes
 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"
+
+async def _gather_tasks(tasks: List[asyncio.Task]) -> List[Any]:
+    """
+    Await a list of asyncio.Tasks and return their results in order.
+    """
+    return await asyncio.gather(*tasks)
+
+
+def _flatten_unique(items: List[Union[List[Any], Any]]) -> List[Any]:
+    """
+    Flatten a list of items where elements may be lists or single values,
+    then deduplicate preserving order.
+    """
+    flat: List[Any] = []
+    seen = set()
+    for elem in items:
+        if isinstance(elem, list):
+            for x in elem:
+                if x not in seen:
+                    seen.add(x)
+                    flat.append(x)
+        elif elem is not None:
+            if elem not in seen:
+                seen.add(elem)
+                flat.append(elem)
+    return flat
+
+
+async def orchestrate_search(
+    query: str,
+    llm: Literal['openai', 'gemini'] = 'openai',
+    max_papers: int = 7,
+    max_trials: int = 10,
+) -> Dict[str, Any]:
+    """
+    Perform a comprehensive biomedical search pipeline:
+      - Literature (PubMed + arXiv)
+      - Entity extraction (UMLS)
+      - Drug safety, gene & variant info, disease-gene mapping
+      - Clinical trials, cBioPortal data
+      - AI-driven summary
+
+    Returns a dict with keys:
+      papers, genes, umls, gene_disease, mesh_defs,
+      drug_safety, clinical_trials, variants, ai_summary, llm_used
+    """
+    # Launch parallel tasks
+    tasks = {
+        'pubmed': asyncio.create_task(fetch_pubmed(query, max_results=max_papers)),
+        'arxiv': asyncio.create_task(fetch_arxiv(query, max_results=max_papers)),
+        'umls': asyncio.create_task(
+            asyncio.to_thread(extract_umls_concepts, query)
+        ),
+        'drug_safety': asyncio.create_task(fetch_drug_safety(query)),
+        'ncbi_gene': asyncio.create_task(search_gene(query)),
+        'mygene': asyncio.create_task(fetch_gene_info(query)),
+        'ensembl': asyncio.create_task(fetch_ensembl(query)),
+        'opentargets': asyncio.create_task(fetch_ot(query)),
+        'mesh': asyncio.create_task(get_mesh_definition(query)),
+        'trials': asyncio.create_task(search_trials(query, max_studies=max_trials)),
+        'cbio': asyncio.create_task(fetch_cbio(query)),
+        'disgenet': asyncio.create_task(disease_to_genes(query)),
+    }
+
+    # Await all
+    results = await _gather_tasks(list(tasks.values()))
+    data = dict(zip(tasks.keys(), results))
+
+    # Process gene sources
+    gene_sources = [data['ncbi_gene'], data['mygene'], data['ensembl'], data['opentargets']]
+    genes = _flatten_unique(gene_sources)
+
+    # Combine literature
+    papers = (data['pubmed'] or []) + (data['arxiv'] or [])
+
+    # AI-driven summary selection
+    summaries = " ".join(p.get('summary', '') for p in papers)
+    if llm == 'gemini':
+        ai_summary = await gemini_summarize(summaries)
+        llm_used = 'gemini'
     else:
-        ai_summary = await ai_summarize(" ".join([p.get("summary", "") for p in papers]))
-        llm_used = "openai"
+        ai_summary = await ai_summarize(summaries)
+        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
+        'papers': papers,
+        'genes': genes,
+        'umls': data['umls'] or [],
+        'gene_disease': data['disgenet'] if isinstance(data['disgenet'], list) else [],
+        'mesh_defs': [data['mesh']] if isinstance(data['mesh'], str) and data['mesh'] else [],
+        'drug_safety': data['drug_safety'] or [],
+        'clinical_trials': data['trials'] or [],
+        'variants': data['cbio'] if isinstance(data['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
+
+async def answer_ai_question(
+    question: str,
+    context: str = "",
+    llm: Literal['openai', 'gemini'] = 'openai',
+) -> Dict[str, str]:
+    """
+    Answer a free-text question using the specified LLM, with fallback.
+
+    Returns a dict {'answer': <text>}.
+    """
     try:
-        if llm == "gemini":
+        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}
+        answer = f"LLM unavailable or quota exceeded: {e}"
+    return {'answer': answer}