app.py

import os
import base64
import tempfile
import gradio as gr
from openai import OpenAI
import fitz  # PyMuPDF

# ── NVIDIA Build client (key stored in HF Secret) ─────────────────────────────
client = OpenAI(
    base_url="https://integrate.api.nvidia.com/v1",
    api_key=os.environ.get("NVIDIA_API_KEY"),
)

MODEL = "openai/gpt-oss-20b"  # GPT via NVIDIA Build — change this to match your NVIDIA model name aka salman's lair

# ── System prompts ────────────────────────────────────────────────────────────
SYSTEM_PROMPTS = {
    "🩺 General Diagnosis": """You are Sehat Guard, an expert AI medical assistant for Pakistan.
Help patients understand symptoms and lab results.

RULES:
- Respond in the SAME LANGUAGE as the user (English or Urdu script).
- When a lab report is shared, analyze ALL values, flag abnormal ones (HIGH/LOW/NORMAL), explain in simple language, give overall interpretation.
- Structure: 📋 Lab Analysis | ⚠️ Abnormal Values | 💡 What This Means | 🩺 Possible Conditions | ✅ Recommendations | 🏥 When to See a Doctor.
- Consider Pakistani diseases: dengue, typhoid, TB, malaria, hepatitis B/C.
- Emergency numbers: 1122 (Rescue) | 115 (Edhi).
- End with: AI guidance only, not a substitute for professional medical care.""",

    "👨‍⚕️ Doctor Mode": """You are Sehat Guard DOCTOR MODE — clinical decision support for licensed physicians in Pakistan.
When lab reports are uploaded:
- Interpret all values with clinical precision
- Flag critical values requiring immediate action
- Differential diagnosis based on the panel
- Suggest additional investigations
- Evidence-based treatment (WHO/DRAP guidelines)
- Drug dosages, interactions, contraindications
- ICD-10 codes where relevant
Respond in the same language as the user.""",

    "🚨 Emergency Triage": """You are Sehat Guard TRIAGE MODULE.
When lab results or symptoms are shared:
- Flag any CRITICAL values immediately (K+ >6.5, Hb <5, Glucose >600, elevated Troponin)
- Assign: 🔴 IMMEDIATE / 🟡 URGENT / 🟢 NON-URGENT
- State exact reason for triage level
- Immediate action steps
- Pakistan emergency: 1122 | 115
Be direct and action-oriented. Respond in user's language.""",

    "💊 Medications": """You are Sehat Guard MEDICATION MODULE.
When lab results or symptoms are shared:
- Identify conditions and recommend medications available in Pakistan
- Generic + brand names, dosage, frequency, duration
- Contraindications based on abnormal values
- Affordable alternatives for Pakistani patients
Prescriptions require a licensed doctor.
Respond in user's language.""",
}

DISCLAIMER = (
    "⚠️ **Medical Disclaimer:** Sehat Guard is for informational guidance only. "
    "Always consult a licensed doctor. "
    "**Emergencies:** 1122 (Rescue) | 115 (Edhi)"
)

# ── Helpers ───────────────────────────────────────────────────────────────────
def encode_image_to_base64(image_path: str) -> str:
    with open(image_path, "rb") as f:
        return base64.b64encode(f.read()).decode("utf-8")

def pdf_to_base64_images(pdf_path: str) -> list:
    doc = fitz.open(pdf_path)
    images = []
    for page in doc:
        pix = page.get_pixmap(dpi=200)
        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
            pix.save(tmp.name)
            images.append(encode_image_to_base64(tmp.name))
    doc.close()
    return images

def build_messages(system, history, user_text, image_b64_list=None):
    messages = [{"role": "system", "content": system}]
    for human, assistant in history:
        if human is not None:
            messages.append({"role": "user",      "content": str(human)})
        if assistant is not None:
            messages.append({"role": "assistant", "content": str(assistant)})

    if image_b64_list:
        content = []
        for b64 in image_b64_list:
            content.append({
                "type": "image_url",
                "image_url": {"url": f"data:image/png;base64,{b64}", "detail": "high"},
            })
        content.append({"type": "text", "text": user_text})
        messages.append({"role": "user", "content": content})
    else:
        messages.append({"role": "user", "content": user_text})

    return messages

