Update app.py

app.py

@@ -68,19 +68,17 @@ 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
     
-    new_documents = []
+    all_data = []
     for file in files:
         data = load_document(file, parser)
         all_data.extend(data)
         total_chunks += len(data)
-        new_documents.append({"name": file.name, "selected": True})
-    
-    uploaded_documents.extend(new_documents)
+        uploaded_documents.append({"name": file.name, "selected": True})
     
     if os.path.exists("faiss_database"):
         database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
@@ -90,8 +88,7 @@ def update_vectors(files, parser):
     
     database.save_local("faiss_database")
     
-    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]
+    return f"Vector store updated successfully. Processed {total_chunks} chunks from {len(files)} files using {parser}."
 
 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")
@@ -413,10 +410,10 @@ css = """
 
 uploaded_documents = []
 
-def display_documents(document_list):
+def display_documents():
     return gr.CheckboxGroup(
-        choices=document_list,
-        value=document_list,
+        choices=[doc["name"] for doc in uploaded_documents],
+        value=[doc["name"] for doc in uploaded_documents if doc["selected"]],
         label="Select documents to query"
     )
 
@@ -426,13 +423,14 @@ document_selector = gr.CheckboxGroup(label="Select documents to query")
 use_web_search = gr.Checkbox(label="Use Web Search", value=False)
 
 demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Dropdown(choices=MODELS, label="Select Model", value=MODELS[0]),
-        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   
-    ],
+       respond,
+       additional_inputs=[
+           gr.Dropdown(choices=MODELS, label="Select Model", value=MODELS[0]),
+           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  # Add this line
+       ],
     title="AI-powered Web Search and PDF Chat Assistant",
     description="Chat with your PDFs or use web search to answer questions.",
     theme=gr.themes.Soft(
@@ -471,16 +469,13 @@ 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")
+        document_selector
     
     update_output = gr.Textbox(label="Update Status")
-    update_button.click(
-        update_vectors, 
-        inputs=[file_input, parser_dropdown], 
-        outputs=[update_output, document_selector]
-    )
+    update_button.click(update_vectors, inputs=[file_input, parser_dropdown], outputs=update_output)
+    update_button.click(display_documents, outputs=document_selector)
 
-gr.Markdown(
+    gr.Markdown(
     """
     ## How to use
     1. Upload PDF documents using the file input at the top.