app.py

"""
app.py — Gradio UI for the Medicine Safety Agent
Run:  python app.py

Changes vs original
───────────────────
* analyse() passes provider/model to agent.analyse() for accurate logging.
* A new "⏱ Performance" tab shows per-phase timing + cache info.
* A small cache-stats panel in the settings accordion.
* Status badge also shows total elapsed time after completion.
"""

from __future__ import annotations
import os
import gradio as gr
from dotenv import load_dotenv

load_dotenv()

from agents.medicine_agent import MedicineAgent, MedicineReport
from llm.connector         import get_llm_connector
from cache.semantic_cache  import get_cache
from utils.logger          import logger


# ─────────────────────────────────────────────
# Helpers
# ─────────────────────────────────────────────

def build_agent(provider: str, model: str) -> MedicineAgent:
    model = model.strip() or None
    llm   = get_llm_connector(provider=provider, model=model)
    return MedicineAgent(llm=llm)


def format_report(report: MedicineReport) -> tuple[str, str, str, str, str]:
    """Return (spell_check, interactions, contraindications, status_html, perf_md)."""
    if report.error:
        err = f"⚠️ **Error:** {report.error}"
        return err, err, err, status_badge("error", 0), ""

    spell = (
        f"**Recognised medicine name:** {report.corrected_name}\n\n"
        f"{report.name_note}"
    )

    perf_md = ""
    total_ms = 0.0
    if report.request_log:
        total_ms = report.request_log.total_ms()
        perf_md  = report.request_log.to_ui_string()

    return (
        spell,
        report.interactions,
        report.contraindications,
        status_badge("done", total_ms),
        perf_md,
    )


def status_badge(state: str, elapsed_ms: float = 0) -> str:
    if state == "idle":
        return '<span style="color:#6b7280">● Idle</span>'
    if state == "running":
        return '<span style="color:#f59e0b">⟳ Analysing…</span>'
    if state == "done":
        t = f"&nbsp;&nbsp;<small style='color:#6b7280'>{elapsed_ms:.0f} ms</small>"
        return f'<span style="color:#10b981">✔ Complete{t}</span>'
    if state == "error":
        return '<span style="color:#ef4444">✘ Error</span>'
    return ""


def cache_stats_md() -> str:
    stats = get_cache().stats()
    return (
        f"**Semantic Cache** — "
        f"{stats['valid_entries']} valid / {stats['total_entries']} total entries &nbsp;|&nbsp; "
        f"TTL {stats['ttl_seconds']//3600} h &nbsp;|&nbsp; "
        f"similarity ≥ {stats['similarity_threshold']}"
    )


def analyse(medicine_input: str, provider: str, model_name: str):
    """Generator: yields intermediate and final UI state."""
    if not medicine_input.strip():
        yield (
            gr.update(value=""),
            gr.update(value=""),
            gr.update(value=""),
            gr.update(value="Please enter a medicine name."),
            gr.update(value=""),
        )
        return

    logger.info(
        "REQUEST START  │  input=%r  provider=%s  model=%s",
        medicine_input.strip(), provider, model_name or "(default)",
    )

    yield (
        gr.update(value="🔍 Checking spelling…"),
        gr.update(value="⏳ Waiting…"),
        gr.update(value="⏳ Waiting…"),
        gr.update(value=status_badge("running")),
        gr.update(value=""),
    )

    try:
        model_str = model_name.strip() or DEFAULT_MODELS.get(provider, "")
        agent     = build_agent(provider, model_name)
        report    = agent.analyse(
            medicine_input,
            provider=provider,
            model=model_str,
        )
        spell, interactions, contraindications, badge, perf = format_report(report)
    except Exception as exc:
        logger.exception("Top-level exception in analyse()")
        spell = f"⚠️ Failed to initialise agent: {exc}"
        interactions = contraindications = perf = ""
        badge = status_badge("error")

    yield (
        gr.update(value=spell),
        gr.update(value=interactions),
        gr.update(value=contraindications),
        gr.update(value=badge),
        gr.update(value=perf),
    )