# ── Streaming chat function ───────────────────────────────────────────────────
def chat(message, history, mode, file_upload, age, gender, weight, conditions, allergies):
    if not message.strip() and file_upload is None:
        yield "", history
        return

    system = SYSTEM_PROMPTS.get(mode, SYSTEM_PROMPTS["🩺 General Diagnosis"])

    # Patient context
    parts = []
    if age:        parts.append(f"Age: {int(age)}y")
    if gender:     parts.append(f"Gender: {gender}")
    if weight:     parts.append(f"Weight: {int(weight)}kg")
    if conditions: parts.append(f"Conditions: {conditions}")
    if allergies:  parts.append(f"Allergies: {allergies}")
    patient_ctx = ("\n\n[Patient Info — " + " | ".join(parts) + "]") if parts else ""

    user_text    = (message.strip() or "Please analyze this uploaded medical document and provide a detailed interpretation.") + patient_ctx
    display_msg  = message.strip() or "📄 File uploaded — please analyze"

    try:
        # NVIDIA Build does not support vision — text only
        if file_upload is not None:
            ext = os.path.splitext(file_upload)[1].lower()
            if ext == ".pdf":
                # Extract text from PDF instead
                doc = fitz.open(file_upload)
                pdf_text = ""
                for page in doc:
                    pdf_text += page.get_text()
                doc.close()
                user_text = f"The patient has uploaded a lab report. Here is the extracted text:\n\n{pdf_text[:4000]}\n\n{user_text}"
                display_msg = "📄 PDF uploaded" + (f" — {message.strip()}" if message.strip() else "")
            elif ext in [".jpg", ".jpeg", ".png", ".webp", ".bmp"]:
                user_text = "The patient has uploaded an image of a medical document (lab report/prescription/X-ray). Unfortunately I cannot read images directly — please type out the key values from the report and I will analyze them for you.\n\n" + user_text
                display_msg = "🖼️ Image uploaded" + (f" — {message.strip()}" if message.strip() else "")
            else:
                history.append((message or "File upload", "❌ Unsupported file type. Please upload PDF, JPG, or PNG."))
                yield "", history
                return

        messages = build_messages(system, history, user_text)

        # ── Streaming response ──
        stream = client.chat.completions.create(
            model=MODEL,
            messages=messages,
            max_tokens=1500,
            temperature=0.3,
            stream=True,
        )

        partial_reply = ""
        history = history + [(display_msg, "")]

        for chunk in stream:
            if not getattr(chunk, "choices", None):
                continue
            delta = chunk.choices[0].delta
            token = getattr(delta, "content", None)
            if token:
                partial_reply += token
                history[-1] = (display_msg, partial_reply)
                yield "", history

    except Exception as e:
        err = f"❌ Error: {str(e)}\n\nCheck your `NVIDIA_API_KEY` secret in Hugging Face Space Settings."
        history.append((display_msg, err))
        yield "", history

# ── Gradio UI ─────────────────────────────────────────────────────────────────
CSS = """
body { background: #f0f7f3; }
.gradio-container { max-width: 1150px !important; margin: 0 auto; }
#header {
    background: linear-gradient(135deg, #0f4230 0%, #1a6b4a 60%, #27a96b 100%);
    padding: 24px 32px; border-radius: 16px; margin-bottom: 16px; text-align: center;
}
#header h1 { color: #fff; font-size: 2rem; margin: 0; }
#header p  { color: #a8dfc2; font-size: 0.9rem; margin: 6px 0 0; }
#disclaimer {
    background: #fff8e1; border: 1.5px solid #f9c84a;
    border-radius: 10px; padding: 10px 16px;
    font-size: 0.82rem; color: #5a4000; margin-bottom: 14px;
}
.message.user > div {
    background: #1a6b4a !important; color: #fff !important;
    border-radius: 18px 18px 4px 18px !important;
}
.message.bot > div {
    background: #e8f5ee !important; color: #1c2c24 !important;
    border-radius: 18px 18px 18px 4px !important;
    border: 1px solid #d0e4d8 !important;
}
"""

