sound added to the chat and new user interface

app.py

@@ -1,145 +1,671 @@
+import os
 from huggingface_hub import InferenceClient
 from typing import List, Tuple
 import gradio as gr
 
-# Ensure you have the required libraries installed
-
-# Chatbot configuration
+# --- Configuration ---
 class ChatConfig:
+    # Switch back to Gemma 3 if the 503 errors are resolved and you prefer it
     # MODEL = "google/gemma-3-4b-it"
-    MODEL = "mistralai/Mistral-7B-Instruct-v0.1"
-    DEFAULT_SYSTEM_MSG = "You are a extremely smart and useful Chatbot."
+    MODEL = "mistralai/Mistral-7B-Instruct-v0.1" # Keep Mistral for stability if needed
+    DEFAULT_SYSTEM_MSG = "You are a highly intelligent and helpful AI assistant."
     DEFAULT_MAX_TOKENS = 1024
-    DEFAULT_TEMP = 0.4
+    DEFAULT_TEMP = 0.6 # Slightly increased default temp for potentially more engaging responses
     DEFAULT_TOP_P = 0.95
 
-import os
+# --- Hugging Face Client Initialization ---
 HF_TOKEN = os.getenv("HUGGING_FACE_HUB_TOKEN")
+if not HF_TOKEN:
+    print("Warning: HUGGING_FACE_HUB_TOKEN environment variable not set. Inference might fail.")
+    # Optionally, raise an error or provide a default (not recommended for production)
+    # raise ValueError("Missing HUGGING_FACE_HUB_TOKEN for Inference API authentication")
 
-# Not recommended to hardcode your token in the script for security reasons.
-# HF_TOKEN = "hf_YOUR_TOKEN_HERE"
-
+# Initialize client, handle potential missing token gracefully if needed elsewhere
 client = InferenceClient(ChatConfig.MODEL, token=HF_TOKEN)
 
+# --- Core Chat Logic ---
 def generate_response(
     message: str,
-    history: List[Tuple[str, str]],
-    system_message: str = ChatConfig.DEFAULT_SYSTEM_MSG,
-    max_tokens: int = ChatConfig.DEFAULT_MAX_TOKENS,
-    temperature: float = ChatConfig.DEFAULT_TEMP,
-    top_p: float = ChatConfig.DEFAULT_TOP_P
+    history_with_sys_prompt: list, # History now includes the system prompt
+    max_tokens: int,
+    temperature: float,
+    top_p: float
 ) -> str:
