Update app.py

app.py

@@ -409,6 +409,36 @@ def get_context_for_query(query, selected_docs):
     else:
         return "No documents available to answer the query."
 
+def validate_response(initial_response, context, query, model, temperature=0.1):
+    validation_prompt = f"""Given the following context and initial response to the query "{query}":
+
+Context:
+{context}
+
+Initial Response:
+{initial_response}
+
+Please validate the initial response against the provided context. Remove any hallucinations, irrelevant, or factually incorrect information. Generate a revised response that is accurate and directly supported by the context. If any information cannot be verified from the context, state that explicitly.
+
+Revised Response:
+"""
+
+    if model == "@cf/meta/llama-3.1-8b-instruct":
+        return get_response_from_cloudflare(prompt=validation_prompt, context="", query="", num_calls=1, temperature=temperature, search_type="validation")
+    else:
+        client = InferenceClient(model, token=huggingface_token)
+        revised_response = ""
+        for message in client.chat_completion(
+            messages=[{"role": "user", "content": validation_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
+                revised_response += chunk
+                yield revised_response
+
 def get_response_from_cloudflare(prompt, context, query, num_calls=3, temperature=0.2, search_type="pdf"):
     headers = {
         "Authorization": f"Bearer {API_TOKEN}",
@@ -420,15 +450,19 @@ def get_response_from_cloudflare(prompt, context, query, num_calls=3, temperatur
         instruction = f"""Using the following context from the PDF documents:
 {context}
 Write a detailed and complete response that answers the following user question: '{query}'"""
-    else:  # web search
+    elif search_type == "web":
         instruction = f"""Using the following context:
 {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."""
+    elif search_type == "validation":
+        instruction = prompt  # For validation, use the provided prompt directly
+    else:
+        raise ValueError("Invalid search_type")
 
     inputs = [
         {"role": "system", "content": instruction},
-        {"role": "user", "content": query}
+        {"role": "user", "content": query if search_type != "validation" else ""}
     ]
 
     payload = {
@@ -492,18 +526,17 @@ def get_response_with_search(query, model, num_calls=3, temperature=0.1):
 {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.
-Importantly, only include information that is directly supported by the provided context. If you're unsure about any information, state that it couldn't be verified from the given context.
-After writing the document, please provide a list of sources used in your response."""
+Importantly, only include information that is directly supported by the provided context. If you're unsure about any information, state that it couldn't be verified from the given context."""
 
+    initial_response = ""
     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
+            initial_response = response
     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}],
@@ -513,8 +546,11 @@ After writing the document, please provide a list of sources used in your respon
             ):
                 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
+                    initial_response += chunk
+
+    # Validation step
+    for revised_response in validate_response(initial_response, context, query, model, temperature):
+        yield revised_response, ""  # Yield streaming revised response without sources
 
 
 INSTRUCTION_PROMPTS = {