app.py

import os
os.environ["HF_HUB_TIMEOUT"] = "60"  # Increase timeout for Hugging Face Hub downloads

import gradio as gr
from fastapi import FastAPI
from fastapi.responses import HTMLResponse
import uvicorn

from backend import process_medical_query, get_graph_html

app = FastAPI()

#########################
# 1) FastAPI Route for Graph
#########################
@app.get("/graph")
def serve_pyvis_graph(query: str):
    """
    Returns the Pyvis HTML as a normal HTML response.
    This bypasses inline script sanitization in Gradio.
    """
    html_content = get_graph_html(query)
    return HTMLResponse(content=html_content, status_code=200)

#########################
# 2) Gradio Interface
#########################
def run_query(user_query: str):
    """
    Called when user submits the query in Gradio.
    Returns a textual answer + an iframe referencing /graph?query=...
    """
    answer = process_medical_query(user_query)
    # Build an iframe that points to the /graph route
    iframe_html = f"""
    <iframe src="/graph?query={user_query}"
            width="100%" height="600" frameborder="0"></iframe>
    """
    return answer, iframe_html

with gr.Blocks() as demo:
    gr.Markdown("## AI-Powered Medical Knowledge Graph Assistant")
    gr.Markdown("**Using BioGPT-Large-PubMedQA + PubMed + Chroma** for advanced retrieval-augmented generation.")

    with gr.Row():
        user_query = gr.Textbox(
            label="Enter biomedical/medical query",
            placeholder="e.g., 'Fever and cough treatment'",
            lines=1
        )
        submit_btn = gr.Button("Submit")

    answer_box = gr.Textbox(label="AI Response")
    graph_box = gr.HTML(label="Knowledge Graph")

    submit_btn.click(fn=run_query, inputs=user_query, outputs=[answer_box, graph_box])

# Mount the Gradio app at path "/"
app = gr.mount_gradio_app(app, demo, path="/")

#########################
# 3) Entry Point (Local)
#########################
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)