Update app.py

app.py

@@ -14,7 +14,6 @@ MODELS = {
 }
 
 DATASETS = ["The Stack", "CodeXGLUE"]  # Dropdown for dataset selection
-
 HF_TOKEN = os.getenv("HF_TOKEN")  # Set in your Space's Secrets
 DB_PATH = "history.db"
 
@@ -91,7 +90,8 @@ def get_messages(session_id: int):
     conn = db()
     cur = conn.cursor()
     cur.execute("""
-        SELECT role, content FROM messages
+        SELECT role, content
+        FROM messages
         WHERE session_id = ?
         ORDER BY id ASC
     """, (session_id,))
@@ -184,6 +184,7 @@ FRONTEND_KEYWORDS = [
 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)
@@ -214,6 +215,7 @@ def send_cb(user_text, selected_label, chatbot_msgs, system_message, max_tokens,
     client = get_client(model_choice)
     api_messages = build_api_messages(sid, system_message)
     partial = ""
+
     try:
         for chunk in client.chat_completion(
             messages=api_messages,
@@ -222,12 +224,23 @@ def send_cb(user_text, selected_label, chatbot_msgs, system_message, max_tokens,
             top_p=float(top_p),
             stream=True,
         ):
-            delta = chunk.choices[0].delta.content or ""
+            # --- 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)
@@ -241,13 +254,16 @@ def regenerate_cb(selected_label, system_message, max_tokens, temperature, top_p
     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
+                SELECT id FROM messages
+                WHERE session_id=?
+                ORDER BY id DESC LIMIT 1
             )
         """, (sid,))
         conn.commit()
@@ -260,6 +276,7 @@ def regenerate_cb(selected_label, system_message, max_tokens, temperature, top_p
 
     client = get_client(model_choice)
     partial = ""
+
     try:
         for chunk in client.chat_completion(
             messages=api_messages,
@@ -268,12 +285,23 @@ def regenerate_cb(selected_label, system_message, max_tokens, temperature, top_p
             top_p=float(top_p),
             stream=True,
         ):
-            delta = chunk.choices[0].delta.content or ""
+            # --- 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
@@ -286,14 +314,24 @@ 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; }
+    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>
     """)
 
@@ -311,7 +349,7 @@ with gr.Blocks(title="Backend-Focused LLaMA/Mistral CRUD Assistant", theme=gr.th
             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")
+                refresh_btn = gr.Button("🔄 Refresh", variant="secondary")
 
             gr.Markdown("### 🤖 Model Selection")
             model_choice = gr.Dropdown(