Update mcp/orchestrator.py

mcp/orchestrator.py

@@ -1,9 +1,11 @@
 # mcp/orchestrator.py
 import asyncio
-from typing import Dict, Any
+from typing import Any, Dict, List
+
 from mcp.arxiv          import fetch_arxiv
 from mcp.pubmed         import fetch_pubmed
 from mcp.nlp            import extract_umls_concepts
+from mcp.umls           import lookup_umls
 from mcp.umls_rel       import fetch_relations
 from mcp.openfda        import fetch_drug_safety
 from mcp.ncbi           import search_gene, get_mesh_definition
@@ -14,37 +16,49 @@ from mcp.opentargets    import ot
 from mcp.cbio           import cbio
 from mcp.openai_utils   import ai_summarize, ai_qa
 from mcp.gemini         import gemini_summarize, gemini_qa
+from mcp.embeddings     import embed_texts, cluster_embeddings
+
 
 def _get_llm(llm: str):
-    return (gemini_summarize, gemini_qa) if llm.lower() == "gemini" else (ai_summarize, ai_qa)
+    """
+    Router for LLM engines: returns (summarize_fn, qa_fn).
+    """
+    if llm.lower() == "gemini":
+        return gemini_summarize, gemini_qa
+    return ai_summarize, ai_qa
+
 
 async def orchestrate_search(query: str, llm: str = "openai") -> Dict[str, Any]:
-    # 1) Parallel literature pulls
-    arxiv_t, pubmed_t = fetch_arxiv(query), fetch_pubmed(query)
-    papers = []
-    for res in await asyncio.gather(arxiv_t, pubmed_t, return_exceptions=True):
+    """
+    Main orchestrator: fetch literature, concepts, enrichments,
+    embeddings, clusters, and AI summary.
+    """
+    # 1) Literature fetch
+    arxiv_task  = fetch_arxiv(query)
+    pubmed_task = fetch_pubmed(query)
+    results     = await asyncio.gather(arxiv_task, pubmed_task, return_exceptions=True)
+    papers: List[Dict] = []
+    for res in results:
         if isinstance(res, list):
             papers.extend(res)
 
-    # 2) SpaCy→UMLS concept linking
-    blob = " ".join(p.get("summary","") for p in papers)
+    # 2) UMLS concept linking via spaCy
+    blob = " ".join(p.get("summary", "") for p in papers)
     umls = await extract_umls_concepts(blob)
 
-    # 3) Fetch UMLS relations in parallel
-    rels = await asyncio.gather(
-        *[fetch_relations(c["cui"]) for c in umls],
-        return_exceptions=True
-    )
+    # 3) Fetch UMLS relations
+    rels_tasks = [fetch_relations(c["cui"]) for c in umls]
+    umls_relations = await asyncio.gather(*rels_tasks, return_exceptions=True)
 
-    # 4) Enrich: OpenFDA, NCBI, DisGeNET, Trials, OpenTargets, cBioPortal
-    keys = [c["name"] for c in umls]
-    fda_tasks    = [fetch_drug_safety(k) for k in keys]
-    gene_task    = search_gene(keys[0])            if keys else asyncio.sleep(0, result=[])
-    mesh_task    = get_mesh_definition(keys[0])    if keys else asyncio.sleep(0, result="")
-    dis_task     = disease_to_genes(keys[0])       if keys else asyncio.sleep(0, result=[])
-    trials_task  = search_trials(query)
-    ot_task      = ot.fetch(keys[0])               if keys else asyncio.sleep(0, result=[])
-    cbio_task    = cbio.fetch_variants(keys[0])    if keys else asyncio.sleep(0, result=[])
+    # 4) Data enrichment
+    names = [c["name"] for c in umls]
+    fda_tasks   = [fetch_drug_safety(n) for n in names]
+    gene_task   = search_gene(names[0]) if names else asyncio.sleep(0, result=[])
+    mesh_task   = get_mesh_definition(names[0]) if names else asyncio.sleep(0, result="")
+    dis_task    = disease_to_genes(names[0]) if names else asyncio.sleep(0, result=[])
+    trials_task = search_trials(query)
+    ot_task     = ot.fetch(names[0]) if names else asyncio.sleep(0, result=[])
+    cbio_task   = cbio.fetch_variants(names[0]) if names else asyncio.sleep(0, result=[])
 
     fda, gene, mesh, dis, trials, ot_assoc, variants = await asyncio.gather(
         asyncio.gather(*fda_tasks, return_exceptions=True),
@@ -53,17 +67,26 @@ async def orchestrate_search(query: str, llm: str = "openai") -> Dict[str, Any]:
         return_exceptions=False
     )
 
-    # 5) AI summary
-    summarize, _ = _get_llm(llm)
+    # 5) Embeddings & clustering
+    summaries = [p.get("summary", "") for p in papers]
+    if summaries:
+        embs = await embed_texts(summaries)
+        clusters = await cluster_embeddings(embs, n_clusters=max(2, min(10, len(embs)//2)))
+    else:
+        embs = []
+        clusters = []
+
+    # 6) AI summary
+    summarize_fn, _ = _get_llm(llm)
     try:
-        ai_summary = await summarize(blob)
+        ai_summary = await summarize_fn(blob)
     except Exception:
         ai_summary = "LLM summary failed."
 
     return {
         "papers": papers,
         "umls": umls,
-        "umls_relations": rels,
+        "umls_relations": umls_relations,
         "drug_safety": fda,
         "genes": [gene],
         "mesh_defs": [mesh],
@@ -71,11 +94,17 @@ async def orchestrate_search(query: str, llm: str = "openai") -> Dict[str, Any]:
         "clinical_trials": trials,
         "ot_associations": ot_assoc,
         "variants": variants,
+        "embeddings": embs,
+        "clusters": clusters,
         "ai_summary": ai_summary,
         "llm_used": llm.lower()
     }
 
-async def answer_ai_question(question: str, context: str = "", llm: str = "openai"):
+
+async def answer_ai_question(question: str, context: str = "", llm: str = "openai") -> Dict[str, str]:
+    """
+    Follow-up Q&A using chosen LLM engine.
+    """
     _, qa_fn = _get_llm(llm)
     try:
         answer = await qa_fn(question, context)