# app.py — Hugging Face Space entrypoint # Adds: Login/Register (first/last name + Lab dropdown), thumbs up/down feedback, Admin CSV export. import os os.environ.setdefault("OMP_NUM_THREADS", "1") os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") import hashlib, hmac, secrets, sqlite3, time, csv from datetime import datetime from typing import Optional, Tuple import gradio as gr from First_Pass import ask # your RAG/LLM core TITLE = "Askstein — CT Rigidity / FE Q&A" DESC = ( "Login or register to use the chatbot. Choose your lab, then ask. After each answer, leave a 👍 or 👎 — feedback is saved for review." ) # ====== Lab options =========================================================== LAB_LABELS = ["Nazarian Lab", "Freedman Lab", "Alboro Lab"] LAB_TO_CODE = { "Nazarian Lab": "nazarian", "Freedman Lab": "freedman", "Alboro Lab": "alboro", } CODE_TO_LAB = {v: k for k, v in LAB_TO_CODE.items()} def to_lab_code(label: str) -> str: return LAB_TO_CODE.get(label, "nazarian") def to_lab_label(code: str) -> str: return CODE_TO_LAB.get((code or "").lower(), "Nazarian Lab") # ====== Simple SQLite storage (users + feedback) ============================= DATA_DIR = os.path.abspath(os.getenv("DATA_DIR", "./data")) DB_PATH = os.path.join(DATA_DIR, "askstein.db") os.makedirs(DATA_DIR, exist_ok=True) def _db(): # check_same_thread=False lets us call from Gradio event threads conn = sqlite3.connect(DB_PATH, check_same_thread=False) conn.execute("PRAGMA journal_mode=WAL;") conn.execute("PRAGMA synchronous=NORMAL;") return conn def _ensure_user_columns(conn: sqlite3.Connection): cur = conn.cursor() cur.execute("PRAGMA table_info(users)") cols = {row[1] for row in cur.fetchall()} # set of column names # Migrate if missing if "first_name" not in cols: conn.execute("ALTER TABLE users ADD COLUMN first_name TEXT DEFAULT ''") if "last_name" not in cols: conn.execute("ALTER TABLE users ADD COLUMN last_name TEXT DEFAULT ''") if "lab_choice" not in cols: # store canonical code: 'nazarian'|'freedman'|'alboro' conn.execute("ALTER TABLE users ADD COLUMN lab_choice TEXT DEFAULT 'nazarian'") conn.commit() def init_db(): conn = _db() cur = conn.cursor() cur.execute(""" CREATE TABLE IF NOT EXISTS users ( id INTEGER PRIMARY KEY AUTOINCREMENT, email TEXT UNIQUE NOT NULL, password_hash TEXT NOT NULL, salt TEXT NOT NULL, created_at TEXT NOT NULL -- first_name, last_name, lab_choice added via migration ) """) cur.execute(""" CREATE TABLE IF NOT EXISTS feedback ( id INTEGER PRIMARY KEY AUTOINCREMENT, user_id INTEGER, question TEXT NOT NULL, answer_preview TEXT NOT NULL, rating INTEGER NOT NULL, -- +1 = thumbs up, -1 = thumbs down created_at TEXT NOT NULL, FOREIGN KEY (user_id) REFERENCES users(id) ) """) conn.commit() # Add new columns if the DB existed before _ensure_user_columns(conn) conn.close() # ====== Password hashing (PBKDF2-HMAC-SHA256) ================================ def _hash_password(password: str, salt_hex: Optional[str]=None) -> Tuple[str, str]: if not salt_hex: salt = secrets.token_bytes(16) else: salt = bytes.fromhex(salt_hex) dk = hashlib.pbkdf2_hmac("sha256", password.encode("utf-8"), salt, 100_000, dklen=32) return dk.hex(), salt.hex() def _verify_password(password: str, salt_hex: str, expected_hash_hex: str) -> bool: actual, _ = _hash_password(password, salt_hex) return hmac.compare_digest(actual, expected_hash_hex) # ====== User & Feedback helpers ============================================== def register_user(first: str, last: str, email: str, password: str, lab_label: str) -> Tuple[bool, str]: first = (first or "").strip() last = (last or "").strip() email = (email or "").strip().lower() lab_code = to_lab_code((lab_label or "").strip()) if not first or len(first) > 80: return False, "Please enter a valid first name." if not last or len(last) > 80: return False, "Please enter a valid last name." if not email or "@" not in email or len(email) > 200: return False, "Please enter a valid email." if not password or len(password) < 8: return False, "Password must be at least 8 characters." if lab_code not in CODE_TO_LAB: return False, "Please select a lab." pw_hash, salt = _hash_password(password) try: conn = _db() conn.execute( "INSERT INTO users (email, password_hash, salt, created_at, first_name, last_name, lab_choice) VALUES (?, ?, ?, ?, ?, ?, ?)", (email, pw_hash, salt, datetime.utcnow().isoformat(), first, last, lab_code) ) conn.commit() return True, "Registration successful! You can log in now." except sqlite3.IntegrityError: return False, "This email is already registered." except Exception as e: return False, f"Registration failed: {e}" finally: conn.close() def login_user(email: str, password: str) -> Tuple[bool, str, Optional[int], Optional[str], Optional[str], Optional[str]]: email = (email or "").strip().lower() if not email or not password: return False, "Missing email or password.", None, None, None, None try: conn = _db() cur = conn.cursor() cur.execute("SELECT id, password_hash, salt, first_name, last_name, lab_choice FROM users WHERE email = ?", (email,)) row = cur.fetchone() if not row: return False, "Invalid email or password.", None, None, None, None uid, pw_hash, salt, first, last, lab_code = row if _verify_password(password, salt, pw_hash): return True, "Login successful.", uid, first or "", last or "", (lab_code or "nazarian") return False, "Invalid email or password.", None, None, None, None except Exception as e: return False, f"Login failed: {e}", None, None, None, None finally: conn.close() def save_feedback(user_id: Optional[int], question: str, answer: str, rating: int) -> Tuple[bool, str]: preview = (answer or "").replace("\n", " ").strip() if len(preview) > 350: preview = preview[:350] + "…" try: conn = _db() conn.execute( "INSERT INTO feedback (user_id, question, answer_preview, rating, created_at) VALUES (?, ?, ?, ?, ?)", (user_id, question, preview, int(rating), datetime.utcnow().isoformat()) ) conn.commit() return True, "Thanks for your feedback!" except Exception as e: return False, f"Could not save feedback: {e}" finally: conn.close() def export_feedback_csv(all_users: bool, user_id: Optional[int]=None) -> str: # returns path to a CSV file ready for download ts = int(time.time()) out_path = os.path.join(DATA_DIR, f"feedback_export_{'all' if all_users else user_id}_{ts}.csv") conn = _db() cur = conn.cursor() if all_users: cur.execute(""" SELECT f.id, f.created_at, f.rating, f.question, f.answer_preview, u.email, u.first_name, u.last_name, u.lab_choice FROM feedback f LEFT JOIN users u ON f.user_id = u.id ORDER BY f.id DESC """) else: cur.execute(""" SELECT f.id, f.created_at, f.rating, f.question, f.answer_preview, u.email, u.first_name, u.last_name, u.lab_choice FROM feedback f LEFT JOIN users u ON f.user_id = u.id WHERE f.user_id = ? ORDER BY f.id DESC """, (user_id,)) rows = cur.fetchall() conn.close() with open(out_path, "w", newline="", encoding="utf-8") as f: w = csv.writer(f) w.writerow(["id","created_at","rating","question","answer_preview","user_email","first_name","last_name","lab"]) for r in rows: # map lab code to label for readability r = list(r) r[-1] = to_lab_label(r[-1] or "nazarian") w.writerow(r) return out_path def export_users_csv() -> str: ts = int(time.time()) out_path = os.path.join(DATA_DIR, f"users_export_{ts}.csv") conn = _db() cur = conn.cursor() cur.execute("SELECT id, email, first_name, last_name, lab_choice, created_at FROM users ORDER BY id DESC") rows = cur.fetchall() conn.close() with open(out_path, "w", newline="", encoding="utf-8") as f: w = csv.writer(f) w.writerow(["id","email","first_name","last_name","lab","created_at"]) for r in rows: r = list(r) r[4] = to_lab_label(r[4] or "nazarian") w.writerow(r) return out_path # Initialize DB on import init_db() ADMIN_TOKEN = os.getenv("ADMIN_TOKEN", "").strip() # ====== Gradio UI ============================================================= with gr.Blocks(title=TITLE) as demo: gr.Markdown(f"## {TITLE}\n{DESC}") # Session state st_user_id = gr.State(value=None) # int or None st_user_email = gr.State(value=None) # str or None st_user_fname = gr.State(value="") # first name st_user_lname = gr.State(value="") # last name st_user_lab = gr.State(value="nazarian") # lab code st_last_q = gr.State(value="") # last question st_last_a = gr.State(value="") # last answer st_can_fb = gr.State(value=False) # whether feedback buttons should be enabled st_admin_ok = gr.State(value=False) # admin unlocked? # ---------- Auth View ---------- auth_view = gr.Column(visible=True) with auth_view: with gr.Tabs(): with gr.Tab("Login"): login_email = gr.Textbox(label="Email", placeholder="you@example.com") login_pass = gr.Textbox(label="Password", type="password", placeholder="Your password") login_btn = gr.Button("Log in", variant="primary") login_msg = gr.Markdown("") with gr.Tab("Register"): reg_first = gr.Textbox(label="First name", placeholder="Ada") reg_last = gr.Textbox(label="Last name", placeholder="Lovelace") reg_email = gr.Textbox(label="Email", placeholder="you@example.com") reg_pass = gr.Textbox(label="Password (min 8 chars)", type="password") reg_lab = gr.Dropdown(choices=LAB_LABELS, value="Nazarian Lab", label="Lab", multiselect=False) reg_btn = gr.Button("Create account", variant="primary") reg_msg = gr.Markdown("") # ---------- Chat View ---------- chat_view = gr.Column(visible=False) with chat_view: welcome_md = gr.Markdown("### Chat") with gr.Row(): inp = gr.Textbox( label="Your question", placeholder="e.g., How is bending rigidity (EI) computed from a cortical cross-section?", lines=3, ) out = gr.Textbox(label="Askstein", lines=12) with gr.Row(): btn_submit = gr.Button("Ask", variant="primary") btn_logout = gr.Button("Log out") # Feedback row with gr.Row(): fb_up = gr.Button("👍 Helpful", visible=True) fb_down = gr.Button("👎 Not helpful", visible=True) fb_status = gr.Markdown("", visible=True) # User export (their own feedback) with gr.Row(): my_export_btn = gr.Button("Download my feedback (CSV)") my_export_file= gr.File(label="Your feedback CSV", visible=False) # ---------- Admin View ---------- admin_view = gr.Column(visible=True) # visible; gated by token field with admin_view: gr.Markdown("### Admin (enter token to unlock)") admin_token_in = gr.Textbox(label="Admin token", type="password", placeholder="Set ADMIN_TOKEN env to use") admin_unlock = gr.Button("Unlock Admin") admin_status = gr.Markdown("") with gr.Group(visible=False) as admin_controls: all_export_btn = gr.Button("Export ALL feedback (CSV)") users_export_btn= gr.Button("Export users (CSV)") all_export_file = gr.File(label="All feedback CSV", visible=False) users_export_file= gr.File(label="Users CSV", visible=False) # ---------- Examples ---------- examples = gr.Examples( examples=[ "Define axial rigidity (EA) and how it is estimated from CT-derived cortical masks.", "How does torsional rigidity (GJ) relate to polar moment of area in FE pre-processing?", "What are typical pitfalls when mapping Hounsfield units to elastic modulus?", "What boundary conditions are common in long-bone FE bending simulations?", ], inputs=inp, ) # ====== Handlers ========================================================== # Register def on_register(first, last, email, password, lab_label): ok, msg = register_user(first, last, email, password, lab_label) return gr.Markdown.update(value=msg), None, None # message only reg_btn.click(on_register, [reg_first, reg_last, reg_email, reg_pass, reg_lab], [reg_msg]) # Login def on_login(email, password): ok, msg, uid, first, last, lab_code = login_user(email, password) if ok: label = to_lab_label(lab_code) welcome = f"Welcome, **{first} {last}** ({label}) — you are logged in as **{email.strip().lower()}**." return ( gr.Markdown.update(value=welcome), gr.update(visible=False), # hide auth_view gr.update(visible=True), # show chat_view uid, email.strip().lower(), first, last, lab_code ) else: return gr.Markdown.update(value=msg), gr.update(), gr.update(), None, None, "", "", "nazarian" login_btn.click( on_login, [login_email, login_pass], [login_msg, auth_view, chat_view, st_user_id, st_user_email, st_user_fname, st_user_lname, st_user_lab] ) # Logout def on_logout(): # clear session state and toggle views return ( None, None, "", "", "nazarian", "", "", False, gr.update(visible=True), # show auth gr.update(visible=False), # hide chat gr.Markdown.update(value=""), # clear feedback status ) btn_logout.click( on_logout, [], [st_user_id, st_user_email, st_user_fname, st_user_lname, st_user_lab, st_last_q, st_last_a, st_can_fb, auth_view, chat_view, fb_status] ) # Run question def on_ask(user_id, q, lab_code): q = (q or "").strip() if not q: return gr.Textbox.update(value="Please enter a question."), q, "", False try: # Light hint to retrieval: include chosen lab keyword (helps detect_lab) lab_hint = to_lab_label(lab_code) # "Freedman Lab" etc. q_eff = f"{q} ({lab_hint})" a = ask(q_eff) except Exception as e: a = f"[runtime error] {e}" # enable feedback for this QA return gr.Textbox.update(value=a), q, a, True btn_submit.click( on_ask, [st_user_id, inp, st_user_lab], [out, st_last_q, st_last_a, st_can_fb] ) inp.submit( on_ask, [st_user_id, inp, st_user_lab], [out, st_last_q, st_last_a, st_can_fb] ) # Feedback buttons def on_feedback(user_id, can_fb, last_q, last_a, rating): if not can_fb or not last_q or not last_a: return gr.Markdown.update(value="No recent answer to rate."), False ok, msg = save_feedback(user_id, last_q, last_a, rating) return gr.Markdown.update(value=msg), False fb_up.click( lambda uid, can, q, a: on_feedback(uid, can, q, a, +1), [st_user_id, st_can_fb, st_last_q, st_last_a], [fb_status, st_can_fb] ) fb_down.click( lambda uid, can, q, a: on_feedback(uid, can, q, a, -1), [st_user_id, st_can_fb, st_last_q, st_last_a], [fb_status, st_can_fb] ) # User export (their own feedback) def on_my_export(user_id): if not user_id: return gr.File.update(visible=False), gr.Markdown.update(value="Please log in first.") path = export_feedback_csv(all_users=False, user_id=user_id) return gr.File.update(value=path, visible=True), gr.Markdown.update(value="") my_export_btn.click( on_my_export, [st_user_id], [my_export_file, fb_status] ) # ---------- Admin View ---------- ADMIN_TOKEN = os.getenv("ADMIN_TOKEN", "").strip() def on_admin_unlock(token): ok = bool(ADMIN_TOKEN) and (token.strip() == ADMIN_TOKEN) if ok: return True, gr.Markdown.update(value="Admin unlocked ✅"), gr.update(visible=True) msg = "Invalid token or ADMIN_TOKEN not set." return False, gr.Markdown.update(value=msg), gr.update(visible=False) admin_unlock.click( on_admin_unlock, [admin_token_in], [st_admin_ok, admin_status, admin_controls] ) # Admin exports def on_export_all(admin_ok): if not admin_ok: return gr.File.update(visible=False), gr.Markdown.update(value="Admin is locked.") path = export_feedback_csv(all_users=True) return gr.File.update(value=path, visible=True), gr.Markdown.update(value="") def on_export_users(admin_ok): if not admin_ok: return gr.File.update(visible=False), gr.Markdown.update(value="Admin is locked.") path = export_users_csv() return gr.File.update(value=path, visible=True), gr.Markdown.update(value="") all_export_btn.click( on_export_all, [st_admin_ok], [all_export_file, admin_status] ) users_export_btn.click( on_export_users, [st_admin_ok], [users_export_file, admin_status] ) # Launch (HF Spaces will call `demo` automatically if not run as __main__) if __name__ == "__main__": demo.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", "7860")))