Update app.py

app.py

@@ -87,29 +87,19 @@ def generate_chunked_response(prompt, model, max_tokens=1000, max_chunks=5, temp
     full_response = ""
     messages = [{"role": "user", "content": prompt}]
     
-    for _ in range(max_chunks):
-        try:
-            chunk_response = ""
-            for message in client.chat_completion(
-                messages=messages,
-                max_tokens=max_tokens,
-                temperature=temperature,
-                stream=True,
-            ):
-                chunk = message.choices[0].delta.content
-                if chunk:
-                    chunk_response += chunk
-                    full_response += chunk
-            
-            if not chunk_response or chunk_response.endswith((".", "!", "?", "</s>", "[/INST]")):
-                break
-            
-            messages.append({"role": "assistant", "content": chunk_response})
-            messages.append({"role": "user", "content": "Continue"})
+    try:
+        for message in client.chat_completion(
+            messages=messages,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            stream=True,
+        ):
+            chunk = message.choices[0].delta.content
+            if chunk:
+                full_response += chunk
         
-        except Exception as e:
-            print(f"Error in generating response: {str(e)}")
-            break
+    except Exception as e:
+        print(f"Error in generating response: {str(e)}")
     
     # Clean up the response
     clean_response = re.sub(r'<s>\[INST\].*?\[/INST\]\s*', '', full_response, flags=re.DOTALL)
@@ -177,6 +167,9 @@ After writing the document, please provide a list of sources used in your respon
     return main_content, sources
 
 def chatbot_interface(message, history, use_web_search, model, temperature):
+    if not message.strip():  # Check if the message is empty or just whitespace
+        return history
+
     if use_web_search:
         main_content, sources = get_response_with_search(message, model, temperature)
         formatted_response = f"{main_content}\n\nSources:\n{sources}"
@@ -184,11 +177,34 @@ def chatbot_interface(message, history, use_web_search, model, temperature):
         response = get_response_from_pdf(message, model, temperature)
         formatted_response = response
 
-    history.append((message, formatted_response))
+    # Check if the last message in history is the same as the current message
+    if history and history[-1][0] == message:
+        # Replace the last response instead of adding a new one
+        history[-1] = (message, formatted_response)
+    else:
+        # Add the new message-response pair
+        history.append((message, formatted_response))
+
     return history
 
+
+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
 with gr.Blocks() as demo:
+    
+    is_generating = 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():
@@ -219,12 +235,12 @@ with gr.Blocks() as demo:
         inputs=msg,
     )
 
-    submit.click(chatbot_interface, 
-             inputs=[msg, chatbot, use_web_search, model_dropdown, temperature_slider], 
-             outputs=[chatbot])
-    msg.submit(chatbot_interface, 
-               inputs=[msg, chatbot, use_web_search, model_dropdown, temperature_slider], 
-               outputs=[chatbot])
+    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(
     """