# ─────────────────────────────────────────────
# Custom CSS  (original + minor additions)
# ─────────────────────────────────────────────
CSS = """
@import url('https://fonts.googleapis.com/css2?family=DM+Serif+Display:ital@0;1&family=DM+Sans:wght@300;400;500;600&display=swap');

:root {
    --primary: #0f4c81;
    --accent:  #00b4d8;
    --danger:  #ef233c;
    --success: #06d6a0;
    --bg:      #f0f4f8;
    --card:    #ffffff;
    --text:    #1a2332;
    --muted:   #64748b;
    --border:  #dde3ec;
    --radius:  14px;
}

body, .gradio-container {
    font-family: 'DM Sans', sans-serif;
    background: var(--bg) !important;
    color: var(--text) !important;
}

#header {
    background: linear-gradient(135deg, var(--primary) 0%, #1565c0 60%, #0288d1 100%);
    padding: 36px 40px 28px;
    border-radius: var(--radius);
    margin-bottom: 24px;
    box-shadow: 0 8px 32px rgba(15,76,129,0.18);
    position: relative;
    overflow: hidden;
}
#header::before {
    content: '';
    position: absolute;
    width: 340px; height: 340px;
    border-radius: 50%;
    background: rgba(255,255,255,0.05);
    top: -120px; right: -80px;
}
#header h1 {
    font-family: 'DM Serif Display', serif;
    font-size: 2.4rem;
    color: #fff;
    margin: 0 0 6px;
    letter-spacing: -0.5px;
}
#header p { color: rgba(255,255,255,0.8); font-size: 1rem; margin: 0; }

.card {
    background: var(--card);
    border: 1px solid var(--border);
    border-radius: var(--radius);
    padding: 24px;
    box-shadow: 0 2px 12px rgba(0,0,0,0.05);
}

#medicine-input textarea {
    font-family: 'DM Sans', sans-serif;
    font-size: 1.1rem;
    border: 2px solid var(--border);
    border-radius: 10px;
    padding: 14px 18px;
    transition: border-color 0.2s;
}
#medicine-input textarea:focus {
    border-color: var(--accent);
    box-shadow: 0 0 0 3px rgba(0,180,216,0.15);
}

#search-btn {
    background: linear-gradient(135deg, var(--primary), var(--accent)) !important;
    border: none !important;
    border-radius: 10px !important;
    font-family: 'DM Sans', sans-serif !important;
    font-weight: 600 !important;
    font-size: 1rem !important;
    color: #fff !important;
    padding: 14px 32px !important;
    cursor: pointer !important;
    transition: opacity 0.2s, transform 0.1s !important;
    box-shadow: 0 4px 14px rgba(15,76,129,0.25) !important;
}
#search-btn:hover  { opacity: 0.9 !important; transform: translateY(-1px) !important; }
#search-btn:active { transform: translateY(0) !important; }

.tab-nav button { font-family: 'DM Sans', sans-serif; font-weight: 500; }
.tab-nav button.selected {
    border-bottom: 3px solid var(--primary) !important;
    color: var(--primary) !important;
}

.result-box {
    border: 1px solid var(--border);
    border-radius: 10px;
    padding: 18px 22px;
    min-height: 160px;
    font-size: 0.96rem;
    line-height: 1.7;
}
.result-box * { color: #111111 !important; }
.result-box .prose { color: #111111 !important; }
.result-box .prose * { color: #111111 !important; }
.result-box h1, .result-box h2, .result-box h3 { color: var(--primary); }

.gradio-container .tabs button { color: #111111 !important; }
.gradio-container .tabs button.selected { color: #1a73e8 !important; }

#settings summary { font-weight: 600; color: var(--muted); cursor: pointer; }
#status { font-family: 'DM Sans', sans-serif; font-size: 0.88rem; margin-top: 8px; }

#disclaimer {
    background: #fff8e1;
    border-left: 4px solid #f59e0b;
    padding: 14px 18px;
    border-radius: 8px;
    font-size: 0.87rem;
    color: #78350f;
    margin-top: 18px;
}

/* Performance tab table */
.perf-box table { width: 100%; border-collapse: collapse; font-size: 0.9rem; }
.perf-box th    { background: #f1f5f9; padding: 8px 12px; text-align: left; }
.perf-box td    { padding: 7px 12px; border-top: 1px solid var(--border); }
"""


# ─────────────────────────────────────────────
# Gradio UI
# ─────────────────────────────────────────────

PROVIDERS = ["openai", "gemini", "claude"]

DEFAULT_MODELS = {
    "openai":  "gpt-4o",
    "gemini":  "gemini-1.5-pro",
    "claude":  "claude-sonnet-4-6",
}

