Update app.py

app.py

@@ -188,15 +188,31 @@ 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):
-    if not message.strip():
+def chatbot_interface(message, history, use_web_search, model, temperature, num_calls, continue_generation=False):
+    if not message.strip() and not continue_generation:
         return "", history
 
-    history = history + [(message, "")]
+    if continue_generation and history:
+        last_message, last_response = history[-1]
+        continuation_prompt = f"""
+        Original query: {last_message}
+        
+        Previously generated response:
+        {last_response}
+        
+        Please continue the response from where it left off, maintaining coherence and relevance to the original query.
+        """
+        history = history[:-1]  # Remove the last response to replace it with the continued one
+        message = continuation_prompt
+    else:
+        history = history + [(message, "")]
 
     try:
         for response in respond(message, history, model, temperature, num_calls, use_web_search):
-            history[-1] = (message, response)
+            if continue_generation and history:
+                history.append((history[-1][0], history[-1][1] + response))
+            else:
+                history[-1] = (message, response)
             yield history
     except gr.CancelledError:
         yield history
@@ -425,13 +441,15 @@ css = """
 # Define the checkbox outside the demo block
 use_web_search = gr.Checkbox(label="Use Web Search", value=False)
 
+# Modify your main chat interface to use the same function
 demo = gr.ChatInterface(
-    respond,
+    fn=chatbot_interface,
     additional_inputs=[
+        use_web_search,
         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"),
         gr.Slider(minimum=1, maximum=5, value=1, step=1, label="Number of API Calls"),
-        use_web_search  # Add this line to include the checkbox
+        gr.State(False)  # Flag for continue_generation, default False for normal chat
     ],
     title="AI-powered Web Search and PDF Chat Assistant",
     description="Chat with your PDFs or use web search to answer questions.",
@@ -464,16 +482,22 @@ demo = gr.ChatInterface(
     analytics_enabled=False,
 )
 
-# In the Gradio interface
+# In your Gradio interface setup
 with demo:
     continue_button = gr.Button("Continue Generation")
-    continue_button.click(continue_generation, 
-                          inputs=[gr.State(lambda: history[-1][0] if history else ""), 
-                                  gr.State(lambda: history[-1][1] if history else ""),
-                                  gr.Dropdown(choices=MODELS, label="Select Model"),
-                                  gr.Slider(label="Temperature"),
-                                  gr.Slider(label="Number of API Calls")],
-                          outputs=gr.Chatbot())
+    continue_button.click(
+        chatbot_interface,
+        inputs=[
+            gr.Textbox(value="", visible=False),  # Hidden textbox for the message
+            gr.State([]),  # Chat history
+            use_web_search,
+            gr.Dropdown(choices=MODELS, label="Select Model"),
+            gr.Slider(label="Temperature"),
+            gr.Slider(label="Number of API Calls"),
+            gr.State(True)  # Flag for continue_generation
+        ],
+        outputs=gr.Chatbot()
+    )
     
 # Add file upload functionality
 with demo: