Update app.py

app.py

@@ -20,6 +20,7 @@ from huggingface_hub import InferenceClient
 import inspect
 import logging
 import shutil
+import tempfile
 
 
 # Set up basic configuration for logging
@@ -274,22 +275,19 @@ def generate_chunked_response(prompt, model, max_tokens=10000, num_calls=3, temp
     return final_response
 
 def get_web_search_database():
-    if not hasattr(get_web_search_database, "database"):
-        embed = get_embeddings()
-        if os.path.exists("web_search_database"):
-            get_web_search_database.database = WebSearchFAISS.load_local("web_search_database", embed, allow_dangerous_deserialization=True)
-        else:
-            get_web_search_database.database = WebSearchFAISS.from_documents([], embed)
-    return get_web_search_database.database
+    embed = get_embeddings()
+    temp_dir = tempfile.mkdtemp()
+    database = WebSearchFAISS.from_documents([], embed)
+    return database, temp_dir
 
-def save_web_search_database(database):
-    database.save_local("web_search_database")
+def cleanup_web_search_database(temp_dir):
+    shutil.rmtree(temp_dir)
 
 def duckduckgo_search(query):
     with DDGS() as ddgs:
         results = list(ddgs.text(query, max_results=5))
     
-    database = get_web_search_database()
+    database, temp_dir = get_web_search_database()
     documents = []
     for result in results:
         content = f"{result['title']}\n{result['body']}"
@@ -297,12 +295,9 @@ def duckduckgo_search(query):
         documents.append(doc)
     
     database.add_documents(documents)
-    save_web_search_database(database)
-    
-    return results
+    return database, temp_dir, results
 
-def retrieve_web_search_results(query):
-    database = get_web_search_database()
+def retrieve_web_search_results(database, query):
     retriever = database.as_retriever(search_kwargs={"k": 5})
     relevant_docs = retriever.get_relevant_documents(query)
     
@@ -453,34 +448,38 @@ After writing the document, please provide a list of sources used in your respon
         yield "I apologize, but I couldn't generate a response at this time. Please try again later."
 
 def get_response_with_search(query, model, num_calls=3, temperature=0.2):
-    search_results = duckduckgo_search(query)
-    context = retrieve_web_search_results(query)
+    database, temp_dir, search_results = duckduckgo_search(query)
+    context = retrieve_web_search_results(database, query)
     
     prompt = f"""Using the following context from web search results:
 {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."""
 
-    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
-    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}],
-                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
-                    main_content += chunk
-                    yield main_content, ""  # Yield partial main content without sources
+    try:
+        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
+        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}],
+                    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
+                        main_content += chunk
+                        yield main_content, ""  # Yield partial main content without sources
+    finally:
+        # Clean up the temporary database
+        cleanup_web_search_database(temp_dir)
 
 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}")