Spaces:
Running
Running
| import os | |
| import time | |
| from openai import OpenAI | |
| import gradio as gr | |
| from auth import register_user, login_user, logout_user, save_chats, delete_account | |
| from response_modes import chat_stream | |
| import mymsg | |
| from retry import retry # Add to requirements.txt: retry | |
| # ================== SETTINGS ================== | |
| API_KEY = os.environ.get("API_KEY") | |
| HF_TOKEN = os.environ.get("HF_TOKEN2") | |
| username = os.environ.get("username") | |
| password = os.environ.get("password") | |
| USER_DATA_REPO = "qwendashcoder/fluade_1" | |
| auth_globals = {"hf_token": HF_TOKEN, "user_repo": USER_DATA_REPO} | |
| if not HF_TOKEN: | |
| raise ValueError("HF_TOKEN secret not set in HF Spaces settings!") | |
| if not API_KEY: | |
| print("WARNING: API_KEY not set—model loading will fail!") # Log for debugging | |
| client = OpenAI(base_url="https://integrate.api.nvidia.com/v1", api_key=API_KEY) | |
| # Model list (deduplicated, sorted, removed potential invalids) | |
| model_choices = sorted(set([ | |
| "meta/llama-3.1-405b-instruct", "meta/llama-3.1-70b-instruct", "meta/llama-3.1-8b-instruct", | |
| "google/gemma-2-27b-it", "google/gemma-2-9b-it", "google/codegemma-7b-it", | |
| "mistralai/mistral-large-3-675b-instruct-2512", | |
| "mistralai/mixtral-8x22b-v0.1", "mistralai/devstral-2-123b-instruct-2512", | |
| "qwen/qwen3-coder-480b-a35b-instruct", "qwen/qwen3.5-397b-a17b", "qwen/qwen3-235b-a22b" | |
| ])) | |
| # Retry on transient failures | |
| def load_model(selected_model): | |
| print(f"[DEBUG] Loading model: {selected_model}") # Log to confirm function call | |
| try: | |
| # Test with timeout and stream=False | |
| test = client.chat.completions.create( | |
| model=selected_model, | |
| messages=[{"role": "user", "content": "Say 'test'"}], | |
| max_tokens=5, | |
| stream=False, | |
| timeout=15 # Added: Prevent hangs | |
| ) | |
| response = test.choices[0].message.content.strip() | |
| print(f"[DEBUG] Model response: {response}") # Log success | |
| if response: | |
| gr.Info("✅ Model loaded successfully!", duration=5) # Longer duration | |
| return selected_model, f"**✅ Loaded & Ready: {selected_model}**" | |
| except Exception as e: # Catch all, but log specifics | |
| error_msg = str(e) | |
| print(f"[ERROR] Model load failed: {error_msg}") # Log for HF Spaces debugging | |
| if "authentication" in error_msg.lower(): | |
| error_msg = "Invalid or missing API_KEY—check HF Secrets!" | |
| elif "timeout" in error_msg.lower(): | |
| error_msg = "API timeout—try again or check NVIDIA status." | |
| elif "not found" in error_msg.lower(): | |
| error_msg = "Model not available on NVIDIA—select another." | |
| gr.Warning(f"❌ Failed to load model: {error_msg}", duration=5) | |
| return None, f"**❌ Unavailable: {error_msg}**" | |
| # ================== MULTI-CHAT HELPERS ================== | |
| def create_new_chat(user_chats, logged_in_username, coding_files): | |
| if not logged_in_username: | |
| return user_chats, None, [], gr.update(choices=[], value=None), "⚠️ Please login first!", "" | |
| new_num = len(user_chats) + 1 | |
| new_title = f"Chat {new_num}" | |
| new_chat = {"title": new_title, "history": [], "context": ""} # Added: Initialize empty context | |
| new_user_chats = user_chats + [new_chat] | |
| choices = [chat["title"] for chat in new_user_chats] | |
| save_chats(logged_in_username, new_user_chats, auth_globals) # Persist changes | |
| return new_user_chats, new_title, [], gr.update(choices=choices, value=new_title), f"✅ New chat created: {new_title}", "" | |
| def switch_chat(selected_title, user_chats, coding_files): | |
| if not selected_title or not user_chats: | |
| return [], None, "" | |
| for chat in user_chats: | |
| if chat["title"] == selected_title: | |
| return chat["history"], selected_title, chat.get("context", "") | |
| return [], None, "" | |
| def delete_current_chat(current_title, user_chats, logged_in_username, coding_files): | |
| if not logged_in_username or not current_title: | |
| return user_chats, current_title, [], gr.update(), "⚠️ Nothing to delete!", "" | |
| new_user_chats = [chat for chat in user_chats if chat["title"] != current_title] | |
| save_chats(logged_in_username, new_user_chats, auth_globals) # Persist changes | |
| choices = [chat["title"] for chat in new_user_chats] | |
| # Inside delete_current_chat, replace ONLY the new_history part: | |
| if new_user_chats: | |
| new_title = new_user_chats[-1]["title"] | |
| new_history = new_user_chats[-1]["history"] # ← direct dict history | |
| new_context = new_user_chats[-1].get("context", "") | |
| else: | |
| new_title = None | |
| new_history = [] | |
| new_context = "" | |
| return new_user_chats, new_title, new_history, gr.update(choices=choices, value=new_title), f"🗑️ Deleted {current_title}", new_context | |
| def save_chat_context(new_context, user_chats, current_chat_title, logged_in_username): | |
| if not logged_in_username or not current_chat_title: | |
| return user_chats, "⚠️ No chat selected!" | |
| if len(new_context) > 2000: | |
| return user_chats, "⚠️ Context too long (max 2000 chars)!" | |
| for chat in user_chats: | |
| if chat["title"] == current_chat_title: | |
| chat["context"] = new_context | |
| break | |
| save_chats(logged_in_username, user_chats, auth_globals) | |
| return user_chats, "✅ Context saved!" | |
| def enforce_cooldown(last_time): | |
| now = time.time() | |
| if now - last_time < 3: | |
| gr.Warning("⏳ Please wait 3 seconds before clicking again.") | |
| return False, last_time | |
| return True, now | |
| def delete_last_message(user_chats, current_title, logged_in_username, auth_globals): | |
| if not current_title or not user_chats: | |
| return user_chats, "Nothing to delete" | |
| for chat in user_chats: | |
| if chat["title"] == current_title and len(chat["history"]) >= 2: | |
| chat["history"] = chat["history"][:-2] # remove last user + assistant | |
| break | |
| save_chats(logged_in_username, user_chats, auth_globals) | |
| return user_chats, "🗑️ Last message deleted" | |
| def edit_last_prompt(user_chats, current_title): | |
| if not current_title or not user_chats: | |
| return "", "No previous prompt" | |
| for chat in user_chats: | |
| if chat["title"] == current_title and len(chat["history"]) >= 2: | |
| last_user_msg = chat["history"][-2]["content"] | |
| return last_user_msg, "Ready to edit & re-run" | |
| return "", "No previous prompt" | |
| def update_token_display(user_chats, current_title): | |
| total = 0 | |
| for chat in user_chats or []: | |
| if chat["title"] == current_title: | |
| for msg in chat.get("history", []): | |
| total += msg.get("tokens", 0) | |
| return f"**Tokens used this chat:** {total:,}" | |
| # ================== HELPER: No conversion needed (Gradio expects messages format) ================== | |
| def format_history_for_chatbot(history): | |
| """Return history as-is (already dict format)""" | |
| return history or [] | |
| def refresh_file_list(user_chats, current_title): | |
| if not current_title or not user_chats: | |
| return gr.update(choices=[], value=None) | |
| for chat in user_chats: | |
| if chat["title"] == current_title: | |
| files = list(chat.get("coding_files", {}).keys()) | |
| return gr.update(choices=files, value=files[0] if files else None) | |
| return gr.update(choices=[], value=None) | |
| def load_file_to_editor(selected_file, user_chats, current_title): | |
| if not selected_file or not current_title: | |
| return "", "python" | |
| for chat in user_chats: | |
| if chat["title"] == current_title: | |
| f = chat.get("coding_files", {}).get(selected_file, {}) | |
| return f.get("content", ""), f.get("language", "python") | |
| return "", "python" | |
| def save_file_content(selected_file, new_content, new_language, user_chats, current_title, logged_in_username, auth_globals): | |
| if not selected_file or not current_title or not logged_in_username: | |
| return user_chats, "❌ No file selected" | |
| for chat in user_chats: | |
| if chat["title"] == current_title: | |
| if "coding_files" not in chat: | |
| chat["coding_files"] = {} | |
| chat["coding_files"][selected_file] = {"content": new_content, "language": new_language} | |
| break | |
| save_chats(logged_in_username, user_chats, auth_globals) | |
| return user_chats, f"✅ Saved {selected_file}" | |
| def add_uploaded_files(uploaded_files, user_chats, current_title, logged_in_username, auth_globals): | |
| if not uploaded_files or not current_title: | |
| return user_chats, gr.update(), "❌ No chat selected" | |
| for chat in user_chats: | |
| if chat["title"] == current_title: | |
| if "coding_files" not in chat: | |
| chat["coding_files"] = {} | |
| for path in uploaded_files if isinstance(uploaded_files, list) else [uploaded_files]: | |
| name = os.path.basename(path) | |
| with open(path, "r", encoding="utf-8") as f: | |
| content = f.read() | |
| lang = "python" if name.endswith(".py") else "javascript" if name.endswith(".js") else "html" if name.endswith(".html") else "text" | |
| chat["coding_files"][name] = {"content": content, "language": lang} | |
| break | |
| save_chats(logged_in_username, user_chats, auth_globals) | |
| return user_chats, refresh_file_list(user_chats, current_title), "✅ Files added to persistent storage" | |
| def delete_selected_file(selected_file, user_chats, current_title, logged_in_username, auth_globals): | |
| if not selected_file or not current_title: | |
| return user_chats, gr.update(), "❌ Nothing to delete" | |
| for chat in user_chats: | |
| if chat["title"] == current_title and "coding_files" in chat: | |
| chat["coding_files"].pop(selected_file, None) | |
| break | |
| save_chats(logged_in_username, user_chats, auth_globals) | |
| return user_chats, refresh_file_list(user_chats, current_title), f"🗑️ Deleted {selected_file}" | |
| # ================== CHAT HANDLER WRAPPER (with file limit) ================== | |
| def handle_chat(msg, chatbot, current_model, logged_in_username, user_chats_state, current_chat_title, mode_radio, auth_globals_state, file_upload, current_chat_context): | |
| if not current_model: | |
| gr.Warning("⚠️ Load a model first!") | |
| return chatbot | |
| # 3-second cooldown | |
| ok, new_time = enforce_cooldown(last_interaction_time.value) | |
| if not ok: | |
| return chatbot | |
| last_interaction_time.value = new_time | |
| is_generating.value = True | |
| stop_btn.visible = True | |
| internal_history = chatbot[:] if chatbot else [] | |
| try: | |
| stream_gen = chat_stream( | |
| msg, internal_history, current_model, logged_in_username, | |
| user_chats_state, current_chat_title, mode_radio, | |
| auth_globals_state, file_upload, current_chat_context | |
| ) | |
| for updated_history in stream_gen: | |
| yield updated_history # ← This makes streaming work | |
| except Exception as e: | |
| gr.Warning(f"⚠️ Error: {str(e)[:250]}") | |
| yield chatbot | |
| finally: | |
| is_generating.value = False | |
| stop_btn.visible = False | |
| # Must return the chatbot at the end (Gradio requirement) | |
| return chatbot | |
| # ================== GUI ================== | |
| with gr.Blocks(title="NVIDIA AI Chat") as demo: | |
| gr.Markdown("# 🟢 NVIDIA AI Chat\n**School Project • Pre-approved usernames only • Max 20 Accounts globally**") | |
| gr.Markdown(f"{mymsg.welcome_caution()}") | |
| with gr.Accordion("📖 How to Use", open=True): | |
| gr.Markdown(f"{mymsg.welcome_message(0)}") | |
| current_model = gr.State(None) | |
| logged_in_username = gr.State(None) | |
| user_chats_state = gr.State([]) | |
| current_chat_title = gr.State(None) | |
| current_chat_context = gr.State("") # Added: State for current chat's context | |
| auth_globals_state = gr.State(auth_globals) | |
| coding_files_state = gr.State({}) # { "filename.py": {"content": "...", "language": "python"} } | |
| current_file_name = gr.State(None) | |
| last_interaction_time = gr.State(0.0) # for 3-second cooldown | |
| is_generating = gr.State(False) # for Stop button | |
| status_label = gr.Markdown("⚠️ No model loaded. Login + create/select a chat to begin.", elem_id="status-bold") # Made more prominent | |
| with gr.Row(): | |
| with gr.Column(scale=1, min_width=200): | |
| gr.Markdown("### 📋 Your Chats") | |
| new_chat_btn = gr.Button("➕ New Chat", variant="primary", size="large") | |
| delete_chat_btn = gr.Button("🗑️ Delete Current", variant="stop", size="large") | |
| chat_dropdown = gr.Dropdown(choices=[], label="Switch Chat", interactive=True, value=None) | |
| with gr.Column(scale=4): | |
| with gr.Row(): | |
| model_dropdown = gr.Dropdown(choices=model_choices, value=model_choices[0], label="🎯 Select Model", interactive=True, scale=4) | |
| load_btn = gr.Button("Load & Test Model", variant="primary", scale=1) | |
| mode_radio = gr.Radio( | |
| choices=["Default Chat", "Fast Search", "Deep Think", "💻 Coding Mode"], | |
| value="Default Chat", | |
| label="🧠 Response Mode", | |
| interactive=True | |
| ) | |
| chatbot = gr.Chatbot(height=480, buttons=["copy", "copy_all"]) | |
| with gr.Row(): # Added: Context editor section | |
| chat_context_tb = gr.Textbox(label="Instructions - Additional instruction for each prompt (saved per chat)", | |
| lines=3, max_lines=10, | |
| placeholder="e.g. 'Use Filipino and be funny' or 'Always answer in bullet points and be super concise'", | |
| interactive=True, visible=False) | |
| save_context_btn = gr.Button("💾 Save Context", variant="secondary", visible=False) | |
| # ================== CODING MODE FILE MANAGER (visible only in Coding Mode) ================== | |
| with gr.Row(visible=False) as coding_manager_row: | |
| with gr.Column(scale=1, min_width=280): | |
| gr.Markdown("### 📁 Persistent Project Files\n(saved per chat • only in Coding Mode)") | |
| file_upload_coding = gr.File( | |
| label="Upload files", | |
| file_types=[".py",".js",".html",".css",".txt",".md",".json",".ts"], | |
| file_count="multiple", | |
| type="filepath" | |
| ) | |
| add_files_btn = gr.Button("➕ Add / Replace files", variant="primary") | |
| file_list_dropdown = gr.Dropdown(choices=[], label="Select file to edit", interactive=True) | |
| delete_file_btn = gr.Button("🗑️ Delete selected file", variant="stop", size="sm") | |
| with gr.Column(scale=4): | |
| file_editor = gr.Code( | |
| label="Editor", | |
| value="", | |
| language="python", | |
| lines=24, | |
| interactive=True, | |
| show_label=True | |
| ) | |
| with gr.Row(): | |
| language_selector = gr.Dropdown( | |
| choices=["python", "javascript", "typescript", "html", "css", "markdown", "json", "text"], | |
| value="python", | |
| label="Language", | |
| scale=1 | |
| ) | |
| save_file_btn = gr.Button("💾 Save file", variant="primary", scale=2) | |
| with gr.Row(): | |
| msg = gr.Textbox(placeholder="Ask anything...", show_label=False, scale=9) | |
| file_upload = gr.File( | |
| label="Attach up to 5 files (.txt, .pdf, .docx)", | |
| file_types=[".txt", ".pdf", ".docx"], | |
| file_count="multiple", | |
| type="filepath" | |
| ) | |
| send_btn = gr.Button("➤", variant="primary", scale=1) | |
| with gr.Row(): | |
| delete_last_btn = gr.Button("🗑️ Delete Last Message", variant="stop", size="sm") | |
| edit_last_btn = gr.Button("✏️ Edit & Re-run Last Prompt", variant="secondary", size="sm") | |
| stop_btn = gr.Button("⏹️ Stop Generation", variant="stop", visible=False) | |
| # Token display (optional but useful) | |
| token_display = gr.Markdown("**Tokens used this chat:** 0", visible=False) | |
| # Authentication (clean version - no phone/OTP) | |
| with gr.Row(): | |
| username_input = gr.Textbox(label="Username (will be forced lowercase)", scale=2, placeholder="admin") | |
| password_input = gr.Textbox(label="Password", type="password", scale=2) | |
| with gr.Row(): | |
| register_btn = gr.Button("Register", variant="secondary") | |
| login_btn = gr.Button("Login", variant="primary") | |
| logout_btn = gr.Button("Logout", visible=False, variant="stop") | |
| with gr.Accordion("⚠️ Danger Zone", open=False): | |
| delete_acc_btn = gr.Button("🗑️ Delete My Account", variant="stop", size="large", visible=False) | |
| with gr.Row(visible=False) as confirm_row: | |
| confirm_username_input = gr.Textbox(label="Type exact username to confirm", placeholder="yourusername") | |
| confirm_delete_btn = gr.Button("YES, DELETE FOREVER", variant="stop") | |
| # ================== EVENTS ================== | |
| load_btn.click(load_model, model_dropdown, [current_model, status_label]) | |
| # New Chat | |
| new_chat_btn.click( | |
| create_new_chat, | |
| [user_chats_state, logged_in_username], | |
| [user_chats_state, current_chat_title, chatbot, chat_dropdown, status_label, current_chat_context] | |
| ) | |
| # Delete Current Chat | |
| delete_chat_btn.click( | |
| delete_current_chat, | |
| [current_chat_title, user_chats_state, logged_in_username], | |
| [user_chats_state, current_chat_title, chatbot, chat_dropdown, status_label, current_chat_context] | |
| ) | |
| # Switch Chat | |
| chat_dropdown.change( | |
| switch_chat, | |
| [chat_dropdown, user_chats_state], | |
| [chatbot, current_chat_title, current_chat_context] | |
| ) | |
| register_btn.click( | |
| register_user, | |
| [username_input, password_input, logged_in_username, auth_globals_state], | |
| [user_chats_state, current_chat_title, chatbot, chat_dropdown, logged_in_username, status_label, logout_btn, delete_acc_btn, confirm_row, chat_context_tb, save_context_btn] | |
| ).then(lambda: gr.Info("✅ Registration successful!", duration=4), None, None).then(lambda: ("",""), None, [username_input, password_input]) | |
| login_btn.click( | |
| login_user, | |
| [username_input, password_input, logged_in_username, auth_globals_state], | |
| [user_chats_state, current_chat_title, chatbot, chat_dropdown, logged_in_username, status_label, logout_btn, delete_acc_btn, confirm_row, chat_context_tb, save_context_btn] | |
| ).then(lambda: gr.Info("✅ Login successful!", duration=4), None, None).then(lambda: ("",""), None, [username_input, password_input]) | |
| logout_btn.click(logout_user, [], [user_chats_state, current_chat_title, chatbot, chat_dropdown, logged_in_username, status_label, logout_btn, delete_acc_btn, confirm_row, chat_context_tb, save_context_btn]) | |
| delete_acc_btn.click(lambda: gr.update(visible=True), None, confirm_row) | |
| confirm_delete_btn.click(delete_account, [confirm_username_input, logged_in_username, auth_globals_state], [status_label, logged_in_username, logout_btn, delete_acc_btn, confirm_row, user_chats_state, current_chat_title, chatbot, chat_dropdown, chat_context_tb, save_context_btn]) | |
| save_context_btn.click(save_chat_context, [chat_context_tb, user_chats_state, current_chat_title, logged_in_username], [user_chats_state, status_label]) | |
| # Cooldown + action buttons | |
| delete_last_btn.click( | |
| delete_last_message, | |
| [user_chats_state, current_chat_title, logged_in_username, auth_globals_state], | |
| [user_chats_state, status_label] | |
| ).then( | |
| lambda: gr.update(visible=True), None, token_display | |
| ).then(update_token_display, [user_chats_state, current_chat_title], token_display) | |
| edit_last_btn.click( | |
| edit_last_prompt, | |
| [user_chats_state, current_chat_title], | |
| [msg, status_label] | |
| ) | |
| stop_btn.click( | |
| lambda: (False, gr.update(visible=False)), | |
| None, | |
| [is_generating, stop_btn] | |
| ) | |
| # Show token display in Coding Mode and after first message | |
| mode_radio.change( | |
| lambda m: gr.update(visible=(m == "💻 Coding Mode")), | |
| mode_radio, | |
| token_display | |
| ) | |
| # Send button + Enter | |
| # Enter key | |
| msg.submit( | |
| handle_chat, | |
| [msg, chatbot, current_model, logged_in_username, user_chats_state, | |
| current_chat_title, mode_radio, auth_globals_state, file_upload, current_chat_context], | |
| chatbot | |
| ).then(lambda: "", None, msg) \ | |
| .then(update_token_display, [user_chats_state, current_chat_title], token_display) | |
| # Send button | |
| send_btn.click( | |
| handle_chat, | |
| [msg, chatbot, current_model, logged_in_username, user_chats_state, | |
| current_chat_title, mode_radio, auth_globals_state, file_upload, current_chat_context], | |
| chatbot | |
| ).then(lambda: "", None, msg) \ | |
| .then(update_token_display, [user_chats_state, current_chat_title], token_display) | |
| # File manager events | |
| add_files_btn.click( | |
| add_uploaded_files, | |
| [file_upload_coding, user_chats_state, current_chat_title, logged_in_username, auth_globals_state], | |
| [user_chats_state, file_list_dropdown, status_label] | |
| ) | |
| file_list_dropdown.change( | |
| load_file_to_editor, | |
| [file_list_dropdown, user_chats_state, current_chat_title], | |
| [file_editor, language_selector] | |
| ) | |
| save_file_btn.click( | |
| save_file_content, | |
| [file_list_dropdown, file_editor, language_selector, user_chats_state, current_chat_title, logged_in_username, auth_globals_state], | |
| [user_chats_state, status_label] | |
| ).then(refresh_file_list, [user_chats_state, current_chat_title], file_list_dropdown) | |
| delete_file_btn.click( | |
| delete_selected_file, | |
| [file_list_dropdown, user_chats_state, current_chat_title, logged_in_username, auth_globals_state], | |
| [user_chats_state, file_list_dropdown, status_label] | |
| ) | |
| language_selector.change( | |
| lambda lang: gr.update(language=lang), | |
| language_selector, | |
| file_editor | |
| ) | |
| # Show/hide file manager when mode changes | |
| mode_radio.change( | |
| lambda m: gr.update(visible=(m == "💻 Coding Mode")), | |
| mode_radio, | |
| coding_manager_row | |
| ).then( | |
| refresh_file_list, | |
| [user_chats_state, current_chat_title], | |
| file_list_dropdown | |
| ) | |
| demo.launch(auth=(username, password), theme=gr.themes.Soft(), server_name="0.0.0.0", debug=True, show_error=True) |