app.py

import os
import sqlite3
from datetime import datetime
import gradio as gr
from huggingface_hub import InferenceClient
from datasets import load_dataset

# ---------------------------
# Config
# ---------------------------
MODELS = {
    "Meta LLaMA 3.1 (8B Instruct)": "meta-llama/Llama-3.1-8B-Instruct",
    "Mistral 7B Instruct": "mistralai/Mistral-7B-Instruct-v0.3",
}

DATASETS = ["The Stack", "CodeXGLUE"]  # Dropdown for dataset selection
HF_TOKEN = os.getenv("HF_TOKEN")  # Set in your Space's Secrets
DB_PATH = "history.db"

SYSTEM_DEFAULT = (
    "Specializes in databases, APIs, auth, CRUD."
    "Provides complete backend code scaffolds."
    "Provide full backend code scaffolds with files, paths, and commands. "
    "Declines frontend-heavy requests."
)

# ---------------------------
# DB Setup
# ---------------------------
def db():
    conn = sqlite3.connect(DB_PATH)
    conn.execute("PRAGMA journal_mode=WAL;")
    return conn

def init_db():
    conn = db()
    cur = conn.cursor()
    cur.execute("""
        CREATE TABLE IF NOT EXISTS sessions (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            title TEXT NOT NULL,
            created_at TEXT NOT NULL
        )
    """)
    cur.execute("""
        CREATE TABLE IF NOT EXISTS messages (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            session_id INTEGER NOT NULL,
            role TEXT NOT NULL,
            content TEXT NOT NULL,
            created_at TEXT NOT NULL,
            FOREIGN KEY(session_id) REFERENCES sessions(id) ON DELETE CASCADE
        )
    """)
    conn.commit()
    conn.close()

def create_session(title: str = "New chat") -> int:
    conn = db()
    cur = conn.cursor()
    cur.execute(
        "INSERT INTO sessions (title, created_at) VALUES (?, ?)",
        (title, datetime.utcnow().isoformat())
    )
    session_id = cur.lastrowid
    conn.commit()
    conn.close()
    return session_id

def delete_session(session_id: int):
    conn = db()
    cur = conn.cursor()
    cur.execute("DELETE FROM messages WHERE session_id = ?", (session_id,))
    cur.execute("DELETE FROM sessions WHERE id = ?", (session_id,))
    conn.commit()
    conn.close()

def list_sessions():
    conn = db()
    cur = conn.cursor()
    cur.execute("SELECT id, title FROM sessions ORDER BY id DESC")
    rows = cur.fetchall()
    conn.close()
    labels = [f"{r[0]} • {r[1]}" for r in rows]
    return labels, rows

def get_messages(session_id: int):
    conn = db()
    cur = conn.cursor()
    cur.execute("""
        SELECT role, content
        FROM messages
        WHERE session_id = ?
        ORDER BY id ASC
    """, (session_id,))
    rows = cur.fetchall()
    conn.close()
    msgs = [{"role": role, "content": content} for (role, content) in rows]
    return msgs

def add_message(session_id: int, role: str, content: str):
    conn = db()
    cur = conn.cursor()
    cur.execute(
        "INSERT INTO messages (session_id, role, content, created_at) VALUES (?, ?, ?, ?)",
        (session_id, role, content, datetime.utcnow().isoformat())
    )
    conn.commit()
    conn.close()

def update_session_title_if_needed(session_id: int, first_user_text: str):
    conn = db()
    cur = conn.cursor()
    cur.execute("SELECT COUNT(*) FROM messages WHERE session_id=? AND role='user'", (session_id,))
    count_users = cur.fetchone()[0]
    if count_users == 1:
        title = first_user_text.strip().split("\n")[0]
        title = (title[:50] + "…") if len(title) > 50 else title
        cur.execute("UPDATE sessions SET title=? WHERE id=?", (title or "New chat", session_id))
        conn.commit()
    conn.close()

# ---------------------------
# Helpers
# ---------------------------
def label_to_id(label: str | None) -> int | None:
    if not label:
        return None
    try:
        return int(label.split("•", 1)[0].strip())
    except Exception:
        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):
    model_id = MODELS.get(model_choice, list(MODELS.values())[0])
    return InferenceClient(model_id, token=HF_TOKEN)

def load_dataset_by_name(name: str):
    if name == "The Stack":
        return load_dataset("bigcode/the-stack", split="train")
    elif name == "CodeXGLUE":
        return load_dataset("google/code_x_glue_cc_code_to_code_trans", split="train")
    return None

