gradio_server.py

import os
import uuid
import json
import asyncio
from dotenv import load_dotenv

import gradio as gr
from mcp.server.fastmcp import FastMCP, Context

# ---------------------------------------------------------------------------
# Environment & MCP initialisation
# ---------------------------------------------------------------------------
load_dotenv()

mcp = FastMCP(
    name="healthcare-mcp",
    version="1.0.0",
    description="Healthcare MCP server for medical information access (Gradio edition)"
)

# ---------------------------------------------------------------------------
# Internal tools & services
# ---------------------------------------------------------------------------
from src.tools.fda_tool import FDATool
from src.tools.pubmed_tool import PubMedTool
from src.tools.healthfinder_tool import HealthFinderTool
from src.tools.clinical_trials_tool import ClinicalTrialsTool
from src.tools.medical_terminology_tool import MedicalTerminologyTool
from src.services.usage_service import UsageService

fda_tool = FDATool()
pubmed_tool = PubMedTool()
healthfinder_tool = HealthFinderTool()
clinical_trials_tool = ClinicalTrialsTool()
medical_terminology_tool = MedicalTerminologyTool()

usage_service = UsageService(db_path="healthcare_usage.db")

# Unique session – ensures all usage is grouped for the current interface user
session_id = str(uuid.uuid4())

# ---------------------------------------------------------------------------
# FastMCP tool definitions – these can be reached programmatically by MCP
# clients (and are also reused by the Gradio UI wrappers below).
# ---------------------------------------------------------------------------

@mcp.tool()
async def fda_drug_lookup(ctx: Context, drug_name: str, search_type: str = "general"):
    """Look up drug information from the FDA database."""
    usage_service.record_usage(session_id, "fda_drug_lookup")
    return await fda_tool.lookup_drug(drug_name, search_type)


@mcp.tool()
async def pubmed_search(ctx: Context, query: str, max_results: int = 5, date_range: str = ""):
    """Search for medical literature in the PubMed database."""
    usage_service.record_usage(session_id, "pubmed_search")
    return await pubmed_tool.search_literature(query, max_results, date_range)


@mcp.tool()
async def health_topics(ctx: Context, topic: str, language: str = "en"):
    """Retrieve evidence‑based information for a specific health topic."""
    usage_service.record_usage(session_id, "health_topics")
    return await healthfinder_tool.get_health_topics(topic, language)


@mcp.tool()
async def clinical_trials_search(
    ctx: Context,
    condition: str,
    status: str = "recruiting",
    max_results: int = 10,
):
    """Search for clinical trials by condition and status."""
    usage_service.record_usage(session_id, "clinical_trials_search")
    return await clinical_trials_tool.search_trials(condition, status, max_results)


@mcp.tool()
async def lookup_icd_code(
    ctx: Context,
    code: str | None = None,
    description: str | None = None,
    max_results: int = 10,
):
    """Look up ICD‑10 codes by code or text description."""
    usage_service.record_usage(session_id, "lookup_icd_code")
    return await medical_terminology_tool.lookup_icd_code(code, description, max_results)


@mcp.tool()
async def get_usage_stats(ctx: Context):
    """Return per‑session usage statistics."""
    return usage_service.get_monthly_usage(session_id)


@mcp.tool()
async def get_all_usage_stats(ctx: Context):
    """Return aggregate usage statistics across all sessions."""
    return usage_service.get_usage_stats()


# ---------------------------------------------------------------------------
# Helper – convert coroutine results to nicely‑formatted JSON for UI display
# ---------------------------------------------------------------------------

def pretty_json(data):
    try:
        return json.dumps(data, indent=2, ensure_ascii=False)
    except TypeError:
        # Fallback if the object isn’t serialisable
        return str(data)


# ---------------------------------------------------------------------------
# Gradio UI definitions
# ---------------------------------------------------------------------------

async def fda_drug_lookup_gr(drug_name: str, search_type: str):
    """Gradio wrapper for :pyfunc:`fda_drug_lookup`.

    Parameters
    ----------
    drug_name : str
        Drug identifier to pass through to :pyfunc:`FDATool.lookup_drug`.
    search_type : str
        See :pydata:`search_type` in :pyfunc:`fda_drug_lookup`.

    Returns
    -------
    str
        JSON‑formatted response suitable for a ``gr.Textbox``.
    """
    usage_service.record_usage(session_id, "fda_drug_lookup")
    res = await fda_tool.lookup_drug(drug_name, search_type)
    return pretty_json(res)


async def pubmed_search_gr(query: str, max_results: int, date_range: str):
    """Gradio wrapper for PubMed literature search.

    Parameters
    ----------
    query : str
        PubMed query string.
    max_results : int
        Maximum number of articles requested from the backend tool.
    date_range : str
        If non‑empty, restricts search to the last *N* years.

    Returns
    -------
    str
        Pretty‑printed JSON with article details.
    """
    usage_service.record_usage(session_id, "pubmed_search")
    res = await pubmed_tool.search_literature(query, int(max_results), str(date_range or ""))
    return pretty_json(res)


