Update mcp/orchestrator.py

mcp/orchestrator.py

@@ -1,90 +1,92 @@
-# mcp/orchestrator.py
 """
-Orchestrates retrieval, enrichment, and AI synthesis for a user query.
+MedGenesis – dual-LLM orchestrator
+----------------------------------
+• Accepts `llm` arg ("openai" | "gemini")
+• Defaults to "openai" if arg omitted
 """
 
-import asyncio
+import asyncio, httpx
 from typing import Dict, Any, List
 
-from mcp.arxiv        import fetch_arxiv
-from mcp.pubmed       import fetch_pubmed
-from mcp.nlp          import extract_keywords
-from mcp.umls         import lookup_umls
-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.arxiv          import fetch_arxiv
+from mcp.pubmed         import fetch_pubmed
+from mcp.nlp            import extract_keywords
+from mcp.umls           import lookup_umls
+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.openai_utils import ai_summarize, ai_qa
+from mcp.openai_utils   import ai_summarize, ai_qa
+from mcp.gemini         import gemini_summarize, gemini_qa   # make sure gemini.py exists
 
-# ---------------------------------------------------------------------
-async def _gene_and_mesh_enrichment(keywords: List[str]) -> Dict[str, Any]:
-    """Run NCBI and DisGeNET on keywords in parallel."""
-    tasks = []
-    for kw in keywords:
-        tasks.append(search_gene(kw))
-        tasks.append(get_mesh_definition(kw))
-        tasks.append(disease_to_genes(kw))
-    results = await asyncio.gather(*tasks, return_exceptions=True)
+# ---------------- LLM router ----------------
+def _get_llm(llm: str):
+    if llm.lower() == "gemini":
+        return gemini_summarize, gemini_qa
+    return ai_summarize, ai_qa                         # default → OpenAI
 
-    genes, meshes, disgen = [], [], []
-    for i, res in enumerate(results):
-        if isinstance(res, Exception):
+
+async def _enrich_genes_mesh_disg(keys: List[str]) -> Dict[str, Any]:
+    jobs = []
+    for k in keys:
+        jobs += [search_gene(k), get_mesh_definition(k), disease_to_genes(k)]
+    res = await asyncio.gather(*jobs, return_exceptions=True)
+
+    genes, meshes, disg = [], [], []
+    for i, r in enumerate(res):
+        if isinstance(r, Exception):  # skip failures quietly
             continue
-        # Cycle: 0 gene, 1 mesh, 2 disgenet, repeat …
-        mod = i % 3
-        if mod == 0:
-            genes.extend(res)
-        elif mod == 1:
-            meshes.append(res)
-        else:
-            disgen.extend(res)
-    return {"genes": genes, "meshes": meshes, "disgenet": disgen}
+        if   i % 3 == 0: genes.extend(r)
+        elif i % 3 == 1: meshes.append(r)
+        else:            disg.extend(r)
+    return {"genes": genes, "meshes": meshes, "disgenet": disg}
 
-# ---------------------------------------------------------------------
-async def orchestrate_search(query: str) -> Dict[str, Any]:
-    """Main entry—returns a rich result dict for app UI."""
-    # -------- literature retrieval in parallel --------
-    arxiv_task  = asyncio.create_task(fetch_arxiv(query))
-    pubmed_task = asyncio.create_task(fetch_pubmed(query))
-    arxiv_results, pubmed_results = await asyncio.gather(arxiv_task, pubmed_task)
-    papers = arxiv_results + pubmed_results
 
-    # -------- keyword extraction --------
-    paper_text = " ".join(p["summary"] for p in papers)
-    keywords   = extract_keywords(paper_text)[:8]
+# ------------------------------------------------------------------
+async def orchestrate_search(query: str, llm: str = "openai") -> Dict[str, Any]:
+    """
+    Main orchestrator – returns dict for UI.
+    """
+    # 1) Literature ---------------------------------------------------
+    arxiv_f  = asyncio.create_task(fetch_arxiv(query))
+    pubmed_f = asyncio.create_task(fetch_pubmed(query))
+    papers   = sum(await asyncio.gather(arxiv_f, pubmed_f), [])
 
-    # -------- enrichment tasks in parallel --------
-    umls_tasks   = [lookup_umls(k)       for k in keywords]
-    fda_tasks    = [fetch_drug_safety(k) for k in keywords]
-    enrich_task  = asyncio.create_task(_gene_and_mesh_enrichment(keywords))
-    trials_task  = asyncio.create_task(search_trials(query, max_studies=10))
+    # 2) Keywords -----------------------------------------------------
+    blob = " ".join(p["summary"] for p in papers)
+    keys = extract_keywords(blob)[:8]
 
-    umls, fda, enrich, trials = await asyncio.gather(
-        asyncio.gather(*umls_tasks),
-        asyncio.gather(*fda_tasks),
-        enrich_task,
-        trials_task,
+    # 3) Enrichment ---------------------------------------------------
+    umls_f   = [lookup_umls(k)       for k in keys]
+    fda_f    = [fetch_drug_safety(k) for k in keys]
+    genes_f  = asyncio.create_task(_enrich_genes_mesh_disg(keys))
+    trials_f = asyncio.create_task(search_trials(query, max_studies=10))
+
+    umls, fda, genes, trials = await asyncio.gather(
+        asyncio.gather(*umls_f,  return_exceptions=True),
+        asyncio.gather(*fda_f,   return_exceptions=True),
+        genes_f,
+        trials_f,
     )
 
-    # -------- AI summary --------
-    summary = await ai_summarize(paper_text)
-    links   = [p["link"] for p in papers[:3]]
+    # 4) AI summary ---------------------------------------------------
+    summarize, _ = _get_llm(llm)
+    summary = await summarize(blob)
 
     return {
         "papers"          : papers,
         "umls"            : umls,
         "drug_safety"     : fda,
         "ai_summary"      : summary,
-        "suggested_reading": links,
-        # new fields
-        "genes"           : enrich["genes"],
-        "mesh_definitions": enrich["meshes"],
-        "gene_disease"    : enrich["disgenet"],
+        "llm_used"        : llm.lower(),
+        "genes"           : genes["genes"],
+        "mesh_defs"       : genes["meshes"],
+        "gene_disease"    : genes["disgenet"],
         "clinical_trials" : trials,
     }
 
-# ---------------------------------------------------------------------
-async def answer_ai_question(question: str, context: str = "") -> Dict[str, str]:
-    """Free-form Q&A using OpenAI."""
-    answer = await ai_qa(question, context)
-    return {"answer": answer}
+
+async def answer_ai_question(question: str, context: str, llm: str = "openai") -> Dict[str, str]:
+    """One-shot follow-up Q-A via chosen engine."""
+    _, qa = _get_llm(llm)
+    return {"answer": await qa(question, context)}