import os import sqlite3 import time from datetime import datetime import gradio as gr from huggingface_hub import InferenceClient # --------------------------- # Config # --------------------------- MODELS = { "Meta LLaMA 3.1 (8B Instruct)": "meta-llama/Llama-3.1-8B-Instruct", "Mistral 7B Instruct": "mistralai/Mistral-7B-Instruct-v0.3", } HF_TOKEN = os.getenv("HF_TOKEN") # set this in your Space's Secrets DB_PATH = "history.db" SYSTEM_DEFAULT = ( "You are a coding assistant. Respond only with clean and complete code " "unless explanation is explicitly requested. Prefer full CRUD scaffolds, " "with files, paths, and commands when asked." ) # --------------------------- # 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): """Return the right InferenceClient for the chosen model.""" model_id = MODELS.get(model_choice, list(MODELS.values())[0]) return InferenceClient(model_id, token=HF_TOKEN) # --------------------------- # Gradio Callbacks # --------------------------- def refresh_sessions_cb(): labels, _ = list_sessions() selected = labels[0] if labels else None return gr.update(choices=labels, value=selected) 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), [], "") def load_session_cb(selected_label): sid = label_to_id(selected_label) if not sid: return [] return get_messages(sid) 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 return gr.update(choices=labels, value=selected), [] def send_cb(user_text, selected_label, chatbot_msgs, system_message, max_tokens, temperature, top_p, model_choice): 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) add_message(sid, "user", user_text) update_session_title_if_needed(sid, user_text) api_messages = build_api_messages(sid, system_message) display_msgs = get_messages(sid) display_msgs.append({"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, ): delta = chunk.choices[0].delta.content or "" if delta: partial += delta display_msgs[-1]["content"] = partial yield (display_msgs, "", selected_label) add_message(sid, "assistant", partial) except Exception as e: err = f"⚠️ Error: {str(e)}" display_msgs[-1]["content"] = err yield (display_msgs, "", selected_label) def regenerate_cb(selected_label, system_message, max_tokens, temperature, top_p, model_choice): sid = label_to_id(selected_label) if sid is None: return [], "" msgs = get_messages(sid) if not msgs: return [], "" 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) 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, ): delta = chunk.choices[0].delta.content or "" 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 # --------------------------- # 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="LLaMA/Mistral CRUD Automation (with History)", theme=gr.themes.Soft()) as demo: # --- Updated CSS to make ALL buttons green --- gr.HTML(""" """) gr.Markdown("## 🦙🤖 LLaMA & Mistral CRUD Automation — with Persistent History") with gr.Row(): with gr.Column(scale=1, min_width=260): gr.Markdown("### 📁 Sessions") session_list = gr.Radio( choices=labels, value=default_selected, label="Your chats", interactive=True ) with gr.Row(): new_btn = gr.Button("➕ New Chat", variant="primary") del_btn = gr.Button("🗑️ Delete", variant="stop") refresh_btn = gr.Button("🔄 Refresh", variant="secondary") gr.Markdown("### 🤖 Model Selection") model_choice = gr.Dropdown( choices=list(MODELS.keys()), value=list(MODELS.keys())[0], label="Choose a model", interactive=True ) gr.Markdown("### ⚙️ Generation Settings") system_box = gr.Textbox( value=SYSTEM_DEFAULT, label="System message", lines=4 ) 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=520, type="messages") with gr.Row(): user_box = gr.Textbox(placeholder="Describe your CRUD task…", lines=3, scale=5) with gr.Row(): send_btn = gr.Button("Send ▶️", variant="primary") regen_btn = gr.Button("Regenerate 🔁", variant="secondary") # Interactions refresh_btn.click(refresh_sessions_cb, outputs=session_list) new_btn.click(new_chat_cb, outputs=[session_list, chatbot, user_box]) del_btn.click(delete_chat_cb, inputs=session_list, outputs=[session_list, chatbot]) session_list.change(load_session_cb, inputs=session_list, outputs=chatbot) send_btn.click( send_cb, inputs=[user_box, session_list, chatbot, system_box, max_tokens, temperature, top_p, model_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], outputs=[chatbot, user_box, session_list] ) regen_btn.click( regenerate_cb, inputs=[session_list, system_box, max_tokens, temperature, top_p, model_choice], outputs=chatbot ) if __name__ == "__main__": demo.launch()