Update app.py

app.py

@@ -1,119 +1,69 @@
-import os
-from fastapi import FastAPI, Query, HTTPException
-from typing import List, Dict, Optional
-import httpx
-import openai
-
-# ENVIRONMENT VARIABLES for API keys (set these securely in your deployment)
-ARXIV_API = "http://export.arxiv.org/api/query?search_query="
-PUBMED_API = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi"
-PUBMED_FETCH = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi"
-OPENFDA_ENDPOINT = "https://api.fda.gov/drug/event.json"
-UMLS_AUTH = "https://utslogin.nlm.nih.gov/cas/v1/api-key"
-UMLS_SEARCH = "https://uts-ws.nlm.nih.gov/rest/search/current"
-OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
-UMLS_API_KEY = os.environ.get("UMLS_KEY")
-PUB_KEY = os.environ.get("PUB_KEY")
-
-app = FastAPI(
-    title="Ultimate AI-Powered Research MCP Server",
-    description="Integrates ArXiv, PubMed, OpenFDA, UMLS, and OpenAI for world-class research.",
-    version="1.0.0",
-)
+# app.py
 
-# -- Utility Functions --
-
-async def fetch_arxiv(query: str) -> List[Dict]:
-    # Use httpx or requests to fetch arXiv results
-    # Parse and return relevant metadata (simplified for brevity)
-    url = f"{ARXIV_API}{query}&max_results=5"
-    async with httpx.AsyncClient() as client:
-        r = await client.get(url)
-        # TODO: Parse arXiv Atom XML and extract paper metadata
-    return []
+import os
+import streamlit as st
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from mcp.orchestrator import orchestrate_search, answer_ai_question
+from mcp.schemas import UnifiedSearchInput, UnifiedSearchResult
 
-async def fetch_pubmed(query: str) -> List[Dict]:
-    # Use PubMed e-utilities for search and fetch
-    params = {"db": "pubmed", "term": query, "retmode": "json", "api_key": PUB_KEY, "retmax": 5}
-    async with httpx.AsyncClient() as client:
-        r = await client.get(PUBMED_API, params=params)
-        ids = r.json()["esearchresult"]["idlist"]
-        fetch_params = {"db": "pubmed", "id": ",".join(ids), "retmode": "xml", "api_key": PUB_KEY}
-        rf = await client.get(PUBMED_FETCH, params=fetch_params)
-        # TODO: Parse PubMed XML to get metadata
-    return []
+# Initialize FastAPI app for API users
+api = FastAPI(title="MCP Research Server", version="2.0")
+api.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"])
 
-async def fetch_umls(term: str) -> Dict:
-    # Authenticate with UMLS and fetch CUIs for term
-    async with httpx.AsyncClient() as client:
-        auth = await client.post(UMLS_AUTH, data={"apikey": UMLS_API_KEY})
-        ticket = auth.text  # TODO: parse ticket
-        params = {"string": term, "ticket": ticket}
-        search = await client.get(UMLS_SEARCH, params=params)
-        # TODO: Extract CUIs and definitions
-    return {}
+@api.post("/unified_search", response_model=UnifiedSearchResult)
+async def unified_search_endpoint(data: UnifiedSearchInput):
+    return await orchestrate_search(data.query)
 
-async def fetch_openfda(drug: str) -> Dict:
-    params = {"search": f"patient.drug.medicinalproduct:{drug}", "limit": 3}
-    async with httpx.AsyncClient() as client:
-        r = await client.get(OPENFDA_ENDPOINT, params=params)
-        # TODO: Parse adverse event data
-    return {}
+@api.post("/ask_ai")
+async def ask_ai_endpoint(question: str, context: str = ""):
+    return await answer_ai_question(question, context)
 