-    
-    messages = [{"role": "system", "content": system_message}]
-    
-    # Conversation history
-    for user_msg, bot_msg in history:
-        if user_msg:
-            messages.append({"role": "user", "content": user_msg})
-        if bot_msg:
-            messages.append({"role": "assistant", "content": bot_msg})
-    
-    messages.append({"role": "user", "content": message})
-    
-    response = ""
-
-    # Stream the response
-    # The stream=True parameter allows for real-time response generation
-    for chunk in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = chunk.choices[0].delta.content or ""
-        response += token
-        yield response
-
-
-def create_interface() -> gr.ChatInterface:
-    """Create and configure the chat interface."""
-    # Custom CSS for a modern look
-    custom_css = """
-    .chatbot .message { 
-        border-radius: 12px; 
-        margin: 5px; 
-        padding: 10px; 
+    """Generates a streaming response using the Hugging Face Inference API."""
+
+    # The history_with_sys_prompt already contains the system message as the first element
+    messages = history_with_sys_prompt + [{"role": "user", "content": message}]
+
+    # Check if client initialization failed earlier due to missing token
+    if not client.token:
+         yield "Error: Hugging Face API token is missing. Please set the HUGGING_FACE_HUB_TOKEN secret."
+         return # Stop generation
+
+    try:
+        response_stream = client.chat_completion(
+            messages,
+            max_tokens=max_tokens,
+            stream=True,
+            temperature=temperature,
+            top_p=top_p,
+        )
+
+        response = ""
+        for chunk in response_stream:
+            token = chunk.choices[0].delta.content or ""
+            response += token
+            yield response # Yield the accumulating response
+
+    except Exception as e:
+        # Basic error handling - yield the error message to the UI
+        print(f"Error during API call: {e}") # Log the full error server-side
+        # Check for specific common errors (like auth or model loading)
+        if "authorization" in str(e).lower() or "401" in str(e):
+             yield f"Error: Authentication failed. Check your Hugging Face token. Details: {e}"
+        elif "service unavailable" in str(e).lower() or "503" in str(e):
+             yield f"Error: The model ({ChatConfig.MODEL}) seems busy or unavailable (503). Please try again later. Details: {e}"
+        else:
+             yield f"An error occurred: {e}"
+
+
+# --- Gradio Interface Creation (Using gr.Blocks) ---
+
+# 1. JavaScript for Sound Effects
+# NOTE: Make sure 'typing_sound.mp3' and 'submit_sound.mp3' are in the root directory of your Space
+js_code = """
+() => {
+  // Ensure audio context is resumed on user interaction
+  let audioContextResumed = false;
+  function resumeAudioContext() {
+    if (!audioContextResumed) {
+      // Create a dummy AudioContext (or use existing if available)
+      // Playing a short silent sound on user interaction is a common workaround
+      // for browsers blocking autoplay audio.
+      const AudioContext = window.AudioContext || window.webkitAudioContext;
+      if (AudioContext) {
+        const dummyContext = new AudioContext();
+        const buffer = dummyContext.createBuffer(1, 1, 22050);
+        const source = dummyContext.createBufferSource();
+        source.buffer = buffer;
+        source.connect(dummyContext.destination);
+        source.start(0);
+        if (dummyContext.state === 'suspended') {
+          dummyContext.resume();
+        }
+      }
+      audioContextResumed = true;
+      // Remove the listener after the first interaction
+      document.removeEventListener('click', resumeAudioContext);
+      document.removeEventListener('keydown', resumeAudioContext);
+      console.log("Audio context potentially resumed by user interaction.");
     }
-    .chatbot .user-message { 
-        background-color: #e3f2fd; 
+  }
+  // Add listeners for the first user interaction
+  document.addEventListener('click', resumeAudioContext, { once: true });
+  document.addEventListener('keydown', resumeAudioContext, { once: true });
+
+
+  // Preload sounds
+  const submitSound = new Audio('file/submit_sound.mp3'); // Gradio serves files from '/file/' route
+  const typingSound = new Audio('file/typing_sound.mp3'); // Make sure these files exist!
+  submitSound.preload = 'auto';
+  typingSound.preload = 'auto';
+  typingSound.loop = true; // Loop typing sound if desired
+
+  let isTyping = false; // Flag to track typing sound status
+
+  // Function to play submit sound
+  function playSubmitSound() {
+    resumeAudioContext(); // Ensure context is active
+    submitSound.currentTime = 0; // Rewind to start
+    submitSound.play().catch(e => console.error("Error playing submit sound:", e));
+    console.log("Submit sound played");
+  }
+
+  // Function to start typing sound
+  function startTypingSound() {
+    if (!isTyping) {
+        resumeAudioContext(); // Ensure context is active
+        typingSound.currentTime = 0; // Rewind
+        typingSound.play().catch(e => console.error("Error playing typing sound:", e));
+        isTyping = true;
+        console.log("Typing sound started");
     }
-    .chatbot .bot-message { 
-        background-color: #f5f5f5; 
+  }
+
+  // Function to stop typing sound
+  function stopTypingSound() {
+    if (isTyping) {
+        typingSound.pause();
+        typingSound.currentTime = 0; // Rewind
+        isTyping = false;
+        console.log("Typing sound stopped");
     }
-    .gr-button { 
-        border-radius: 8px; 
-        padding: 8px 16px; 
+  }
+
+ // --- Event Listeners ---
+
+  // Use MutationObserver to detect when bot starts/stops responding
+  const observer = new MutationObserver((mutationsList, observer) => {
+    for(const mutation of mutationsList) {
+      if (mutation.type === 'childList' || mutation.type === 'characterData') {
+        // Check if a bot message is being added or actively updated
+        // This selector might need adjustment based on Gradio's internal structure.
+        // Inspect the chatbot HTML to confirm the classes for bot messages.
+        const botMessages = document.querySelectorAll('.chatbot .message-wrap.bot'); // More specific selector
+        if (botMessages.length > 0) {
+            const lastBotMessage = botMessages[botMessages.length - 1];
+            // Check if the last message is potentially being streamed into
+            // This is heuristic: check if it was recently added or text changed
+            // A better check might involve looking for specific gradio streaming classes if they exist
+            if (lastBotMessage.classList.contains('generating')) { // Gradio might add a class like this
+                 startTypingSound();
+            } else {
+                 // Simple check: if bot messages exist and the sound isn't playing, start it.
+                 // This might start the sound even when just loading history initially.
+                 // Refinement needed for perfect start/stop based on actual streaming state.
+                 // A simple approach: Start sound when bot message appears/updates, stop when input is enabled again?
+                 startTypingSound(); // Tentative start
+            }
+        }
+        // We need a reliable way to know when the *streaming* specifically stops.
+        // Often, the input textbox becomes enabled again. Let's try monitoring that.
+        const textbox = document.querySelector('#chat-input textarea'); // Find input textbox (check ID/class)
+        if (textbox && !textbox.disabled && isTyping) {
+             stopTypingSound();
+        }
+
+      }
     }
-    """
-
-    # Custom chatbot 
-    chatbot = gr.Chatbot(
-        label="Gemma Chat",
-        avatar_images=("./user.png", "./gemma.png"),
-        height=450,
-        show_copy_button=True
-    )
+  });
+
+  // Target the chatbot output area for observation
+  // Need to find a stable parent element containing the messages.
+  // Use a timeout to ensure the element exists after Gradio renders.
+  setTimeout(() => {
+      const chatArea = document.querySelector('.chatbot'); // Adjust selector if needed
+      if (chatArea) {
+          observer.observe(chatArea, { childList: true, subtree: true, characterData: true });
+          console.log("MutationObserver attached to chatbot area.");
+      } else {
+          console.error("Could not find chatbot area to observe.");
+      }
+
+      // Attach listener to the SUBMIT button
+      // Find the submit button - its selector might change. Inspect the element!
+      // Often it's a button next to the input textbox.
+      // Let's try finding it by looking for a button within the input row.
+      const inputRow = document.querySelector('#chat-input'); // Assuming input textbox container has this ID
+       if (inputRow) {
+            const submitButton = inputRow.querySelector('button'); // Find the first button within
+            if (submitButton) {
+                submitButton.addEventListener('click', () => {
+                    console.log("Submit button clicked via JS.");
+                    playSubmitSound();
+                    // Also stop typing sound immediately on submit
+                    stopTypingSound();
+                });
+                 console.log("Submit button listener attached.");
+            } else {
+                console.error("Could not find the submit button within #chat-input.");
+            }
+       } else {
+            console.error("Could not find the chat input container #chat-input.");
+       }
+
+
+  }, 2000); // Wait 2 seconds for Gradio UI to likely be ready
+
+
+ // Monitor the input textbox enable/disable state as another way to stop sound
+  setInterval(() => {
+    const textbox = document.querySelector('#chat-input textarea');
+    if (textbox && textbox.disabled && !isTyping) {
+        // If textbox becomes disabled (likely bot is thinking/typing), start sound
+        // This might conflict with MutationObserver, needs careful testing
+        // startTypingSound(); // Let's rely on MutationObserver primarily for starting
+    } else if (textbox && !textbox.disabled && isTyping) {
+        // If textbox becomes enabled (bot finished), stop sound
+        stopTypingSound();
+    }
+  }, 500); // Check every 500ms
+
+
+}
+"""
+
+# 2. Custom CSS for Gemini/Gemma Theme
+# (Adjust colors and gradients as desired)
+css = """
+:root {
+    /* Core Colors */
+    --color-background-dark: #131314; /* Very dark grey/black */
+    --color-background-secondary: #202124; /* Slightly lighter dark grey */
+    --color-text-light: #E8EAED; /* Light grey/off-white text */
+    --color-text-medium: #BDC1C6; /* Medium grey text */
+    --color-accent-blue: #89b4f8; /* Google's lighter blue */
+    --color-accent-blue-darker: #1a73e8; /* Google's darker blue */
+    --color-border: #5f6368; /* Border color */
+
+    /* Gradients */
+    --gradient-blue-white: linear-gradient(to right, var(--color-accent-blue), #ffffff);
+    --gradient-blue-dark: linear-gradient(to right, var(--color-accent-blue-darker), var(--color-accent-blue));
+
+    /* Component Specifics */
+    --input-background: var(--color-background-secondary);
+    --button-background: var(--color-accent-blue-darker);
+    --button-text-color: var(--color-background-dark);
+    --slider-track-color: var(--color-background-secondary);
+    --slider-thumb-color: var(--color-accent-blue);
+
+    /* Font */
+    --font-family: 'Roboto', sans-serif; /* Example: Using Roboto */
+    --font-family-mono: 'Roboto Mono', monospace;
+}
+
+/* --- Global Styles --- */
+body, .gradio-container {
+    background-color: var(--color-background-dark) !important;
+    color: var(--color-text-light) !important;
+    font-family: var(--font-family) !important;
+    height: 100vh;
+    display: flex;
+    flex-direction: column;
+}
+
+/* Ensure container takes full height */
+.gradio-container {
+    flex-grow: 1;
+    display: flex;
+    flex-direction: column;
+}
+
+/* Adjust main content area if needed */
+.main, .contain {
+    flex-grow: 1;
+    display: flex;
+    flex-direction: column;
+    overflow: hidden; /* Prevent double scrollbars */
+}
+
+
+/* --- Titles & Labels --- */
+.gr-panel > h1, .gr-panel > .label, .gr-box > .label, .gr-form > .label { /* Adjust selectors as needed */
+    color: var(--color-text-light) !important;
+    font-weight: bold;
+    margin-bottom: 10px;
+}
+.gr-panel > h1 { /* Main Title */
+   text-align: center;
+   font-size: 2em;
+   margin-bottom: 20px;
+   color: var(--color-accent-blue) !important; /* Blue title */
+}
+
+/* --- Chatbot Area --- */
+.chatbot {
+    background-color: var(--color-background-secondary);
+    border-radius: 12px;
+    border: 1px solid var(--color-border);
+    flex-grow: 1; /* Make chatbot take available space */
+    overflow-y: auto; /* Ensure scrolling within chatbot */
+    display: flex;
+    flex-direction: column-reverse; /* New messages at the bottom */
+    padding: 10px;
+    margin-bottom: 15px; /* Space before input */
+}
+
+.message-wrap { /* Container for avatar + message */
+   display: flex;
+   margin-bottom: 10px;
+   width: 100%;
+}
+
+.message-wrap.user {
+    justify-content: flex-end; /* Align user messages to the right */
+}
+.message-wrap.bot {
+   justify-content: flex-start; /* Align bot messages to the left */
+}
+
+.message {
+    padding: 10px 15px;
+    border-radius: 18px;
+    max-width: 75%; /* Limit message width */
+    word-wrap: break-word;
+    box-shadow: 0 1px 2px rgba(0, 0, 0, 0.1);
+}
 
-    # Chat interface
-    interface = gr.ChatInterface(
-        fn=generate_response,
-        chatbot=chatbot,
-        title="Your Gemma 3-(iend)",
-        theme=gr.themes.Soft(),
-        css=custom_css,
-        additional_inputs=[
-            gr.Textbox(
+/* User Message Bubble */
+.message.user {
+    background-color: var(--color-accent-blue-darker); /* Dark blue background */
+    color: var(--color-text-light); /* Light text */
+    border-bottom-right-radius: 4px; /* Slightly shape the bubble */
+}
+
+/* Bot Message Bubble */
+.message.bot {
+    background-color: #3c4043; /* Darker grey for bot */
+    color: var(--color-text-light);
+    border-bottom-left-radius: 4px;
+}
+
+/* Avatar Styling (Optional) */
+.avatar-container img {
+    width: 35px;
+    height: 35px;
+    border-radius: 50%;
+    margin: 0 8px;
+    align-self: flex-end; /* Align avatar with bottom of message */
+}
+.message-wrap.user .avatar-container { order: 1; margin-left: 8px; } /* Avatar after message */
+.message-wrap.bot .avatar-container { order: -1; margin-right: 8px; } /* Avatar before message */
+
+
+/* --- Input Area --- */
+#chat-input, .gr-form { /* Target input row/form */
+   background-color: var(--color-background-dark);
+   border-top: 1px solid var(--color-border);
+   padding: 10px 15px;
+   margin-top: auto; /* Push input to bottom if container allows */
+}
+
+textarea {
+    background-color: var(--input-background) !important;
+    color: var(--color-text-light) !important;
+    border: 1px solid var(--color-border) !important;
+    border-radius: 8px !important;
+    padding: 10px !important;
+}
+textarea:focus {
+    border-color: var(--color-accent-blue) !important;
+    box-shadow: 0 0 0 2px rgba(137, 180, 248, 0.3) !important; /* Subtle focus ring */
+}
+
+/* Submit Button */
+.gr-button {
+    background: var(--gradient-blue-dark) !important; /* Blue Gradient */
+    color: white !important; /* White text on button */
+    border: none !important;
+    border-radius: 8px !important;
+    padding: 10px 20px !important;
+    font-weight: bold !important;
+    transition: filter 0.2s ease;
+}
+.gr-button:hover {
+    filter: brightness(1.1); /* Slightly brighter on hover */
+}
+.gr-button:disabled {
+    background: var(--color-border) !important;
+    cursor: not-allowed;
+}
+
+
+/* --- Advanced Settings Accordion --- */
+.gr-accordion {
+    background-color: transparent !important; /* Make accordion background transparent */
+    border: 1px solid var(--color-border) !important;
+    border-radius: 8px !important;
+    margin-top: 15px;
+}
+.gr-accordion > .label, details > summary { /* Accordion Header */
+    color: var(--color-text-medium) !important;
+    padding: 10px 15px !important;
+    background-color: var(--color-background-secondary) !important;
+    border-radius: 8px 8px 0 0 !important; /* Round top corners */
+    border-bottom: 1px solid var(--color-border);
+    cursor: pointer;
+}
+details[open] > summary {
+    border-radius: 8px 8px 0 0 !important;
+}
+.gr-accordion .gr-form { /* Content inside accordion */
+   background-color: var(--color-background-secondary) !important;
+   border: none !important;
+   padding: 15px !important;
+   border-radius: 0 0 8px 8px !important; /* Round bottom corners */
+}
+
+/* Sliders */
+.gr-slider {
+    background: transparent !important; /* Remove default slider background */
+}
+
+input[type="range"] {
+  -webkit-appearance: none; /* Override default look */
+  appearance: none;
+  width: 100%;
+  height: 8px;
+  background: var(--slider-track-color); /* Track color */
+  border-radius: 5px;
+  cursor: pointer;
+}
+
+/* Thumb (the draggable circle) */
+input[type="range"]::-webkit-slider-thumb {
+  -webkit-appearance: none;
+  appearance: none;
+  width: 18px;
+  height: 18px;
+  background: var(--slider-thumb-color); /* Thumb color */
+  border-radius: 50%;
+  cursor: pointer;
+}
+
+input[type="range"]::-moz-range-thumb {
+  width: 18px;
+  height: 18px;
+  background: var(--slider-thumb-color);
+  border-radius: 50%;
+  cursor: pointer;
+  border: none; /* Remove default border in Firefox */
+}
+
+/* Slider Value/Label */
+.gr-slider .number { /* Adjust selector if needed */
+   color: var(--color-text-medium) !important;
+}
+.gr-slider .label, .gr-slider .info {
+    color: var(--color-text-light) !important;
+}
+
+/* Textbox within Accordion */
+.gr-accordion .gr-form textarea {
+    background-color: var(--color-background-dark) !important; /* Even darker input bg */
+     border-color: #444 !important;
+}
+
+/* --- Scrollbars (Optional, for Webkit browsers) --- */
+::-webkit-scrollbar {
+  width: 8px;
+  height: 8px;
+}
+::-webkit-scrollbar-track {
+  background: var(--color-background-secondary);
+  border-radius: 4px;
+}
+::-webkit-scrollbar-thumb {
+  background: var(--color-border);
+  border-radius: 4px;
+}
+::-webkit-scrollbar-thumb:hover {
+  background: var(--color-text-medium);
+}
+
+/* Make interface take full height */
+div[data-testid="block-label"] { /* This targets the container Gradio adds */
+    height: 100%;
+    display: flex;
+    flex-direction: column;
+}
+"""
+
+# 3. Build the UI with gr.Blocks
+def create_interface() -> gr.Blocks:
+    """Creates the Gradio interface using gr.Blocks."""
+
+    with gr.Blocks(theme=gr.themes.Base(), # Use Base and rely on CSS overrides
+                   css=css,
+                   title="Gemma/Gemini Style Chat",
+                   js=js_code) as interface:
+
+        gr.Markdown("# ✨ Gemini/Gemma Style Chat ✨", elem_classes="title") # Main title
+
+        # Store conversation history, including the initial system prompt
+        # The first element will always be the system prompt dict
+        # We use gr.State to manage history between interactions
+        initial_history = [{"role": "system", "content": ChatConfig.DEFAULT_SYSTEM_MSG}]
+        chat_history_state = gr.State(initial_history)
+
+        # Chatbot display - Use elem_id for easier CSS targeting if needed
+        chatbot_display = gr.Chatbot(
+            value=lambda hist: hist[1:], # Display history *without* the system prompt
+            label="Chat",
+            elem_id="chatbot",
+            elem_classes="chatbot", # CSS class
+            show_label=False,
+            avatar_images=("./user.png", "./gemma.png"), # Make sure these exist!
+            height=550, # Adjust height as needed
+            show_copy_button=True,
+            bubble_full_width=False # Bubbles don't take full width
+        )
+
+        # Input Textbox area
+        with gr.Row(elem_id="chat-input-row"): # Use elem_id for JS targeting
+            chat_input_box = gr.Textbox(
+                show_label=False,
+                placeholder="Type your message here...",
+                container=False, # Remove default container styling/padding
+                scale=7, # Take more horizontal space
+                elem_id="chat-input" # Use elem_id for JS targeting
+            )
+            submit_button = gr.Button("Send", scale=1, variant="primary") # Use variant for potential base styling
+
+        # Accordion for Advanced Settings
+        with gr.Accordion("Advanced Settings", open=False, elem_classes="gr-accordion"):
+            system_prompt_input = gr.Textbox(
                 value=ChatConfig.DEFAULT_SYSTEM_MSG,
-                label="System Prompt (You can change the text below)",
-                lines=2,
-                placeholder="Enter system message..."
-            ),
-            gr.Slider(
-                minimum=1,
-                maximum=8192,
-                value=ChatConfig.DEFAULT_MAX_TOKENS,
-                step=1,
-                label="Max Tokens",
-                info="Controls response length"
-            ),
-            gr.Slider(
-                minimum=0.1,
-                maximum=1.0,
-                value=ChatConfig.DEFAULT_TEMP,
-                step=0.1,
-                label="Temperature",
-                info="Controls randomness"
-            ),
-            gr.Slider(
-                minimum=0.1,
-                maximum=1.0,
-                value=ChatConfig.DEFAULT_TOP_P,
-                step=0.05,
-                label="Top-P",
-                info="Controls diversity"
+                label="System Prompt",
+                lines=3, # More lines for system prompt
+                # elem_id="system-prompt" # Optional ID
             )
-        ],
-        additional_inputs_accordion=gr.Accordion(label="Advanced Settings", open=False)
-    )
-    
+            max_tokens_slider = gr.Slider(
+                minimum=256, maximum=8192, value=ChatConfig.DEFAULT_MAX_TOKENS,
+                step=128, label="Max New Tokens", info="Max length of the bot's response."
+            )
+            temperature_slider = gr.Slider(
+                minimum=0.1, maximum=1.5, value=ChatConfig.DEFAULT_TEMP, # Wider temp range
+                step=0.1, label="Temperature", info="Controls randomness (higher = more random)."
+            )
+            top_p_slider = gr.Slider(
+                minimum=0.1, maximum=1.0, value=ChatConfig.DEFAULT_TOP_P,
+                step=0.05, label="Top-P", info="Controls diversity via nucleus sampling."
+            )
+            clear_button = gr.Button("Clear Chat History")
+
+
+        # --- Event Handling ---
+
+        def handle_submit(message: str, history_state: list, max_tokens: int, temp: float, top_p: float):
+            """Called when user sends message or presses Enter."""
+            # Add user message to state (history includes system prompt)
+            updated_history = history_state + [{"role": "user", "content": message}]
+            # Immediately update chatbot display to show user message
+            # Yield history state and empty textbox
+            yield updated_history, gr.update(value="")
+
+            # Stream response
+            # Pass the updated history directly to the generator
+            bot_stream = generate_response(message, history_state, max_tokens, temp, top_p)
+
+            # Accumulate the full response while streaming to the UI
+            full_response = ""
+            for partial_response in bot_stream:
+                full_response = partial_response
+                 # Update history state with the partial bot response for display
+                display_history = updated_history + [{"role": "assistant", "content": full_response}]
+                yield display_history, gr.update(value="") # Keep textbox empty
+
+            # Once streaming is done, update the *actual* history state with the final bot message
+            final_history = updated_history + [{"role": "assistant", "content": full_response}]
+            # Update the state, no need to update chatbot display again if streaming worked
+            yield final_history, gr.update(value="")
+
+
+        def update_history_and_chatbot(new_history_state):
+            """Helper to update both state and chatbot display."""
+            # Chatbot display needs history *without* the system prompt
+            return new_history_state, new_history_state[1:]
+
+        # Connect Submit Button click
+        submit_button.click(
+            fn=handle_submit,
+            inputs=[chat_input_box, chat_history_state, max_tokens_slider, temperature_slider, top_p_slider],
+            outputs=[chat_history_state, chat_input_box] # Update history state and clear textbox
+        ).then(
+            fn=lambda hist: hist[1:], # Get history without sys prompt for display
+            inputs=[chat_history_state],
+            outputs=[chatbot_display] # Update chatbot display
+        )
+
+        # Connect Textbox Enter key press
+        chat_input_box.submit(
+             fn=handle_submit,
+             inputs=[chat_input_box, chat_history_state, max_tokens_slider, temperature_slider, top_p_slider],
+             outputs=[chat_history_state, chat_input_box] # Update history state and clear textbox
+         ).then(
+             fn=lambda hist: hist[1:], # Get history without sys prompt for display
+             inputs=[chat_history_state],
+             outputs=[chatbot_display] # Update chatbot display
+         )
+
+        def update_system_prompt(new_prompt: str, history_state: list):
+            """Update the system prompt in the history state."""
+            if history_state and history_state[0]["role"] == "system":
+                history_state[0]["content"] = new_prompt
+            else: # Should not happen with initial_history, but safety check
+                history_state.insert(0, {"role": "system", "content": new_prompt})
+            # No UI update needed here, just update the state
+            return history_state
+
+        # Update system prompt in state when textbox changes
+        system_prompt_input.change(
+            fn=update_system_prompt,
+            inputs=[system_prompt_input, chat_history_state],
+            outputs=[chat_history_state] # Only update the state
+        )
+
+        def clear_chat(current_system_prompt: str):
+            """Clears chat history, keeping the current system prompt."""
+            new_history = [{"role": "system", "content": current_system_prompt}]
+            return new_history, [] # Return new state and empty list for chatbot display
+
+        # Connect Clear Button
+        clear_button.click(
+            fn=clear_chat,
+            inputs=[system_prompt_input], # Get the current system prompt from its input box
+            outputs=[chat_history_state, chatbot_display] # Update state and clear chatbot
+        )
+
     return interface
 
+# --- Main Execution ---
 def main():
     app = create_interface()
     app.launch(
-        server_name="0.0.0.0",
+        server_name="0.0.0.0", # Important for Docker/Spaces
         server_port=7860,
-        share=False,
         show_api=False,
         show_error=True,
-        debug=True
+        debug=True # Enable debug for easier troubleshooting
     )
 
 if __name__ == "__main__":
+    # Ensure avatar files exist or handle gracefully
+    if not os.path.exists("./user.png"):
+        print("Warning: ./user.png not found. Using default avatar.")
+        # You might want to download a default one here if needed
+    if not os.path.exists("./gemma.png"):
+        print("Warning: ./gemma.png not found. Using default avatar.")
+
     main()