Update mcp/orchestrator.py

mcp/orchestrator.py

@@ -21,7 +21,7 @@ from mcp.embeddings     import embed_texts, cluster_embeddings
 
 def _get_llm(llm: str):
     """
-    Router for LLM engines: returns (summarize_fn, qa_fn).
+    Route summarization and QA to the chosen engine.
     """
     if llm.lower() == "gemini":
         return gemini_summarize, gemini_qa
@@ -30,27 +30,27 @@ def _get_llm(llm: str):
 
 async def orchestrate_search(query: str, llm: str = "openai") -> Dict[str, Any]:
     """
-    Main orchestrator: fetch literature, concepts, enrichments,
-    embeddings, clusters, and AI summary.
+    Fetch papers, extract concepts & relations, enrich data,
+    compute embeddings+clusters, and run LLM summary.
     """
-    # 1) Literature fetch
+    # Gather literature
     arxiv_task  = fetch_arxiv(query)
     pubmed_task = fetch_pubmed(query)
-    results     = await asyncio.gather(arxiv_task, pubmed_task, return_exceptions=True)
+    lit_results = await asyncio.gather(arxiv_task, pubmed_task, return_exceptions=True)
     papers: List[Dict] = []
-    for res in results:
+    for res in lit_results:
         if isinstance(res, list):
             papers.extend(res)
 
-    # 2) UMLS concept linking via spaCy
+    # Concept extraction
     blob = " ".join(p.get("summary", "") for p in papers)
     umls = await extract_umls_concepts(blob)
 
-    # 3) Fetch UMLS relations
-    rels_tasks = [fetch_relations(c["cui"]) for c in umls]
-    umls_relations = await asyncio.gather(*rels_tasks, return_exceptions=True)
+    # Fetch UMLS relations
+    rel_tasks = [fetch_relations(c["cui"]) for c in umls]
+    umls_relations = await asyncio.gather(*rel_tasks, return_exceptions=True)
 
-    # 4) Data enrichment
+    # Data enrichment tasks
     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=[])
@@ -60,6 +60,7 @@ async def orchestrate_search(query: str, llm: str = "openai") -> Dict[str, Any]:
     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=[])
 
+    # Run enrichment
     fda, gene, mesh, dis, trials, ot_assoc, variants = await asyncio.gather(
         asyncio.gather(*fda_tasks, return_exceptions=True),
         gene_task, mesh_task, dis_task,
@@ -67,16 +68,17 @@ async def orchestrate_search(query: str, llm: str = "openai") -> Dict[str, Any]:
         return_exceptions=False
     )
 
-    # 5) Embeddings & clustering
+    # 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)))
+        embeddings = await embed_texts(summaries)
+        clusters = await cluster_embeddings(
+            embeddings, n_clusters = max(2, min(10, len(embeddings)//2))
+        )
     else:
-        embs = []
-        clusters = []
+        embeddings, clusters = [], []
 
-    # 6) AI summary
+    # LLM summary
     summarize_fn, _ = _get_llm(llm)
     try:
         ai_summary = await summarize_fn(blob)
@@ -94,7 +96,7 @@ async def orchestrate_search(query: str, llm: str = "openai") -> Dict[str, Any]:
         "clinical_trials": trials,
         "ot_associations": ot_assoc,
         "variants": variants,
-        "embeddings": embs,
+        "embeddings": embeddings,
         "clusters": clusters,
         "ai_summary": ai_summary,
         "llm_used": llm.lower()
@@ -103,11 +105,11 @@ async def orchestrate_search(query: str, llm: str = "openai") -> Dict[str, Any]:
 
 async def answer_ai_question(question: str, context: str = "", llm: str = "openai") -> Dict[str, str]:
     """
-    Follow-up Q&A using chosen LLM engine.
+    Follow-up Q&A via chosen LLM.
     """
     _, qa_fn = _get_llm(llm)
     try:
-        answer = await qa_fn(question, context)
+        ans = await qa_fn(question, context)
     except Exception:
-        answer = "LLM follow-up failed."
-    return {"answer": answer}
+        ans = "LLM follow-up failed."
+    return {"answer": ans}