# ---------------------------
# Gradio Callbacks
# ---------------------------
def refresh_sessions_cb():
    labels, _ = list_sessions()
    selected = labels[0] if labels else None
    visible = True if selected else False
    return (
        gr.update(choices=labels, value=selected),
        gr.update(visible=visible),
        gr.update(visible=visible),
        gr.update(visible=visible),
        gr.update(visible=visible),
    )

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: session_list, chatbot (empty), user_box cleared, then make edit/save/delete/refresh visible
    return (
        gr.update(choices=labels, value=selected),
        [],  # chatbot messages
        "",  # user_box cleared
        gr.update(visible=True),  # edit_title_box visible
        gr.update(visible=True),  # save_title_btn visible
        gr.update(visible=True),  # del_btn visible
        gr.update(visible=True),  # refresh_btn visible
    )

def load_session_cb(selected_label):
    sid = label_to_id(selected_label)
    if not sid:
        # no valid selection -> hide controls
        return [], gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
    msgs = get_messages(sid)
    return (
        msgs,
        gr.update(visible=True),
        gr.update(visible=True),
        gr.update(visible=True),
        gr.update(visible=True),
    )

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 = True if selected else False
    return (
        gr.update(choices=labels, value=selected),
        [],  # chatbot cleared
        gr.update(visible=visible),
        gr.update(visible=visible),
        gr.update(visible=visible),
        gr.update(visible=visible),
    )

FRONTEND_KEYWORDS = [
    "react", "vue", "angular", "html", "css", "javascript", "tailwind", "recharts", "typescript"
]

def is_frontend_request(user_text: str) -> bool:
    text_lower = user_text.lower()
    return any(kw in text_lower for kw in FRONTEND_KEYWORDS)

# --- Fixed send_cb to show user message ---
def send_cb(user_text, selected_label, chatbot_msgs, system_message, max_tokens, temperature, top_p, model_choice, dataset_choice, *args):
    sid = label_to_id(selected_label)
    if sid is None:
        sid = create_session("New chat")
        labels, _ = list_sessions()
        selected_label = next((lbl for lbl in labels if lbl.startswith(f"{sid} ")), None)

    # Save user message
    add_message(sid, "user", user_text)
    update_session_title_if_needed(sid, user_text)

    display_msgs = chatbot_msgs[:]
    display_msgs.append({"role": "user", "content": user_text})

    # Check for frontend-heavy 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)
        return

    # Normal backend response
    display_msgs.append({"role": "assistant", "content": "…"})
    yield (display_msgs, "", selected_label)

    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,
        ):
            # --- FIX: handle models that send empty chunks or use message instead of delta ---
            if not hasattr(chunk, "choices") or not chunk.choices:
                continue
            choice = chunk.choices[0]
            delta = ""
            if hasattr(choice, "delta") and choice.delta and getattr(choice.delta, "content", None) is not None:
                delta = choice.delta.content
            elif hasattr(choice, "message") and getattr(choice.message, "content", None) is not None:
                delta = choice.message.content

            if delta:
                partial += delta
                display_msgs[-1]["content"] = partial
                yield (display_msgs, "", selected_label)

        add_message(sid, "assistant", partial)

    except Exception as e:
        display_msgs[-1]["content"] = f"⚠️ Error: {str(e)}"
        yield (display_msgs, "", selected_label)

def regenerate_cb(selected_label, system_message, max_tokens, temperature, top_p, model_choice, dataset_choice):
    sid = label_to_id(selected_label)
    if sid is None:
        return [], ""

    msgs = get_messages(sid)
    if not msgs:
        return [], ""

    # Remove the last assistant message if it exists (to regenerate it)
    if msgs and msgs[-1]["role"] == "assistant":
        conn = db()
        cur = conn.cursor()
        cur.execute("""
            DELETE FROM messages
            WHERE id = (
                SELECT id FROM messages
                WHERE session_id=?
                ORDER BY id DESC LIMIT 1
            )
        """, (sid,))
        conn.commit()
        conn.close()
        msgs = get_messages(sid)

    dataset = load_dataset_by_name(dataset_choice)
    api_messages = [{"role": "system", "content": system_message.strip()}] + msgs
    display_msgs = msgs + [{"role": "assistant", "content": ""}]

    client = get_client(model_choice)
    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,
        ):
            # --- FIX: handle models that send empty chunks or use message instead of delta ---
            if not hasattr(chunk, "choices") or not chunk.choices:
                continue
            choice = chunk.choices[0]
            delta = ""
            if hasattr(choice, "delta") and choice.delta and getattr(choice.delta, "content", None) is not None:
                delta = choice.delta.content
            elif hasattr(choice, "message") and getattr(choice.message, "content", None) is not None:
                delta = choice.message.content

            if delta:
                partial += delta
                display_msgs[-1]["content"] = partial
                yield display_msgs

        add_message(sid, "assistant", partial)

    except Exception as e:
        display_msgs[-1]["content"] = f"⚠️ Error: {str(e)}"
        yield (display_msgs, "", selected_label)

# ---------------------------
# App UI
# ---------------------------
init_db()
labels, _ = list_sessions()
if not labels:
    first_sid = create_session("New chat")
    labels, _ = list_sessions()

default_selected = labels[0] if labels else None