async def health_topics_gr(topic: str, language: str):
    """Gradio wrapper for HealthFinder topic retrieval.

    Parameters
    ----------
    topic : str
        Health topic slug or phrase.
    language : str
        Output language (``"en"`` / ``"es"``).

    Returns
    -------
    str
        JSON‑formatted HealthFinder result.
    """
    usage_service.record_usage(session_id, "health_topics")
    res = await healthfinder_tool.get_health_topics(topic, language)
    return pretty_json(res)


async def clinical_trials_search_gr(condition: str, status: str, max_results: int):
    """Gradio wrapper for :pyfunc:`clinical_trials_search`.

    Parameters
    ----------
    condition : str
        Condition keyword, e.g. "melanoma".
    status : str
        Recruitment status filter.
    max_results : int
        Cap on trial records returned.

    Returns
    -------
    str
        Pretty‑printed JSON list of trials.
    """
    usage_service.record_usage(session_id, "clinical_trials_search")
    res = await clinical_trials_tool.search_trials(condition, status, int(max_results))
    return pretty_json(res)


async def lookup_icd_code_gr(code: str, description: str, max_results: int):
    """Gradio wrapper for ICD‑10 code search.

    The user may supply *either* an ICD‑10 code fragment *or* a free‑text
    description. Empty strings are converted to ``None`` so the underlying tool
    can decide which parameter to use.

    Parameters
    ----------
    code : str
        ICD‑10 code or prefix (optional).
    description : str
        Natural‑language description (optional).
    max_results : int
        Maximum number of matches.

    Returns
    -------
    str
        JSON‑formatted list of matching codes.
    """
    usage_service.record_usage(session_id, "lookup_icd_code")
    code = code or None
    description = description or None
    res = await medical_terminology_tool.lookup_icd_code(code, description, int(max_results))
    return pretty_json(res)


# ---------------------------------------------------------------------------
# Usage‑statistics helpers (synchronous; no external I/O)
# ---------------------------------------------------------------------------

def get_usage_stats_gr() -> str:
    """Return **session‑specific** usage statistics as pretty JSON.

    Returns
    -------
    str
        JSON rendition of the current session’s monthly usage counts.
    """
    return pretty_json(usage_service.get_monthly_usage(session_id))


def get_all_usage_stats_gr() -> str:
    """Return **global** usage statistics across *all* sessions.

    Returns
    -------
    str
        JSON rendition of cumulative endpoint invocations.
    """
    return pretty_json(usage_service.get_usage_stats())


# ---------------------------------------------------------------------------
# Build the Gradio Blocks interface
# ---------------------------------------------------------------------------

with gr.Blocks(title="Healthcare MCP Explorer") as demo:
    gr.Markdown("""# 🩺 Healthcare MCP Dashboard\nInteractively query FDA, PubMed, HealthFinder, Clinical‑Trials, and ICD‑10 data — all from one screen.""")

    # ---------------- Drug lookup ----------------
    with gr.Group():
        gr.Markdown("## 💊 FDA Drug Lookup")
        with gr.Row():
            drug_name_tb = gr.Textbox(label="Drug name", placeholder="e.g. atorvastatin")
            search_type_dd = gr.Dropdown(
                label="Search type",
                choices=["general", "label", "adverse_events"],
                value="general",
            )
            fda_button = gr.Button("Search")
        fda_output = gr.Textbox(label="Result", lines=12)

    # ---------------- PubMed search -------------
    with gr.Group():
        gr.Markdown("## 🔍 PubMed Literature Search")
        query_tb = gr.Textbox(label="Query", placeholder="e.g. type 2 diabetes glycemic control")
        with gr.Row():
            max_results_slider = gr.Slider(1, 20, value=5, step=1, label="Max articles")
            date_range_tb = gr.Textbox(label="Published within last N years (optional)")
            pubmed_button = gr.Button("Search PubMed")
        pubmed_output = gr.Textbox(label="Result", lines=12)

    # ---------------- Health topics -------------
    with gr.Group():
        gr.Markdown("## 📚 HealthFinder Topics")
        topic_tb = gr.Textbox(label="Health topic", placeholder="e.g. hypertension")
        language_dd = gr.Dropdown(label="Language", choices=["en", "es"], value="en")
        hf_button = gr.Button("Get Topic Info")
        hf_output = gr.Textbox(label="Result", lines=12)


    # -------------------------------------------------------------------
    # Wire‑up callbacks
    # -------------------------------------------------------------------
    fda_button.click(fda_drug_lookup_gr, inputs=[drug_name_tb, search_type_dd], outputs=fda_output)
    pubmed_button.click(pubmed_search_gr, inputs=[query_tb, max_results_slider, date_range_tb], outputs=pubmed_output)
    hf_button.click(health_topics_gr, inputs=[topic_tb, language_dd], outputs=hf_output)

# ---------------------------------------------------------------------------
# Entrypoint – run Gradio and MCP in one process
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    demo.launch(mcp_server=True, show_error=True)