Update app.py

app.py

@@ -71,10 +71,9 @@ def scan_faiss_database():
     uploaded_documents = []
     
     if os.path.exists("faiss_database"):
-        embed = HuggingFaceEmbeddings(model_name="sentence-transformers/stsb-roberta-large")
+        embed = get_embeddings()
         database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
         
-        # Assuming each document in FAISS has a 'source' metadata field with the file name
         for doc in database.docstore._dict.values():
             file_name = os.path.basename(doc.metadata['source'])
             if not any(d['name'] == file_name for d in uploaded_documents):
@@ -85,6 +84,10 @@ def scan_faiss_database():
 # Call this function when the application starts
 uploaded_documents = scan_faiss_database()
 
+def initialize_session():
+    global uploaded_documents
+    uploaded_documents = scan_faiss_database()
+    return uploaded_documents
 
 def update_vectors(files, parser):
     global uploaded_documents
@@ -109,12 +112,12 @@ def update_vectors(files, parser):
             logging.info(f"Loaded {len(data)} chunks from {file.name}")
             all_data.extend(data)
             total_chunks += len(data)
-            # 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}")
+            file_name = os.path.basename(file.name)
+            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}")
+                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)}")
     
@@ -131,9 +134,6 @@ def update_vectors(files, parser):
     database.save_local("faiss_database")
     logging.info("FAISS database saved")
     
-    # After processing new files, rescan the database
-    uploaded_documents = scan_faiss_database()
-    
     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"]],
@@ -240,14 +240,14 @@ class CitingSources(BaseModel):
         ...,
         description="List of sources to cite. Should be an URL of the source."
     )
-def chatbot_interface(message, history, use_web_search, model, temperature, num_calls):
+def chatbot_interface(message, history, use_web_search, model, temperature, num_calls, selected_docs):
     if not message.strip():
         return "", history
 
     history = history + [(message, "")]
 
     try:
-        for response in respond(message, history, model, temperature, num_calls, use_web_search):
+        for response in respond(message, history, model, temperature, num_calls, use_web_search, selected_docs):
             history[-1] = (message, response)
             yield history
     except gr.CancelledError:
@@ -270,15 +270,12 @@ 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:
         if use_web_search:
             for main_content, sources in get_response_with_search(message, model, num_calls=num_calls, temperature=temperature):
                 response = f"{main_content}\n\n{sources}"
-                first_line = response.split('\n')[0] if response else ''
-#                logging.info(f"Generated Response (first line): {first_line}")
                 yield response
         else:
             embed = get_embeddings()
@@ -286,9 +283,8 @@ def respond(message, history, model, temperature, num_calls, use_web_search, sel
                 database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
                 retriever = database.as_retriever()
                 
-                # 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 doc.metadata["source"] in selected_docs]
+                relevant_docs = [doc for doc in all_relevant_docs if os.path.basename(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."
@@ -303,14 +299,10 @@ def respond(message, history, model, temperature, num_calls, use_web_search, sel
             if model == "@cf/meta/llama-3.1-8b-instruct":
                 # Use Cloudflare API
                 for partial_response in get_response_from_cloudflare(prompt="", context=context_str, query=message, num_calls=num_calls, temperature=temperature, search_type="pdf"):
-                    first_line = partial_response.split('\n')[0] if partial_response else ''
-                    logging.info(f"Generated Response (first line): {first_line}")
                     yield partial_response
             else:
                 # Use Hugging Face API
                 for partial_response in get_response_from_pdf(message, model, selected_docs, num_calls=num_calls, temperature=temperature):
-                    first_line = partial_response.split('\n')[0] if partial_response else ''
-                    logging.info(f"Generated Response (first line): {first_line}")
                     yield partial_response
     except Exception as e:
         logging.error(f"Error with {model}: {str(e)}")
@@ -563,7 +555,9 @@ demo = gr.ChatInterface(
 )
 
 # Add file upload functionality
-with demo:
+with gr.Blocks() as demo:
+    session_documents = gr.State(initialize_session)
+    
     gr.Markdown("## Upload PDF Documents")
 
     with gr.Row():
@@ -572,11 +566,29 @@ with demo:
         update_button = gr.Button("Upload Document")
     
     update_output = gr.Textbox(label="Update Status")
+    document_selector = 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"
+    )
     
     # Update both the output text and the document selector
-    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, document_selector]
+    )
+
+    # Add a refresh button to update the document selector
+    refresh_button = gr.Button("Refresh Document List")
+    refresh_button.click(
+        lambda: 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"
+        ),
+        outputs=[document_selector]
+    )  
 
     gr.Markdown(
     """