Update app.py

app.py

@@ -15,7 +15,6 @@ from langchain_core.documents import Document
 from huggingface_hub import InferenceClient
 import inspect
 import logging
-import shutil
 
 
 # Set up basic configuration for logging
@@ -48,20 +47,18 @@ llama_parser = LlamaParse(
 )
 
 def load_document(file: NamedTemporaryFile, parser: str = "llamaparse") -> List[Document]:
-    file_path = os.path.join(UPLOAD_FOLDER, os.path.basename(file.name))
-    shutil.copy(file.name, file_path)
-    
+    """Loads and splits the document into pages."""
     if parser == "pypdf":
-        loader = PyPDFLoader(file_path)
+        loader = PyPDFLoader(file.name)
         return loader.load_and_split()
     elif parser == "llamaparse":
         try:
-            documents = llama_parser.load_data(file_path)
-            return [Document(page_content=doc.text, metadata={"source": file_path}) for doc in documents]
+            documents = llama_parser.load_data(file.name)
+            return [Document(page_content=doc.text, metadata={"source": file.name}) for doc in documents]
         except Exception as e:
-            logging.error(f"Error using Llama Parse: {str(e)}")
-            logging.info("Falling back to PyPDF parser")
-            loader = PyPDFLoader(file_path)
+            print(f"Error using Llama Parse: {str(e)}")
+            print("Falling back to PyPDF parser")
+            loader = PyPDFLoader(file.name)
             return loader.load_and_split()
     else:
         raise ValueError("Invalid parser specified. Use 'pypdf' or 'llamaparse'.")
@@ -75,7 +72,11 @@ def update_vectors(files, parser):
     
     if not files:
         logging.warning("No files provided for update_vectors")
-        return "Please upload at least one PDF file.", gr.update(choices=[], value=[])
+        return "Please upload at least one PDF file.", gr.CheckboxGroup(
+            choices=[doc["name"] for doc in uploaded_documents],
+            value=[doc["name"] for doc in uploaded_documents if doc["selected"]],
+            label="Select documents to query"
+        )
     
     embed = get_embeddings()
     total_chunks = 0
@@ -88,8 +89,9 @@ def update_vectors(files, parser):
             logging.info(f"Loaded {len(data)} chunks from {file.name}")
             all_data.extend(data)
             total_chunks += len(data)
-            if not any(doc["name"] == os.path.basename(file.name) for doc in uploaded_documents):
-                uploaded_documents.append({"name": os.path.basename(file.name), "selected": True})
+            # Append new documents instead of replacing
+            if not any(doc["name"] == file.name for doc in uploaded_documents):
+                uploaded_documents.append({"name": file.name, "selected": True})
                 logging.info(f"Added new document to uploaded_documents: {file.name}")
             else:
                 logging.info(f"Document already exists in uploaded_documents: {file.name}")
@@ -98,72 +100,53 @@ def update_vectors(files, parser):
     
     logging.info(f"Total chunks processed: {total_chunks}")
     
-    if all_data:
-        if os.path.exists("faiss_database"):
-            logging.info("Updating existing FAISS database")
-            database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
-            database.add_documents(all_data)
-        else:
-            logging.info("Creating new FAISS database")
-            database = FAISS.from_documents(all_data, embed)
-        
-        database.save_local("faiss_database")
-        logging.info("FAISS database saved")
-        
-        return f"Vector store updated successfully. Processed {total_chunks} chunks from {len(files)} files using {parser}.", gr.update(choices=[doc["name"] for doc in uploaded_documents], value=[doc["name"] for doc in uploaded_documents if doc["selected"]])
-    else:
-        return "No data was processed. Please check your files and try again.", gr.update(choices=[doc["name"] for doc in uploaded_documents], value=[doc["name"] for doc in uploaded_documents if doc["selected"]])
-
-UPLOAD_FOLDER = "uploaded_files"
-if not os.path.exists(UPLOAD_FOLDER):
-    os.makedirs(UPLOAD_FOLDER)
-
-# Add this new function to handle file deletion
-def delete_file(file_name):
-    global uploaded_documents
-    logging.info(f"Attempting to delete file: {file_name}")
-    
-    # Remove the file from uploaded_documents
-    uploaded_documents = [doc for doc in uploaded_documents if doc["name"] != file_name]
-    
-    # Remove the file from the file system if it exists
-    file_path = os.path.join(UPLOAD_FOLDER, file_name)
-    if os.path.exists(file_path):
-        os.remove(file_path)
-        logging.info(f"Deleted file: {file_path}")
+    if os.path.exists("faiss_database"):
+        logging.info("Updating existing FAISS database")
+        database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
+        database.add_documents(all_data)
     else:
