Update app.py

app.py

@@ -27,6 +27,8 @@ import warnings
 warnings.filterwarnings("ignore", category=FutureWarning)
 warnings.filterwarnings("ignore", category=DeprecationWarning)
 
+import glob
+
 # os.environ['HUGGINGFACEHUB_API_TOKEN'] = 'your_api_key' # for using HuggingFace Inference API
 
 
@@ -97,21 +99,16 @@ st.set_page_config(
 )
 
 # Document Config
+if 'countries_override' not in st.session_state: 
+    # countries to override with own documents from uploaded pdf or updated scraped search results
+    # must first scrape or upload own documents to use this
+    st.session_state['countries_override'] = []
+   
 if 'chunk_size' not in st.session_state:
     st.session_state['chunk_size'] = 1000 # choose one of [500, 600, 700, 800, 900, 1000, 1250, 1500, 1750, 2000, 2250, 2500, 2750, 3000]
     
 if 'chunk_overlap' not in st.session_state:
-    st.session_state['chunk_overlap'] = 100 # choose one of [50, 100, 150, 200]
-
-# scraping results using DuckDuckGo
-if 'top_n_results' not in st.session_state:
-    st.session_state['top_n_results'] = 10 # this is for returning top n search results using DuckDuckGo
-    
-if 'countries_to_scrape' not in st.session_state:
-    st.session_state['countries_to_scrape'] = [] # this is for returning top n search results using DuckDuckGo
-
-# in main app, add configuration for user to scrape new data from DuckDuckGo
-# in main app, add configuration for user to upload PDF to override country's existing policies in vectorstore
+    st.session_state['chunk_overlap'] = 100 # choose one of [50, 100, 150, 200]    
 
 # Retriever Config
 if 'chroma_n_similar_documents' not in st.session_state:
@@ -379,15 +376,49 @@ with st.sidebar:
         [
             "Main Chatbot", 
             "View Source Docs for Last Query",      
-            "Scrape or Upload Docs",  
+            "Scrape or Upload Own Docs",  
         ], 
         icons=['house', 'gear', 'gear', 'gear'], 
         menu_icon="", default_index=0)
 
-    with st.container(border = True):
-        st.write("DO NOT NAVIGATE between pages or change when agent is still generating messages in the chat.  Wait for query to complete first.")
-        st.write("")  
     
+    with st.expander("Warning", expanded = True):
+        st.write("⚠️ DO NOT navigate between pages or change config when chat is ongoing.  Wait for query to complete first.")
+         
+    st.container()
+
+    # see if retrievers created by user's own uploaded PDF or newly scraped data is found
+    new_documents_chroma = glob.glob("chromadb/new*")
+    new_documents_bm25 = glob.glob("chromadb/new*")
+    new_countries = [doc.split('_')[1] for doc in new_documents_chroma]
+    if len(new_countries) == 0:
+        info = '(Own documents not found. Must first scrape or upload own PDF (see menu above) to use this.)'
+    else:
+        info = '(⚠️Own documents for the following countries found, select them in the list below to override)'
+
+    with st.expander("Document Config", expanded = True):
+        st.multiselect(
+                'Countries to Override with Own Docs:' + info,
+                new_countries,
+                key="countries_override"           
+            )
+
+        st.selectbox(
+                "Chunk Size", 
+                options=[500, 600, 700, 800, 900, 1000, 1250, 1500, 1750, 2000, 2250, 2500, 2750, 3000], 
+                on_change=update_retrievers,
+                key="chunk_size"
+            )
+        
+        st.selectbox(
+                "Chunk Overlap", 
+                options=[50, 100, 150, 200],
+                on_change=update_retrievers,
+                key="chunk_overlap"
+            )
+    
+    st.write("")
+
     with st.expander("LLM Config", expanded = True):
         
         st.selectbox(
@@ -410,22 +441,9 @@ with st.sidebar:
                 200, 1000,
                 on_change=update_llm,
                 key="max_new_tokens"
-            )     
+            )  
 
-    with st.expander("Document Config", expanded = True):
-        st.selectbox(
-                "Chunk Size", 
-                options=[500, 600, 700, 800, 900, 1000, 1250, 1500, 1750, 2000, 2250, 2500, 2750, 3000], 
-                on_change=update_retrievers,
-                key="chunk_size"
-            )
-        
-        st.selectbox(
-                "Chunk Overlap", 
-                options=[50, 100, 150, 200],
-                on_change=update_retrievers,
-                key="chunk_overlap"
-            )
+    st.write("")
 
     with st.expander("Retriever Config", expanded = True):
 
@@ -517,18 +535,9 @@ if page == "Main Chatbot":
             st.markdown(response)
 
 
-
-################################ Document Page ################################
-# to scrape new documents from DuckDuckGo
-# to chnange paramters like chunk size
-# to upload own PDF
-# to override existing data on new scraped data or new pdf uploaded
-
-
-
 ################################ Source Documents Page ################################
 if page == "View Source Docs for Last Query":
-    st.header("Source Documents for Last Query")
+    st.subheader("Source Documents for Last Query")
     try:
         st.subheader(st.session_state['source_documents'][0])
         for doc in st.session_state['source_documents'][1:]:
@@ -539,6 +548,45 @@ if page == "View Source Docs for Last Query":
 
 
 
-# in main app, add configuration for user to scrape new data from DuckDuckGo
-# in main app, add configuration for user to upload PDF to override country's existing policies in vectorstore
+################################ Scrap or Upload Documents Page ################################
+# to scrape new documents from DuckDuckGo
+# to upload own PDF
+# to override existing data on new scraped data or new pdf uploaded
+if page == "Scrape or Upload Own Docs":
+    st.header("Scrape or Upload Own PDF")
+    st.write("Here you can choose to upload your own PDF or scrape more recent data via DuckDuckGo search for a selected country below.")
+    st.write("You will create new BM2.5 (keyword) and Chroma (semantic) retrievers for it. Note that this can take a very long time.")
+
+    country_scrape_upload = st.selectbox(
+        "Select Country", 
+        options=[
+            "Australia", "Bangladesh", "Brunei", "Cambodia", "China", "India", "Indonesia", "Japan", "Laos", "Macau", "Malaysia", "Myanmar",
+            "Nepal", "Philippines", "Singapore", "South Korea", "Sri Lanka", "Thailand", "Vietnam", "France", "Germany", "Israel", "Poland",
+            "Sweden", "Turkey", "United Kingdom", "United States"
+            ], 
+    )
+    options = [
+                "Upload Own PDF", 
+                "Automatically Scrape Web Data using DuckDuckGo (Will take a long time, 10 mins or more)"
+            ]
+
+    option = st.selectbox(
+        "How Do You Wish To Create New Documents", 
+        options=options 
+    )
+
+    col1, col2 = st.columns(2)
+    with col1:
+        with st.container(border = True):
+            st.write("New Documents Chunk Size: (Can change in sidebar)" )
+            st.text(f"{st.session_state['chunk_size']}" )
+    with col2:
+        with st.container(border = True):
+            st.write("New Documents Chunk Overlap: (Can change in sidebar)" )
+            st.text(f"{st.session_state['chunk_overlap']}")
+
+    st.subheader(f"Selected Option: {option}")
+    
+
 
+    # spinner