Update app.py

app.py

@@ -13,8 +13,9 @@ from streamlit_extras.add_vertical_space import add_vertical_space
 from streamlit_extras.switch_page_button import switch_page
 import json
 import pandas as pd
+from st_aggrid import AgGrid, GridOptionsBuilder, GridUpdateMode, DataReturnMode
 # ---------------------- Configuration ----------------------
-st.set_page_config(page_title="Building Regulations Chatbot", layout="wide", page_icon="🏗️", initial_sidebar_state="expanded")
+st.set_page_config(page_title="Building Regulations Chatbot", layout="wide", initial_sidebar_state="expanded")
 # Load environment variables from .env file
 load_dotenv()
 
@@ -155,22 +156,22 @@ def chat_with_assistant(file_ids, user_message):
 
 # ---------------------- Streamlit App ----------------------
 
-page = st.sidebar.selectbox("Choose a page", ["Home", "Documents", "Admin"])
+page = st.sidebar.selectbox("Choose a page", ["Documents", "Home", "Admin"])
 
 if page == "Home":
-    st.title("🏗️ Building Regulations Chatbot", anchor=False)
+    st.title("Building Regulations Chatbot", anchor=False)
 
     # Sidebar improvements
     with st.sidebar:
-        colored_header("📄 Selected Documents", description="Documents for chat")
-        if 'selected_pdfs' in st.session_state and st.session_state.selected_pdfs:
-            for pdf in st.session_state.selected_pdfs:
-                st.write(f"- {pdf['doc_title']}")
+        colored_header("Selected Documents", description="Documents for chat")
+        if 'selected_pdfs' in st.session_state and not st.session_state.selected_pdfs.empty:
+            for _, pdf in st.session_state.selected_pdfs.iterrows():
+                st.write(f"- {pdf['Doc Title']}")
         else:
             st.write("No documents selected. Please go to the Documents page.")
 
     # Main chat area improvements
-    colored_header("💬 Chat", description="Ask questions about building regulations")
+    colored_header("Chat", description="Ask questions about building regulations")
 
     # Display chat messages with improved styling
     for chat in st.session_state.chat_history:
@@ -225,7 +226,7 @@ if page == "Home":
         st.caption("© 2023 Your Company Name")
 
 elif page == "Documents":
-    st.title("📚 Document Selection")
+    st.title("Document Selection")
 
     city_code_input = st.text_input("Enter city code:", key="city_code_input")
     load_documents_button = st.button("Load Documents", key="load_documents_button")
@@ -258,40 +259,59 @@ elif page == "Documents":
         })
         
         # Add a checkbox column to the DataFrame at the beginning
-        df.insert(0, "Checkbox", False)
+        df.insert(0, "Select", False)
         
-        # Display the DataFrame as an editable table
-        edited_df = st.data_editor(
+        # Configure grid options
+        gb = GridOptionsBuilder.from_dataframe(df)
+        gb.configure_default_column(enablePivot=True, enableValue=True, enableRowGroup=True)
+        gb.configure_column("Select", header_name="Select", cellRenderer='checkboxRenderer')
+        gb.configure_column("File Href", cellRenderer='linkRenderer')
+        gb.configure_selection(selection_mode="multiple", use_checkbox=True)
+        gb.configure_side_bar()
+        gridOptions = gb.build()
+
+        # Display the AgGrid
+        grid_response = AgGrid(
             df,
-            hide_index=True,
-            num_rows="dynamic",
-            column_config={
-                "Checkbox": st.column_config.CheckboxColumn(required=True),
-                "File Href": st.column_config.LinkColumn("File Href")
-            },
-            disabled=["Municipality", "Abbreviation", "Doc Title", "File Title", "File Href", "Enactment Date", "Prio"]
+            gridOptions=gridOptions,
+            enable_enterprise_modules=True,
+            update_mode=GridUpdateMode.MODEL_CHANGED,
+            data_return_mode=DataReturnMode.FILTERED_AND_SORTED,
+            fit_columns_on_grid_load=False,
         )
+
+        # Get the selected rows
+        selected_rows = grid_response['selected_rows']
         
-        # Get the selected PDFs
-        selected_pdfs = edited_df[edited_df["Checkbox"]].to_dict('records')
+        # Debug: Print the structure of selected_rows
+        st.write("Debug - Selected Rows Structure:", selected_rows)
         
         if st.button("Process Selected PDFs"):
-            if selected_pdfs:
-                st.session_state.selected_pdfs = selected_pdfs
+            if len(selected_rows) > 0:  # Check if there are any selected rows
+                # Convert selected_rows to a DataFrame
+                st.session_state.selected_pdfs = pd.DataFrame(selected_rows)
                 st.session_state.assistant_id = create_assistant()
                 with st.spinner("Processing PDFs and creating/updating assistant..."):
                     file_ids = []
-                    for pdf in selected_pdfs:
-                        file_name = download_pdf(pdf['File Href'])
+                    
+                    for _, pdf in st.session_state.selected_pdfs.iterrows():
+                        # Debug: Print each pdf item
+                        st.write("Debug - PDF item:", pdf)
+                        
+                        file_href = pdf['File Href']
+                        doc_title = pdf['Doc Title']
+                        
+                        file_name = download_pdf(file_href)
                         if file_name:
                             file_path = f"./{file_name}"
                             file_id = upload_file_to_openai(file_path)
                             if file_id:
                                 file_ids.append(file_id)
                             else:
-                                st.warning(f"Failed to upload {pdf['Doc Title']}. Skipping this file.")
+                                st.warning(f"Failed to upload {doc_title}. Skipping this file.")
                         else:
-                            st.warning(f"Failed to download {pdf['Doc Title']}. Skipping this file.")
+                            st.warning(f"Failed to download {doc_title}. Skipping this file.")
+                    
                     st.session_state.file_ids = file_ids
                 st.success("PDFs processed successfully. You can now chat on the Home page.")
             else:
@@ -301,8 +321,8 @@ elif page == "Documents":
         switch_page("Home")
 
 elif page == "Admin":
-    st.title("🔧 Admin Panel")
-    colored_header("Vector Stores Information", description="List of all vector stores")
+    st.title("Admin Panel")
+    st.header("Vector Stores Information")
     
     vector_stores = get_vector_stores()
     json_vector_stores = json.dumps([vs.model_dump() for vs in vector_stores])