-        logging.warning(f"File not found: {file_path}")
+        logging.info("Creating new FAISS database")
+        database = FAISS.from_documents(all_data, embed)
     
-    # Rebuild the FAISS database
-    rebuild_faiss_database()
+    database.save_local("faiss_database")
+    logging.info("FAISS database saved")
     
-    return gr.update(value=[doc["name"] for doc in uploaded_documents], choices=[doc["name"] for doc in uploaded_documents])
+    return f"Vector store updated successfully. Processed {total_chunks} chunks from {len(files)} files using {parser}.", gr.CheckboxGroup(
+        choices=[doc["name"] for doc in uploaded_documents],
+        value=[doc["name"] for doc in uploaded_documents if doc["selected"]],
+        label="Select documents to query"
+    )
 
-def rebuild_faiss_database():
-    logging.info("Rebuilding FAISS database")
-    embed = get_embeddings()
-    all_data = []
-    
-    for doc in uploaded_documents:
-        try:
-            file_path = os.path.join(UPLOAD_FOLDER, doc["name"])
-            temp_file = NamedTemporaryFile(delete=False, suffix=".pdf", dir=UPLOAD_FOLDER)
-            temp_file.write(open(file_path, 'rb').read())
-            temp_file.close()
-            data = load_document(temp_file, "llamaparse")
-            all_data.extend(data)
-            os.unlink(temp_file.name)
-        except Exception as e:
-            logging.error(f"Error processing file {doc['name']}: {str(e)}")
-    
-    if all_data:
-        database = FAISS.from_documents(all_data, embed)
+def delete_files(files_to_delete):
+    global uploaded_documents
+    if not files_to_delete:
+        return "No files selected for deletion.", document_selector
+
+    deleted_files = []
+    for file_name in files_to_delete:
+        # Remove the file from uploaded_documents
+        uploaded_documents = [doc for doc in uploaded_documents if doc["name"] != file_name]
+        deleted_files.append(file_name)
+
+    # Update the FAISS database
+    if os.path.exists("faiss_database"):
+        embed = get_embeddings()
+        database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
+        
+        # Remove documents from the database
+        database.delete(lambda doc: doc.metadata["source"] in deleted_files)
+        
+        # Save the updated database
         database.save_local("faiss_database")
-        logging.info("FAISS database rebuilt and saved")
-    else:
-        if os.path.exists("faiss_database"):
-            shutil.rmtree("faiss_database")
-        logging.info("No documents left, removed FAISS database")
+
+    # Update the document selector
+    updated_selector = gr.CheckboxGroup(
+        choices=[doc["name"] for doc in uploaded_documents],
+        value=[doc["name"] for doc in uploaded_documents if doc["selected"]],
+        label="Select documents to query"
+    )
+
+    return f"Deleted files: {', '.join(deleted_files)}", updated_selector
 
 def generate_chunked_response(prompt, model, max_tokens=10000, num_calls=3, temperature=0.2, should_stop=False):
     print(f"Starting generate_chunked_response with {num_calls} calls")
@@ -265,14 +248,14 @@ class CitingSources(BaseModel):
         ...,
         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, selected_docs):
+def chatbot_interface(message, history, use_web_search, model, temperature, num_calls):
     if not message.strip():
         return "", history
 
     history = history + [(message, "")]
 
     try:
-        for response in respond(message, history, model, temperature, num_calls, use_web_search, selected_docs):
+        for response in respond(message, history, model, temperature, num_calls, use_web_search):
             history[-1] = (message, response)
             yield history
     except gr.CancelledError:
@@ -295,67 +278,58 @@ 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:
-            logging.info("Entering web search flow")
             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.info(f"Generated Response (first 100 chars): {response[:100]}...")
+                first_line = response.split('\n')[0] if response else ''
+#                logging.info(f"Generated Response (first line): {first_line}")
                 yield response
         else:
-            logging.info("Entering PDF search flow")
             embed = get_embeddings()
             if os.path.exists("faiss_database"):
-                logging.info("FAISS database exists, loading it")
                 database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
                 retriever = database.as_retriever()
                 
