Update app.py

app.py

@@ -226,11 +226,15 @@ class CitingSources(BaseModel):
         ...,
         description="List of sources to cite. Should be an URL of the source."
     )
-def chatbot_interface(message, history, use_web_search, model, temperature, num_calls):
+def chatbot_interface(message, history, model, temperature, num_calls, use_web_search, selected_docs):
     if not message.strip():
-        return "", history
+        return history
 
-    history = history + [(message, "")]
+    for response in respond(message, history, model, temperature, num_calls, use_web_search, selected_docs):
+        yield response
+
+    # Make the Continue Generation button visible after the response is complete
+    demo.update(visible=True, elem_id="continue_btn")
 
     try:
         for response in respond(message, history, model, temperature, num_calls, use_web_search):
@@ -254,7 +258,7 @@ def retry_last_response(history, use_web_search, model, temperature, num_calls):
 
 def respond(message, history, model, temperature, num_calls, use_web_search, selected_docs):
     if not message.strip():
-        return "", history
+        return history
 
     history = history + [(message, "")]
 
@@ -270,13 +274,14 @@ def respond(message, history, model, temperature, num_calls, use_web_search, sel
                 yield history
     except Exception as e:
         if "cancelled" in str(e).lower():
-            # This is likely a cancellation, so we'll just return the current history
             yield history
         else:
             logging.error(f"Unexpected error in respond: {str(e)}")
             history[-1] = (message, f"An unexpected error occurred: {str(e)}")
             yield history
 
+    return history  # Ensure we always return the history at the end
+
 logging.basicConfig(level=logging.DEBUG)
 
 def get_response_from_cloudflare(prompt, context, query, num_calls=3, temperature=0.2, search_type="pdf"):
@@ -449,7 +454,7 @@ def continue_generation(history, model, temperature, num_calls, use_web_search,
     """
     
     try:
-        for response in respond(continuation_prompt, history[:-1], use_web_search, model, temperature, num_calls, selected_docs, continuation=True):
+        for response in respond(continuation_prompt, history[:-1], model, temperature, num_calls, use_web_search, selected_docs):
             new_response = f"{last_ai_response}\n\n{response[-1][1]}"
             history[-1] = (last_user_msg, new_response)
             yield history
@@ -458,6 +463,8 @@ def continue_generation(history, model, temperature, num_calls, use_web_search,
         history[-1] = (last_user_msg, f"{last_ai_response}\n\nError continuing generation: {str(e)}")
         yield history
 
+    return history
+
 css = """
 /* Add your custom CSS here */
 """
@@ -477,7 +484,7 @@ document_selector = gr.CheckboxGroup(label="Select documents to query")
 use_web_search = gr.Checkbox(label="Use Web Search", value=False)
 
 demo = gr.ChatInterface(
-    respond,
+    chatbot_interface,
     additional_inputs=[
         gr.Dropdown(choices=MODELS, label="Select Model", value=MODELS[0]),
         gr.Slider(minimum=0.1, maximum=1.0, value=0.2, step=0.1, label="Temperature"),
@@ -531,12 +538,9 @@ with demo:
                         inputs=[file_input, parser_dropdown], 
                         outputs=[update_output, document_selector])
 
-    # Add the Continue Generation button
-    continue_btn = gr.Button("Continue Generation", visible=False)
-    
-    # Add the click event for the Continue Generation button
+    continue_btn = gr.Button("Continue Generation", visible=False, elem_id="continue_btn")
     continue_btn.click(continue_generation, 
-                       inputs=[demo.chatbot] + demo.additional_inputs,  # Use existing inputs
+                       inputs=[demo.chatbot] + demo.additional_inputs,
                        outputs=[demo.chatbot])
 
     gr.Markdown(