Update app.py

app.py

@@ -86,7 +86,7 @@ def generate_chunked_response(prompt, model, max_tokens=1000, num_calls=3, tempe
     
     for i in range(num_calls):
         print(f"Starting API call {i+1}")
-        if stop_clicked and stop_clicked.value:  # Check if stop_clicked is not None and its value is True
+        if (isinstance(stop_clicked, gr.State) and stop_clicked.value) or stop_clicked:
             print("Stop clicked, breaking loop")
             break
         try:
@@ -97,7 +97,7 @@ def generate_chunked_response(prompt, model, max_tokens=1000, num_calls=3, tempe
                 temperature=temperature,
                 stream=True,
             ):
-                if stop_clicked and stop_clicked.value:
+                if (isinstance(stop_clicked, gr.State) and stop_clicked.value) or stop_clicked:
                     print("Stop clicked during streaming, breaking")
                     break
                 if message.choices and message.choices[0].delta and message.choices[0].delta.content:
@@ -255,23 +255,29 @@ with gr.Blocks() as demo:
         if is_generating:
             print("Already generating, returning")
             return message, history, is_generating, stop_clicked
+    
         is_generating = True
-        stop_clicked.value = False  # Reset stop_clicked at the start of generation
         
+        # Handle both boolean and Gradio State object
+        if isinstance(stop_clicked, gr.State):
+            stop_clicked.value = False
+        else:
+            stop_clicked = False
+    
         try:
             print(f"Generating response for: {message}")
             if use_web_search:
                 print("Using web search")
-                main_content, sources = get_response_with_search(query, model, num_calls=num_calls, temperature=temperature)
+                main_content, sources = get_response_with_search(message, model, num_calls=num_calls, temperature=temperature, stop_clicked=stop_clicked)
                 formatted_response = f"{main_content}\n\nSources:\n{sources}"
             else:
                 print("Using PDF search")
-                response = get_response_from_pdf(query, model, num_calls=num_calls, temperature=temperature)
+                response = get_response_from_pdf(message, model, num_calls=num_calls, temperature=temperature, stop_clicked=stop_clicked)
                 formatted_response = response
     
             print(f"Generated response: {formatted_response[:100]}...")
     
-            if not stop_clicked.value:
+            if not (isinstance(stop_clicked, gr.State) and stop_clicked.value) and not stop_clicked:
                 print("Appending to history")
                 history.append((message, formatted_response))
             else:
@@ -283,7 +289,7 @@ with gr.Blocks() as demo:
         is_generating = False
         print(f"Returning history with {len(history)} items")
         return "", history, is_generating, stop_clicked
-
+    
     def on_submit(message, history, use_web_search, model, temperature, num_calls, is_generating, stop_clicked):
         print(f"Submit button clicked with message: {message}")
         _, new_history, new_is_generating, new_stop_clicked = protected_generate_response(