Update app.py

app.py

@@ -1,35 +1,54 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
-import time
+import os
 
 MODELS = [
     {"name": "Qwen 2.5 Coder 32B", "id": "Qwen/Qwen2.5-Coder-32B-Instruct"},
-    {"name": "Qwen 2.5 72B Instruct", "id": "Qwen/Qwen2.5-72B-Instruct"},
-    {"name": "Llama 3.3 70B Instruct", "id": "meta-llama/Llama-3.3-70B-Instruct"},
-    {"name": "Mixtral 8x7B Instruct", "id": "mistralai/Mixtral-8x7B-Instruct-v0.1"},
+    {"name": "Qwen 2.5 72B", "id": "Qwen/Qwen2.5-72B-Instruct"},
+    {"name": "Llama 3.3 70B", "id": "meta-llama/Llama-3.3-70B-Instruct"},
+    {"name": "Mixtral 8x7B", "id": "mistralai/Mixtral-8x7B-Instruct-v0.1"},
     {"name": "DeepSeek Coder V2", "id": "deepseek-ai/DeepSeek-Coder-V2-Instruct"},
 ]
 
-def stream_chat(message, history, model_name):
-    model_id = next((m["id"] for m in MODELS if m["name"] == model_name), MODELS[0]["id"])
-    client = InferenceClient(model_id)
+def stream_chat(message, history, model_name, api_key):
+    if not api_key or api_key.strip() == "":
+        yield "⚠️ Please enter your Hugging Face API key in the sidebar first.\n\nGet your free API key at: https://huggingface.co/settings/tokens"
+        return
     
-    messages = []
-    for h in history:
-        if h[0]:
-            messages.append({"role": "user", "content": h[0]})
-        if h[1]:
-            messages.append({"role": "assistant", "content": h[1]})
-    messages.append({"role": "user", "content": message})
+    model_id = next((m["id"] for m in MODELS if m["name"] == model_name), MODELS[0]["id"])
     
-    response = ""
     try:
-        for chunk in client.chat_completion(messages=messages, max_tokens=8192, stream=True, temperature=0.7):
+        client = InferenceClient(token=api_key.strip())
+        
+        messages = []
+        for h in history:
+            if h["role"] == "user":
+                messages.append({"role": "user", "content": h["content"]})
+            elif h["role"] == "assistant":
+                messages.append({"role": "assistant", "content": h["content"]})
+        
+        messages.append({"role": "user", "content": message})
+        
+        response = ""
+        for chunk in client.chat_completion(
+            model=model_id,
+            messages=messages,
+            max_tokens=8192,
+            stream=True,
+            temperature=0.7
+        ):
             if chunk.choices[0].delta.content:
                 response += chunk.choices[0].delta.content
                 yield response
+                
     except Exception as e:
-        yield f"Error: {str(e)}"
+        error_msg = str(e)
+        if "401" in error_msg or "authorization" in error_msg.lower():
+            yield "❌ Invalid API Key. Please check your Hugging Face token.\n\nGet a new one at: https://huggingface.co/settings/tokens"
+        elif "429" in error_msg:
+            yield "⏳ Rate limit exceeded. Please wait a moment and try again."
+        else:
+            yield f"❌ Error: {error_msg}\n\nPlease try again or select a different model."
 
 CSS = """
 @import url('https://fonts.googleapis.com/css2?family=JetBrains+Mono:wght@400;500;600;700&family=Space+Grotesk:wght@400;500;600;700&display=swap');
@@ -51,14 +70,13 @@ CSS = """
     --accent-hover: #79C0FF;
     --success: #3FB950;
     --error: #F85149;
