Update app.py

add region filter

app.py

@@ -4,7 +4,7 @@ from appStore.prep_data import process_giz_worldwide, remove_duplicates
 from appStore.prep_utils import create_documents, get_client
 from appStore.embed import hybrid_embed_chunks
 from appStore.search import hybrid_search
-from appStore.region_utils import load_region_data, get_country_name
+from appStore.region_utils import load_region_data, get_country_name, get_regions
 from appStore.tfidf_extraction import extract_top_keywords 
 from torch import cuda
 import json
@@ -36,45 +36,53 @@ collection_name = "giz_worldwide"
 client = get_client()
 print(client.get_collections())
 
+# Get all unique sub-regions
+_, unique_sub_regions = get_regions(region_df)
+
 # Fetch unique country codes and map to country names
 @st.cache_data
-def get_country_name_mapping(_client, collection_name, region_df):
+def get_country_name_and_region_mapping(_client, collection_name, region_df):
     results = hybrid_search(_client, "", collection_name)
     country_set = set()
     for res in results[0] + results[1]:
         countries = res.payload.get('metadata', {}).get('countries', "[]")
         try:
             country_list = json.loads(countries.replace("'", '"'))
-            # ADD: only add codes of length 2
+            # Only add codes of length 2
             two_digit_codes = [code.upper() for code in country_list if len(code) == 2]
             country_set.update(two_digit_codes)
         except json.JSONDecodeError:
             pass
-    
-    # Create a mapping of {CountryName -> ISO2Code} 
-    # so you can display the name in the selectbox but store the 2-digit code
+
+    # Create a mapping of {CountryName -> ISO2Code} and {ISO2Code -> SubRegion}
     country_name_to_code = {}
+    iso_code_to_sub_region = {}
+
     for code in country_set:
         name = get_country_name(code, region_df)
+        sub_region_row = region_df[region_df['alpha-2'] == code]
+        sub_region = sub_region_row['sub-region'].values[0] if not sub_region_row.empty else "Not allocated"
         country_name_to_code[name] = code
+        iso_code_to_sub_region[code] = sub_region
 
-    return country_name_to_code
+    return country_name_to_code, iso_code_to_sub_region
 
-
-# Get country name mapping
+# Get country name and region mappings
 client = get_client()
-country_name_mapping = get_country_name_mapping(client, collection_name, region_df)
+country_name_mapping, iso_code_to_sub_region = get_country_name_and_region_mapping(client, collection_name, region_df)
 unique_country_names = sorted(country_name_mapping.keys())  # List of country names
 
 # Layout filters in columns
 col1, col2, col3 = st.columns([1, 1, 4])
 
 with col1:
-    country_filter = st.selectbox("Country", ["All/Not allocated"] + unique_country_names)  # Display country names
+    region_filter = st.selectbox("Region", ["All/Not allocated"] + sorted(unique_sub_regions))  # Display region names
 with col2:
+    country_filter = st.selectbox("Country", ["All/Not allocated"] + unique_country_names)  # Display country names
+with col3:
     current_year = datetime.now().year
     default_start_year = current_year - 5  # Default to 5 years ago
-    
+
     end_year_range = st.slider(
         "Project End Year",
         min_value=2010,
@@ -86,7 +94,7 @@ with col2:
 show_exact_matches = st.checkbox("Show only exact matches", value=False)
 button = st.button("Search")
 
-def filter_results(results, country_filter, end_year_range):
+def filter_results(results, country_filter, region_filter, end_year_range):
     filtered = []
     for r in results:
         metadata = r.payload.get('metadata', {})
@@ -103,15 +111,21 @@ def filter_results(results, country_filter, end_year_range):
         # Translate selected country name to iso2
         selected_iso_code = country_name_mapping.get(country_filter, None)
 
+        # Check if any country in the metadata matches the selected region
+        if region_filter != "All/Not allocated":
+            countries_in_region = [code for code in c_list if iso_code_to_sub_region.get(code) == region_filter]
+        else:
+            countries_in_region = c_list
+
         # Filtering
         if (
             (country_filter == "All/Not allocated" or selected_iso_code in c_list)
+            and (region_filter == "All/Not allocated" or countries_in_region)
             and (end_year_range[0] <= end_year_val <= end_year_range[1])
         ):
             filtered.append(r)
     return filtered
 
-
 if button:
     # 1) Use a bigger limit so we get more than 10 results
     #    We'll filter them first, then slice the top 10 from the filtered set.
@@ -122,8 +136,8 @@ if button:
     lexical_all = results[1]
 
     # 2) Filter the entire sets
-    filtered_semantic = filter_results(semantic_all, country_filter, end_year_range)
-    filtered_lexical = filter_results(lexical_all, country_filter, end_year_range)
+    filtered_semantic = filter_results(semantic_all, country_filter, region_filter, end_year_range)
+    filtered_lexical = filter_results(lexical_all, country_filter, region_filter, end_year_range)
 
     filtered_semantic_no_dupe = remove_duplicates(filtered_semantic)
     filtered_lexical_no_dupe = remove_duplicates(filtered_lexical)