Update app.py

app.py

@@ -803,52 +803,57 @@ def side():
             accept_multiple_files=True,
             key="file_uploader"
         )        
+        
         if uploaded_files:
             for uploaded_file in uploaded_files:
                 try:
-                    
-                        upload_to_firebase(st.session_state["wix_user_id"], uploaded_file)
-                        st.sidebar.success(f"File '{uploaded_file.name}' uploaded and converted to TXT!")
-              
+                    upload_to_firebase(st.session_state["wix_user_id"], uploaded_file)
+                    st.sidebar.success(f"File '{uploaded_file.name}' uploaded and converted to TXT!")
                 except Exception as e:
                     st.sidebar.error(f"Error processing file '{uploaded_file.name}': {e}")
-
         else:
-                st.sidebar.warning("Please select a file to upload.")
+            st.sidebar.warning("Please select a file to upload.")
+        
         # Display and delete functionality for documents
         if st.session_state.get("documents"):
             doc_ids = list(st.session_state["documents"].keys())
             doc_options = ["None (use only main knowledge base)"] + doc_ids
-            selected_option = st.selectbox(
-                "Select a document to include in your query:",
+            selected_options = st.multiselect(
+                "Select documents to include in your query:",
                 options=doc_options,
+                default="None (use only main knowledge base)",
                 format_func=lambda x: st.session_state["documents"][x].get("name", f"Document {x}") if x != "None (use only main knowledge base)" else x,
-                key="select_doc"
+                key="select_docs"
             )
-            selected_doc_id = selected_option if selected_option != "None (use only main knowledge base)" else None
-            if selected_doc_id:
-                st.info(f"Selected Document: {st.session_state['documents'][selected_doc_id]['name']}")
+            selected_doc_ids = [doc_id for doc_id in selected_options if doc_id != "None (use only main knowledge base)"]
+        
+            if selected_doc_ids:
+                selected_doc_names = [st.session_state['documents'][doc_id]['name'] for doc_id in selected_doc_ids]
+                st.info(f"Selected Documents: {', '.join(selected_doc_names)}")
             else:
                 st.info("Using only the main knowledge base.")
         else:
-            st.warning("Please upload a document to include in your queries.")
-            selected_doc_id = None
-                
-            # Button to delete the selected document
-            if st.sidebar.button("Delete ", key="delete_button"):
+            st.warning("Please upload documents to include in your queries.")
+            selected_doc_ids = []
+        
+        # Button to delete the selected documents
+        if selected_doc_ids:
+            if st.button("Delete Selected Documents", key="delete_button"):
                 try:
-                    # Remove the document from Firebase
-                    db.child("users").child(st.session_state["wix_user_id"]).child("KnowledgeBase").child(selected_doc_id).remove()
-                    
-                    # Remove the document from session state
-                    fetch_documents()
-
+                    for doc_id in selected_doc_ids:
+                        # Remove the document from Firebase
+                        db.child("users").child(st.session_state["wix_user_id"]).child("KnowledgeBase").child(doc_id).remove()
         
+                        # Remove from session state
+                        st.session_state["vector_store"].pop(doc_id, None)
+                        st.session_state["documents"].pop(doc_id, None)
         
-                    st.sidebar.success("Document deleted successfully!")
-                    st.rerun()
+                    st.success("Selected documents deleted successfully!")
+                    st.experimental_rerun()
                 except Exception as e:
-                    st.sidebar.error(f"Error deleting document: {e}")
+                    st.error(f"Error deleting documents: {e}")
+
+
                             
         st.sidebar.markdown("</div>", unsafe_allow_html=True)
         trust_buckets = ["Any","Stability", "Development", "Relationship", "Benefit", "Vision", "Competence"]
@@ -1120,9 +1125,9 @@ def google_search(query):
 
 
 # RAG response function
-def rag_response(query, selected_doc_id=None):
+def rag_response(query, selected_doc_ids=None):
     """
-    Handle queries by searching both the main knowledge base and the selected document.
+    Handle queries by searching both the main knowledge base and the selected documents.
     """
     try:
         results = []
@@ -1132,16 +1137,21 @@ def rag_response(query, selected_doc_id=None):
             retrieved_docs = search_knowledge_base(query)
             results.extend(retrieved_docs)
 
-        # Search vector store of the selected document
-        if selected_doc_id:
-            vector_store = st.session_state.get("vector_store", {}).get(selected_doc_id)
-            if vector_store:
-                vector_store_results = vector_store.similarity_search(query, k=5)
-                results.extend(vector_store_results)
-            else:
-                st.warning("Selected document's vector store not found.")
+        # If selected_doc_ids is None, try to get it from session state
+        if selected_doc_ids is None:
+            selected_doc_ids = st.session_state.get('selected_doc_ids', [])
+
+        # Search vector stores of the selected documents
+        if selected_doc_ids:
+            for doc_id in selected_doc_ids:
+                vector_store = st.session_state.get("vector_store", {}).get(doc_id)
+                if vector_store:
+                    vector_store_results = vector_store.similarity_search(query, k=5)
+                    results.extend(vector_store_results)
+                else:
+                    st.warning(f"Vector store for document '{st.session_state['documents'][doc_id]['name']}' not found.")
         else:
-            st.warning("No document selected. Proceeding with main knowledge base only.")
+            st.info("No documents selected. Proceeding with main knowledge base only.")
 
         # Combine results into a single context
         context = "\n".join([doc.page_content for doc in results])
@@ -1159,9 +1169,6 @@ def rag_response(query, selected_doc_id=None):
         logger.error(f"Error generating RAG response: {e}")
         return "An error occurred during the RAG response generation process."
 
-
-
-
 # Define tools
 @tool
 def knowledge_base_tool(query: str):