Update app.py

app.py

@@ -66,37 +66,22 @@ def load_document(file: NamedTemporaryFile, parser: str = "llamaparse") -> List[
 def get_embeddings():
     return HuggingFaceEmbeddings(model_name="sentence-transformers/stsb-roberta-large")
 
-# File to store the list of uploaded documents
-# File to store metadata about uploaded documents
-METADATA_FILE = "document_metadata.json"
-
-def load_document_metadata():
-    if os.path.exists(METADATA_FILE):
-        with open(METADATA_FILE, 'r') as f:
-            return json.load(f)
-    return []
-
-def save_document_metadata(metadata):
-    with open(METADATA_FILE, 'w') as f:
-        json.dump(metadata, f)
-
 def update_vectors(files, parser):
+    global uploaded_documents
     logging.info(f"Entering update_vectors with {len(files)} files and parser: {parser}")
     
     if not files:
         logging.warning("No files provided for update_vectors")
         return "Please upload at least one PDF file.", gr.CheckboxGroup(
-            choices=[],
-            value=[],
+            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 = HuggingFaceEmbeddings(model_name="sentence-transformers/stsb-roberta-large")
+    embed = get_embeddings()
     total_chunks = 0
     
     all_data = []
-    metadata = load_document_metadata()
-    
     for file in files:
         logging.info(f"Processing file: {file.name}")
         try:
@@ -104,13 +89,12 @@ def update_vectors(files, parser):
             logging.info(f"Loaded {len(data)} chunks from {file.name}")
             all_data.extend(data)
             total_chunks += len(data)
-            
-            # Update metadata
-            if not any(doc["name"] == file.name for doc in metadata):
-                metadata.append({"name": file.name, "selected": True})
-                logging.info(f"Added new document to metadata: {file.name}")
+            # 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 metadata: {file.name}")
+                logging.info(f"Document already exists in uploaded_documents: {file.name}")
         except Exception as e:
             logging.error(f"Error processing file {file.name}: {str(e)}")
     
@@ -127,25 +111,9 @@ def update_vectors(files, parser):
     database.save_local("faiss_database")
     logging.info("FAISS database saved")
     
-    # Save the updated metadata
-    save_document_metadata(metadata)
-    logging.info("Document metadata saved")
-    
-    completion_message = f"Vector store updated successfully. Processed {total_chunks} chunks from {len(files)} files using {parser}."
-    logging.info(completion_message)
-    
-    updated_choices = [doc["name"] for doc in metadata]
-    updated_values = [doc["name"] for doc in metadata if doc["selected"]]
-    
-    return completion_message, gr.CheckboxGroup.update(choices=updated_choices, value=updated_values, label="Select documents to query")
-
-# Make sure to call this function at the start of your script
-
-def initialize_document_selector():
-    metadata = load_document_metadata()
-    return gr.CheckboxGroup(
-        choices=[doc["name"] for doc in metadata],
-        value=[doc["name"] for doc in metadata if doc["selected"]],
+    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"
     )
 
@@ -520,8 +488,6 @@ def initial_conversation():
 # Define the checkbox outside the demo block
 document_selector = gr.CheckboxGroup(label="Select documents to query")
 
-document_selector = initialize_document_selector()
-
 use_web_search = gr.Checkbox(label="Use Web Search", value=True)
 
 custom_placeholder = "Ask a question (Note: You can toggle between Web Search and PDF Chat in Additional Inputs below)"
@@ -583,7 +549,6 @@ with demo:
         update_button = gr.Button("Upload Document")
     
     update_output = gr.Textbox(label="Update Status")
-    document_selector = gr.CheckboxGroup(label="Select documents to query")
     
     # Update both the output text and the document selector
     update_button.click(update_vectors, 
@@ -604,4 +569,4 @@ with demo:
     )
 
 if __name__ == "__main__":
-    demo.launch(share=True)
+    demo.launch(share=True)