-                logging.info("Attempting to retrieve relevant documents")
-                try:
-                    relevant_docs = retriever.invoke(message)
-                    logging.info(f"Retrieved {len(relevant_docs)} relevant documents")
-                except Exception as e:
-                    logging.error(f"Error retrieving relevant documents: {str(e)}")
-                    yield f"An error occurred while retrieving relevant documents: {str(e)}"
-                    return
-                
                 # Filter relevant documents based on user selection
-                filtered_docs = [doc for doc in relevant_docs if os.path.basename(doc.metadata["source"]) in [os.path.basename(doc) for doc in selected_docs]]
-                logging.info(f"Number of filtered documents: {len(filtered_docs)}")
+                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 filtered_docs:
-                    logging.warning("No relevant information found in the selected documents")
+                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 filtered_docs])
-                logging.info(f"Total context length: {len(context_str)}")
+                context_str = "\n".join([doc.page_content for doc in relevant_docs])
             else:
-                logging.warning("No FAISS database found")
                 context_str = "No documents available."
                 yield "No documents available. Please upload PDF documents to answer questions."
                 return
             
             if model == "@cf/meta/llama-3.1-8b-instruct":
-                logging.info("Using Cloudflare API")
+                # 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.info(f"Generated Response (first 100 chars): {partial_response[:100]}...")
+                    first_line = partial_response.split('\n')[0] if partial_response else ''
+                    logging.info(f"Generated Response (first line): {first_line}")
                     yield partial_response
             else:
-                logging.info("Entering PDF search flow")
+                # Use Hugging Face API
                 for partial_response in get_response_from_pdf(message, model, selected_docs, num_calls=num_calls, temperature=temperature):
-                    logging.info(f"Generated Response (first 100 chars): {partial_response[:100]}...")
+                    first_line = partial_response.split('\n')[0] if partial_response else ''
+                    logging.info(f"Generated Response (first line): {first_line}")
                     yield partial_response
-        
-            logging.info("Finished respond function")
     except Exception as e:
-        logging.error(f"Unexpected error in respond function: {str(e)}")
-        yield f"An unexpected error occurred: {str(e)}"
+        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"
+            yield from respond(message, history, fallback_model, temperature, num_calls, use_web_search, selected_docs)
+        else:
+            yield f"An error occurred with the {model} model: {str(e)}. Please try again or select a different model."
 
