Update app.py

app.py

@@ -381,6 +381,37 @@ Write a detailed and complete response that answers the following user question:
                     response += chunk
                     yield response  # Yield partial response
 
+def continue_generation(message, last_response, model, temperature, num_calls):
+    continue_prompt = f"""
+    Original query: {message}
+    
+    Previously generated response:
+    {last_response}
+    
+    Please continue the response from where it left off, maintaining coherence and relevance to the original query.
+    """
+    
+    if model == "@cf/meta/llama-3.1-8b-instruct":
+        # Use Cloudflare API
+        for response in get_response_from_cloudflare(prompt=continue_prompt, context="", query="", num_calls=num_calls, temperature=temperature, search_type="continue"):
+            yield last_response + response
+    else:
+        # Use Hugging Face API
+        client = InferenceClient(model, token=huggingface_token)
+        
+        continued_response = last_response
+        for i in range(num_calls):
+            for message in client.chat_completion(
+                messages=[{"role": "user", "content": continue_prompt}],
+                max_tokens=1000,
+                temperature=temperature,
+                stream=True,
+            ):
+                if message.choices and message.choices[0].delta and message.choices[0].delta.content:
+                    chunk = message.choices[0].delta.content
+                    continued_response += chunk
+                    yield continued_response
+
 def vote(data: gr.LikeData):
     if data.liked:
         print(f"You upvoted this response: {data.value}")
@@ -433,6 +464,17 @@ demo = gr.ChatInterface(
     analytics_enabled=False,
 )
 
+# In the Gradio interface
+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())
+    
 # Add file upload functionality
 with demo:
     gr.Markdown("## Upload PDF Documents")