Update app.py

app.py

@@ -195,19 +195,28 @@ def clear_and_update_chat(message, history, use_web_search, model, temperature):
     updated_history = chatbot_interface(message, history, use_web_search, model, temperature)
     return "", updated_history  # Return empty string to clear the input
 
-# Gradio interface
+def retry_last_response(history):
+    if history:
+        last_user_message = history[-1][0]
+        return last_user_message, history[:-1]
+    return "", history
+
+def undo_last_interaction(history):
+    if len(history) >= 1:
+        return history[:-1]
+    return history
+
+def clear_conversation():
+    return []
+
+def stop_generation():
+    global is_generating
+    is_generating = False
+
 with gr.Blocks() as demo:
-    
     is_generating = gr.State(False)
+    stop_clicked = gr.State(False)
 
-    def protected_clear_and_update_chat(message, history, use_web_search, model, temperature, is_generating):
-        if is_generating:
-            return message, history, is_generating
-        is_generating = True
-        updated_message, updated_history = clear_and_update_chat(message, history, use_web_search, model, temperature)
-        is_generating = False
-        return updated_message, updated_history, is_generating
-    
     gr.Markdown("# AI-powered Web Search and PDF Chat Assistant")
     
     with gr.Row():
@@ -225,8 +234,71 @@ with gr.Blocks() as demo:
     with gr.Row():
         model_dropdown = gr.Dropdown(choices=MODELS, label="Select Model", value=MODELS[1])
         temperature_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.2, step=0.1, label="Temperature")
+        num_calls_slider = gr.Slider(minimum=1, maximum=5, value=3, step=1, label="Number of API Calls")
     
-    submit = gr.Button("Submit")
+    with gr.Row():
+        submit_btn = gr.Button("Send")
+        stop_btn = gr.Button("Stop", visible=False)
+        retry_btn = gr.Button("Retry")
+        undo_btn = gr.Button("Undo")
+        clear_btn = gr.Button("Clear")
+
+    def protected_generate_response(message, history, use_web_search, model, temperature, num_calls, is_generating, stop_clicked):
+        if is_generating:
+            return message, history, is_generating, stop_clicked
+        is_generating = True
+        stop_clicked = False
+        
+        try:
+            if use_web_search:
+                main_content, sources = get_response_with_search(message, model, num_calls=num_calls, temperature=temperature)
+                formatted_response = f"{main_content}\n\nSources:\n{sources}"
+            else:
+                response = get_response_from_pdf(message, model, num_calls=num_calls, temperature=temperature)
+                formatted_response = response
+
+            if not stop_clicked:
+                history.append((message, formatted_response))
+        except Exception as e:
+            print(f"Error generating response: {str(e)}")
+            history.append((message, "I'm sorry, but I encountered an error while generating the response. Please try again."))
+        
+        is_generating = False
+        return "", history, is_generating, stop_clicked
+
+    submit_btn.click(
+        protected_generate_response,
+        inputs=[msg, chatbot, use_web_search, model_dropdown, temperature_slider, num_calls_slider, is_generating, stop_clicked],
+        outputs=[msg, chatbot, is_generating, stop_clicked],
+        show_progress=True
+    ).then(
+        lambda: gr.update(visible=True),
+        None,
+        stop_btn
+    ).then(
+        lambda: gr.update(visible=False),
+        None,
+        stop_btn
+    )
+
+    stop_btn.click(
+        lambda: (True, gr.update(visible=False)),
+        None,
+        [stop_clicked, stop_btn]
+    )
+
+    retry_btn.click(
+        retry_last_response,
+        inputs=[chatbot],
+        outputs=[msg, chatbot]
+    ).then(
+        protected_generate_response,
+        inputs=[msg, chatbot, use_web_search, model_dropdown, temperature_slider, num_calls_slider, is_generating, stop_clicked],
+        outputs=[msg, chatbot, is_generating, stop_clicked]
+    )
+
+    undo_btn.click(undo_last_interaction, inputs=[chatbot], outputs=[chatbot])
+    clear_btn.click(clear_conversation, outputs=[chatbot])
 
     gr.Examples(
         examples=[
@@ -238,13 +310,6 @@ with gr.Blocks() as demo:
         inputs=msg,
     )
 
-    submit.click(protected_clear_and_update_chat, 
-                 inputs=[msg, chatbot, use_web_search, model_dropdown, temperature_slider, is_generating], 
-                 outputs=[msg, chatbot, is_generating])
-    msg.submit(protected_clear_and_update_chat, 
-               inputs=[msg, chatbot, use_web_search, model_dropdown, temperature_slider, is_generating], 
-               outputs=[msg, chatbot, is_generating])
-
     gr.Markdown(
     """
     ## How to use
@@ -252,8 +317,10 @@ with gr.Blocks() as demo:
     2. Select the PDF parser (pypdf or llamaparse) and click "Upload Document" to update the vector store.
     3. Ask questions in the textbox. 
     4. Toggle "Use Web Search" to switch between PDF chat and web search.
-    5. Adjust Temperature and Repetition Penalty sliders to fine-tune the response generation.
-    6. Click "Submit" or press Enter to get a response.
+    5. Adjust Temperature and Number of API Calls sliders to fine-tune the response generation.
+    6. Click "Send" or press Enter to get a response.
+    7. Use "Retry" to regenerate the last response, "Undo" to remove the last interaction, and "Clear" to reset the conversation.
+    8. Click "Stop" during generation to halt the process.
     """
     )