Update app.py

app.py

@@ -13,25 +13,20 @@ MODELS = {
     "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
+DATASETS = ["The Stack", "CodeXGLUE"]
+HF_TOKEN = os.getenv("HF_TOKEN")
 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. "
+    "Specializes in databases, APIs, auth, CRUD. "
+    "Provides complete backend code scaffolds. "
     "Declines frontend-heavy requests."
 )
 
 # ---------------------------
 # DB Setup
 # ---------------------------
-def db():
-    conn = sqlite3.connect(DB_PATH)
-    conn.execute("PRAGMA journal_mode=WAL;")
-    return conn
-
+def db(): return sqlite3.connect(DB_PATH)
 def init_db():
     conn = db()
     cur = conn.cursor()
@@ -62,10 +57,10 @@ def create_session(title: str = "New chat") -> int:
         "INSERT INTO sessions (title, created_at) VALUES (?, ?)",
         (title, datetime.utcnow().isoformat())
     )
-    session_id = cur.lastrowid
+    sid = cur.lastrowid
     conn.commit()
     conn.close()
-    return session_id
+    return sid
 
 def delete_session(session_id: int):
     conn = db()
@@ -93,9 +88,8 @@ def get_messages(session_id: int):
         WHERE session_id = ?
         ORDER BY id ASC
     """, (session_id,))
-    rows = cur.fetchall()
+    msgs = [{"role": role, "content": content} for (role, content) in 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):
@@ -112,24 +106,17 @@ 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:
+    if cur.fetchone()[0] == 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
+    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()}]
@@ -137,23 +124,24 @@ def build_api_messages(session_id: int, system_message: str):
     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)
+    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")
-    elif name == "CodeXGLUE":
-        return load_dataset("google/code_x_glue_cc_code_to_code_trans", split="train")
+    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)
+
 # ---------------------------
-# Gradio Callbacks
+# Callbacks
 # ---------------------------
 def refresh_sessions_cb():
     labels, _ = list_sessions()
     selected = labels[0] if labels else None
-    visible = True if selected else False
+    visible = bool(selected)
     return (
         gr.update(choices=labels, value=selected),
         gr.update(visible=visible),
@@ -166,350 +154,159 @@ 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
+        [], "",  # chatbot cleared, user box cleared
+        gr.update(visible=True), gr.update(visible=True),
+        gr.update(visible=True), gr.update(visible=True),
     )
 
 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),
-    )
+    if not sid: return [], gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+    return get_messages(sid), 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)
+    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)
+    visible = bool(selected)
+    return gr.update(choices=labels, value=selected), [], gr.update(visible=visible), gr.update(visible=visible), gr.update(visible=visible), gr.update(visible=visible)
 
-# --- Fixed send_cb to show user message ---
+# --- Send ---
 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
+    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[:]
-    display_msgs.append({"role": "user", "content": user_text})
-
-    # Check for frontend-heavy request
+    display_msgs = chatbot_msgs[:] + [{"role": "user", "content": user_text}]
     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)
+        yield display_msgs, "", selected_label
         return
 
-    # Normal backend response
     display_msgs.append({"role": "assistant", "content": "…"})
-    yield (display_msgs, "", selected_label)
+    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
+        for chunk in client.chat_completion(messages=api_messages, max_tokens=int(max_tokens), temperature=float(temperature), top_p=float(top_p), stream=True):
+            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
-
+            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)
-
+                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)
+        display_msgs[-1]["content"] = f"⚠️ Error: {e}"
+        yield display_msgs, "", selected_label
 
+# --- Regenerate ---
 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 [], ""
-
+    if not sid: 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,))
+        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
-
+        for chunk in client.chat_completion(messages=[{"role": "system", "content": system_message.strip()}]+msgs, max_tokens=int(max_tokens), temperature=float(temperature), top_p=float(top_p), stream=True):
+            if not hasattr(chunk, "choices") or not chunk.choices: continue
+            delta = getattr(getattr(chunk.choices[0], "delta", None), "content", None) or getattr(getattr(chunk.choices[0], "message", None), "content", None) 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, "", selected_label)
+        display_msgs[-1]["content"] = f"⚠️ Error: {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()
-
+if not labels: 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
+    # --- Custom CSS ---
     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;
-    }
+    /* Chat scroll */
+    .gr-chatbot { max-height: 500px !important; overflow-y: auto !important; }
+    /* Compact left panel */
+    .compact-sliders .gr-slider, .compact-sliders .gr-number { margin-bottom:0.5rem !important; font-size:0.9rem;}
+    .system-message-compact .gr-textarea { min-height:80px !important; max-height:120px !important; font-size:0.9rem; }
+    /* Action buttons compact */
+    #session-actions-row .gr-button { margin-right: 6px; margin-bottom:3px; font-size:0.9rem;}
+    /* Input + buttons row */
+    .user-input-row .gr-textbox, .user-input-row .gr-button { margin:0 !important; padding:0.4rem 0.6rem !important; }
     </style>
     """)
 
-    gr.Markdown("## 🗄️ LLaMA & Mistral Backend-Focused CRUD Automation — with Persistent History")
+    gr.Markdown("## 🗄️ LLaMA & Mistral Backend-Focused CRUD Automation — 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)
+            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", 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 --------
+                new_btn = gr.Button("New Chat")
+                rename_btn = gr.Button("Rename", visible=False)
+                save_title_btn = gr.Button("Save", visible=False)
+                del_btn = gr.Button("Delete", visible=False)
+                refresh_btn = gr.Button("Refresh", visible=False)
+            edit_title_box = gr.Textbox(label="Edit Chat Name", placeholder="Type new chat name…", visible=False)
+
+            # 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 --------
+            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 --------
+            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
-                )
-
+                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")
         with gr.Column(scale=3):
-            chatbot = gr.Chatbot(label="Assistant", height=500, type="messages")
-            
-            # --- Single row for textbox + buttons ---
-            with gr.Row():
-                user_box = gr.Textbox(
-                    placeholder="Describe your CRUD/backend task…",
-                    lines=3,
-                    scale=5
-                )
-                send_btn = gr.Button("Send ▶️", variant="primary", scale=1)
-                regen_btn = gr.Button("Regenerate 🔁", variant="secondary", scale=1)
-
-
-    # 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
-    )
+            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)
+                regen_btn = gr.Button("🔁", variant="secondary", scale=1)
+
+    # --- Wire callbacks ---
+    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()