Update app.py

app.py

@@ -271,16 +271,45 @@ def generate_chunked_response(prompt, model, max_tokens=10000, num_calls=3, temp
     print(f"Final clean response: {final_response[:100]}...")
     return final_response
 
+def get_web_search_database():
+    embed = get_embeddings()
+    if os.path.exists("web_search_database"):
+        return FAISS.load_local("web_search_database", embed, allow_dangerous_deserialization=True)
+    else:
+        return FAISS.from_documents([], embed)
+
+def save_web_search_database(database):
+    database.save_local("web_search_database")
+
+
+
 def duckduckgo_search(query):
     with DDGS() as ddgs:
-        results = ddgs.text(query, max_results=5)
+        results = list(ddgs.text(query, max_results=5))
+    
+    database = get_web_search_database()
+    documents = []
+    for result in results:
+        content = f"{result['title']}\n{result['body']}"
+        doc = Document(page_content=content, metadata={"source": result['href']})
+        documents.append(doc)
+    
+    database.add_documents(documents)
+    save_web_search_database(database)
+    
     return results
 
+def retrieve_web_search_results(query):
+    database = get_web_search_database()
+    retriever = database.as_retriever(search_kwargs={"k": 5})
+    return retriever.get_relevant_documents(query)
+
 class CitingSources(BaseModel):
     sources: List[str] = Field(
         ...,
         description="List of sources to cite. Should be an URL of the source."
     )
+
 def chatbot_interface(message, history, use_web_search, model, temperature, num_calls):
     if not message.strip():
         return "", history
@@ -318,8 +347,6 @@ def respond(message, history, model, temperature, num_calls, use_web_search, sel
         if use_web_search:
             for main_content, sources in get_response_with_search(message, model, num_calls=num_calls, temperature=temperature):
                 response = f"{main_content}\n\n{sources}"
-                first_line = response.split('\n')[0] if response else ''
-#                logging.info(f"Generated Response (first line): {first_line}")
                 yield response
         else:
             embed = get_embeddings()
@@ -421,8 +448,10 @@ After writing the document, please provide a list of sources used in your respon
 
 def get_response_with_search(query, model, num_calls=3, temperature=0.2):
     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)
+    relevant_docs = retrieve_web_search_results(query)
+    
+    context = "\n".join([doc.page_content for doc in relevant_docs])
+    sources = "\n".join([f"Source: {doc.metadata['source']}" for doc in relevant_docs])
     
     prompt = f"""Using the following context:
 {context}
@@ -432,7 +461,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, search_type="web"):
-            yield response, ""  # Yield streaming response without sources
+            yield response, sources
     else:
         # Use Hugging Face API
         client = InferenceClient(model, token=huggingface_token)
@@ -448,7 +477,7 @@ 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
+                    yield main_content, sources
 
 def get_response_from_pdf(query, model, selected_docs, num_calls=3, temperature=0.2):
     logging.info(f"Entering get_response_from_pdf with query: {query}, model: {model}, selected_docs: {selected_docs}")