Update app.py

app.py

@@ -8,8 +8,6 @@ from typing import List
 from pydantic import BaseModel, Field
 from tempfile import NamedTemporaryFile
 from langchain_community.vectorstores import FAISS
-from langchain_community.vectorstores import FAISS as WebSearchFAISS
-from langchain_core.documents import Document as WebSearchDocument
 from langchain_core.vectorstores import VectorStore
 from langchain_core.documents import Document
 from langchain_community.document_loaders import PyPDFLoader
@@ -20,8 +18,6 @@ from huggingface_hub import InferenceClient
 import inspect
 import logging
 import shutil
-import tempfile
-from typing import List, Tuple
 
 
 # Set up basic configuration for logging
@@ -275,72 +271,10 @@ 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()
-    temp_dir = tempfile.mkdtemp()
-    
-    try:
-        # Create a dummy document to initialize the database
-        dummy_doc = Document(page_content="Dummy content", metadata={"source": "dummy"})
-        database = FAISS.from_documents([dummy_doc], embed)
-        logging.info("Successfully initialized WebSearchFAISS database with dummy document")
-        # Remove the dummy document
-        database.delete(["dummy"])
-        logging.info("Removed dummy document from database")
-    except Exception as e:
-        logging.error(f"Error initializing WebSearchFAISS: {str(e)}", exc_info=True)
-        # If initialization fails, create an empty database
-        database = FAISS(embed, None, {}, {}, None)
-        logging.info("Created empty WebSearchFAISS database manually")
-    
-    return database, temp_dir
-
-def cleanup_web_search_database(temp_dir):
-    shutil.rmtree(temp_dir)
-
 def duckduckgo_search(query):
