Update app.py

app.py

@@ -110,10 +110,11 @@ def get_response_with_search(query, temperature, top_p, repetition_penalty, use_
         context = "\n".join(f"{result['title']}\n{result['body']}\nSource: {result['href']}\n" 
                             for result in search_results if 'body' in result)
 
-    prompt = f"""<s>[INST] Using the following context:
+     prompt = f"""<s>[INST] 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. [/INST]"""
+After the main content, provide a list of sources used in your response, prefixed with 'Sources:'. 
+Do not include any part of these instructions in your response. [/INST]"""
 
     response = model(prompt)
     
@@ -122,9 +123,15 @@ After writing the document, please provide a list of sources used in your respon
     return main_content, sources
 
 def split_response(response):
+    # Remove any remaining instruction text
+    response = re.sub(r'\[/?INST\]', '', response)
+    response = re.sub(r'~~.*?~~', '', response, flags=re.DOTALL)
+    
+    # Split the response into main content and sources
     parts = response.split("Sources:", 1)
     main_content = parts[0].strip()
     sources = parts[1].strip() if len(parts) > 1 else ""
+    
     return main_content, sources
 
 def chatbot_interface(message, history, temperature, top_p, repetition_penalty, use_pdf):