Update app.py

app.py

@@ -238,7 +238,7 @@ def respond(message, history, model, temperature, num_calls, use_web_search):
                 else:
                     context_str = "No documents available."
                 
-                for partial_response in get_response_from_cloudflare(prompt="", context=context_str, query=message, num_calls=num_calls, temperature=temperature):
+                for partial_response in get_response_from_cloudflare(prompt="", context=context_str, query=message, num_calls=num_calls, temperature=temperature, search_type="pdf"):
                     first_line = partial_response.split('\n')[0] if partial_response else ''
                     logging.info(f"Generated Response (first line): {first_line}")
                     yield partial_response
@@ -259,16 +259,26 @@ def respond(message, history, model, temperature, num_calls, use_web_search):
 
 logging.basicConfig(level=logging.DEBUG)
 
-def get_response_from_cloudflare(prompt, context, query, num_calls=3, temperature=0.2):
+def get_response_from_cloudflare(prompt, context, query, num_calls=3, temperature=0.2, search_type="pdf"):
     headers = {
         "Authorization": f"Bearer {API_TOKEN}",
         "Content-Type": "application/json"
     }
     model = "@cf/meta/llama-3.1-8b-instruct"
 
+    if search_type == "pdf":
+        instruction = f"""Using the following context:
+{context}
+Write a detailed and complete research document that fulfills the following user request: '{query}'"""
+    else:  # web search
+        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."""
+
     inputs = [
-        {"role": "system", "content": "Write a detailed and complete research document that fulfills the following user request: '{query}'"},
-        {"role": "user", "content": f"Context: {context}\n\nQuestion: {query}"}
+        {"role": "system", "content": instruction},
+        {"role": "user", "content": query}
     ]
 
     payload = {
@@ -315,7 +325,7 @@ After writing the document, please provide a list of sources used in your respon
 
     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):
+        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
@@ -348,7 +358,7 @@ def get_response_from_pdf(query, model, num_calls=3, temperature=0.2):
 
     if model == "@cf/meta/llama-3.1-8b-instruct":
         # Use Cloudflare API with the retrieved context
-        for response in get_response_from_cloudflare(prompt="", context=context_str, query=query, num_calls=num_calls, temperature=temperature):
+        for response in get_response_from_cloudflare(prompt="", context=context_str, query=query, num_calls=num_calls, temperature=temperature, search_type="pdf"):
             yield response
     else:
         # Use Hugging Face API