app.py

import gradio as gr
import json
import os
import re
import logging
from datetime import datetime
from transformers import pipeline
import google.generativeai as genai

# Configure Gemini API
GEMINI_API_KEY = "AIzaSyBxpTHcJP3dmR9Pqppp4zmc2Tfut6nic6A"
genai.configure(api_key=GEMINI_API_KEY)
model = genai.GenerativeModel(model_name="models/gemini-2.0-flash")

# NER pipeline
NER_MODEL = "raynardj/ner-disease-ncbi-bionlp-bc5cdr-pubmed"
ers = pipeline(task="ner", model=NER_MODEL, tokenizer=NER_MODEL)

# Chat history file
CHAT_RECORD_FILE = "chat_records.json"

def load_records():
    if os.path.exists(CHAT_RECORD_FILE):
        with open(CHAT_RECORD_FILE, "r") as f:
            return json.load(f)
    return []

def save_record(chat_history):
    records = load_records()
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    records.append({"timestamp": timestamp, "history": chat_history})
    with open(CHAT_RECORD_FILE, "w") as f:
        json.dump(records, f, indent=2)

health_keywords = [
    "fever", "cold", "headache", "pain", "diabetes", "pressure", "bp", "covid",
    "infection", "symptom", "cancer", "flu", "aids", "allergy", "disease", "vomit", "asthma",
    "medicine", "tablet", "ill", "sick", "nausea", "health", "injury", "cough", "treatment",
    "doctor", "hospital", "clinic", "vaccine", "antibiotic", "therapy", "mental health", "stress",
    "anxiety", "depression", "diet", "nutrition", "fitness", "exercise", "weight loss", "cholesterol",
    "thyroid", "migraine", "burn", "fracture", "wound", "emergency", "blood sugar", "sugar", "heart", "lungs"
]

def is_health_related(text):
    return any(re.search(rf"\\b{re.escape(word)}\\b", text.lower()) for word in health_keywords)

def extract_diseases(text):
    entities = ers(text)
    return set(ent['word'] for ent in entities if 'disease' in ent.get('entity', '').lower())

def highlight_diseases(text):
    diseases = extract_diseases(text)
    for disease in diseases:
        text = re.sub(fr"\\b({re.escape(disease)})\\b", r"<mark>\\1</mark>", text, flags=re.IGNORECASE)
    return text

def ask_gemini(prompt):
    try:
        response = model.generate_content(prompt)
        return response.text.strip()
    except Exception as e:
        logging.error(e)
        return "⚠️ An unexpected error occurred."

def respond(name, age, gender, topic, user_input, history):
    chat_history = history or []
    chat_history.append(("You", user_input))

    if is_health_related(user_input):
        if not (name and age and gender):
            response = "⚠️ Please fill in your name, age, and gender."
        elif not age.isdigit() or not (0 <= int(age) <= 120):
            response = "⚠️ Please enter a valid age (0-120)."
        else:
            prompt = f"""
You are a helpful AI healthcare assistant.
Provide simple, safe, general health-related answers without diagnoses or prescriptions.

User Info:
Name: {name}
Age: {age}
Gender: {gender}

Topic: {topic}

User's Question: {user_input}
"""
            response = ask_gemini(prompt)
    else:
        prompt = f"""
You are a friendly, polite assistant.
Respond naturally and supportively.

Topic: {topic}

User's Message:
{user_input}
"""
        response = ask_gemini(prompt)

    chat_history.append(("Gemini", response))
    save_record(chat_history)

    chat_display = "\n".join(
        f"<b>{sender}:</b> {highlight_diseases(msg) if sender != 'You' else msg}"
        for sender, msg in chat_history
    )
    return chat_display, chat_history

def export_json(history):
    return gr.File.update(value=json.dumps(history, indent=2).encode("utf-8"), visible=True)

with gr.Blocks(css="mark { background-color: #ffeb3b; font-weight: bold; }") as demo:
    gr.Markdown("# 🩺 Gemini Healthcare Assistant")

    with gr.Accordion("User Info", open=True):
        name = gr.Textbox(label="Name")
        age = gr.Textbox(label="Age")
        gender = gr.Dropdown(["Male", "Female", "Other"], label="Gender")
        topic = gr.Radio(["General", "Mental Health", "Diet", "Fitness", "Stress"], label="Topic", value="General")

    chat = gr.Textbox(label="Ask something", placeholder="Type your question here")
    output = gr.HTML()
    state = gr.State([])

    btn = gr.Button("💬 Send")
    btn.click(fn=respond, inputs=[name, age, gender, topic, chat, state], outputs=[output, state])

    export_btn = gr.Button("⬇️ Export Chat")
    export_file = gr.File(visible=False)
    export_btn.click(fn=export_json, inputs=[state], outputs=[export_file])

demo.launch()