Update app.py

app.py

@@ -78,7 +78,7 @@ def update_vectors(files, parser):
     
     return f"Vector store updated successfully. Processed {total_chunks} chunks from {len(files)} files using {parser}."
 
-def generate_chunked_response(prompt, model, max_tokens=1000, num_calls=3, temperature=0.2, stop_clicked=False):
+def generate_chunked_response(prompt, model, max_tokens=1000, num_calls=3, temperature=0.2, stop_clicked=None):
     print(f"Starting generate_chunked_response with {num_calls} calls")
     client = InferenceClient(model, token=huggingface_token)
     full_responses = []
@@ -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():  # Assume stop_clicked is now a function that returns the current state
+        if stop_clicked and stop_clicked.value:  # Check if stop_clicked is not None and its value is True
             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():
+                if stop_clicked and stop_clicked.value:
                     print("Stop clicked during streaming, breaking")
                     break
                 if message.choices and message.choices[0].delta and message.choices[0].delta.content:
@@ -129,7 +129,7 @@ class CitingSources(BaseModel):
         description="List of sources to cite. Should be an URL of the source."
     )
 
-def get_response_with_search(query, model, num_calls=3, temperature=0.2):
+def get_response_with_search(query, model, num_calls=3, temperature=0.2, stop_clicked=None):
     search_results = duckduckgo_search(query)
     context = "\n".join(f"{result['title']}\n{result['body']}\nSource: {result['href']}\n" 
                         for result in search_results if 'body' in result)
@@ -139,8 +139,8 @@ def get_response_with_search(query, model, num_calls=3, temperature=0.2):
 Write a detailed and complete research document that fulfills the following user request: '{query}'
 After writing the document, please provide a list of sources used in your response. [/INST]"""
     
-    generated_text = generate_chunked_response(prompt, model, num_calls=num_calls, temperature=temperature)
-    
+    generated_text = generate_chunked_response(prompt, model, num_calls=num_calls, temperature=temperature, stop_clicked=stop_clicked)
+   
     # Clean the response
     clean_text = re.sub(r'<s>\[INST\].*?\[/INST\]\s*', '', generated_text, flags=re.DOTALL)
     clean_text = clean_text.replace("Using the following context:", "").strip()
@@ -152,7 +152,7 @@ After writing the document, please provide a list of sources used in your respon
     
     return main_content, sources
 
-def get_response_from_pdf(query, model, num_calls=3, temperature=0.2):
+def get_response_from_pdf(query, model, num_calls=3, temperature=0.2, stop_clicked=None):
     embed = get_embeddings()
     if os.path.exists("faiss_database"):
         database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
@@ -168,7 +168,7 @@ def get_response_from_pdf(query, model, num_calls=3, temperature=0.2):
 Write a detailed and complete response that answers the following user question: '{query}'
 Do not include a list of sources in your response. [/INST]"""
 
-    generated_text = generate_chunked_response(prompt, model, num_calls=num_calls, temperature=temperature)
+    generated_text = generate_chunked_response(prompt, model, num_calls=num_calls, temperature=temperature, stop_clicked=stop_clicked)
 
     # Clean the response
     clean_text = re.sub(r'<s>\[INST\].*?\[/INST\]\s*', '', generated_text, flags=re.DOTALL)
@@ -256,22 +256,22 @@ with gr.Blocks() as demo:
             print("Already generating, returning")
             return message, history, is_generating, stop_clicked
         is_generating = True
-        stop_clicked = False  # Reset stop_clicked at the start of generation
+        stop_clicked.value = False  # Reset stop_clicked at the start of generation
         
         try:
             print(f"Generating response for: {message}")
             if use_web_search:
                 print("Using web search")
-                main_content, sources = get_response_with_search(message, model, num_calls=num_calls, temperature=temperature)
+                main_content, sources = get_response_with_search(query, model, num_calls=num_calls, temperature=temperature)
                 formatted_response = f"{main_content}\n\nSources:\n{sources}"
             else:
                 print("Using PDF search")
-                response = get_response_from_pdf(message, model, num_calls=num_calls, temperature=temperature)
+                response = get_response_from_pdf(query, model, num_calls=num_calls, temperature=temperature)
                 formatted_response = response
     
             print(f"Generated response: {formatted_response[:100]}...")
     
-            if not stop_clicked():
+            if not stop_clicked.value:
                 print("Appending to history")
                 history.append((message, formatted_response))
             else:
@@ -291,7 +291,7 @@ with gr.Blocks() as demo:
         )
         print(f"New history has {len(new_history)} items")
         return "", new_history, new_is_generating, new_stop_clicked
-    
+
     submit_btn.click(
         on_submit,
         inputs=[msg, chatbot, use_web_search, model_dropdown, temperature_slider, num_calls_slider, is_generating, stop_clicked],