-    logging.info("Finished respond function")
-    logging.info(f"Selected docs: {selected_docs}")
+logging.basicConfig(level=logging.DEBUG)
 
 def get_response_from_cloudflare(prompt, context, query, num_calls=3, temperature=0.2, search_type="pdf"):
     headers = {
@@ -457,19 +431,11 @@ def get_response_from_pdf(query, model, selected_docs, num_calls=3, temperature=
 
     retriever = database.as_retriever()
     logging.info(f"Retrieving relevant documents for query: {query}")
-    try:
-        relevant_docs = retriever.invoke(query)
-        logging.info(f"Number of relevant documents retrieved: {len(relevant_docs)}")
-    except Exception as e:
-        logging.error(f"Error retrieving relevant documents: {str(e)}")
-        yield f"An error occurred while retrieving relevant documents: {str(e)}"
-        return
-    
-    # Log the sources of relevant documents
-    logging.info(f"Sources of relevant documents: {[os.path.basename(doc.metadata['source']) for doc in relevant_docs]}")
+    relevant_docs = retriever.get_relevant_documents(query)
+    logging.info(f"Number of relevant documents retrieved: {len(relevant_docs)}")
     
     # Filter relevant_docs based on selected documents
-    filtered_docs = [doc for doc in relevant_docs if os.path.basename(doc.metadata["source"]) in selected_docs]
+    filtered_docs = [doc for doc in relevant_docs if doc.metadata["source"] in selected_docs]
     logging.info(f"Number of filtered documents: {len(filtered_docs)}")
     
     if not filtered_docs:
@@ -478,7 +444,7 @@ def get_response_from_pdf(query, model, selected_docs, num_calls=3, temperature=
         return
 
     for doc in filtered_docs:
-        logging.info(f"Document source: {os.path.basename(doc.metadata['source'])}")
+        logging.info(f"Document source: {doc.metadata['source']}")
         logging.info(f"Document content preview: {doc.page_content[:100]}...")  # Log first 100 characters of each document
 
     context_str = "\n".join([doc.page_content for doc in filtered_docs])
@@ -501,25 +467,18 @@ Write a detailed and complete response that answers the following user question:
         response = ""
         for i in range(num_calls):
             logging.info(f"API call {i+1}/{num_calls}")
-            try:
-                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
-                        response += chunk
-                        yield response  # Yield partial response
-                logging.info(f"API call {i+1} completed successfully")
-            except Exception as e:
-                logging.error(f"Error in API call {i+1}: {str(e)}")
+            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
+                    response += chunk
+                    yield response  # Yield partial response
         
         logging.info("Finished generating response")
-        logging.info(f"Relevant docs: {[doc.metadata['source'] for doc in relevant_docs]}")
-        logging.info(f"Selected docs: {selected_docs}")
-        logging.info(f"Filtered docs: {[doc.metadata['source'] for doc in filtered_docs]}")
 
 def vote(data: gr.LikeData):
     if data.liked:
@@ -543,8 +502,8 @@ uploaded_documents = []
 
 def display_documents():
     return gr.CheckboxGroup(
-        choices=[os.path.basename(doc["name"]) for doc in uploaded_documents],
-        value=[os.path.basename(doc["name"]) for doc in uploaded_documents if doc["selected"]],
+        choices=[doc["name"] for doc in uploaded_documents],
+        value=[doc["name"] for doc in uploaded_documents if doc["selected"]],
         label="Select documents to query"
     )
 
@@ -554,13 +513,13 @@ document_selector = gr.CheckboxGroup(label="Select documents to query")
 use_web_search = gr.Checkbox(label="Use Web Search", value=True)
 
 demo = gr.ChatInterface(
-    chatbot_interface,
+    respond,
     additional_inputs=[
         gr.Dropdown(choices=MODELS, label="Select Model", value=MODELS[3]),
         gr.Slider(minimum=0.1, maximum=1.0, value=0.2, step=0.1, label="Temperature"),
         gr.Slider(minimum=1, maximum=5, value=1, step=1, label="Number of API Calls"),
         use_web_search,
-        document_selector  # This should now correctly pass the selected documents
+        document_selector  # Add the document selector to the chat interface
     ],
     title="AI-powered Web Search and PDF Chat Assistant",
     description="Chat with your PDFs or use web search to answer questions (Please use toggle under Additional Inputs to swithc between PDF and Web Search, Default Value Web Search)",
@@ -603,42 +562,31 @@ with demo:
     
     update_output = gr.Textbox(label="Update Status")
     
-    # Create a new row for displaying uploaded files with delete buttons
-    with gr.Row():
-        uploaded_files = gr.CheckboxGroup(label="Uploaded Documents", interactive=True)
-        delete_button = gr.Button("Delete Selected")
-    
     # Update both the output text and the document selector
-    update_button.click(
-        update_vectors, 
-        inputs=[file_input, parser_dropdown], 
-        outputs=[update_output, uploaded_files]
-    )
-    
-    # Handle file deletion
-    delete_button.click(
-        lambda selected: [delete_file(file) for file in selected],
-        inputs=[uploaded_files],
-        outputs=[uploaded_files]
-    )
+    update_button.click(update_vectors, 
+                        inputs=[file_input, parser_dropdown], 
+                        outputs=[update_output, document_selector])
+
+    # Add delete button
+    delete_button = gr.Button("Delete Selected Files")
+    delete_output = gr.Textbox(label="Delete Status")
+
+    # Connect delete button to the delete_files function
+    delete_button.click(delete_files,
+                        inputs=[document_selector],
+                        outputs=[delete_output, document_selector])
 
-    # Update the document selector in the chat interface
-    uploaded_files.change(
-        lambda x: gr.update(choices=x, value=x),
-        inputs=[uploaded_files],
-        outputs=[document_selector]
-    )
-    
     gr.Markdown(
     """
     ## How to use
     1. Upload PDF documents using the file input at the top.
     2. Select the PDF parser (pypdf or llamaparse) and click "Upload Document" to update the vector store.
     3. Select the documents you want to query using the checkboxes.
-    4. Ask questions in the chat interface. 
-    5. Toggle "Use Web Search" to switch between PDF chat and web search.
-    6. Adjust Temperature and Number of API Calls to fine-tune the response generation.
-    7. Use the provided examples or ask your own questions.
+    4. To delete files, select them in the checkbox and click "Delete Selected Files".
+    5. Ask questions in the chat interface. 
+    6. Toggle "Use Web Search" to switch between PDF chat and web search.
+    7. Adjust Temperature and Number of API Calls to fine-tune the response generation.
+    8. Use the provided examples or ask your own questions.
     """
     )