Spaces:
Running
Running
| import gradio as gr | |
| import sqlite3 | |
| import bcrypt | |
| from datetime import datetime | |
| import re | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from transformers import AutoTokenizer, AutoModel, AutoModelForSequenceClassification | |
| import os | |
| import logging | |
| from openai import OpenAI | |
| import json | |
| from fpdf import FPDF | |
| print("ENV:", os.environ) # 👈 Add this for debugging | |
| api_key = os.getenv("OPENAI_API_KEY") | |
| if not api_key: | |
| raise RuntimeError("OPENAI_API_KEY environment variable not found.") | |
| client = OpenAI(api_key=api_key) | |
| import json | |
| from fpdf import FPDF | |
| # -------------------------- | |
| # Environment Setup | |
| # -------------------------- | |
| logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| print("Using device:", device) | |
| # -------------------------- | |
| # Global Tokenizer and Hybrid Model for Treatment Prediction | |
| # -------------------------- | |
| tokenizer = AutoTokenizer.from_pretrained("emilyalsentzer/Bio_ClinicalBERT") | |
| class HybridMentalHealthModel(nn.Module): | |
| def __init__(self, bert_model, num_genders, num_medications, num_therapies, hidden_size=128): | |
| super(HybridMentalHealthModel, self).__init__() | |
| self.bert = AutoModel.from_pretrained(bert_model) | |
| bert_output_size = self.bert.config.hidden_size | |
| self.age_fc = nn.Linear(1, 16) | |
| self.gender_fc = nn.Embedding(num_genders, 16) | |
| self.fc = nn.Linear(bert_output_size + 32, hidden_size) | |
| self.medication_head = nn.Linear(hidden_size, num_medications) | |
| self.therapy_head = nn.Linear(hidden_size, num_therapies) | |
| def forward(self, input_ids, attention_mask, age, gender): | |
| bert_output = self.bert(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state[:, 0, :] | |
| age_out = self.age_fc(age) | |
| gender_out = self.gender_fc(gender) | |
| combined = torch.cat((bert_output, age_out, gender_out), dim=1) | |
| hidden = torch.relu(self.fc(combined)) | |
| return self.medication_head(hidden), self.therapy_head(hidden) | |
| # -------------------------- | |
| # Global Label Mappings and Age Scaler | |
| # -------------------------- | |
| medication_classes = ["Anxiolytics", "Benzodiazepines", "Antidepressants", "Mood Stabilizers", "Antipsychotics", "Stimulants"] | |
| therapy_classes = ["Cognitive Behavioral Therapy", "Dialectical Behavioral Therapy", "Interpersonal Therapy", "Mindfulness-Based Therapy"] # Update with your types | |
| gender_classes = ["Male", "Female", "Other"] | |
| medication_encoder = {name: idx for idx, name in enumerate(medication_classes)} | |
| inv_medication_encoder = {idx: name for name, idx in medication_encoder.items()} | |
| therapy_encoder = {name: idx for idx, name in enumerate(therapy_classes)} | |
| inv_therapy_encoder = {idx: name for name, idx in therapy_encoder.items()} | |
| gender_encoder = {name: idx for idx, name in enumerate(gender_classes)} | |
| mean_age = 50 | |
| std_age = 10 | |
| def scale_age(age): | |
| return (age - mean_age) / std_age | |
| # -------------------------- | |
| # Load the Hybrid Model (Treatment Prediction) | |
| # -------------------------- | |
| num_genders = len(gender_classes) | |
| num_medications = len(medication_classes) | |
| num_therapies = len(therapy_classes) | |
| MODEL_SAVE_PATH = "22.03.2025-16.02-ML128E10" # Update accordingly | |
| model = HybridMentalHealthModel("emilyalsentzer/Bio_ClinicalBERT", num_genders, num_medications, num_therapies) | |
| state_dict = torch.load(MODEL_SAVE_PATH, map_location=device) | |
| if "gender_fc.weight" in state_dict: | |
| del state_dict["gender_fc.weight"] | |
| model.load_state_dict(state_dict, strict=False) | |
| model.to(device) | |
| model.eval() | |
| # -------------------------- | |
| # Global Diagnosis Model (Mental Health Diagnosis) | |
| # -------------------------- | |
| diagnosis_tokenizer = AutoTokenizer.from_pretrained("ethandavey/mental-health-diagnosis-bert") # Update with your model ID | |
| diagnosis_model = AutoModelForSequenceClassification.from_pretrained("ethandavey/mental-health-diagnosis-bert") # Update with your model ID | |
| diagnosis_model.to(device) | |
| diagnosis_model.eval() | |
| def predict_disease(text): | |
| inputs = diagnosis_tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128) | |
| inputs = {k: v.to(device) for k, v in inputs.items()} | |
| with torch.no_grad(): | |
| outputs = diagnosis_model(**inputs) | |
| probabilities = F.softmax(outputs.logits, dim=1).squeeze() | |
| label_mapping = {0: "Anxiety", 1: "Normal", 2: "Depression", 3: "Suicidal", 4: "Stress"} | |
| topk = torch.topk(probabilities, k=3) | |
| top_preds = [(label_mapping[i.item()], probabilities[i].item()) for i in topk.indices] | |
| return top_preds | |
| def predict_med_therapy(symptoms, age, gender): | |
| encoding = tokenizer(symptoms, return_tensors="pt", truncation=True, padding='max_length', max_length=128) | |
| input_ids = encoding["input_ids"].to(device) | |
| attention_mask = encoding["attention_mask"].to(device) | |
| age_norm = torch.tensor([[scale_age(age)]], dtype=torch.float32).to(device) | |
| gender_idx = gender_encoder.get(gender, 0) | |
| gender_tensor = torch.tensor([gender_idx], dtype=torch.long).to(device) | |
| with torch.no_grad(): | |
| med_logits, therapy_logits = model(input_ids, attention_mask, age_norm, gender_tensor) | |
| med_probabilities = torch.softmax(med_logits, dim=1) | |
| therapy_probabilities = torch.softmax(therapy_logits, dim=1) | |
| med_pred = torch.argmax(med_probabilities, dim=1).item() | |
| therapy_pred = torch.argmax(therapy_probabilities, dim=1).item() | |
| med_confidence = med_probabilities[0][med_pred].item() | |
| therapy_confidence = therapy_probabilities[0][therapy_pred].item() | |
| predicted_med = inv_medication_encoder.get(med_pred, "Unknown") | |
| predicted_therapy = inv_therapy_encoder.get(therapy_pred, "Unknown") | |
| return (predicted_med, med_confidence), (predicted_therapy, therapy_confidence) | |
| # -------------------------- | |
| # OpenAI Functions (Summarization and Explanation) | |
| # -------------------------- | |
| def get_concise_rewrite(text, max_tokens, temperature=0.7): | |
| messages = [ | |
| {"role": "system", "content": "You are an expert rewriting assistant. Rewrite the given statement into a concise version while preserving its tone and vocabulary."}, | |
| {"role": "user", "content": text} | |
| ] | |
| try: | |
| response = client.chat.completions.create(model="gpt-4o-mini", messages=messages, max_tokens=max_tokens, temperature=temperature) | |
| concise_text = response.choices[0].message.content.strip() | |
| except Exception as e: | |
| concise_text = f"API call failed: {e}" | |
| return concise_text | |
| def get_explanation(patient_statement, predicted_diagnosis): | |
| messages = [ | |
| {"role": "system", "content": "You are an expert mental health assistant. Provide a concise, evidence-based explanation of how the patient's statement supports the diagnosis."}, | |
| {"role": "user", "content": f"Patient statement: {patient_statement}\nPredicted diagnosis: {predicted_diagnosis}\nExplain briefly."} | |
| ] | |
| try: | |
| response = client.chat.completions.create(model="gpt-4o-mini", messages=messages, max_tokens=256) | |
| explanation = response.choices[0].message.content.strip() | |
| except Exception as e: | |
| explanation = "API call failed." | |
| return explanation | |
| # -------------------------- | |
| # Database Functions | |
| # -------------------------- | |
| def init_db(): | |
| conn = sqlite3.connect("users.db") | |
| c = conn.cursor() | |
| c.execute(""" | |
| CREATE TABLE IF NOT EXISTS users ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| username TEXT UNIQUE NOT NULL, | |
| password TEXT NOT NULL, | |
| full_name TEXT, | |
| email TEXT | |
| ) | |
| """) | |
| c.execute(""" | |
| CREATE TABLE IF NOT EXISTS chat_history ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| username TEXT NOT NULL, | |
| message TEXT NOT NULL, | |
| response TEXT NOT NULL, | |
| timestamp DATETIME DEFAULT CURRENT_TIMESTAMP | |
| ) | |
| """) | |
| c.execute(""" | |
| CREATE TABLE IF NOT EXISTS patient_sessions ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| username TEXT, | |
| patient_name TEXT, | |
| age REAL, | |
| gender TEXT, | |
| symptoms TEXT, | |
| diagnosis TEXT, | |
| medication TEXT, | |
| therapy TEXT, | |
| summary TEXT, | |
| explanation TEXT, | |
| pdf_report TEXT, | |
| session_timestamp DATETIME DEFAULT CURRENT_TIMESTAMP, | |
| appointment_date DATE | |
| ) | |
| """) | |
| conn.commit() | |
| conn.close() | |
| def register_user(username, password, full_name, email): | |
| if not re.fullmatch(r"[^@]+@[^@]+\.[^@]+", email): | |
| return "Invalid email format." | |
| if len(password) <= 8: | |
| return "Password must be more than 8 characters." | |
| conn = sqlite3.connect("users.db") | |
| c = conn.cursor() | |
| hashed = bcrypt.hashpw(password.encode(), bcrypt.gensalt()) | |
| try: | |
| c.execute("INSERT INTO users (username, password, full_name, email) VALUES (?, ?, ?, ?)", (username, hashed, full_name, email)) | |
| conn.commit() | |
| return "User registered successfully." | |
| except sqlite3.IntegrityError: | |
| return "Username already exists." | |
| finally: | |
| conn.close() | |
| def login_user(username, password): | |
| conn = sqlite3.connect("users.db") | |
| c = conn.cursor() | |
| c.execute("SELECT password FROM users WHERE username = ?", (username,)) | |
| user = c.fetchone() | |
| conn.close() | |
| if user and bcrypt.checkpw(password.encode(), user[0]): | |
| return True | |
| return False | |
| def get_user_info(username): | |
| conn = sqlite3.connect("users.db") | |
| c = conn.cursor() | |
| c.execute("SELECT username, email, full_name FROM users WHERE username = ?", (username,)) | |
| user = c.fetchone() | |
| conn.close() | |
| if user: | |
| return f"Username: {user[0]}\nFull Name: {user[2]}\nEmail: {user[1]}" | |
| else: | |
| return "User info not found." | |
| def get_chat_history(username): | |
| conn = sqlite3.connect("users.db") | |
| c = conn.cursor() | |
| c.execute("SELECT message, response, timestamp FROM chat_history WHERE username = ? ORDER BY timestamp DESC", (username,)) | |
| history = c.fetchall() | |
| conn.close() | |
| return history | |
| def get_patient_sessions(filter_name="", filter_date=""): | |
| conn = sqlite3.connect("users.db") | |
| c = conn.cursor() | |
| query = "SELECT patient_name, age, gender, symptoms, diagnosis, medication, therapy, summary, explanation, pdf_report, session_timestamp FROM patient_sessions WHERE 1=1" | |
| params = [] | |
| if filter_name: | |
| query += " AND patient_name LIKE ?" | |
| params.append(f"%{filter_name}%") | |
| if filter_date: | |
| query += " AND DATE(session_timestamp)=?" | |
| params.append(filter_date) | |
| c.execute(query, params) | |
| sessions = c.fetchall() | |
| conn.close() | |
| return sessions | |
| def insert_patient_session(session_data): | |
| conn = sqlite3.connect("users.db") | |
| c = conn.cursor() | |
| c.execute(""" | |
| INSERT INTO patient_sessions (username, patient_name, age, gender, symptoms, diagnosis, medication, therapy, summary, explanation, pdf_report, appointment_date) | |
| VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) | |
| """, ( | |
| session_data.get("username"), session_data.get("patient_name"), session_data.get("age"), session_data.get("gender"), | |
| session_data.get("symptoms"), session_data.get("diagnosis"), session_data.get("medication"), | |
| session_data.get("therapy"), session_data.get("summary"), session_data.get("explanation"), | |
| session_data.get("pdf_report"), session_data.get("appointment_date"))) | |
| conn.commit() | |
| conn.close() | |
| # -------------------------- | |
| # PDF Report Generation Function | |
| # -------------------------- | |
| def generate_pdf_report(session_data): | |
| pdf = FPDF() | |
| pdf.add_page() | |
| pdf.set_font("Arial", size=12) | |
| # Use safe_text to ensure the title is safe for latin-1 encoding | |
| pdf.cell(200, 10, txt=safe_text("Patient Session Report"), ln=True, align='C') | |
| pdf.ln(10) | |
| for key, value in session_data.items(): | |
| # Convert each line to a safe text version before writing it | |
| pdf.multi_cell(0, 10, txt=safe_text(f"{key.capitalize()}: {value}")) | |
| reports_dir = "reports" | |
| os.makedirs(reports_dir, exist_ok=True) | |
| timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") | |
| filename = f"{reports_dir}/{session_data.get('patient_name')}_{timestamp}.pdf" | |
| pdf.output(filename) | |
| return filename | |
| def safe_text(txt): | |
| # Encode the text to latin-1, replacing characters that can't be encoded | |
| return txt.encode("latin-1", "replace").decode("latin-1") | |
| # -------------------------- | |
| # Helper: Autofill Previous Patient Info | |
| # -------------------------- | |
| def get_previous_patients(): | |
| # Use the current logged-in user from user_session.value | |
| username = user_session.value | |
| conn = sqlite3.connect("users.db") | |
| c = conn.cursor() | |
| c.execute("SELECT DISTINCT patient_name FROM patient_sessions WHERE username=?", (username,)) | |
| records = c.fetchall() | |
| conn.close() | |
| return [r[0] for r in records] | |
| def get_previous_patient_info(selected_patient): | |
| # Use the current logged-in user from user_session.value | |
| username = user_session.value | |
| conn = sqlite3.connect("users.db") | |
| c = conn.cursor() | |
| c.execute( | |
| "SELECT patient_name, age, gender FROM patient_sessions WHERE username=? AND patient_name=? ORDER BY session_timestamp DESC LIMIT 1", | |
| (username, selected_patient) | |
| ) | |
| record = c.fetchone() | |
| conn.close() | |
| if record: | |
| return record[0], record[1], record[2] | |
| else: | |
| return "", None, "" | |
| # -------------------------- | |
| # Gradio UI Setup with External CSS | |
| # -------------------------- | |
| with gr.Blocks(css=open("styles.css", "r").read(), theme="soft") as app: | |
| user_session = gr.State(value="") | |
| profile_visible = gr.State(value=False) | |
| session_data_state = gr.State(value="") | |
| with gr.Row(elem_id="header") as header_row: | |
| with gr.Column(scale=8): | |
| gr.Markdown("## Mental Health Chatbot") | |
| with gr.Column(visible=False, elem_id="profile_container") as profile_container: | |
| profile_button = gr.Button("👤", elem_id="profile_button", variant="secondary") | |
| with gr.Column(visible=False, elem_id="profile_info_box") as profile_info_box: | |
| profile_info = gr.HTML() | |
| logout_button = gr.Button("Logout", elem_id="logout_button") | |
| with gr.Column(visible=True, elem_id="login_page") as login_page: | |
| gr.Markdown("## Login") | |
| with gr.Row(): | |
| username_login = gr.Textbox(label="Username") | |
| password_login = gr.Textbox(label="Password", type="password") | |
| login_btn = gr.Button("Login") | |
| login_output = gr.Textbox(label="Login Status", interactive=False) | |
| gr.Markdown("New user? Click below to register.") | |
| go_to_register = gr.Button("Go to Register") | |
| with gr.Column(visible=False, elem_id="register_page") as register_page: | |
| gr.Markdown("## Register") | |
| new_username = gr.Textbox(label="New Username") | |
| new_password = gr.Textbox(label="New Password", type="password") | |
| full_name = gr.Textbox(label="Full Name") | |
| email = gr.Textbox(label="Email") | |
| register_btn = gr.Button("Register") | |
| register_output = gr.Textbox(label="Registration Status", interactive=False) | |
| gr.Markdown("Already have an account?") | |
| back_to_login = gr.Button("Back to Login") | |
| with gr.Tabs(visible=False, elem_id="main_panel") as main_panel: | |
| with gr.Tab("Chatbot"): | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| previous_patient = gr.Dropdown(label="Previous Patients", choices=[], interactive=True) | |
| patient_name_input = gr.Textbox(placeholder="Enter patient name", label="Patient Name") | |
| gender_input = gr.Dropdown(choices=list(gender_encoder.keys()), label="Gender") | |
| age_input = gr.Number(label="Age") | |
| symptoms_input = gr.Textbox(placeholder="Describe symptoms", label="Symptoms", lines=4) | |
| submit = gr.Button("Submit") | |
| generate_report_btn = gr.Button("Generate Report", visible=False) | |
| with gr.Column(scale=1): | |
| with gr.Row(): | |
| with gr.Column(scale=4, min_width=240): # Textbox column | |
| diagnosis_textbox = gr.Textbox(label="Diagnosis", | |
| interactive=False) | |
| with gr.Column(scale=1, min_width=120): # Confidence column | |
| diagnosis_conf_html = gr.HTML(elem_classes=["confidence-container"]) | |
| with gr.Row(): | |
| with gr.Column(scale=4, min_width=240): | |
| medication_textbox = gr.Textbox(label="Medication", | |
| interactive=False) | |
| with gr.Column(scale=1, min_width=120): | |
| medication_conf_html = gr.HTML(elem_classes=["confidence-container"]) | |
| with gr.Row(): | |
| with gr.Column(scale=4, min_width=240): | |
| therapy_textbox = gr.Textbox(label="Therapy", | |
| interactive=False) | |
| with gr.Column(scale=1, min_width=120): | |
| therapy_conf_html = gr.HTML(elem_classes=["confidence-container"]) | |
| summary_textbox = gr.Textbox(label="Concise Summary", interactive=False) | |
| explanation_textbox = gr.Textbox(label="Explanation", interactive=False) | |
| with gr.Row(): | |
| report_download = gr.File(label="Download Report", interactive=False) | |
| def handle_chat_extended(patient_name, gender, age, symptoms): | |
| if age is None or age <= 0: | |
| error_msg = "Age must be greater than 0." | |
| return (error_msg, "", error_msg, "", error_msg, "", error_msg, error_msg, gr.update(visible=False)) | |
| if age > 150: | |
| error_msg2 = "Age must be lower than 150" | |
| return (error_msg2, "", error_msg2, "", error_msg2, "", error_msg2, error_msg2, gr.update(visible=False)) | |
| if len(symptoms.split()) > 512: | |
| msg = "Input exceeds maximum allowed length of 512 words." | |
| return (msg, "", msg, "", msg, "", msg, msg, gr.update(visible=False)) | |
| full_statement = f"Patient Name: {patient_name}, Gender: {gender}, Age: {age}, Symptoms: {symptoms}" | |
| summary = get_concise_rewrite(full_statement, max_tokens=150, temperature=0.7) | |
| # Predict top 3 diagnoses | |
| diagnosis_preds = predict_disease(full_statement) # Now returns list of (label, confidence) | |
| diagnosis_display = "\n".join([f"{label}" for label, _ in diagnosis_preds]) | |
| def get_confidence_class(percentage): | |
| if percentage <= 50: | |
| return "confidence-low" | |
| elif percentage <= 74: | |
| return "confidence-medium" | |
| else: | |
| return "confidence-high" | |
| diagnosis_conf_html_val = "<div class='confidence-multi'>" + "<br>".join([ | |
| f"<div class='confidence-display'><span class='confidence-value {get_confidence_class(conf * 100)}'>{conf * 100:.1f}% confidence</span></div>" | |
| for _, conf in diagnosis_preds | |
| ]) + "</div>" | |
| # Predict medication and therapy | |
| (med_pred, med_conf), (therapy_pred, therapy_conf) = predict_med_therapy(symptoms, age, gender) | |
| med_percentage = med_conf * 100 | |
| therapy_percentage = therapy_conf * 100 | |
| def get_conf_html(percentage): | |
| return f""" | |
| <div class="confidence-display"> | |
| <span class="confidence-value {get_confidence_class(percentage)}"> | |
| {percentage:.1f}% confidence | |
| </span> | |
| </div> | |
| """ | |
| medication_conf_html_val = get_conf_html(med_percentage) | |
| therapy_conf_html_val = get_conf_html(therapy_percentage) | |
| # Explanation | |
| top_diag_labels = ", ".join([label for label, _ in diagnosis_preds]) | |
| explanation = get_explanation(full_statement, f"{top_diag_labels}, {med_pred} and {therapy_pred}") | |
| # Prepare session data | |
| top_diag_with_conf = ", ".join([f"{label} ({conf * 100:.1f}%)" for label, conf in diagnosis_preds]) | |
| session_data = { | |
| "patient_name": patient_name, | |
| "age": age, | |
| "gender": gender, | |
| "symptoms": symptoms, | |
| "diagnosis": top_diag_with_conf, | |
| "medication": f"{med_pred} ({med_percentage:.1f}% confidence)", | |
| "therapy": f"{therapy_pred} ({therapy_percentage:.1f}% confidence)", | |
| "summary": summary, | |
| "explanation": explanation, | |
| "session_timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S") | |
| } | |
| session_data_state.value = json.dumps(session_data) | |
| # Save to chat history | |
| conn = sqlite3.connect("users.db") | |
| c = conn.cursor() | |
| if user_session.value: | |
| c.execute("INSERT INTO chat_history (username, message, response) VALUES (?, ?, ?)", | |
| (user_session.value, full_statement, top_diag_with_conf)) | |
| conn.commit() | |
| conn.close() | |
| return ( | |
| diagnosis_display, diagnosis_conf_html_val, | |
| med_pred, medication_conf_html_val, | |
| therapy_pred, therapy_conf_html_val, | |
| summary, explanation, | |
| gr.update(visible=True) | |
| ) | |
| submit.click(handle_chat_extended, | |
| inputs=[patient_name_input, gender_input, age_input, symptoms_input], | |
| outputs=[diagnosis_textbox, diagnosis_conf_html, medication_textbox, medication_conf_html, | |
| therapy_textbox, therapy_conf_html, summary_textbox, explanation_textbox, | |
| generate_report_btn]) | |
| def handle_generate_report(): | |
| try: | |
| # Try to load session data and generate the PDF report. | |
| data = json.loads(session_data_state.value) | |
| pdf_file = generate_pdf_report(data) | |
| data["username"] = user_session.value | |
| data["appointment_date"] = "" | |
| data["pdf_report"] = pdf_file | |
| insert_patient_session(data) | |
| return pdf_file | |
| except Exception as e: | |
| # Create an error file that contains the error message. | |
| error_msg = f"Error generating PDF report: {str(e)}" | |
| reports_dir = "reports" | |
| os.makedirs(reports_dir, exist_ok=True) | |
| error_filename = f"{reports_dir}/error_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt" | |
| with open(error_filename, "w", encoding="utf-8") as f: | |
| f.write(error_msg) | |
| return error_filename | |
| generate_report_btn.click(handle_generate_report, outputs=report_download) | |
| def autofill_previous(selected_patient): | |
| name, age_val, gender_val = get_previous_patient_info(selected_patient) | |
| return name, age_val, gender_val | |
| previous_patient.change(autofill_previous, | |
| inputs=[previous_patient], | |
| outputs=[patient_name_input, age_input, gender_input]) | |
| with gr.Tab("Chat History"): | |
| history_output = gr.Textbox(label="Chat History", interactive=False) | |
| load_history_btn = gr.Button("Load History") | |
| def load_history(): | |
| if user_session.value: | |
| history = get_chat_history(user_session.value) | |
| chat_history_text = "\n".join([f"[{h[2]}] {h[0]}\nBot: {h[1]}" for h in history]) | |
| return f"Username: {user_session.value}\n\n{chat_history_text}" | |
| else: | |
| return "Please log in to view history." | |
| load_history_btn.click(load_history, outputs=history_output) | |
| with gr.Tab("Book an Appointment"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| patient_name_appt = gr.Textbox(label="Patient Name", placeholder="Enter your name") | |
| doctor_name = gr.Dropdown(choices=["Dr. Smith", "Dr. Johnson", "Dr. Lee"], label="Select Doctor") | |
| appointment_date = gr.Textbox(label="Appointment Date", placeholder="YYYY-MM-DD") | |
| appointment_time = gr.Textbox(label="Appointment Time", placeholder="HH:MM (24-hour format)") | |
| reason = gr.TextArea(label="Reason for Visit", placeholder="Describe your symptoms or reason for the visit") | |
| book_button = gr.Button("Book Appointment") | |
| with gr.Column(): | |
| booking_output = gr.Textbox(label="Booking Confirmation", interactive=False) | |
| def book_appointment(patient_name, doctor_name, appointment_date, appointment_time, reason): | |
| if not user_session.value: | |
| return "Please log in to book an appointment." | |
| patient_name = (patient_name or "").strip() | |
| doctor_name = (doctor_name or "").strip() | |
| appointment_date = (appointment_date or "").strip() | |
| appointment_time = (appointment_time or "").strip() | |
| reason = (reason or "").strip() | |
| if not (patient_name and doctor_name and appointment_date and appointment_time and reason): | |
| return "Please fill in all the fields." | |
| if not re.fullmatch(r"[A-Za-z ]+", patient_name): | |
| return "Patient name should contain only letters and spaces." | |
| try: | |
| # Parse the appointment date and time strings | |
| appointment_date_obj = datetime.strptime(appointment_date, "%Y-%m-%d") | |
| except ValueError: | |
| return "Appointment date must be in YYYY-MM-DD format." | |
| try: | |
| appointment_time_obj = datetime.strptime(appointment_time, "%H:%M") | |
| except ValueError: | |
| return "Appointment time must be in HH:MM (24-hour) format." | |
| # Combine date and time into a single datetime object | |
| appointment_datetime = datetime.combine(appointment_date_obj.date(), appointment_time_obj.time()) | |
| now = datetime.now() | |
| if appointment_datetime <= now: | |
| return "Appointment date/time has already passed. Please select a future date and time." | |
| confirmation = (f"Appointment booked for {patient_name} with {doctor_name} on {appointment_date} at {appointment_time}.\n\n" | |
| f"Reason: {reason}") | |
| return confirmation | |
| book_button.click(book_appointment, | |
| inputs=[patient_name_appt, doctor_name, appointment_date, appointment_time, reason], | |
| outputs=booking_output) | |
| with gr.Tab("Patient Sessions"): | |
| gr.Markdown("### Search Patient Sessions") | |
| search_name = gr.Textbox(label="Patient Name (optional)") | |
| search_date = gr.Textbox(label="Date (YYYY-MM-DD, optional)") | |
| search_button = gr.Button("Search") | |
| sessions_output = gr.Textbox(label="Sessions", interactive=False) | |
| def search_sessions(name, date): | |
| sessions = get_patient_sessions(filter_name=name, filter_date=date) | |
| if sessions: | |
| output = "\n\n".join([f"Patient: {s[0]}\nAge: {s[1]}\nGender: {s[2]}\nSymptoms: {s[3]}\nDiagnosis: {s[4]}\nMedication: {s[5]}\nTherapy: {s[6]}\nSummary: {s[7]}\nExplanation: {s[8]}\nReport: {s[9]}\nSession Time: {s[10]}" for s in sessions]) | |
| return output | |
| else: | |
| return "No sessions found." | |
| search_button.click(search_sessions, inputs=[search_name, search_date], outputs=sessions_output) | |
| def handle_register(username, password, full_name, email): | |
| return register_user(username, password, full_name, email) | |
| def go_to_register_page(): | |
| return gr.update(visible=False), gr.update(visible=True) | |
| def back_to_login_page(): | |
| return gr.update(visible=True), gr.update(visible=False) | |
| go_to_register.click(go_to_register_page, outputs=[login_page, register_page]) | |
| register_btn.click(handle_register, | |
| inputs=[new_username, new_password, full_name, email], | |
| outputs=register_output) | |
| back_to_login.click(back_to_login_page, outputs=[login_page, register_page]) | |
| def toggle_profile(current_visible): | |
| #print("toggle_profile called with user:", user_session.value) # Debug print | |
| if not user_session.value: | |
| return gr.update(visible=False), False, "" | |
| new_visible = not current_visible | |
| info = get_user_info(user_session.value) if new_visible else "" | |
| return gr.update(visible=new_visible), new_visible, info | |
| # Connect profile button click with correct input order: | |
| profile_button.click( | |
| toggle_profile, | |
| inputs=[profile_visible], | |
| outputs=[profile_info_box, profile_visible, profile_info] | |
| ) | |
| # Handle login: update previous patients dropdown | |
| def handle_login(username, password): | |
| prev_choices = [] | |
| if login_user(username, password): | |
| user_session.value = username | |
| prev_choices = get_previous_patients() | |
| return ( | |
| f"Welcome, {username}!", # login_output | |
| gr.update(visible=True), # main_panel | |
| gr.update(visible=False), # login_page | |
| gr.update(visible=True), # header_row | |
| gr.update(choices=prev_choices, value=None), # previous_patient | |
| "", # patient_name_input | |
| None, # age_input | |
| None, # gender_input | |
| "", # symptoms_input | |
| "", # diagnosis_textbox | |
| "", # diagnosis_conf_html | |
| "", # medication_textbox | |
| "", # medication_conf_html | |
| "", # therapy_textbox | |
| "", # therapy_conf_html | |
| "", # summary_textbox | |
| "", # explanation_textbox | |
| gr.update(visible=False), # generate_report_btn | |
| None, # report_download | |
| "", # session_data_state | |
| "", # search_name (Patient Sessions tab) | |
| "", # search_date (Patient Sessions tab) | |
| "", # booking_output (Book an Appointment tab) | |
| "", # patient_name_appt (Booking tab field) | |
| "", # appointment_date (Booking tab field) | |
| "", # appointment_time (Booking tab field) | |
| "", # reason (Booking tab field) | |
| gr.update(visible=True) # profile_container: show profile icon now | |
| ) | |
| else: | |
| return ( | |
| "Invalid credentials.", # login_output | |
| gr.update(), # main_panel | |
| gr.update(), # login_page | |
| gr.update(), # header_row | |
| gr.update(choices=[], value=None), # previous_patient (cleared) | |
| "", # patient_name_input | |
| None, # age_input | |
| None, # gender_input | |
| "", # symptoms_input | |
| "", # diagnosis_textbox | |
| "", # diagnosis_conf_html | |
| "", # medication_textbox | |
| "", # medication_conf_html | |
| "", # therapy_textbox | |
| "", # therapy_conf_html | |
| "", # summary_textbox | |
| "", # explanation_textbox | |
| gr.update(visible=False), # generate_report_btn | |
| None, # report_download | |
| "", # session_data_state | |
| "", # search_name | |
| "", # search_date | |
| "", # booking_output | |
| "", # patient_name_appt | |
| "", # appointment_date | |
| "", # appointment_time | |
| "", # reason | |
| gr.update(visible=False) # profile_container: hide profile icon on failure | |
| ) | |
| login_btn.click( | |
| handle_login, | |
| inputs=[username_login, password_login], | |
| outputs=[ | |
| login_output, main_panel, login_page, header_row, previous_patient, | |
| patient_name_input, age_input, gender_input, symptoms_input, | |
| diagnosis_textbox, diagnosis_conf_html, | |
| medication_textbox, medication_conf_html, | |
| therapy_textbox, therapy_conf_html, | |
| summary_textbox, explanation_textbox, | |
| generate_report_btn, report_download, session_data_state, | |
| search_name, search_date, booking_output, patient_name_appt, appointment_date, appointment_time, reason, | |
| profile_container # new output for profile container | |
| ] | |
| ) | |
| def handle_logout(): | |
| user_session.value = "" | |
| return ( | |
| gr.update(visible=False), # Hide main_panel | |
| gr.update(visible=True), # Show login_page | |
| gr.update(visible=False), # Hide header_row | |
| gr.update(visible=False), # Hide profile_info_box | |
| False, # Reset profile_visible | |
| "", # Clear profile_info | |
| "", # Clear login_output | |
| "", # Clear history_output | |
| "", # Clear username_login textbox | |
| "", # Clear password_login textbox | |
| "", # Clear new_username textbox (register page) | |
| "", # Clear new_password textbox (register page) | |
| "", # Clear full_name textbox (register page) | |
| "", # Clear email textbox (register page) | |
| gr.update(visible=False) # profile_container: hide profile icon | |
| ) | |
| logout_button.click( | |
| handle_logout, | |
| outputs=[ | |
| main_panel, | |
| login_page, | |
| header_row, | |
| profile_info_box, | |
| profile_visible, | |
| profile_info, | |
| login_output, | |
| history_output, | |
| username_login, | |
| password_login, | |
| new_username, | |
| new_password, | |
| full_name, | |
| ] | |
| ) | |
| def main(): | |
| init_db() | |
| app.launch() | |
| if __name__ == "__main__": | |
| main() | |