Update app.py

app.py

@@ -110,28 +110,6 @@ class CitingSources(BaseModel):
         description="List of sources to cite. Should be an URL of the source."
     )
 
-def get_response_with_search(query):
-    search_results = duckduckgo_search(query)
-    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:
-{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]"""
-    
-    generated_text = generate_chunked_response(prompt)
-    
-    content_start = generated_text.find("[/INST]")
-    if content_start != -1:
-        generated_text = generated_text[content_start + 7:].strip()
-    
-    parts = generated_text.split("Sources:", 1)
-    main_content = parts[0].strip()
-    sources = parts[1].strip() if len(parts) > 1 else ""
-    
-    return main_content, sources
-
 def get_response_from_pdf(query):
     embed = get_embeddings()
     if os.path.exists("faiss_database"):
@@ -150,18 +128,40 @@ After writing the response, please provide a list of sources used (document name
 
     generated_text = generate_chunked_response(prompt)
 
-    # Remove the instruction part from the response
-    content_start = generated_text.find("[/INST]")
-    if content_start != -1:
-        generated_text = generated_text[content_start + 7:].strip()
+    # Remove all instruction-related content
+    clean_text = re.sub(r'<s>\[INST\].*?\[/INST\]\s*', '', generated_text, flags=re.DOTALL)
+    clean_text = clean_text.replace("Using the following context from the PDF documents:", "").strip()
 
     # Split the content and sources
-    parts = generated_text.split("Sources:", 1)
+    parts = clean_text.split("Sources:", 1)
     main_content = parts[0].strip()
     sources = parts[1].strip() if len(parts) > 1 else ""
 
     return main_content, sources
 
+def get_response_with_search(query):
+    search_results = duckduckgo_search(query)
+    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:
+{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]"""
+    
+    generated_text = generate_chunked_response(prompt)
+    
+    # Remove all instruction-related content
+    clean_text = re.sub(r'<s>\[INST\].*?\[/INST\]\s*', '', generated_text, flags=re.DOTALL)
+    clean_text = clean_text.replace("Using the following context:", "").strip()
+    
+    # Split the content and sources
+    parts = clean_text.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, use_web_search):
     if use_web_search:
         main_content, sources = get_response_with_search(message)