Update app.py

app.py

@@ -66,6 +66,21 @@ 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
+DOCUMENTS_FILE = "uploaded_documents.json"
+
+def load_uploaded_documents():
+    global uploaded_documents
+    if os.path.exists(DOCUMENTS_FILE):
+        with open(DOCUMENTS_FILE, 'r') as f:
+            uploaded_documents = json.load(f)
+    else:
+        uploaded_documents = []
+
+def save_uploaded_documents():
+    with open(DOCUMENTS_FILE, 'w') as f:
+        json.dump(uploaded_documents, f)
+
 def update_vectors(files, parser):
     global uploaded_documents
     logging.info(f"Entering update_vectors with {len(files)} files and parser: {parser}")
@@ -78,7 +93,7 @@ def update_vectors(files, parser):
             label="Select documents to query"
         )
     
-    embed = get_embeddings()
+    embed = HuggingFaceEmbeddings(model_name="sentence-transformers/stsb-roberta-large")
     total_chunks = 0
     
     all_data = []
@@ -111,12 +126,19 @@ def update_vectors(files, parser):
     database.save_local("faiss_database")
     logging.info("FAISS database saved")
     
+    # Save the updated list of documents
+    save_uploaded_documents()
+    logging.info("Uploaded documents list 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],
         value=[doc["name"] for doc in uploaded_documents if doc["selected"]],
         label="Select documents to query"
     )
 
+# Make sure to call this function at the start of your script
+load_uploaded_documents()
+
 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 = ""