Update app.py

app.py

@@ -517,6 +517,26 @@ def send_message(message, history):
             yield history + [{"role": "assistant", "content": partial}]    
     
 
+def start_generation(message, history):
+    history.append({"role": "user", "content": message})
+    return message, history  # Pass message along with updated history
+
+
+def continue_generation(message, history):
+    if message.lower() in database:
+        context = random.choice(database[message.lower()])
+        message += " \n" + context
+    response_generator = wrapper.generate(message, history)
+    partial = ""
+    idx = 0
+    for t in response_generator:
+        if idx <= 3:
+            idx += 1
+            continue
+        partial += t
+        yield history + [{"role": "assistant", "content": partial}]
+
+
 # Create the dashboard
 with gr.Blocks(theme=ArtemisiaTheme(), css=custom_css) as demo:
         # Header
@@ -568,6 +588,7 @@ with gr.Blocks(theme=ArtemisiaTheme(), css=custom_css) as demo:
     #    inputs=[chat_input, chatbot],
     #    outputs=[chatbot]
     #).then(lambda: "", None, chat_input)
+    """
     send_button.click(
         lambda msg: (gr.update(value=""), msg),  # clears textbox immediately
         inputs=[chat_input],
@@ -587,6 +608,34 @@ with gr.Blocks(theme=ArtemisiaTheme(), css=custom_css) as demo:
         inputs=[chat_input, chatbot],
         outputs=[chatbot]
     )
+    """
+    send_button.click(
+        fn=lambda msg: (msg, gr.update(value="")),
+        inputs=[chat_input],
+        outputs=[gr.State(), chat_input]  # Store original msg in State
+    ).then(
+        fn=start_generation,
+        inputs=[gr.State(), chatbot],     # Use the stored message
+        outputs=[gr.State(), chatbot]     # Store again for streaming
+    ).then(
+        fn=continue_generation,
+        inputs=[gr.State(), chatbot],     # Stream using same message
+        outputs=[chatbot]
+    )
+    
+    chat_input.submit(
+        fn=lambda msg: (msg, gr.update(value="")),
+        inputs=[chat_input],
+        outputs=[gr.State(), chat_input]
+    ).then(
+        fn=start_generation,
+        inputs=[gr.State(), chatbot],
+        outputs=[gr.State(), chatbot]
+    ).then(
+        fn=continue_generation,
+        inputs=[gr.State(), chatbot],
+        outputs=[chatbot]
+    )
     # JavaScript for UI enhancements
     gr.HTML("""
     <script>