Update app.py

app.py

@@ -68,7 +68,7 @@ def get_embeddings():
 def update_vectors(files, parser):
     global uploaded_documents
     if not files:
-        return "Please upload at least one PDF file."
+        return "Please upload at least one PDF file.", []
     
     embed = get_embeddings()
     total_chunks = 0
@@ -88,7 +88,8 @@ def update_vectors(files, parser):
     
     database.save_local("faiss_database")
     
-    return f"Vector store updated successfully. Processed {total_chunks} chunks from {len(files)} files using {parser}."
+    status_message = f"Vector store updated successfully. Processed {total_chunks} chunks from {len(files)} files using {parser}."
+    return status_message, [doc["name"] for doc in uploaded_documents]
 
 def generate_chunked_response(prompt, model, max_tokens=1000, num_calls=3, temperature=0.2, should_stop=False):
     print(f"Starting generate_chunked_response with {num_calls} calls")
@@ -410,10 +411,10 @@ css = """
 
 uploaded_documents = []
 
-def display_documents():
+def display_documents(document_list):
     return gr.CheckboxGroup(
-        choices=[doc["name"] for doc in uploaded_documents],
-        value=[doc["name"] for doc in uploaded_documents if doc["selected"]],
+        choices=document_list,
+        value=document_list,
         label="Select documents to query"
     )
 
@@ -469,12 +470,16 @@ with demo:
         file_input = gr.Files(label="Upload your PDF documents", file_types=[".pdf"])
         parser_dropdown = gr.Dropdown(choices=["pypdf", "llamaparse"], label="Select PDF Parser", value="llamaparse")
         update_button = gr.Button("Upload Document")
+        document_selector = gr.CheckboxGroup(label="Select documents to query")
     
     update_output = gr.Textbox(label="Update Status")
-    update_button.click(update_vectors, inputs=[file_input, parser_dropdown], outputs=update_output)
-    update_button.click(display_documents, outputs=document_selector)
+    update_button.click(
+        update_vectors, 
+        inputs=[file_input, parser_dropdown], 
+        outputs=[update_output, document_selector]
+    )
 
-    gr.Markdown(
+gr.Markdown(
     """
     ## How to use
     1. Upload PDF documents using the file input at the top.