Update mcp/orchestrator.py

mcp/orchestrator.py

@@ -1,4 +1,5 @@
 import asyncio
+import httpx
 from typing import Any, Dict, List, Literal, Union
 
 from mcp.pubmed import fetch_pubmed
@@ -16,6 +17,23 @@ from mcp.gemini import gemini_summarize, gemini_qa
 from mcp.openai_utils import ai_summarize, ai_qa
 
 
+def _safe_call(
+    func: Any,
+    *args,
+    default: Any = None,
+    **kwargs,
+) -> Any:
+    """
+    Safely call an async function, returning a default on HTTP or other failures.
+    """
+    try:
+        return await func(*args, **kwargs)  # type: ignore
+    except httpx.HTTPStatusError:
+        return default
+    except Exception:
+        return default
+
+
 async def _gather_tasks(tasks: List[asyncio.Task]) -> List[Any]:
     """
     Await a list of asyncio.Tasks and return their results in order.
@@ -50,47 +68,46 @@ async def orchestrate_search(
     max_trials: int = 10,
 ) -> Dict[str, Any]:
     """
-    Perform a comprehensive biomedical search pipeline:
+    Perform a comprehensive biomedical search pipeline with fault tolerance:
       - 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
+    Individual fetch functions that fail with an HTTP error will return an empty default,
+    ensuring the pipeline always completes.
     """
-    # Launch parallel tasks
+    # Launch parallel tasks with safe wrapper for potential HTTP errors
     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)),
+        'drug_safety': asyncio.create_task(_safe_call(fetch_drug_safety, query, default=[])),
+        'ncbi_gene': asyncio.create_task(_safe_call(search_gene, query, default=[])),
+        'mygene': asyncio.create_task(_safe_call(fetch_gene_info, query, default=[])),
+        'ensembl': asyncio.create_task(_safe_call(fetch_ensembl, query, default=[])),
+        'opentargets': asyncio.create_task(_safe_call(fetch_ot, query, default=[])),
+        'mesh': asyncio.create_task(_safe_call(get_mesh_definition, query, default="")),
+        'trials': asyncio.create_task(_safe_call(search_trials, query, default=[], max_studies=max_trials)),
+        'cbio': asyncio.create_task(_safe_call(fetch_cbio, query, default=[])),
+        'disgenet': asyncio.create_task(_safe_call(disease_to_genes, query, default=[])),
     }
 
-    # Await all
+    # Await all tasks
     results = await _gather_tasks(list(tasks.values()))
     data = dict(zip(tasks.keys(), results))
 
-    # Process gene sources
+    # Consolidate gene sources
     gene_sources = [data['ncbi_gene'], data['mygene'], data['ensembl'], data['opentargets']]
     genes = _flatten_unique(gene_sources)
 
-    # Combine literature
+    # Merge literature results
     papers = (data['pubmed'] or []) + (data['arxiv'] or [])
 
-    # AI-driven summary selection
+    # AI-driven summary
     summaries = " ".join(p.get('summary', '') for p in papers)
     if llm == 'gemini':
         ai_summary = await gemini_summarize(summaries)
@@ -103,11 +120,11 @@ async def orchestrate_search(
         '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 [],
+        'gene_disease': data['disgenet'] or [],
+        'mesh_defs': [data['mesh']] if data['mesh'] else [],
         'drug_safety': data['drug_safety'] or [],
         'clinical_trials': data['trials'] or [],
-        'variants': data['cbio'] if isinstance(data['cbio'], list) else [],
+        'variants': data['cbio'] or [],
         'ai_summary': ai_summary,
         'llm_used': llm_used,
     }
@@ -120,8 +137,7 @@ async def answer_ai_question(
 ) -> Dict[str, str]:
     """
     Answer a free-text question using the specified LLM, with fallback.
-
-    Returns a dict {'answer': <text>}.
+    Returns {'answer': text}.
     """
     try:
         if llm == 'gemini':