Update app.py

app.py

@@ -12,8 +12,10 @@ 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"]
-HF_TOKEN = os.getenv("HF_TOKEN")
+
+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 = (
@@ -90,7 +92,8 @@ def get_messages(session_id: int):
     cur = conn.cursor()
     cur.execute("""
         SELECT role, content FROM messages
-        WHERE session_id = ? ORDER BY id ASC
+        WHERE session_id = ?
+        ORDER BY id ASC
     """, (session_id,))
     rows = cur.fetchall()
     conn.close()
@@ -116,7 +119,7 @@ def update_session_title_if_needed(session_id: int, first_user_text: str):
         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.commit()
     conn.close()
 
 # ---------------------------
@@ -175,13 +178,13 @@ def delete_chat_cb(selected_label):
     return gr.update(choices=labels, value=selected), []
 
 FRONTEND_KEYWORDS = [
-    "react","vue","angular","html","css","javascript","tailwind","recharts","typescript"
+    "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:
@@ -189,12 +192,14 @@ def send_cb(user_text, selected_label, chatbot_msgs, system_message, max_tokens,
         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})
@@ -202,6 +207,7 @@ def send_cb(user_text, selected_label, chatbot_msgs, system_message, max_tokens,
         yield (display_msgs, "", selected_label)
         return
 
+    # Normal backend response
     display_msgs.append({"role": "assistant", "content": "…"})
     yield (display_msgs, "", selected_label)
 
@@ -216,20 +222,11 @@ def send_cb(user_text, selected_label, chatbot_msgs, system_message, max_tokens,
             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 = 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:
         display_msgs[-1]["content"] = f"⚠️ Error: {str(e)}"
@@ -248,7 +245,8 @@ def regenerate_cb(selected_label, system_message, max_tokens, temperature, top_p
         conn = db()
         cur = conn.cursor()
         cur.execute("""
-            DELETE FROM messages WHERE id = (
+            DELETE FROM messages
+            WHERE id = (
                 SELECT id FROM messages WHERE session_id=? ORDER BY id DESC LIMIT 1
             )
         """, (sid,))
@@ -270,20 +268,11 @@ def regenerate_cb(selected_label, system_message, max_tokens, temperature, top_p
             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 = 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)}"
@@ -297,29 +286,61 @@ 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:
+    gr.HTML("""
+    <style>
+      button { background-color: #22c55e !important; color: #ffffff !important; border: none !important; }
+      button:hover { background-color: #16a34a !important; }
+      button:focus { outline: 2px solid #166534 !important; outline-offset: 2px; }
+    </style>
+    """)
+
     gr.Markdown("## 🗄️ LLaMA & Mistral Backend-Focused CRUD Automation — with Persistent History")
 
     with gr.Row():
         with gr.Column(scale=1, min_width=260):
-            session_list = gr.Radio(choices=labels, value=default_selected, label="Your chats", interactive=True)
-            new_btn = gr.Button("➕ New Chat", variant="primary")
-            del_btn = gr.Button("🗑️ Delete", variant="stop")
+            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")
 
-            model_choice = gr.Dropdown(choices=list(MODELS.keys()), value=list(MODELS.keys())[0], label="Choose a model", interactive=True)
-            dataset_choice = gr.Dropdown(choices=DATASETS, value=DATASETS[0], label="Select a dataset", interactive=True)
+            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("### 📚 Dataset Selection")
+            dataset_choice = gr.Dropdown(
+                choices=DATASETS,
+                value=DATASETS[0],
+                label="Select a dataset",
+                interactive=True
+            )
 
-            system_box = gr.Textbox(value=SYSTEM_DEFAULT, label="System message", lines=5)
+            gr.Markdown("### ⚙️ Generation Settings")
+            system_box = gr.Textbox(
+                value=SYSTEM_DEFAULT,
+                label="System message",
+                lines=5
+            )
             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=720, type="messages")
+            chatbot = gr.Chatbot(label="Assistant", height=520, type="messages")
             with gr.Row():
                 user_box = gr.Textbox(placeholder="Describe your CRUD/backend task…", lines=3, scale=5)
             with gr.Row():
@@ -350,4 +371,4 @@ with gr.Blocks(title="Backend-Focused LLaMA/Mistral CRUD Assistant", theme=gr.th
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()