-    --shadow: rgba(0,0,0,0.4);
+    --warning: #F0883E;
 }
 
 body {
     font-family: 'JetBrains Mono', monospace !important;
     background: var(--bg-primary) !important;
     color: var(--text-primary) !important;
-    overflow-x: hidden;
 }
 
 .gradio-container {
@@ -68,7 +86,6 @@ body {
     background: var(--bg-primary) !important;
 }
 
-/* Header */
 #header {
     background: linear-gradient(135deg, var(--bg-secondary) 0%, var(--bg-tertiary) 100%);
     border-bottom: 2px solid var(--border);
@@ -76,7 +93,6 @@ body {
     position: sticky;
     top: 0;
     z-index: 100;
-    backdrop-filter: blur(10px);
 }
 
 #header h1 {
@@ -100,40 +116,12 @@ body {
     text-transform: uppercase;
 }
 
-/* Main Layout */
-#main-wrapper {
-    display: grid;
-    grid-template-columns: 320px 1fr;
-    min-height: calc(100vh - 100px);
-    gap: 0;
-}
-
-/* Sidebar */
-#sidebar {
-    background: var(--bg-secondary);
-    border-right: 1px solid var(--border);
-    padding: 32px 24px;
-    overflow-y: auto;
-    position: sticky;
-    top: 100px;
-    height: calc(100vh - 100px);
-}
-
-#sidebar::-webkit-scrollbar {
-    width: 6px;
-}
-
-#sidebar::-webkit-scrollbar-track {
-    background: transparent;
-}
-
-#sidebar::-webkit-scrollbar-thumb {
-    background: var(--border);
-    border-radius: 3px;
-}
-
 .sidebar-section {
-    margin-bottom: 32px;
+    background: var(--bg-secondary);
+    border: 1px solid var(--border);
+    border-radius: 12px;
+    padding: 20px;
+    margin-bottom: 20px;
 }
 
 .sidebar-title {
@@ -145,47 +133,43 @@ body {
     margin-bottom: 12px;
 }
 
-.model-card {
-    background: var(--bg-tertiary);
-    border: 1px solid var(--border);
-    border-radius: 8px;
-    padding: 16px;
-    margin-bottom: 12px;
-    cursor: pointer;
-    transition: all 0.2s cubic-bezier(0.4, 0, 0.2, 1);
-}
-
-.model-card:hover {
-    border-color: var(--accent);
-    transform: translateX(4px);
-    box-shadow: 0 4px 12px var(--shadow);
-}
-
-.model-card.active {
-    background: linear-gradient(135deg, var(--accent) 0%, var(--success) 100%);
-    border-color: var(--success);
+.api-key-input input {
+    background: var(--bg-tertiary) !important;
+    border: 1px solid var(--border) !important;
+    color: var(--text-primary) !important;
+    font-family: 'JetBrains Mono', monospace !important;
+    font-size: 12px !important;
+    padding: 12px !important;
+    border-radius: 8px !important;
 }
 
-.model-card.active .model-name {
-    color: var(--bg-primary);
-    font-weight: 700;
+.api-key-input input:focus {
+    border-color: var(--accent) !important;
+    box-shadow: 0 0 0 3px rgba(88, 166, 255, 0.1) !important;
 }
 
-.model-name {
-    font-size: 13px;
-    font-weight: 600;
-    color: var(--text-primary);
-    margin-bottom: 4px;
+.security-note {
+    background: rgba(63, 185, 80, 0.1);
+    border: 1px solid var(--success);
+    border-radius: 8px;
+    padding: 12px;
+    font-size: 11px;
+    color: var(--success);
+    line-height: 1.6;
+    margin-top: 12px;
 }
 
-.model-info {
-    font-size: 10px;
-    color: var(--text-secondary);
-    text-transform: uppercase;
-    letter-spacing: 0.05em;
+.warning-note {
+    background: rgba(240, 136, 62, 0.1);
+    border: 1px solid var(--warning);
+    border-radius: 8px;
+    padding: 12px;
+    font-size: 11px;
+    color: var(--warning);
+    line-height: 1.6;
+    margin-top: 12px;
 }
 
-/* Stats */
 .stat-grid {
     display: grid;
     grid-template-columns: 1fr 1fr;
@@ -196,13 +180,13 @@ body {
 .stat-box {
     background: var(--bg-tertiary);
     border: 1px solid var(--border);
-    border-radius: 6px;
+    border-radius: 8px;
     padding: 12px;
     text-align: center;
 }
 
 .stat-value {
-    font-size: 18px;
+    font-size: 16px;
     font-weight: 700;
     color: var(--success);
     margin-bottom: 4px;
@@ -215,125 +199,70 @@ body {
     letter-spacing: 0.1em;
 }
 
-/* Chat Area */
-#chat-area {
-    background: var(--bg-primary);
-    padding: 32px 48px;
-    display: flex;
-    flex-direction: column;
-    gap: 32px;
-}
-
-#chatbot {
+.chatbot-container {
     background: transparent !important;
     border: none !important;
-    min-height: 500px;
-    padding: 0 !important;
 }
 
 .message {
-    margin-bottom: 24px !important;
-    animation: slideUp 0.3s cubic-bezier(0.4, 0, 0.2, 1);
+    margin-bottom: 16px !important;
 }
 
-@keyframes slideUp {
-    from {
-        opacity: 0;
-        transform: translateY(20px);
-    }
-    to {
-        opacity: 1;
-        transform: translateY(0);
-    }
+.user .message-row {
+    justify-content: flex-end !important;
 }
 
-.message.user .message-content {
-    background: linear-gradient(135deg, var(--accent) 0%, var(--accent-hover) 100%) !important;
-    color: var(--bg-primary) !important;
-    border-radius: 16px 16px 4px 16px !important;
-    padding: 16px 20px !important;
-    font-size: 14px !important;
-    line-height: 1.6 !important;
-    box-shadow: 0 4px 12px rgba(88, 166, 255, 0.2);
+.bot .message-row {
+    justify-content: flex-start !important;
 }
 
-.message.bot .message-content {
-    background: var(--bg-tertiary) !important;
-    color: var(--text-primary) !important;
-    border: 1px solid var(--border) !important;
-    border-radius: 16px 16px 16px 4px !important;
+.message-content {
     padding: 16px 20px !important;
+    border-radius: 16px !important;
     font-size: 14px !important;
     line-height: 1.6 !important;
+    max-width: 70% !important;
 }
 
-.message-content code {
-    background: rgba(255, 255, 255, 0.08) !important;
-    padding: 2px 6px !important;
-    border-radius: 4px !important;
-    font-size: 13px !important;
-    font-family: 'JetBrains Mono', monospace !important;
+.user .message-content {
+    background: linear-gradient(135deg, var(--accent) 0%, var(--accent-hover) 100%) !important;
+    color: var(--bg-primary) !important;
+    border-radius: 16px 16px 4px 16px !important;
 }
 
-.message-content pre {
-    background: var(--bg-secondary) !important;
+.bot .message-content {
+    background: var(--bg-tertiary) !important;
+    color: var(--text-primary) !important;
     border: 1px solid var(--border) !important;
-    border-radius: 8px !important;
-    padding: 16px !important;
-    overflow-x: auto !important;
-    margin: 12px 0 !important;
-}
-
-.message-content pre code {
-    background: none !important;
-    padding: 0 !important;
-}
-
-/* Input Area */
-#input-wrapper {
-    background: var(--bg-secondary);
-    border: 2px solid var(--border);
-    border-radius: 16px;
-    padding: 8px;
-    transition: all 0.2s cubic-bezier(0.4, 0, 0.2, 1);
-}
-
-#input-wrapper:focus-within {
-    border-color: var(--accent);
-    box-shadow: 0 0 0 4px rgba(88, 166, 255, 0.1);
+    border-radius: 16px 16px 16px 4px !important;
 }
 
-#message-input {
-    background: transparent !important;
-    border: none !important;
+#message-input textarea {
+    background: var(--bg-tertiary) !important;
+    border: 2px solid var(--border) !important;
     color: var(--text-primary) !important;
-    font-size: 14px !important;
-    padding: 16px 20px !important;
-    resize: none !important;
     font-family: 'JetBrains Mono', monospace !important;
+    font-size: 14px !important;
+    padding: 16px !important;
+    border-radius: 12px !important;
 }
 
-#message-input::placeholder {
-    color: var(--text-secondary) !important;
+#message-input textarea:focus {
+    border-color: var(--accent) !important;
+    box-shadow: 0 0 0 4px rgba(88, 166, 255, 0.1) !important;
 }
 
-/* Buttons */
-.btn {
+.btn-primary {
+    background: linear-gradient(135deg, var(--accent) 0%, var(--success) 100%) !important;
+    color: var(--bg-primary) !important;
     font-family: 'Space Grotesk', sans-serif !important;
     font-weight: 600 !important;
     text-transform: uppercase !important;
     letter-spacing: 0.05em !important;
-    border-radius: 8px !important;
-    transition: all 0.2s cubic-bezier(0.4, 0, 0.2, 1) !important;
-    cursor: pointer !important;
     border: none !important;
-}
-
-.btn-primary {
-    background: linear-gradient(135deg, var(--accent) 0%, var(--success) 100%) !important;
-    color: var(--bg-primary) !important;
     padding: 14px 28px !important;
-    font-size: 13px !important;
+    border-radius: 8px !important;
+    cursor: pointer !important;
 }
 
 .btn-primary:hover {
@@ -344,9 +273,14 @@ body {
 .btn-secondary {
     background: var(--bg-tertiary) !important;
     color: var(--text-primary) !important;
+    font-family: 'Space Grotesk', sans-serif !important;
+    font-weight: 600 !important;
+    text-transform: uppercase !important;
+    letter-spacing: 0.05em !important;
     border: 1px solid var(--border) !important;
     padding: 12px 24px !important;
-    font-size: 12px !important;
+    border-radius: 8px !important;
+    cursor: pointer !important;
 }
 
 .btn-secondary:hover {
@@ -354,120 +288,29 @@ body {
     border-color: var(--accent) !important;
 }
 
-/* Examples */
-.examples-grid {
-    display: grid;
-    grid-template-columns: repeat(auto-fit, minmax(250px, 1fr));
-    gap: 16px;
-    margin-top: 24px;
-}
-
-.example-card {
-    background: var(--bg-tertiary);
-    border: 1px solid var(--border);
-    border-radius: 10px;
-    padding: 16px 20px;
-    cursor: pointer;
-    transition: all 0.2s cubic-bezier(0.4, 0, 0.2, 1);
-    font-size: 13px;
-    color: var(--text-secondary);
-}
-
-.example-card:hover {
-    background: var(--bg-secondary);
-    border-color: var(--accent);
-    color: var(--text-primary);
-    transform: translateY(-3px);
-    box-shadow: 0 6px 20px var(--shadow);
-}
-
-/* Action Bar */
-.action-bar {
-    display: flex;
-    gap: 12px;
-    margin-top: 16px;
-}
-
-/* Loading State */
-.loading-indicator {
-    display: inline-flex;
-    gap: 4px;
-}
-
-.loading-dot {
-    width: 6px;
-    height: 6px;
-    border-radius: 50%;
-    background: var(--accent);
-    animation: bounce 1.4s infinite ease-in-out;
-}
-
-.loading-dot:nth-child(1) {
-    animation-delay: -0.32s;
-}
-
-.loading-dot:nth-child(2) {
-    animation-delay: -0.16s;
-}
-
-@keyframes bounce {
-    0%, 80%, 100% {
-        transform: scale(0);
-    }
-    40% {
-        transform: scale(1);
-    }
-}
-
-/* Responsive */
-@media (max-width: 1024px) {
-    #main-wrapper {
-        grid-template-columns: 1fr;
-    }
-    
-    #sidebar {
-        position: relative;
-        height: auto;
-        border-right: none;
-        border-bottom: 1px solid var(--border);
-    }
-    
-    #header {
-        padding: 20px 24px;
-    }
-    
-    #chat-area {
-        padding: 24px;
-    }
-}
-
-/* Gradio Overrides */
-.gr-button {
+label {
     font-family: 'Space Grotesk', sans-serif !important;
+    color: var(--text-secondary) !important;
+    font-size: 11px !important;
+    font-weight: 600 !important;
+    text-transform: uppercase !important;
+    letter-spacing: 0.05em !important;
 }
 
-.gr-input {
-    font-family: 'JetBrains Mono', monospace !important;
-}
-
-.gr-box {
+.gr-dropdown {
     background: var(--bg-tertiary) !important;
-    border-color: var(--border) !important;
+    border: 1px solid var(--border) !important;
+    color: var(--text-primary) !important;
+    border-radius: 8px !important;
 }
 
-label {
-    font-family: 'Space Grotesk', sans-serif !important;
-    color: var(--text-secondary) !important;
-    font-size: 11px !important;
-    font-weight: 700 !important;
-    text-transform: uppercase !important;
-    letter-spacing: 0.1em !important;
+::selection {
+    background: var(--accent);
+    color: var(--bg-primary);
 }
 
-/* Scrollbar */
 ::-webkit-scrollbar {
     width: 10px;
-    height: 10px;
 }
 
 ::-webkit-scrollbar-track {
@@ -482,12 +325,6 @@ label {
 ::-webkit-scrollbar-thumb:hover {
     background: var(--text-secondary);
 }
-
-/* Selection */
-::selection {
-    background: var(--accent);
-    color: var(--bg-primary);
-}
 """
 
 with gr.Blocks(css=CSS, theme=gr.themes.Base(), title="Neural Chat") as demo:
@@ -495,118 +332,133 @@ with gr.Blocks(css=CSS, theme=gr.themes.Base(), title="Neural Chat") as demo:
     gr.HTML("""
     <div id="header">
         <h1>⬡ NEURAL CHAT</h1>
-        <div class="subtitle">Advanced AI Communication Interface</div>
+        <div class="subtitle">Secure AI Interface with Personal API Keys</div>
     </div>
     """)
     
-    with gr.Row(elem_id="main-wrapper"):
+    with gr.Row():
         
-        with gr.Column(elem_id="sidebar", scale=1):
+        with gr.Column(scale=1, min_width=320):
             
-            gr.HTML('<div class="sidebar-title">MODEL SELECTION</div>')
-            
-            model_selector = gr.Dropdown(
-                choices=[m["name"] for m in MODELS],
-                value=MODELS[0]["name"],
-                label="Active Model",
-                interactive=True,
-                elem_classes=["model-dropdown"]
-            )
-            
-            gr.HTML('<div class="sidebar-title" style="margin-top: 24px;">SYSTEM STATUS</div>')
-            
-            gr.HTML("""
-            <div class="stat-grid">
-                <div class="stat-box">
-                    <div class="stat-value">● </div>
-                    <div class="stat-label">ONLINE</div>
-                </div>
-                <div class="stat-box">
-                    <div class="stat-value">32K</div>
-                    <div class="stat-label">CONTEXT</div>
-                </div>
-                <div class="stat-box">
-                    <div class="stat-value">~2.1s</div>
-                    <div class="stat-label">LATENCY</div>
+            with gr.Group(elem_classes=["sidebar-section"]):
+                gr.HTML('<div class="sidebar-title">🔐 API AUTHENTICATION</div>')
+                
+                api_key_input = gr.Textbox(
+                    label="Hugging Face API Key",
+                    type="password",
+                    placeholder="hf_...",
+                    info="Your key is stored locally in your browser session",
+                    elem_classes=["api-key-input"]
+                )
+                
+                gr.HTML("""
+                <div class="security-note">
+                    ✓ Your API key is never stored on our servers<br>
+                    ✓ Each user uses their own key<br>
+                    ✓ Completely secure and private
                 </div>
-                <div class="stat-box">
-                    <div class="stat-value">99.8%</div>
-                    <div class="stat-label">UPTIME</div>
+                """)
+                
+                gr.HTML("""
+                <div class="warning-note">
+                    📌 Get your free API key:<br>
+                    <a href="https://huggingface.co/settings/tokens" target="_blank" style="color: #F0883E;">
+                    https://huggingface.co/settings/tokens
+                    </a>
                 </div>
-            </div>
-            """)
-            
-            gr.HTML('<div class="sidebar-title" style="margin-top: 24px;">QUICK ACTIONS</div>')
+                """)
             
-            clear_btn = gr.Button("⟳ NEW CONVERSATION", elem_classes=["btn", "btn-secondary"], size="sm")
-            export_btn = gr.Button("↓ EXPORT CHAT", elem_classes=["btn", "btn-secondary"], size="sm")
+            with gr.Group(elem_classes=["sidebar-section"]):
+                gr.HTML('<div class="sidebar-title">🤖 MODEL SELECTION</div>')
+                
+                model_selector = gr.Dropdown(
+                    choices=[m["name"] for m in MODELS],
+                    value=MODELS[0]["name"],
+                    label="Active Model",
+                    interactive=True,
+                )
             
-            gr.HTML("""
-            <div style="margin-top: 32px; padding-top: 24px; border-top: 1px solid var(--border);">
-                <div style="font-size: 10px; color: var(--text-secondary); line-height: 1.6;">
-                    <strong>NEURAL CHAT v2.0</strong><br>
-                    Powered by Hugging Face<br>
-                    React + Python Architecture<br><br>
-                    All models are free and open-source
+            with gr.Group(elem_classes=["sidebar-section"]):
+                gr.HTML('<div class="sidebar-title">📊 STATISTICS</div>')
+                
+                gr.HTML("""
+                <div class="stat-grid">
+                    <div class="stat-box">
+                        <div class="stat-value">●</div>
+                        <div class="stat-label">SECURE</div>
+                    </div>
+                    <div class="stat-box">
+                        <div class="stat-value">32K</div>
+                        <div class="stat-label">CONTEXT</div>
+                    </div>
+                    <div class="stat-box">
+                        <div class="stat-value">FREE</div>
+                        <div class="stat-label">TIER</div>
+                    </div>
+                    <div class="stat-box">
+                        <div class="stat-value">5</div>
+                        <div class="stat-label">MODELS</div>
+                    </div>
                 </div>
-            </div>
-            """)
+                """)
+            
+            with gr.Group(elem_classes=["sidebar-section"]):
+                gr.HTML('<div class="sidebar-title">⚡ ACTIONS</div>')
+                
+                clear_btn = gr.Button("⟳ NEW CHAT", elem_classes=["btn-secondary"], size="sm")
+                retry_btn = gr.Button("🔄 RETRY LAST", elem_classes=["btn-secondary"], size="sm")
         
-        with gr.Column(elem_id="chat-area", scale=3):
+        with gr.Column(scale=3):
             
             chatbot = gr.Chatbot(
                 value=[],
-                elem_id="chatbot",
-                height=550,
+                height=600,
                 show_label=False,
+                type="messages",
                 avatar_images=(None, "⬡"),
-                bubble_full_width=False,
-                render_markdown=True,
-                show_copy_button=True,
+                elem_classes=["chatbot-container"]
             )
             
-            with gr.Group(elem_id="input-wrapper"):
-                msg_input = gr.Textbox(
-                    placeholder="Enter your message here... (Shift+Enter for new line)",
-                    show_label=False,
-                    lines=1,
-                    max_lines=10,
-                    elem_id="message-input",
-                    container=False,
-                )
+            msg_input = gr.Textbox(
+                placeholder="Type your message... (Shift+Enter for new line)",
+                show_label=False,
+                lines=3,
+                max_lines=10,
+                elem_id="message-input"
+            )
             
-            with gr.Row(elem_classes=["action-bar"]):
-                send_btn = gr.Button("SEND MESSAGE", elem_classes=["btn", "btn-primary"], scale=2)
-                retry_btn = gr.Button("⟲ RETRY", elem_classes=["btn", "btn-secondary"], scale=1)
-                stop_btn = gr.Button("⏹ STOP", elem_classes=["btn", "btn-secondary"], scale=1)
+            with gr.Row():
+                send_btn = gr.Button("SEND MESSAGE", elem_classes=["btn-primary"], scale=3)
+                stop_btn = gr.Button("⏹ STOP", elem_classes=["btn-secondary"], scale=1)
             
-            gr.HTML('<div class="sidebar-title" style="margin-top: 32px;">EXAMPLE PROMPTS</div>')
+            gr.HTML('<div class="sidebar-title" style="margin-top: 24px;">💡 EXAMPLE PROMPTS</div>')
             
             gr.Examples(
                 examples=[
-                    "Explain quantum entanglement in simple terms",
-                    "Write a Python script to analyze CSV data",
-                    "What are the key principles of clean code?",
-                    "Create a React component for a dropdown menu",
-                    "How does machine learning differ from deep learning?",
-                    "Explain the SOLID principles with code examples",
+                    "Explain quantum computing in simple terms",
+                    "Write a Python function for fibonacci sequence",
+                    "What are the SOLID principles with examples?",
+                    "Create a React component for a modal dialog",
+                    "Help me optimize this SQL query",
+                    "Explain async/await in JavaScript",
                 ],
                 inputs=msg_input,
-                elem_classes=["examples-grid"]
             )
     
     def user_msg(message, history):
-        return "", history + [[message, None]]
+        if not message.strip():
+            return "", history
+        return "", history + [{"role": "user", "content": message}]
     
-    def bot_msg(history, model):
-        if not history or history[-1][1] is not None:
+    def bot_msg(history, model, api_key):
+        if not history or history[-1]["role"] != "user":
             return history
         
-        user_message = history[-1][0]
-        history[-1][1] = ""
+        user_message = history[-1]["content"]
+        history.append({"role": "assistant", "content": ""})
         
-        for response in stream_chat(user_message, history[:-1], model):
-            history[-1][1] = response
+        for response in stream_chat(user_message, history[:-1], model, api_key):
+            history[-1]["content"] = response
             yield history
     
     msg_input.submit(
@@ -616,7 +468,7 @@ with gr.Blocks(css=CSS, theme=gr.themes.Base(), title="Neural Chat") as demo:
         queue=False
     ).then(
         bot_msg,
-        [chatbot, model_selector],
+        [chatbot, model_selector, api_key_input],
         chatbot
     )
     
@@ -627,7 +479,7 @@ with gr.Blocks(css=CSS, theme=gr.themes.Base(), title="Neural Chat") as demo:
         queue=False
     ).then(
         bot_msg,
-        [chatbot, model_selector],
+        [chatbot, model_selector, api_key_input],
         chatbot
     )
     
@@ -636,8 +488,15 @@ with gr.Blocks(css=CSS, theme=gr.themes.Base(), title="Neural Chat") as demo:
     def retry_last_msg(history):
         if not history:
             return history
-        if history[-1][1] is not None:
-            history[-1][1] = None
+        
+        if history[-1]["role"] == "assistant":
+            history = history[:-1]
+        
+        if history and history[-1]["role"] == "user":
+            last_user_msg = history[-1]
+            history = history[:-1]
+            return history + [last_user_msg]
+        
         return history
     
     retry_btn.click(
@@ -647,7 +506,7 @@ with gr.Blocks(css=CSS, theme=gr.themes.Base(), title="Neural Chat") as demo:
         queue=False
     ).then(
         bot_msg,
-        [chatbot, model_selector],
+        [chatbot, model_selector, api_key_input],
         chatbot
     )