-    logging.info(f"Performing DuckDuckGo search for query: {query}")
-    try:
-        with DDGS() as ddgs:
-            results = list(ddgs.text(query, max_results=5))
-        
-        logging.info(f"Number of search results: {len(results)}")
-        
-        database, temp_dir = 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)
-        
-        logging.info(f"Number of documents created: {len(documents)}")
-        if documents:
-            try:
-                database.add_documents(documents)
-                logging.info(f"Successfully added {len(documents)} documents to the database")
-            except Exception as e:
-                logging.error(f"Error adding documents to database: {str(e)}", exc_info=True)
-                # If adding documents fails, create a new database with these documents
-                database = FAISS.from_documents(documents, get_embeddings())
-                logging.info("Created new WebSearchFAISS database with search results")
-        return database, temp_dir, results
-    except Exception as e:
-        logging.error(f"Error in duckduckgo_search: {str(e)}", exc_info=True)
-        return None, None, []
-
-def retrieve_web_search_results(database, query):
-    logging.info(f"Retrieving web search results for query: {query}")
-    retriever = database.as_retriever(search_kwargs={"k": 5})
-    relevant_docs = retriever.get_relevant_documents(query)
-    
-    logging.info(f"Number of relevant documents retrieved: {len(relevant_docs)}")
-    
-    if not relevant_docs:
-        logging.warning("No relevant documents found in the database")
-        return "No relevant information found."
-    
-    context = "\n".join([f"{doc.page_content}\nSource: {doc.metadata['source']}" for doc in relevant_docs])
-    return context
+    with DDGS() as ddgs:
+        results = ddgs.text(query, max_results=5)
+    return results
 
 class CitingSources(BaseModel):
     sources: List[str] = Field(
@@ -377,66 +311,50 @@ def respond(message, history, model, temperature, num_calls, use_web_search, sel
     logging.info(f"User Query: {message}")
     logging.info(f"Model Used: {model}")
     logging.info(f"Search Type: {'Web Search' if use_web_search else 'PDF Search'}")
+
     logging.info(f"Selected Documents: {selected_docs}")
 
     try:
         if use_web_search:
-            database, temp_dir, search_results = duckduckgo_search(message)
-            if database is None or not search_results:
-                yield "I'm sorry, but I couldn't find any search results for your query. Could you please rephrase or ask a different question?"
-                return
-            
-            try:
-                context = retrieve_web_search_results(database, message)
-                logging.info(f"Retrieved context length: {len(context)}")
-                
-                if model == "@cf/meta/llama-3.1-8b-instruct":
-                    # Use Cloudflare API
-                    for partial_response in get_response_from_cloudflare(prompt="", context=context, query=message, num_calls=num_calls, temperature=temperature, search_type="web"):
-                        logging.debug(f"Partial response: {partial_response[:100]}...")  # Log first 100 chars
-                        yield partial_response
-                else:
-                    # Use Hugging Face API
-                    for main_content, sources in get_response_with_search(message, model, num_calls=num_calls, temperature=temperature):
-                        response = f"{main_content}\n\n{sources}"
-                        logging.debug(f"Response: {response[:100]}...")  # Log first 100 chars
-                        yield response
-            finally:
-                # Clean up the temporary database
-                cleanup_web_search_database(temp_dir)
+            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:
-            # PDF search logic
             embed = get_embeddings()
-            if not os.path.exists("faiss_database"):
+            if os.path.exists("faiss_database"):
+                database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
+                retriever = database.as_retriever(search_kwargs={"k": 20})
+                
+                # Filter relevant documents based on user selection
+                all_relevant_docs = retriever.get_relevant_documents(message)
+                relevant_docs = [doc for doc in all_relevant_docs if doc.metadata["source"] in selected_docs]
+                
+                if not relevant_docs:
+                    yield "No relevant information found in the selected documents. Please try selecting different documents or rephrasing your query."
+                    return
+
+                context_str = "\n".join([doc.page_content for doc in relevant_docs])
+            else:
+                context_str = "No documents available."
                 yield "No documents available. Please upload PDF documents to answer questions."
                 return
             
-            database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
-            retriever = database.as_retriever(search_kwargs={"k": 20})
-            
-            all_relevant_docs = retriever.get_relevant_documents(message)
-            relevant_docs = [doc for doc in all_relevant_docs if doc.metadata["source"] in selected_docs]
-            
-            if not relevant_docs:
-                yield "No relevant information found in the selected documents. Please try selecting different documents or rephrasing your query."
-                return
-
-            context_str = "\n".join([doc.page_content for doc in relevant_docs])
-            logging.info(f"Context length for PDF search: {len(context_str)}")
-            
             if model == "@cf/meta/llama-3.1-8b-instruct":
                 # Use Cloudflare API
                 for partial_response in get_response_from_cloudflare(prompt="", context=context_str, query=message, num_calls=num_calls, temperature=temperature, search_type="pdf"):
-                    logging.debug(f"Partial response: {partial_response[:100]}...")  # Log first 100 chars
+                    first_line = partial_response.split('\n')[0] if partial_response else ''
+#                   logging.info(f"Generated Response (first line): {first_line}")
                     yield partial_response
             else:
                 # Use Hugging Face API
                 for partial_response in get_response_from_pdf(message, model, selected_docs, num_calls=num_calls, temperature=temperature):
-                    logging.debug(f"Partial response: {partial_response[:100]}...")  # Log first 100 chars
+                    first_line = partial_response.split('\n')[0] if partial_response else ''
+#                    logging.info(f"Generated Response (first line): {first_line}")
                     yield partial_response
-    
     except Exception as e:
-        logging.error(f"Error in respond function with {model}: {str(e)}", exc_info=True)
+        logging.error(f"Error with {model}: {str(e)}")
         if "microsoft/Phi-3-mini-4k-instruct" in model:
             logging.info("Falling back to Mistral model due to Phi-3 error")
             fallback_model = "mistralai/Mistral-7B-Instruct-v0.3"
@@ -501,42 +419,65 @@ After writing the document, please provide a list of sources used in your respon
     if not full_response:
         yield "I apologize, but I couldn't generate a response at this time. Please try again later."
 
+# New global variable for web search database
+web_search_database = None
+
+def update_web_search_vectors(search_results):
+    global web_search_database
+    embed = get_embeddings()
+    
+    documents = []
+    for result in search_results:
+        if 'body' in result:
+            content = f"{result['title']}\n{result['body']}\nSource: {result['href']}"
+            documents.append(Document(page_content=content, metadata={"source": result['href']}))
+    
+    if web_search_database is None:
+        web_search_database = FAISS.from_documents(documents, embed)
+    else:
+        web_search_database.add_documents(documents)
+
 def get_response_with_search(query, model, num_calls=3, temperature=0.2):
-    database, temp_dir, search_results = duckduckgo_search(query)
-    context = retrieve_web_search_results(database, query)
+    global web_search_database
+    
+    search_results = duckduckgo_search(query)
+    update_web_search_vectors(search_results)
+    
+    if web_search_database is None:
+        yield "No web search results available. Please try again.", ""
+        return
+    
+    retriever = web_search_database.as_retriever(search_kwargs={"k": 5})
+    relevant_docs = retriever.get_relevant_documents(query)
+    
+    context = "\n".join([doc.page_content for doc in relevant_docs])
     
     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."""
+{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 len(message.choices) > 0 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
-    except Exception as e:
-        logging.error(f"Error in get_response_with_search: {str(e)}", exc_info=True)
-        yield f"An error occurred while processing the search results: {str(e)}", ""
-    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}")