Spaces:
Runtime error
Runtime error
| 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(""" | |
| <style> | |
| /* ========== Compact section headings ========== */ | |
| .compact-heading h3 { | |
| font-size: 0.9rem !important; | |
| font-weight: 600 !important; | |
| margin: 0.3rem 0 !important; | |
| padding: 0 !important; | |
| } | |
| /* ========== Main top title smaller ========== */ | |
| .main-title h2 { | |
| font-size: 1.2rem !important; | |
| font-weight: 700 !important; | |
| margin-bottom: 0.6rem !important; | |
| } | |
| /* ========== Reduce spacing below sliders & dropdowns ========== */ | |
| .compact-sliders .gr-slider, | |
| .compact-sliders .gr-number, | |
| .compact-sliders .gr-dropdown { | |
| margin-bottom: 0.4rem !important; | |
| } | |
| /* ========== Compact the left panel column ========== */ | |
| .left-panel .gr-column { | |
| gap: 0.3rem !important; | |
| } | |
| /* ========== Tiny / pressed buttons ========== */ | |
| .tiny-btn .gr-button { | |
| font-size: 0.75rem !important; /* smaller text */ | |
| padding: 2px 6px !important; /* tight padding */ | |
| min-height: 24px !important; /* shorter button */ | |
| line-height: 1.1 !important; /* compact vertical spacing */ | |
| border-radius: 4px !important; /* slightly rounded */ | |
| } | |
| /* ========== Remove extra top/left page gaps ========== */ | |
| body, html { | |
| margin: 0 !important; | |
| padding: 0 !important; | |
| } | |
| .gradio-container { | |
| margin-top: -40px !important; /* was -10px, move up more */ | |
| margin-left: -40px !important; | |
| padding-top: 0 !important; | |
| padding-left: 0 !important; | |
| } | |
| .main-title, | |
| .gr-block:first-of-type { | |
| margin-top: -10px !important; /* was -5px, pull up further */ | |
| margin-left: -10 !important; | |
| padding-top: 0 !important; | |
| padding-left: 0 !important; | |
| } | |
| /* ========== Compact sliders: reduce vertical height and spacing ========== */ | |
| .compact-sliders .gr-slider { | |
| height: 28px !important; /* shorter track */ | |
| padding: 0 !important; /* remove inner padding */ | |
| margin-bottom: 0.3rem !important; /* less spacing below each slider */ | |
| } | |
| /* Compact slider label and value text */ | |
| .compact-sliders .gr-slider label, | |
| .compact-sliders .gr-slider span { | |
| font-size: 0.8rem !important; /* smaller text */ | |
| } | |
| /* Tighten the numeric input box on sliders */ | |
| .compact-sliders input[type="number"] { | |
| max-width: 70px !important; /* smaller numeric input width */ | |
| font-size: 0.8rem !important; | |
| padding: 2px 4px !important; | |
| } | |
| /* Stop button styling */ | |
| .stop-btn { | |
| background-color: #ff4444 !important; | |
| color: white !important; | |
| } | |
| .stop-btn:hover { | |
| background-color: #cc0000 !important; | |
| } | |
| </style> | |
| """) | |
| 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() |