Update app.py

app.py

@@ -188,14 +188,16 @@ 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):
+Here's a modified version of the `chatbot_interface` function that includes the "Continue Generation" button logic:
+
+```
+def chatbot_interface(message, history, use_web_search, model, temperature, num_calls, is_interrupted=False, partial_response=""):
     if not message.strip():
         return "", history
 
     history = history + [(message, "")]
-
     try:
-        for response in respond(message, history, model, temperature, num_calls, use_web_search):
+        for response in respond(message, history, model, temperature, num_calls, use_web_search, is_interrupted, partial_response):
             history[-1] = (message, response)
             yield history
     except gr.CancelledError:
@@ -214,7 +216,7 @@ def retry_last_response(history, use_web_search, model, temperature, num_calls):
     
     return chatbot_interface(last_user_msg, history, use_web_search, model, temperature, num_calls)
 
-def respond(message, history, model, temperature, num_calls, use_web_search):
+def respond(message, history, model, temperature, num_calls, use_web_search, is_interrupted=False, partial_response=""):
     logging.info(f"User Query: {message}")
     logging.info(f"Model Used: {model}")
     logging.info(f"Search Type: {'Web Search' if use_web_search else 'PDF Search'}")
@@ -223,9 +225,11 @@ def respond(message, history, model, temperature, num_calls, use_web_search):
         if use_web_search:
             for main_content, sources in get_response_with_search(message, model, num_calls=num_calls, temperature=temperature):
                 response = f"{main_content}\n\n{sources}"
-                first_line = response.split('\n')[0] if response else ''
-                logging.info(f"Generated Response (first line): {first_line}")
-                yield response
+                if is_interrupted:
+                    partial_response += response
+                    yield partial_response
+                else:
+                    yield response
         else:
             embed = get_embeddings()
             if os.path.exists("faiss_database"):
@@ -235,19 +239,21 @@ def respond(message, history, model, temperature, num_calls, use_web_search):
                 context_str = "\n".join([doc.page_content for doc in relevant_docs])
             else:
                 context_str = "No documents available."
-            
+
             if model == "@cf/meta/llama-3.1-8b-instruct":
-                # Use Cloudflare API
-                for partial_response in get_response_from_cloudflare(prompt="", context=context_str, query=message, num_calls=num_calls, temperature=temperature, search_type="pdf"):
-                    first_line = partial_response.split('\n')[0] if partial_response else ''
-                    logging.info(f"Generated Response (first line): {first_line}")
-                    yield partial_response
+                for response in get_response_from_cloudflare(prompt="", context=context_str, query=message, num_calls=num_calls, temperature=temperature, search_type="pdf"):
+                    if is_interrupted:
+                        partial_response += response
+                        yield partial_response
+                    else:
+                        yield response
             else:
-                # Use Hugging Face API
-                for partial_response in get_response_from_pdf(message, model, num_calls=num_calls, temperature=temperature):
-                    first_line = partial_response.split('\n')[0] if partial_response else ''
-                    logging.info(f"Generated Response (first line): {first_line}")
-                    yield partial_response
+                for response in get_response_from_pdf(message, model, num_calls=num_calls, temperature=temperature):
+                    if is_interrupted:
+                        partial_response += response
+                        yield partial_response
+                    else:
+                        yield response
     except Exception as e:
         logging.error(f"Error with {model}: {str(e)}")
         if "microsoft/Phi-3-mini-4k-instruct" in model:
@@ -255,8 +261,7 @@ def respond(message, history, model, temperature, num_calls, use_web_search):
             fallback_model = "mistralai/Mistral-7B-Instruct-v0.3"
             yield from respond(message, history, fallback_model, temperature, num_calls, use_web_search)
         else:
-            yield f"An error occurred with the {model} model: {str(e)}. Please try again or select a different model."
-
+            yield f"An error occurred with the {model} model: {str(e)}. Please try again or select a different model.
 logging.basicConfig(level=logging.DEBUG)
 
 def get_response_from_cloudflare(prompt, context, query, num_calls=3, temperature=0.2, search_type="pdf"):
@@ -395,12 +400,14 @@ css = """
 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"),
         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
+        use_web_search,
+        gr.Boolean(value=False, label="Is Interrupted"),  # Flag to indicate interruption
+        gr.Textbox(label="Partial Response"),  # Store partial response
     ],
     title="AI-powered Web Search and PDF Chat Assistant",
     description="Chat with your PDFs or use web search to answer questions.",
@@ -422,6 +429,30 @@ demo = gr.ChatInterface(
         color_accent_soft_dark="transparent",
         code_background_fill_dark="#140b0b"
     ),
+    css=css,
+    examples=[
+        ["Tell me about the contents of the uploaded PDFs."],
+        ["What are the main topics discussed in the documents?"],
+        ["Can you summarize the key points from the PDFs?"],
+    ],
+    cache_examples=False,
+    analytics_enabled=False,
+)
+
+continue_button = gr.Button("Continue Generation")
+continue_output = gr.Textbox(label="Response")
+
+continue_button.click(
+    chatbot_interface,
+    inputs=[
+        message_input,
+        history,
+        use_web_search,
+        model_dropdown,
+        temperature_slider,
+        num_calls_slider,
+        gr.Boolean(value=True, label="Is Interrupted"),  # Set is_interrupted to True
+        gr.Textbox(label="Partial Response"),
     
     css=css,
     examples=[