Update app.py

app.py

@@ -475,7 +475,25 @@ def create_web_search_vectors(search_results):
     
     return FAISS.from_documents(documents, embed)
 
-def get_response_with_search(query, model, num_calls=3, temperature=0.2):
+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)
     
@@ -483,36 +501,47 @@ def get_response_with_search(query, model, num_calls=3, temperature=0.2):
         yield "No web search results available. Please try again.", ""
         return
     
-    retriever = web_search_database.as_retriever(search_kwargs={"k": 5})
+    retriever = web_search_database.as_retriever(search_kwargs={"k": 10})
     relevant_docs = retriever.get_relevant_documents(query)
     
-    context = "\n".join([doc.page_content for doc in relevant_docs])
+    context = "\n".join([doc.page_content for doc in relevant_docs[:5]])
     
-    prompt = f"""Using the following context from web search results:
+    initial_prompt = f"""Using the following context from web search results:
 {context}
-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."""
+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."""
 
-    if model == "@cf/meta/llama-3.1-8b-instruct":
-        # Use Cloudflare API
-        for response in get_response_from_cloudflare(prompt="", context=context, query=query, num_calls=num_calls, temperature=temperature, search_type="web"):
-            yield response, ""  # Yield streaming response without sources
-    else:
-        # Use Hugging Face API
-        client = InferenceClient(model, token=huggingface_token)
-        
-        main_content = ""
-        for i in range(num_calls):
-            for message in client.chat_completion(
-                messages=[{"role": "user", "content": prompt}],
-                max_tokens=10000,
-                temperature=temperature,
-                stream=True,
-            ):
-                if message.choices and message.choices[0].delta and message.choices[0].delta.content:
-                    chunk = message.choices[0].delta.content
-                    main_content += chunk
-                    yield main_content, ""  # Yield partial main content without sources
+    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.
+
+User Query: {query}
+
+Initial Response:
+{initial_response}
+
+Context:
+{context}
+
+Critique:
+{critique}
+
+Revised Response:"""
+
+    # 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, ""
 
 
 INSTRUCTION_PROMPTS = {