Update mcp/orchestrator.py

mcp/orchestrator.py

@@ -1,161 +1,127 @@
-import asyncio
-import httpx
-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, 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
-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
-
-
-async 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.
-    """
-    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 insertion order.
-    """
-    flat: List[Any] = []
-    for elem in items:
-        if isinstance(elem, list):
-            for x in elem:
-                if x not in flat:
-                    flat.append(x)
-        else:
-            if elem is not None and elem not in flat:
-                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 with fault tolerance:
-      - Extract UMLS concepts and fetch definitions
-      - Literature (PubMed + arXiv)
-      - Drug safety, gene & variant info, disease-gene mapping
-      - Clinical trials, cBioPortal data
-      - AI-driven summary
-
-    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)),
-        '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=[])),
+"""
+MedGenesis – dual-LLM orchestrator
+----------------------------------
+• Accepts llm = "openai" | "gemini"   (falls back to OpenAI)
+• Returns one unified dict the UI can rely on.
+"""
+from __future__ import annotations
+import asyncio, itertools, logging
+from typing import Dict, Any, List, Tuple
+
+from mcp.arxiv            import fetch_arxiv
+from mcp.pubmed           import fetch_pubmed
+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.umls             import lookup_umls
+from mcp.openfda          import fetch_drug_safety
+from mcp.disgenet         import disease_to_genes
+from mcp.clinicaltrials   import search_trials
+from mcp.cbio             import fetch_cbio
+from mcp.openai_utils     import ai_summarize, ai_qa
+from mcp.gemini           import gemini_summarize, gemini_qa
+
+log = logging.getLogger(__name__)
+_DEF = "openai"                                # default engine
+
+
+# ─────────────────────────────────── helpers ───────────────────────────────────
+def _llm_router(engine: str = _DEF) -> Tuple:
+    if engine.lower() == "gemini":
+        return gemini_summarize, gemini_qa, "gemini"
+    return ai_summarize, ai_qa, "openai"
+
+async def _gather_safely(*aws, as_list: bool = True):
+    """await gather() that converts Exception → RuntimeError placeholder"""
+    out = await asyncio.gather(*aws, return_exceptions=True)
+    if as_list:
+        # filter exceptions – keep structure but drop failures
+        return [x for x in out if not isinstance(x, Exception)]
+    return out
+
+async def _gene_enrichment(keys: List[str]) -> Dict[str, Any]:
+    jobs = []
+    for k in keys:
+        jobs += [
+            search_gene(k),                    # basic gene info
+            get_mesh_definition(k),            # MeSH definitions
+            fetch_gene_info(k),                # MyGene
+            fetch_ensembl(k),                  # Ensembl x-refs
+            fetch_ot(k),                       # Open Targets associations
+        ]
+    res = await _gather_safely(*jobs, as_list=False)
+
+    # slice & compress five-way fan-out
+    combo = lambda idx: [r for i, r in enumerate(res) if i % 5 == idx and r]
+    return {
+        "ncbi"     : combo(0),
+        "mesh"     : combo(1),
+        "mygene"   : combo(2),
+        "ensembl"  : combo(3),
+        "ot_assoc" : combo(4),
     }
 
-    # 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)
-        llm_used = 'gemini'
-    else:
-        ai_summary = await ai_summarize(summaries)
-        llm_used = 'openai'
+# ───────────────────────────────── orchestrator ────────────────────────────────
+async def orchestrate_search(query: str, *, llm: str = _DEF) -> Dict[str, Any]:
+    """Main entry – returns dict for the Streamlit UI"""
+    # 1  Literature – run in parallel
+    arxiv_task  = asyncio.create_task(fetch_arxiv(query))
+    pubmed_task = asyncio.create_task(fetch_pubmed(query))
+    papers_raw  = await _gather_safely(arxiv_task, pubmed_task)
+    papers      = list(itertools.chain.from_iterable(papers_raw))[:30]   # keep ≤30
+
+    # 2  Keyword extraction (very light – only from abstracts)
+    kws = {w for p in papers for w in (p["summary"][:500].split()) if w.isalpha()}
+    kws = list(kws)[:10]                           # coarse, fast -> 10 seeds
+
+    # 3  Bio-enrichment fan-out
+    umls_f       = [_safe_task(lookup_umls, k) for k in kws]
+    fda_f        = [_safe_task(fetch_drug_safety, k) for k in kws]
+    gene_bundle  = asyncio.create_task(_gene_enrichment(kws))
+    trials_task  = asyncio.create_task(search_trials(query, max_studies=20))
+    cbio_task    = asyncio.create_task(fetch_cbio(kws[0] if kws else ""))
+
+    umls, fda, gene_dat, trials, variants = await asyncio.gather(
+        _gather_safely(*umls_f),
+        _gather_safely(*fda_f),
+        gene_bundle,
+        trials_task,
+        cbio_task,
+    )
+
+    # 4  LLM summary
+    summarise_fn, _, engine = _llm_router(llm)
+    summary = await summarise_fn(" ".join(p["summary"] for p in papers)[:12000])
 
     return {
-        'papers': papers,
-        'genes': genes,
-        'umls': umls_results,
-        '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'] or [],
-        'ai_summary': ai_summary,
-        'llm_used': llm_used,
+        "papers"          : papers,
+        "umls"            : umls,
+        "drug_safety"     : fda,
+        "ai_summary"      : summary,
+        "llm_used"        : engine,
+        "genes"           : gene_dat["ncbi"] + gene_dat["ensembl"] + gene_dat["mygene"],
+        "mesh_defs"       : gene_dat["mesh"],
+        "gene_disease"    : gene_dat["ot_assoc"],
+        "clinical_trials" : trials,
+        "variants"        : variants or [],
     }
 
-
-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 {'answer': text}.
-    """
-    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}
+# ─────────────────────────────── follow-up QA ─────────────────────────────────
+async def answer_ai_question(question: str, *, context: str, llm: str = _DEF) -> Dict[str, str]:
+    """Follow-up QA using chosen LLM."""
+    _, qa_fn, _ = _llm_router(llm)
+    return {"answer": await qa_fn(f"Q: {question}\nContext: {context}\nA:")}
+
+
+# ─────────────────────────── internal util  ───────────────────────────────────
+def _safe_task(fn, *args):
+    """Helper to wrap callable → Task returning RuntimeError on exception."""
+    async def _wrapper():
+        try:
+            return await fn(*args)
+        except Exception as exc:
+            log.warning("background task %s failed: %s", fn.__name__, exc)
+            return RuntimeError(str(exc))
+    return asyncio.create_task(_wrapper())