Update mcp/orchestrator.py

mcp/orchestrator.py

@@ -4,7 +4,7 @@ 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
+from mcp.umls import extract_umls_concepts, lookup_umls
 from mcp.openfda import fetch_drug_safety
 from mcp.ncbi import search_gene, get_mesh_definition
 from mcp.mygene import fetch_gene_info
@@ -44,7 +44,7 @@ async def _gather_tasks(tasks: List[asyncio.Task]) -> List[Any]:
 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.
+    then deduplicate preserving insertion order.
     """
     flat: List[Any] = []
     for elem in items:
@@ -66,21 +66,36 @@ async def orchestrate_search(
 ) -> Dict[str, Any]:
     """
     Perform a comprehensive biomedical search pipeline with fault tolerance:
+      - Extract UMLS concepts and fetch definitions
       - Literature (PubMed + arXiv)
-      - Entity extraction (UMLS)
       - Drug safety, gene & variant info, disease-gene mapping
       - Clinical trials, cBioPortal data
       - AI-driven summary
 
-    Individual fetch functions that fail with an HTTP error will return an empty default,
-    ensuring the pipeline always completes.
+    Returns a dict with structured results ready for UI/graph building.
     """
+    # 1) Extract concepts and perform UMLS lookups
+    raw_concepts = await asyncio.to_thread(extract_umls_concepts, query)
+    umls_tasks = [
+        asyncio.create_task(
+            _safe_call(
+                lookup_umls,
+                term,
+                default={
+                    'term': term,
+                    'cui': None,
+                    'name': None,
+                    'definition': None,
+                },
+            )
+        )
+        for term in raw_concepts
+    ]
+
+    # 2) Launch parallel data-fetch tasks (excluding UMLS)
     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(_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=[])),
@@ -92,14 +107,19 @@ async def orchestrate_search(
         'disgenet': asyncio.create_task(_safe_call(disease_to_genes, query, default=[])),
     }
 
+    # 3) Await all tasks
     results = await _gather_tasks(list(tasks.values()))
     data = dict(zip(tasks.keys(), results))
+    umls_results = await asyncio.gather(*umls_tasks)
 
+    # 4) Consolidate gene sources
     gene_sources = [data['ncbi_gene'], data['mygene'], data['ensembl'], data['opentargets']]
     genes = _flatten_unique(gene_sources)
 
+    # 5) Merge literature
     papers = (data['pubmed'] or []) + (data['arxiv'] or [])
 
+    # 6) AI-driven summary
     summaries = " ".join(p.get('summary', '') for p in papers)
     if llm == 'gemini':
         ai_summary = await gemini_summarize(summaries)
@@ -111,7 +131,7 @@ async def orchestrate_search(
     return {
         'papers': papers,
         'genes': genes,
-        'umls': data['umls'] or [],
+        'umls': umls_results,
         'gene_disease': data['disgenet'] or [],
         'mesh_defs': [data['mesh']] if data['mesh'] else [],
         'drug_safety': data['drug_safety'] or [],