import os from datetime import datetime import gradio as gr from huggingface_hub import InferenceClient from datasets import load_dataset from supabase import create_client import threading import time # =========================== # Supabase Setup # =========================== SUPABASE_URL = os.getenv("SUPABASE_URL") SUPABASE_KEY = os.getenv("SUPABASE_KEY") if not SUPABASE_URL or not SUPABASE_KEY: raise ValueError("SUPABASE_URL and SUPABASE_KEY must be set in HF Secrets") supabase = create_client(SUPABASE_URL, SUPABASE_KEY) # =========================== # Global stop flag # =========================== stop_generation = False # =========================== # Config # =========================== MODELS = { "Meta LLaMA 3.1 (8B Instruct)": "meta-llama/Llama-3.1-8B-Instruct", "Mistral 7B Instruct": "mistralai/Mistral-7B-Instruct-v0.2", } DATASETS = ["The Stack", "CodeXGLUE"] HF_TOKEN = os.getenv("HF_TOKEN") SYSTEM_DEFAULT = ( "Specializes in databases, APIs, auth, CRUD. " "Provides complete backend code scaffolds. " "Declines frontend-heavy requests." ) # =========================== # Database Functions (Supabase) # =========================== def init_db(): pass # Tables must be created in Supabase dashboard def create_session(title: str = "New chat") -> int: res = supabase.table("sessions").insert({ "title": title, "created_at": datetime.utcnow().isoformat() }).execute() return res.data[0]["id"] def delete_session(session_id: int): supabase.table("messages").delete().eq("session_id", session_id).execute() supabase.table("sessions").delete().eq("id", session_id).execute() def list_sessions(): res = supabase.table("sessions").select("id,title").order("id", desc=True).execute() rows = res.data or [] labels = [f"{r['id']} • {r['title']}" for r in rows] return labels, rows def get_messages(session_id: int): res = supabase.table("messages") \ .select("role,content") \ .eq("session_id", session_id) \ .order("id", desc=False) \ .execute() rows = res.data or [] return [{"role": r["role"], "content": r["content"]} for r in rows] def add_message(session_id: int, role: str, content: str): print(f"🔍 ADDING MESSAGE - Session: {session_id}, Role: {role}") print(f"🔍 Content preview: {content[:50]}..." if len(content) > 50 else f"🔍 Content: {content}") supabase.table("messages").insert({ "session_id": session_id, "role": role, "content": content, "created_at": datetime.utcnow().isoformat() }).execute() print(f"✅ Message added successfully to session {session_id}") def update_session_title_if_needed(session_id: int, first_user_text: str): res = supabase.table("messages") \ .select("id") \ .eq("session_id", session_id) \ .eq("role", "user") \ .execute() if len(res.data) == 1: title = first_user_text.strip().split("\n")[0] title = (title[:50] + "…") if len(title) > 50 else title supabase.table("sessions") \ .update({"title": title or "New chat"}) \ .eq("id", session_id) \ .execute() def label_to_id(label: str | None) -> int | None: if not label: return None try: return int(label.split("•", 1)[0].strip()) except: return None def build_api_messages(session_id: int, system_message: str): msgs = [{"role": "system", "content": system_message.strip()}] msgs.extend(get_messages(session_id)) return msgs def get_client(model_choice: str): return InferenceClient(MODELS.get(model_choice, list(MODELS.values())[0]), token=HF_TOKEN) def load_dataset_by_name(name: str): if name == "The Stack": return load_dataset("bigcode/the-stack", split="train") if name == "CodeXGLUE": return load_dataset("google/code_x_glue_cc_code_to_code_trans", split="train") return None FRONTEND_KEYWORDS = ["react", "vue", "angular", "html", "css", "javascript", "tailwind", "recharts", "typescript"] def is_frontend_request(user_text: str) -> bool: return any(kw in user_text.lower() for kw in FRONTEND_KEYWORDS) # =========================== # Callbacks # =========================== def stop_generation_cb(): """Stop the current generation""" global stop_generation stop_generation = True return gr.update(visible=True), gr.update(visible=False) # Show send, hide stop def refresh_sessions_cb(): labels, _ = list_sessions() selected = labels[0] if labels else None visible = bool(selected) return ( gr.update(choices=labels, value=selected), gr.update(visible=False), # edit_title_box hidden gr.update(visible=visible), # rename_btn visible if session exists gr.update(visible=False), # save_title_btn hidden gr.update(visible=visible), # del_btn visible if session exists gr.update(visible=visible), # refresh_btn visible if session exists ) def new_chat_cb(): sid = create_session("New chat") labels, _ = list_sessions() selected = next((lbl for lbl in labels if lbl.startswith(f"{sid} ")), None) return ( gr.update(choices=labels, value=selected), [], "", # chatbot cleared, user box cleared gr.update(visible=False), # edit_title_box hidden gr.update(visible=True), # rename_btn visible gr.update(visible=False), # save_title_btn hidden gr.update(visible=True), # del_btn visible gr.update(visible=True), # refresh_btn visible gr.update(visible=True), # send_btn visible gr.update(visible=False) # stop_btn hidden ) def load_session_cb(selected_label): sid = label_to_id(selected_label) if not sid: return [], gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=True), gr.update(visible=False) return ( get_messages(sid), gr.update(visible=False), # edit_title_box hidden initially gr.update(visible=True), # rename_btn visible gr.update(visible=False), # save_title_btn hidden gr.update(visible=True), # del_btn visible gr.update(visible=True), # refresh_btn visible gr.update(visible=True), # send_btn visible gr.update(visible=False) # stop_btn hidden ) def delete_chat_cb(selected_label): sid = label_to_id(selected_label) if sid: delete_session(sid) labels, _ = list_sessions() selected = labels[0] if labels else None visible = bool(selected) return ( gr.update(choices=labels, value=selected), [], # chatbot cleared gr.update(visible=False), # edit_title_box hidden gr.update(visible=visible), # rename_btn visible if sessions exist gr.update(visible=False), # save_title_btn hidden gr.update(visible=visible), # del_btn visible if sessions exist gr.update(visible=visible), # refresh_btn visible if sessions exist gr.update(visible=True), # send_btn visible gr.update(visible=False) # stop_btn hidden ) # --- New Rename Callbacks --- def start_rename_cb(): """Show the rename input box and save button, hide rename button""" return ( gr.update(visible=True), # edit_title_box visible gr.update(visible=False), # rename_btn hidden gr.update(visible=True) # save_title_btn visible ) def save_rename_cb(selected_label, new_title): """Save the new chat title to database""" sid = label_to_id(selected_label) if not sid or not new_title or not new_title.strip(): # If no valid title, just revert UI return ( gr.update(visible=False), # edit_title_box hidden gr.update(visible=True), # rename_btn visible gr.update(visible=False), # save_title_btn hidden gr.update() # session_list unchanged ) # Update the title in database supabase.table("sessions") \ .update({"title": new_title.strip()}) \ .eq("id", sid) \ .execute() # Refresh the session list labels, _ = list_sessions() selected = next((lbl for lbl in labels if lbl.startswith(f"{sid} ")), None) return ( gr.update(visible=False), # edit_title_box hidden gr.update(visible=True), # rename_btn visible gr.update(visible=False), # save_title_btn hidden gr.update(choices=labels, value=selected) # updated session list ) def cancel_rename_cb(): """Cancel rename operation (if user clicks away)""" return ( gr.update(visible=False), # edit_title_box hidden gr.update(visible=True), # rename_btn visible gr.update(visible=False) # save_title_btn hidden ) # --- Send with stop capability --- def send_cb(user_text, selected_label, chatbot_msgs, system_message, max_tokens, temperature, top_p, model_choice, dataset_choice, *args): global stop_generation stop_generation = False # Reset stop flag at start sid = label_to_id(selected_label) or create_session("New chat") add_message(sid, "user", user_text) update_session_title_if_needed(sid, user_text) display_msgs = chatbot_msgs[:] + [{"role": "user", "content": user_text}] # Check for frontend request if is_frontend_request(user_text): apology = "⚠️ I'm a backend-focused assistant and cannot provide frontend code." display_msgs.append({"role": "assistant", "content": apology}) add_message(sid, "assistant", apology) yield display_msgs, "", selected_label, gr.update(visible=True), gr.update(visible=False) return display_msgs.append({"role": "assistant", "content": "…"}) # Show stop button, hide send button yield display_msgs, "", selected_label, gr.update(visible=False), gr.update(visible=True) client = get_client(model_choice) api_messages = build_api_messages(sid, system_message) partial = "" try: for chunk in client.chat_completion( messages=api_messages, max_tokens=int(max_tokens), temperature=float(temperature), top_p=float(top_p), stream=True ): # Check if stop was requested if stop_generation: display_msgs[-1]["content"] = partial + "\n\n[Generation stopped by user]" yield display_msgs, "", selected_label, gr.update(visible=True), gr.update(visible=False) add_message(sid, "assistant", partial + "\n\n[Generation stopped by user]") return if not hasattr(chunk, "choices") or not chunk.choices: continue choice = chunk.choices[0] delta = getattr(getattr(choice, "delta", None), "content", None) or getattr(getattr(choice, "message", None), "content", None) or "" if delta: partial += delta display_msgs[-1]["content"] = partial yield display_msgs, "", selected_label, gr.update(visible=False), gr.update(visible=True) # Only add message if not stopped if not stop_generation: add_message(sid, "assistant", partial) yield display_msgs, "", selected_label, gr.update(visible=True), gr.update(visible=False) else: yield display_msgs, "", selected_label, gr.update(visible=True), gr.update(visible=False) except Exception as e: display_msgs[-1]["content"] = f"⚠️ Error: {e}" yield display_msgs, "", selected_label, gr.update(visible=True), gr.update(visible=False) # =========================== # Init # =========================== init_db() labels, _ = list_sessions() if not labels: create_session("New chat") labels, _ = list_sessions() default_selected = labels[0] if labels else None # =========================== # App UI (YOUR ORIGINAL UI WITH STOP BUTTON) # =========================== with gr.Blocks(title="Backend-Focused LLaMA/Mistral CRUD Assistant", theme=gr.themes.Soft()) as demo: # --- Custom CSS (your original CSS) --- gr.HTML(""" """) with gr.Row(equal_height=True): with gr.Column(scale=1, min_width=260): session_list = gr.Radio(choices=labels, value=default_selected, label="Your Chats", interactive=True) with gr.Row(elem_id="session-actions-row"): new_btn = gr.Button("New Chat") rename_btn = gr.Button("Rename", visible=bool(default_selected)) save_title_btn = gr.Button("Save", visible=False) del_btn = gr.Button("Delete", visible=bool(default_selected)) refresh_btn = gr.Button("Refresh", visible=bool(default_selected)) edit_title_box = gr.Textbox(label="Edit Chat Name", placeholder="Type new chat name…", visible=False) # Model selection model_choice = gr.Dropdown(choices=list(MODELS.keys()), value=list(MODELS.keys())[0], label="Choose a model", interactive=True) # Dataset selection dataset_choice = gr.Dropdown(choices=DATASETS, value=DATASETS[0], label="Select a dataset", interactive=True) # Generation settings with gr.Group(elem_classes="system-message-compact"): system_box = gr.Textbox(value=SYSTEM_DEFAULT, label="System message", lines=1) with gr.Group(elem_classes="compact-sliders"): max_tokens = gr.Slider(256, 4096, value=1200, step=16, label="Max tokens") temperature = gr.Slider(0.0, 2.0, value=0.25, step=0.05, label="Temperature") top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p") with gr.Column(scale=3): chatbot = gr.Chatbot(label="Assistant", type="messages") with gr.Row(elem_classes="user-input-row"): user_box = gr.Textbox(placeholder="Describe your CRUD/backend task…", lines=3, scale=5) send_btn = gr.Button("▶️", variant="primary", scale=1, visible=True) stop_btn = gr.Button("⏹️", variant="stop", scale=1, visible=False, elem_classes="stop-btn") # --- Wire callbacks --- # Stop button stop_btn.click( stop_generation_cb, outputs=[send_btn, stop_btn] ) # Refresh button refresh_btn.click( refresh_sessions_cb, outputs=[session_list, edit_title_box, rename_btn, save_title_btn, del_btn, refresh_btn] ) # New chat button new_btn.click( new_chat_cb, outputs=[session_list, chatbot, user_box, edit_title_box, rename_btn, save_title_btn, del_btn, refresh_btn, send_btn, stop_btn] ) # Delete button del_btn.click( delete_chat_cb, inputs=session_list, outputs=[session_list, chatbot, edit_title_box, rename_btn, save_title_btn, del_btn, refresh_btn, send_btn, stop_btn] ) # Session selection change session_list.change( load_session_cb, inputs=session_list, outputs=[chatbot, edit_title_box, rename_btn, save_title_btn, del_btn, refresh_btn, send_btn, stop_btn] ) # Rename button - starts the rename process rename_btn.click( start_rename_cb, outputs=[edit_title_box, rename_btn, save_title_btn] ) # Save button - saves the new title save_title_btn.click( save_rename_cb, inputs=[session_list, edit_title_box], outputs=[edit_title_box, rename_btn, save_title_btn, session_list] ) # Pressing Enter in the rename box also saves edit_title_box.submit( save_rename_cb, inputs=[session_list, edit_title_box], outputs=[edit_title_box, rename_btn, save_title_btn, session_list] ) # Clicking away from the rename box (blur) cancels edit_title_box.blur( cancel_rename_cb, outputs=[edit_title_box, rename_btn, save_title_btn] ) # Send message send_btn.click( send_cb, inputs=[user_box, session_list, chatbot, system_box, max_tokens, temperature, top_p, model_choice, dataset_choice], outputs=[chatbot, user_box, session_list, send_btn, stop_btn] ) # Enter key in user box user_box.submit( send_cb, inputs=[user_box, session_list, chatbot, system_box, max_tokens, temperature, top_p, model_choice, dataset_choice], outputs=[chatbot, user_box, session_list, send_btn, stop_btn] ) if __name__ == "__main__": demo.launch()