Update app.py

app.py

@@ -11,22 +11,23 @@ def respond(
     max_tokens: int,
     temperature: float,
     top_p: float,
+    stop_event: gr.EventData,
 ) -> Iterator[str]:
     messages = [{"role": "system", "content": system_message}]
-    
+
     # Add history to messages
     for user_msg, assistant_msg in history:
         if user_msg:
             messages.append({"role": "user", "content": user_msg})
         if assistant_msg:
             messages.append({"role": "assistant", "content": assistant_msg})
-    
+
     # Add current message
     messages.append({"role": "user", "content": message})
-    
+
     # Initialize response
     response = ""
-    
+
     # Stream the response
     try:
         for chunk in client.chat_completion(
@@ -36,6 +37,8 @@ def respond(
             temperature=temperature,
             top_p=top_p,
         ):
+            if stop_event.originator.get("clicked"):
+                break
             if chunk.choices[0].delta.content is not None:
                 token = chunk.choices[0].delta.content
                 response += token
@@ -81,15 +84,12 @@ details[open] summary:after {
 
 # Create Gradio interface
 with gr.Blocks(css=css) as demo:
-    gr.Markdown("## Chat with <span style='color: #d026ff'>Superthoughts</span> lite! (1.7B)")
+    gr.Markdown("## Chat with Superthoughts lite! (1.7B)")
     gr.Markdown("**Note:** First response may take a moment to initialize. Subsequent responses will be faster.")
-    
-    chatbot = gr.Chatbot(height=600, type='messages')
-    with gr.Row():
-        msg = gr.Textbox(label="Your message", placeholder="Type your message here...", scale=9)
-        stop = gr.Button("Stop", scale=1)
-    
-    
+
+    chatbot = gr.Chatbot(height=600)
+    msg = gr.Textbox(label="Your message", placeholder="Type your message here...")
+
     with gr.Accordion("Advanced Settings", open=False):
         system_message = gr.Textbox(
             value="You must act in a conversational matter and always include <think> ... </think> <output> </output> tokens.",
@@ -119,39 +119,42 @@ with gr.Blocks(css=css) as demo:
 
     def user(user_message: str, history: list) -> tuple[str, list]:
         """Add user message to history"""
-        return "", history + [{"role": "user", "content": user_message}, {"role": "assistant", "content": None}]
+        return "", history + [[user_message, None]]
 
-    def bot(history: list, system_message: str, max_tokens: int, temperature: float, top_p: float) -> Iterator[list]:
+    def bot(history: list, system_message: str, max_tokens: int, temperature: float, top_p: float, stop_event: gr.EventData) -> Iterator[list]:
         """Generate and stream bot responses"""
-        user_message = history[-2]["content"]  # Get the last user message
-        for partial_response in respond(user_message, history[:-2], system_message, max_tokens, temperature, top_p):
-            history[-1]["content"] = partial_response
+        user_message, _ = history[-1]
+        history[-1][1] = ""  # Initialize bot's response
+
+        for partial_response in respond(user_message, history[:-1], system_message, max_tokens, temperature, top_p, stop_event):
+            history[-1][1] = partial_response
             yield history
 
-    # Create event handlers
-    submit_event = msg.submit(
+    # Set up chat message handling
+    msg.submit(
         user,
         [msg, chatbot],
         [msg, chatbot],
         queue=False
     ).then(
         bot,
-        [chatbot, system_message, max_tokens, temperature, top_p],
+        [chatbot, system_message, max_tokens, temperature, top_p, gr.EventData()],
         chatbot
     )
-    
-    # Set up stop button
-    stop.click(fn=None, inputs=None, outputs=None, cancels=[submit_event])
 
     # Add a clear button
     clear = gr.Button("Clear Conversation")
     clear.click(lambda: None, None, chatbot, queue=False)
-    
+
+    # Add a stop button
+    stop_button = gr.Button("Stop")
+    stop_button.click(lambda: gr.EventData(clicked=True), outputs=None, queue=False)
+
     # Add disclaimer
     gr.Markdown(
         """
         ---
-        ⚠️ **Disclaimer:** Superthoughts may make mistakes. Always verify important information. 
+        ⚠️ **Disclaimer:** Superthoughts may make mistakes. Always verify important information.
         This chat interface is intended for testing and experimentation purposes only.
         """
     )
@@ -159,4 +162,4 @@ with gr.Blocks(css=css) as demo:
 # Launch the interface
 if __name__ == "__main__":
     demo.queue()
-    demo.launch(share=True)
+    demo.launch(share=True)