Update mcp/orchestrator.py

mcp/orchestrator.py

@@ -7,22 +7,22 @@ from mcp.umls import lookup_umls
 from mcp.openfda import fetch_drug_safety
 from mcp.openai_utils import ai_summarize, ai_qa
 
-async def orchestrate_search(query: str) -> dict:
-    # Fetch results in parallel (use asyncio.gather for speed)
-    arxiv_results = await fetch_arxiv(query)
-    pubmed_results = await fetch_pubmed(query)
+import asyncio
+
+async def orchestrate_search(query: str):
+    # Fetch from arXiv and PubMed 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)
     all_papers = arxiv_results + pubmed_results
-    # Semantic ranking (use OpenAI embeddings or similar)
-    # ...
-    # NLP: extract keywords/drugs
-    keywords = extract_keywords(" ".join([p['summary'] for p in all_papers]))
-    # UMLS enrichment
-    umls_results = [await lookup_umls(k) for k in keywords]
-    # Drug safety
-    drug_data = [await fetch_drug_safety(k) for k in keywords]
-    # Summarization
-    summary = await ai_summarize(" ".join([p['summary'] for p in all_papers]))
-    # Suggest reading (top links)
+    paper_text = " ".join([p['summary'] for p in all_papers])
+    keywords = extract_keywords(paper_text)[:8]  # Limit for speed
+    # UMLS and OpenFDA in parallel
+    umls_tasks = [lookup_umls(k) for k in keywords]
+    drug_tasks = [fetch_drug_safety(k) for k in keywords]
+    umls_results = await asyncio.gather(*umls_tasks)
+    drug_data = await asyncio.gather(*drug_tasks)
+    summary = await ai_summarize(paper_text)
     links = [p['link'] for p in all_papers[:3]]
     return {
         "papers": all_papers,
@@ -32,6 +32,6 @@ async def orchestrate_search(query: str) -> dict:
         "suggested_reading": links,
     }
 
-async def answer_ai_question(question: str, context: str = "") -> dict:
+async def answer_ai_question(question: str, context: str = ""):
     answer = await ai_qa(question, context)
     return {"answer": answer}