-async def ai_summarize(text: str) -> str:
-    openai.api_key = OPENAI_API_KEY
-    response = openai.ChatCompletion.create(
-        model="gpt-4o",
-        messages=[
-            {"role": "system", "content": "You are an expert biomedical research assistant."},
-            {"role": "user", "content": f"Summarize and synthesize: {text}"}
-        ],
-        max_tokens=256
+# Streamlit UI for Hugging Face Space
+def render_ui():
+    st.set_page_config(page_title="Ultimate Research Assistant", page_icon=":microscope:", layout="wide")
+    st.image("assets/logo.png", width=100)
+    st.title("🔬 Next-Gen AI-Powered Biomedical Research Assistant")
+    st.markdown(
+        """
+        *Combine the power of ArXiv, PubMed, UMLS, OpenFDA, and OpenAI.  
+        Get instant, unified, semantically-ranked answers—plus drug safety, concept enrichment, and expert Q&A!*
+        """
     )
-    return response['choices'][0]['message']['content']
 
-# -- Endpoints --
+    query = st.text_input("Enter your research question or topic:", value="What are the latest treatments for Alzheimer's disease?")
+    if st.button("Run Unified Search 🚀"):
+        with st.spinner("Retrieving and synthesizing knowledge..."):
+            results = orchestrate_search(query)
+            st.success("Here are the results!")
+            for i, paper in enumerate(results['papers'], 1):
+                st.markdown(f"**{i}. [{paper['title']}]({paper['link']})**  \n*{paper['authors']}*")
+                st.write(paper['summary'])
+            st.subheader("UMLS Concept Enrichment")
+            for c in results['umls']:
+                st.write(f"**{c['name']}** (CUI: {c['cui']}): {c['definition']}")
+            st.subheader("Drug & Safety Insights")
+            for d in results['drug_safety']:
+                st.write(d)
+            st.subheader("AI-Generated Synthesis")
+            st.info(results['ai_summary'])
+            st.markdown("#### 📚 Suggested Reading")
+            for link in results['suggested_reading']:
+                st.write(f"- {link}")
 
-@app.post("/unified_search")
-async def unified_search(query: str):
-    arxiv = await fetch_arxiv(query)
-    pubmed = await fetch_pubmed(query)
-    all_papers = arxiv + pubmed
-    # TODO: Semantic ranking (e.g., use OpenAI embeddings for similarity)
-    # Extract keywords for UMLS/Drug enrichment
-    keywords = []  # TODO: implement keyword extraction (e.g., spaCy)
-    umls_results = [await fetch_umls(k) for k in keywords]
-    drug_data = [await fetch_openfda(k) for k in keywords]
-    summary = await ai_summarize(str(all_papers))
-    return {
-        "papers": all_papers,
-        "umls": umls_results,
-        "drug_safety": drug_data,
-        "ai_summary": summary,
-        "suggested_reading": [p.get('link') for p in all_papers[:3]]
-    }
-
-@app.get("/get_umls_info")
-async def get_umls_info(term: str):
-    result = await fetch_umls(term)
-    return result
-
-@app.get("/fetch_drug_safety")
-async def fetch_drug_safety(drug: str):
-    result = await fetch_openfda(drug)
-    return result
-
-@app.post("/ask_ai")
-async def ask_ai(question: str, context: Optional[str] = None):
-    openai.api_key = OPENAI_API_KEY
-    response = openai.ChatCompletion.create(
-        model="gpt-4o",
-        messages=[
-            {"role": "system", "content": "You are an advanced multi-domain biomedical research agent. Use information from ArXiv, PubMed, OpenFDA, and UMLS when answering."},
-            {"role": "user", "content": f"Question: {question}\nContext: {context or ''}"}
-        ],
-        max_tokens=512
-    )
-    return {"answer": response['choices'][0]['message']['content']}
+    st.markdown("---")
+    st.subheader("🤖 Ask a follow-up (AI Q&A):")
+    follow_up = st.text_input("Type your question here:")
+    if st.button("Ask AI"):
+        with st.spinner("AI is thinking..."):
+            answer = answer_ai_question(follow_up, context=query)
+            st.success("AI says:")
+            st.write(answer['answer'])
 
-# -- Extensibility: Plug-and-play new tools via simple Python modules and endpoint registration. --
+if __name__ == "__main__":
+    import sys
+    if "runserver" in sys.argv:
+        import uvicorn
+        uvicorn.run(api, host="0.0.0.0", port=7860)
+    else:
+        render_ui()