with gr.Blocks(css=CSS, title="💊 MedCheck Agent", theme=gr.themes.Base()) as demo:

    gr.HTML("""
    <div id="header">
      <h1>💊 MedCheck Agent</h1>
      <p>AI-powered drug information · interactions · contraindications</p>
    </div>
    """)

    with gr.Row():
        # ── Left column ───────────────────────
        with gr.Column(scale=4):
            with gr.Group(elem_classes="card"):
                gr.Markdown("### 🔎 Enter Medicine Name")
                medicine_input = gr.Textbox(
                    label="",
                    placeholder="e.g.  Asprin,  Metformin,  Lisinopril …",
                    lines=2,
                    elem_id="medicine-input",
                )
                with gr.Row():
                    search_btn = gr.Button("🔍  Analyse Medicine", elem_id="search-btn", scale=2)
                    clear_btn  = gr.Button("✕ Clear", variant="secondary", scale=1)

                status_html = gr.HTML(value=status_badge("idle"), elem_id="status")

            with gr.Accordion("⚙️ LLM Settings", open=False, elem_id="settings"):
                provider_radio = gr.Radio(
                    choices=PROVIDERS,
                    value="claude",
                    label="LLM Provider",
                    elem_id="provider-radio",
                )
                model_input = gr.Textbox(
                    label="Model name (leave blank for default)",
                    placeholder="e.g. gpt-4o / gemini-1.5-flash / claude-sonnet-4-6",
                    value="",
                )
                gr.Markdown(
                    "_Set your API keys via environment variables: "
                    "`OPENAI_API_KEY`, `GEMINI_API_KEY`, `ANTHROPIC_API_KEY`_",
                )
                cache_stats_box = gr.Markdown(value=cache_stats_md())
                refresh_cache_btn = gr.Button("🔄 Refresh cache stats", size="sm")

        # ── Right column ──────────────────────
        with gr.Column(scale=6):
            with gr.Tabs():
                with gr.Tab("✏️ Spell Check & Name"):
                    spell_out = gr.Markdown(
                        value="Results will appear here after analysis.",
                        elem_classes="result-box",
                    )
                with gr.Tab("⚠️ Drug Interactions"):
                    inter_out = gr.Markdown(
                        value="Results will appear here after analysis.",
                        elem_classes="result-box",
                    )
                with gr.Tab("🚫 Contraindications"):
                    contra_out = gr.Markdown(
                        value="Results will appear here after analysis.",
                        elem_classes="result-box",
                    )
                with gr.Tab("⏱ Performance"):
                    perf_out = gr.Markdown(
                        value="_Run an analysis to see per-phase timing and cache info._",
                        elem_classes="result-box perf-box",
                    )

    gr.HTML("""
    <div id="disclaimer">
      ⚕️ <strong>Medical Disclaimer:</strong> This tool provides general information only and
      is <em>not</em> a substitute for professional medical advice. Always consult a qualified
      healthcare provider or pharmacist before making any medication decisions.
    </div>
    """)

    gr.Examples(
        examples=[["Asprin"], ["Metformin"], ["Lisinopril"], ["Warfrin"], ["Amoxicilin"]],
        inputs=medicine_input,
        label="💡 Try an example",
        cache_examples=False,
    )

    # ── event wiring ──────────────────────────

    def update_model_placeholder(provider):
        return gr.update(placeholder=f"default: {DEFAULT_MODELS.get(provider, '')}")

    provider_radio.change(update_model_placeholder, provider_radio, model_input)

    OUTPUTS = [spell_out, inter_out, contra_out, status_html, perf_out]

    search_btn.click(
        fn=analyse,
        inputs=[medicine_input, provider_radio, model_input],
        outputs=OUTPUTS,
    )

    medicine_input.submit(
        fn=analyse,
        inputs=[medicine_input, provider_radio, model_input],
        outputs=OUTPUTS,
    )

    def clear():
        return (
            gr.update(value=""),
            gr.update(value="Results will appear here after analysis."),
            gr.update(value="Results will appear here after analysis."),
            gr.update(value="Results will appear here after analysis."),
            gr.update(value=status_badge("idle")),
            gr.update(value="_Run an analysis to see per-phase timing and cache info._"),
        )

    clear_btn.click(
        fn=clear,
        inputs=[],
        outputs=[medicine_input, spell_out, inter_out, contra_out, status_html, perf_out],
    )

    refresh_cache_btn.click(
        fn=lambda: gr.update(value=cache_stats_md()),
        inputs=[],
        outputs=[cache_stats_box],
    )


if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=int(os.getenv("PORT", 7860)),
        share=True,
    )