WELCOME = (
    "👋 Hello! I'm **Sehat Guard** 🛡️ — your AI medical assistant for Pakistan.\n\n"
    "**I can help with:**\n"
    "• 📋 **Lab report analysis** — upload PDF or photo\n"
    "• 🩺 Symptom assessment & diagnosis\n"
    "• 🚨 Emergency triage (🔴🟡🟢)\n"
    "• 💊 Medication information\n"
    "• 👨‍⚕️ Clinical support for doctors\n\n"
    "Type in **English** or **Urdu** (اردو). Upload a lab report on the left."
)

with gr.Blocks(css=CSS, title="Sehat Guard — AI Medical Assistant") as demo:

    gr.HTML("""
    <div id="header">
      <h1>🛡️ Sehat Guard &nbsp;|&nbsp; <span style="font-family:serif;font-size:1.5rem">صحت گارڈ</span></h1>
      <p>AI Medical Assistant · Pakistan · Lab Reports · English & Urdu · Powered by NVIDIA Build</p>
    </div>
    """)
    gr.HTML(f'<div id="disclaimer">{DISCLAIMER}</div>')

    with gr.Row():

        # ── Sidebar ───────────────────────────────────────────────────────────
        with gr.Column(scale=1, min_width=270):

            mode = gr.Radio(
                choices=list(SYSTEM_PROMPTS.keys()),
                value="🩺 General Diagnosis",
                label="Consultation Mode",
            )

            gr.Markdown("---\n### 📁 Upload Lab Report / X-Ray / Prescription")
            file_upload = gr.File(
                label="Upload PDF or Image",
                file_types=[".pdf", ".jpg", ".jpeg", ".png", ".webp"],
                file_count="single",
            )
            gr.Markdown("_Supports: CBC, LFT, RFT, HbA1c, X-ray, ECG, ultrasound, prescriptions_")

            gr.Markdown("---\n### 👤 Patient Info *(optional)*")
            age        = gr.Number(label="Age (years)", precision=0, minimum=0, maximum=120)
            gender     = gr.Dropdown(["", "Male", "Female", "Other"], label="Gender", value="")
            weight     = gr.Number(label="Weight (kg)", precision=0, minimum=0)
            conditions = gr.Textbox(label="Known Conditions", placeholder="e.g. Diabetes, Hypertension")
            allergies  = gr.Textbox(label="Allergies",         placeholder="e.g. Penicillin")

            gr.Markdown("""
---
🇵🇰 **Emergency Numbers**
- **1122** — Rescue / Ambulance
- **115** — Edhi Foundation
- **115** — Chippa Welfare
""")

        # ── Chat ──────────────────────────────────────────────────────────────
        with gr.Column(scale=3):

            chatbot = gr.Chatbot(
                value=[[None, WELCOME]],
                label="Sehat Guard Chat",
                height=520,
            )

            with gr.Row():
                user_input = gr.Textbox(
                    placeholder="Describe symptoms or ask a question... / علامات بتائیں یا سوال پوچھیں...",
                    label="Your Message",
                    scale=5,
                    lines=2,
                )
                send_btn = gr.Button("Send ➤", scale=1, variant="primary")

            gr.Examples(
                examples=[
                    ["I have fever 38.5°C, body aches and headache for 2 days. Could it be dengue?"],
                    ["مجھے 2 دن سے بخار ہے اور جسم میں درد ہے، کیا یہ ڈینگی ہو سکتا ہے؟"],
                    ["Please analyze my uploaded CBC report and explain what's abnormal."],
                    ["My HbA1c is 9.2%, fasting glucose 210 mg/dL. What does this mean?"],
                    ["میرا بلڈ پریشر 160/100 ہے، کیا کرنا چاہیے؟"],
                    ["CBC: Hb 7.2, WBC 14,000, Platelets 45,000. What could this indicate?"],
                    ["Triage: 70yo male, sudden severe headache, vomiting, neck stiffness."],
                    ["TSH 12.5 mIU/L, T4 low. What medication is needed in Pakistan?"],
                ],
                inputs=user_input,
                label="💡 Example Questions",
            )

    inputs  = [user_input, chatbot, mode, file_upload, age, gender, weight, conditions, allergies]
    outputs = [user_input, chatbot]

    send_btn.click(fn=chat, inputs=inputs, outputs=outputs)
    user_input.submit(fn=chat, inputs=inputs, outputs=outputs)

demo.launch()