Blame for app.py

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(tries=2, delay=1)  # 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)