with gr.Blocks(title="Backend-Focused LLaMA/Mistral CRUD Assistant", theme=gr.themes.Soft()) as demo:
    # Minimal CSS for compact, professional look
    gr.HTML("""
    <style>
    /* compact action row below chat list */
    #session-actions-row .gr-button {
        margin-right: 6px;
    }
    /* keep settings compact */
    .compact-sliders .gr-slider,
    .compact-sliders .gr-number { margin-bottom: 0.5rem !important; }
    .system-message-compact .gr-textarea {
        min-height: 80px !important;
        max-height: 120px !important;
    }
    </style>
    """)

    gr.Markdown("## 🗄️ LLaMA & Mistral Backend-Focused CRUD Automation — with Persistent History")

    with gr.Row(equal_height=True):
        with gr.Column(scale=1, min_width=260):
            gr.Markdown("### 📁 Sessions")

            # Main session list
            session_list = gr.Radio(
                choices=labels,
                value=default_selected,
                label="Your Chats",
                interactive=True
            )

            # Row for actions (visible/hidden controlled via callbacks)
            with gr.Row(elem_id="session-actions-row"):
                new_btn = gr.Button("New Chat", elem_id="new-chat-btn")
                rename_btn = gr.Button("Rename", elem_id="rename-chat-btn", visible=False)
                save_title_btn = gr.Button("Save", elem_id="save-title-btn", visible=False)
                del_btn = gr.Button("Delete", elem_id="delete-chat-btn", visible=False)
                refresh_btn = gr.Button("Refresh", elem_id="refresh-btn", visible=False)

            # Rename textbox (hidden until Rename is clicked)
            edit_title_box = gr.Textbox(
                label="Edit Chat Name",
                placeholder="Type new chat name here…",
                visible=False
            )

            # Callback to open the rename box (show edit box + save button)
            def show_rename_ui():
                return gr.update(visible=True), gr.update(visible=True)

            rename_btn.click(fn=show_rename_ui, outputs=[edit_title_box, save_title_btn])

            # Callback for renaming -> updates DB and refreshes list, clears edit box
            def rename_session_cb(new_title, selected_label):
                sid = label_to_id(selected_label)
                if sid and new_title.strip():
                    conn = db()
                    cur = conn.cursor()
                    cur.execute("UPDATE sessions SET title=? WHERE id=?", (new_title.strip(), sid))
                    conn.commit()
                    conn.close()
                labels, _ = list_sessions()
                new_selected = next((lbl for lbl in labels if lbl.startswith(f"{sid} ")), None)
                # return updated session list selection and clear the edit box value
                return gr.update(choices=labels, value=new_selected), gr.update(value="")

            save_title_btn.click(
                rename_session_cb,
                inputs=[edit_title_box, session_list],
                outputs=[session_list, edit_title_box]
            )

            # -------- Model selection --------
            gr.Markdown("### 🤖 Model Selection")
            model_choice = gr.Dropdown(
                choices=list(MODELS.keys()),
                value=list(MODELS.keys())[0],
                label="Choose a model",
                interactive=True
            )

            # -------- Dataset selection --------
            gr.Markdown("### 📚 Dataset Selection")
            dataset_choice = gr.Dropdown(
                choices=DATASETS,
                value=DATASETS[0],
                label="Select a dataset",
                interactive=True
            )

            # -------- Generation settings --------
            gr.Markdown("### ⚙️ Generation Settings")
            with gr.Group(elem_classes="system-message-compact"):
                system_box = gr.Textbox(
                    value=SYSTEM_DEFAULT,
                    label="System message",
                    lines=3
                )

            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", height=500, type="messages")
            with gr.Row():
                user_box = gr.Textbox(placeholder="Describe your CRUD/backend task…", lines=3, scale=5)
            with gr.Row():
                send_btn = gr.Button("Send ▶️", variant="primary")
                regen_btn = gr.Button("Regenerate 🔁", variant="secondary")

    # Wire callbacks — updated outputs so visibility changes flow correctly
    refresh_btn.click(
        refresh_sessions_cb,
        outputs=[session_list, edit_title_box, save_title_btn, del_btn, refresh_btn],
    )
    new_btn.click(
        new_chat_cb,
        outputs=[session_list, chatbot, user_box, edit_title_box, save_title_btn, del_btn, refresh_btn],
    )
    del_btn.click(
        delete_chat_cb,
        inputs=session_list,
        outputs=[session_list, chatbot, edit_title_box, save_title_btn, del_btn, refresh_btn],
    )
    session_list.change(
        load_session_cb,
        inputs=session_list,
        outputs=[chatbot, edit_title_box, save_title_btn, del_btn, refresh_btn],
    )

    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]
    )

    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]
    )

    regen_btn.click(
        regenerate_cb,
        inputs=[session_list, system_box, max_tokens, temperature, top_p, model_choice, dataset_choice],
        outputs=chatbot
    )

if __name__ == "__main__":
    demo.launch()