Update app.py

app.py

@@ -477,24 +477,6 @@ def create_web_search_vectors(search_results):
     
     return FAISS.from_documents(documents, embed)
 
-def critique_response(response, context, query):
-    critique_prompt = f"""Given the following response, original context, and user query, identify any statements that might be inaccurate, unsupported by the context, or irrelevant to the query. Be specific about which parts may be hallucinations or extrapolations beyond the given information.
-
-User Query: {query}
-
-Response:
-{response}
-
-Original Context:
-{context}
-
-Critique:"""
-
-    client = InferenceClient(model, token=huggingface_token)
-    critique = client.text_generation(critique_prompt, max_new_tokens=500, temperature=0.2)
-    
-    return critique
-
 def get_response_with_search(query, model, num_calls=3, temperature=0.1):
     search_results = duckduckgo_search(query)
     web_search_database = create_web_search_vectors(search_results)
@@ -513,15 +495,16 @@ def get_response_with_search(query, model, num_calls=3, temperature=0.1):
 Write a detailed and complete response that answers the following user query: '{query}'
 Stick closely to the information provided in the context and avoid making unsupported claims."""
 
-    client = InferenceClient(model, token=huggingface_token)
-    
-    # Generate initial response
-    initial_response = client.text_generation(initial_prompt, max_new_tokens=1000, temperature=temperature)
-    
-    # Generate critique
-    critique = critique_response(initial_response, context, query)
-    
-    final_prompt = f"""Given the following initial response, context, critique, and original query, provide a revised response that addresses the identified issues and sticks closely to the information provided in the context while fully answering the user's query.
+    try:
+        client = InferenceClient(model, token=huggingface_token)
+        
+        # Generate initial response
+        initial_response = client.text_generation(initial_prompt, max_new_tokens=1000, temperature=temperature)
+        
+        # Generate critique
+        critique = critique_response(initial_response, context, query, model)
+        
+        final_prompt = f"""Given the following initial response, context, critique, and original query, provide a revised response that addresses the identified issues and sticks closely to the information provided in the context while fully answering the user's query.
 
 User Query: {query}
 
@@ -536,14 +519,36 @@ Critique:
 
 Revised Response:"""
 
-    # Generate final response
-    for chunk in client.text_generation(final_prompt, max_new_tokens=1500, temperature=temperature, stream=True):
-        yield chunk, ""
+        # Generate final response
+        for chunk in client.text_generation(final_prompt, max_new_tokens=1500, temperature=temperature, stream=True):
+            yield chunk, ""
+
+        # Add a disclaimer
+        disclaimer = ("\nNote: This response was generated by an AI model based on web search results. "
+                      "While efforts have been made to ensure accuracy, please verify important information from authoritative sources.")
+        yield disclaimer, ""
+    
+    except Exception as e:
+        logging.error(f"Error in multi-step generation process: {str(e)}")
+        yield f"An error occurred during the response generation process: {str(e)}", ""
+
+def critique_response(response, context, query, model):
+    critique_prompt = f"""Given the following response, original context, and user query, identify any statements that might be inaccurate, unsupported by the context, or irrelevant to the query. Be specific about which parts may be hallucinations or extrapolations beyond the given information.
+
+User Query: {query}
+
+Response:
+{response}
+
+Original Context:
+{context}
 
-    # Add a disclaimer
-    disclaimer = ("\nNote: This response was generated by an AI model based on web search results. "
-                  "While efforts have been made to ensure accuracy, please verify important information from authoritative sources.")
-    yield disclaimer, ""
+Critique:"""
+
+    client = InferenceClient(model, token=huggingface_token)
+    critique = client.text_generation(critique_prompt, max_new_tokens=500, temperature=0.2)
+    
+    return critique
 
 
 INSTRUCTION_PROMPTS = {