Update app.py

app.py

@@ -15,8 +15,6 @@ from langchain_core.documents import Document
 from huggingface_hub import InferenceClient
 import inspect
 import logging
-import tempfile
-import shutil
 
 
 # Set up basic configuration for logging
@@ -48,19 +46,19 @@ llama_parser = LlamaParse(
     language="en",
 )
 
-def load_document(file_path: str, parser: str = "llamaparse") -> List[Document]:
+def load_document(file: NamedTemporaryFile, parser: str = "llamaparse") -> List[Document]:
     """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:
             print(f"Error using Llama Parse: {str(e)}")
             print("Falling back to PyPDF parser")
-            loader = PyPDFLoader(file_path)
+            loader = PyPDFLoader(file.name)
             return loader.load_and_split()
     else:
         raise ValueError("Invalid parser specified. Use 'pypdf' or 'llamaparse'.")
@@ -87,29 +85,13 @@ def update_vectors(files, parser):
     for file in files:
         logging.info(f"Processing file: {file.name}")
         try:
-            # Create a temporary file
-            with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf', mode='wb') as temp_file:
-                # Write the content to the temporary file
-                temp_file.write(file.name.encode('utf-8'))  # Write the file path
-                temp_file_path = temp_file.name
-
-            # Save the uploaded file
-            os.makedirs("uploaded_files", exist_ok=True)
-            file_path = os.path.join("uploaded_files", os.path.basename(file.name))
-            shutil.copy(file.name, file_path)  # Copy the actual file
-            
-            # Remove the temporary file
-            os.unlink(temp_file_path)
-            
-            data = load_document(file_path, parser)
+            data = load_document(file, parser)
             logging.info(f"Loaded {len(data)} chunks from {file.name}")
-            for chunk in data:
-                logging.info(f"Chunk content preview: {chunk.page_content[:100]}...")  # Log first 100 characters of each chunk
             all_data.extend(data)
             total_chunks += len(data)
             # Append new documents instead of replacing
-            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})
+            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}")
@@ -118,10 +100,16 @@ 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")
+    
+    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.CheckboxGroup(
         choices=[doc["name"] for doc in uploaded_documents],
@@ -129,52 +117,6 @@ def update_vectors(files, parser):
         label="Select documents to query"
     )
 
-def delete_selected_files(selected_files):
-    global uploaded_documents
-    if not selected_files:
-        return "No files selected for deletion.", document_selector
-
-    deleted_files = []
-    for file_name in selected_files:
-        # Remove the file from uploaded_documents
-        uploaded_documents = [doc for doc in uploaded_documents if doc["name"] != file_name]
-        
-        # Delete the file from the file system if it exists
-        file_path = os.path.join("uploaded_files", file_name)
-        if os.path.exists(file_path):
-            os.remove(file_path)
-        
-        deleted_files.append(file_name)
-
-    # Update the FAISS database
-    update_faiss_database()
-
-    remaining_files = [doc["name"] for doc in uploaded_documents]
-    return f"Deleted files: {', '.join(deleted_files)}", gr.CheckboxGroup(choices=remaining_files, value=remaining_files, label="Select documents to query")
-
-def update_faiss_database():
-    global uploaded_documents
-    
-    embed = get_embeddings()
-    all_data = []
-
-    for doc in uploaded_documents:
-        file_path = os.path.join("uploaded_files", doc["name"])
-        if os.path.exists(file_path):
-            data = load_document(file_path, parser="llamaparse")  # or use your preferred parser
-            all_data.extend(data)
-
-    if all_data:
-        database = FAISS.from_documents(all_data, embed)
-        database.save_local("faiss_database")
-        logging.info("FAISS database updated after deletion")
-    else:
-        # If no documents left, remove the FAISS database
-        if os.path.exists("faiss_database"):
-            import shutil
-            shutil.rmtree("faiss_database")
-        logging.info("All documents deleted, FAISS database removed")
-
 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")
     full_response = ""
@@ -305,6 +247,7 @@ 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:
@@ -322,7 +265,7 @@ def respond(message, history, model, temperature, num_calls, use_web_search, sel
                 
                 # 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 os.path.basename(doc.metadata["source"]) in selected_docs]
+                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."
@@ -460,12 +403,8 @@ def get_response_from_pdf(query, model, selected_docs, num_calls=3, temperature=
     relevant_docs = retriever.get_relevant_documents(query)
     logging.info(f"Number of relevant documents retrieved: {len(relevant_docs)}")
     
-    # Log the sources of all retrieved documents
-    for doc in relevant_docs:
-        logging.info(f"Retrieved document source: {doc.metadata['source']}")
-    
     # 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: