Refactor Streamlit app to support jurisdiction-specific data handling. Updated functions to accept jurisdiction as a parameter, modified file paths for user data and grading templates, and added jurisdiction selection in the login process. This enhances the app's flexibility for different jurisdictions.

src/streamlit_app.py

@@ -15,9 +15,12 @@ HF_TOKEN = os.getenv("HF_TOKEN")
 # Fallback to st.secrets for local development (if not found in environment)
 if not HF_DATASET_REPO:
     try:
-        HF_DATASET_REPO = st.secrets.get("HF_DATASET_REPO", "TransLegal/en-us-grading-answers")
+        HF_DATASET_REPO = st.secrets.get("HF_DATASET_REPO", "TransLegal/grading-answers")
     except Exception:
-        HF_DATASET_REPO = "TransLegal/en-us-grading-answers"
+        HF_DATASET_REPO = "TransLegal/grading-answers"
+
+# Available jurisdictions
+AVAILABLE_JURISDICTIONS = ["en-us", "hr-hr", "sv-se"]
 
 if not HF_TOKEN:
     try:
@@ -33,35 +36,32 @@ if not HF_TOKEN:
 @st.cache_resource
 def get_hf_api():
     """Get cached Hugging Face API client - only initializes once per session"""
-    try:
-        login(token=HF_TOKEN)
+try:
+    login(token=HF_TOKEN)
         return HfApi(token=HF_TOKEN)
-    except Exception as e:
-        st.error(f"❌ **Error initializing Hugging Face API**: {str(e)}")
-        st.stop()
+except Exception as e:
+    st.error(f"❌ **Error initializing Hugging Face API**: {str(e)}")
+    st.stop()
 
 # Initialize HF API - cached to avoid re-initialization on every rerun
 hf_api = get_hf_api()
 
 @st.cache_data
-def load_grading_template():
-    """Load grading template from Hugging Face Dataset"""
+def load_grading_template(jurisdiction):
+    """Load grading template from Hugging Face Dataset for the specified jurisdiction"""
     try:
         file_path = hf_hub_download(
             repo_id=HF_DATASET_REPO,
-            filename="grading_template.parquet",
+            filename=f"{jurisdiction}/grading_template.parquet",
             repo_type="dataset",
             token=HF_TOKEN
         )
         return pd.read_parquet(file_path)
     except Exception as e:
         st.error(f"❌ **Error loading grading template from Hugging Face Dataset**: {str(e)}")
-        st.error(f"Please ensure the file `grading_template.parquet` exists in the dataset repository: {HF_DATASET_REPO}")
+        st.error(f"Please ensure the file `{jurisdiction}/grading_template.parquet` exists in the dataset repository: {HF_DATASET_REPO}")
         st.stop()
 
-# Load data from the grading template
-df = load_grading_template()
-
 # Assessment options with descriptive labels
 ASSESSMENT_OPTIONS = [
     "Perfect",
@@ -138,12 +138,12 @@ def hash_password(password):
     return hashlib.sha256(password.encode()).hexdigest()
 
 @st.cache_data
-def load_users():
-    """Load user credentials from Hugging Face Dataset"""
+def load_users(jurisdiction):
+    """Load user credentials from Hugging Face Dataset for the specified jurisdiction"""
     try:
         file_path = hf_hub_download(
             repo_id=HF_DATASET_REPO,
-            filename="users/users.json",
+            filename=f"{jurisdiction}/users/users.json",
             repo_type="dataset",
             token=HF_TOKEN
         )
@@ -153,8 +153,8 @@ def load_users():
         # File doesn't exist yet (first run), return empty dict
         return {}
 
-def save_users(users):
-    """Save user credentials to Hugging Face Dataset"""
+def save_users(users, jurisdiction):
+    """Save user credentials to Hugging Face Dataset for the specified jurisdiction"""
     try:
         with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f:
             json.dump(users, f, indent=2)
@@ -162,7 +162,7 @@ def save_users(users):
         
         hf_api.upload_file(
             path_or_fileobj=temp_path,
-            path_in_repo="users/users.json",
+            path_in_repo=f"{jurisdiction}/users/users.json",
             repo_id=HF_DATASET_REPO,
             repo_type="dataset",
             token=HF_TOKEN
@@ -170,7 +170,7 @@ def save_users(users):
         os.unlink(temp_path)
         
         # Clear cache for users to ensure fresh data on next load
-        load_users.clear()
+        load_users.clear(jurisdiction)
         
         return True
     except Exception as e:
@@ -178,25 +178,26 @@ def save_users(users):
         raise
 
 @st.cache_data(ttl=3600)  # Cache for 1 hour as safety measure
-def load_user_data(username):
-    """Load user's answer data from Hugging Face Dataset"""
+def load_user_data(username, jurisdiction):
+    """Load user's answer data from Hugging Face Dataset for the specified jurisdiction"""
     try:
         file_path = hf_hub_download(
             repo_id=HF_DATASET_REPO,
-            filename=f"users/{username}_answers.parquet",
+            filename=f"{jurisdiction}/users/{username}_answers.parquet",
             repo_type="dataset",
             token=HF_TOKEN
         )
         return pd.read_parquet(file_path)
     except Exception:
-        # File doesn't exist yet (new user), create new dataframe
+        # File doesn't exist yet (new user), create new dataframe from grading template
+        df = load_grading_template(jurisdiction)
         user_df = df.copy()
         user_df['legal_accuracy_score'] = None
         user_df['time_stamp'] = None
         return user_df
 
-def save_user_data(username, user_df, commit_message=None):
-    """Save user's answer data to Hugging Face Dataset"""
+def save_user_data(username, user_df, jurisdiction, commit_message=None):
+    """Save user's answer data to Hugging Face Dataset for the specified jurisdiction"""
     try:
         with tempfile.NamedTemporaryFile(suffix='.parquet', delete=False) as f:
             user_df.to_parquet(f.name, index=False)
@@ -204,7 +205,7 @@ def save_user_data(username, user_df, commit_message=None):
         
         upload_kwargs = {
             'path_or_fileobj': temp_path,
-            'path_in_repo': f"users/{username}_answers.parquet",
+            'path_in_repo': f"{jurisdiction}/users/{username}_answers.parquet",
             'repo_id': HF_DATASET_REPO,
             'repo_type': "dataset",
             'token': HF_TOKEN
@@ -217,18 +218,18 @@ def save_user_data(username, user_df, commit_message=None):
         hf_api.upload_file(**upload_kwargs)
         os.unlink(temp_path)
         
-        # Clear cache for this user to ensure fresh data on next load
-        load_user_data.clear(username)
+        # Clear cache for this user/jurisdiction to ensure fresh data on next load
+        load_user_data.clear(username, jurisdiction)
         
         return True
     except Exception as e:
         st.error(f"❌ **Error saving user data to Hugging Face Dataset**: {str(e)}")
         raise
 
-def update_user_answer(username, term, category, subcategory, question, answer, score):
+def update_user_answer(username, term, category, subcategory, question, answer, score, jurisdiction, df):
     """Update a specific answer in the user's data (deprecated - use update_category_answers for bulk updates)"""
     try:
-        user_df = load_user_data(username)
+        user_df = load_user_data(username, jurisdiction)
         
         # Find the matching row
         mask = (
@@ -241,7 +242,7 @@ def update_user_answer(username, term, category, subcategory, question, answer,
         
         if mask.any():
             user_df.loc[mask, 'legal_accuracy_score'] = score
-            save_user_data(username, user_df)
+            save_user_data(username, user_df, jurisdiction)
             return True
         else:
             print(f"Warning: Could not find matching row for: {term}, {category}, {subcategory}, {question}")
@@ -273,7 +274,7 @@ def auto_score_unknown_answers(username, term, category, df):
     return [(row['subcategory'], row['question'], row['answer'], "NA") 
             for _, row in unknown_rows.iterrows()]
 
-def auto_score_all_unknown_answers_for_new_user(username):
+def auto_score_all_unknown_answers_for_new_user(username, jurisdiction, df):
     """
     Automatically score all Unknown answers for all categories when a new user is created.
     This runs in the background during account creation.
@@ -297,7 +298,7 @@ def auto_score_all_unknown_answers_for_new_user(username):
             return True
         
         # Load user dataframe once
-        user_df = load_user_data(username)
+        user_df = load_user_data(username, jurisdiction)
         
         # Get current timestamp once for all updates
         current_timestamp = pd.Timestamp.now()
@@ -327,14 +328,14 @@ def auto_score_all_unknown_answers_for_new_user(username):
         
         # Save once with a single commit message
         commit_message = f"Auto-score all Unknown answers for new user {username}"
-        save_user_data(username, user_df, commit_message=commit_message)
+        save_user_data(username, user_df, jurisdiction, commit_message=commit_message)
         
         return True
     except Exception as e:
         print(f"Error auto-scoring Unknown answers for new user {username}: {str(e)}")
         return False
 
-def update_category_answers(username, term, category, answers_list, commit_message=None):
+def update_category_answers(username, term, category, answers_list, jurisdiction, commit_message=None):
     """
     Update all answers for a category in a single commit.
     
@@ -343,6 +344,7 @@ def update_category_answers(username, term, category, answers_list, commit_messa
         term: Term name
         category: Category name
         answers_list: List of tuples (subcategory, question, answer, score)
+        jurisdiction: Jurisdiction identifier
         commit_message: Optional commit message (auto-generated if None)
     
     Returns:
@@ -350,7 +352,7 @@ def update_category_answers(username, term, category, answers_list, commit_messa
     """
     try:
         # Load user dataframe once
-        user_df = load_user_data(username)
+        user_df = load_user_data(username, jurisdiction)
         
         # Get current timestamp once for all updates in this category
         current_timestamp = pd.Timestamp.now()
@@ -383,15 +385,15 @@ def update_category_answers(username, term, category, answers_list, commit_messa
             commit_message = f"Update answers for {username} - {term} - {category}"
         
         # Save once with commit message
-        save_user_data(username, user_df, commit_message=commit_message)
+        save_user_data(username, user_df, jurisdiction, commit_message=commit_message)
         return True
     except Exception as e:
         print(f"Error updating category answers: {str(e)}")
         return False
 
-def get_user_answer(username, term, category, subcategory, question, answer):
+def get_user_answer(username, term, category, subcategory, question, answer, jurisdiction):
     """Get user's answer for a specific question"""
-    user_df = load_user_data(username)
+    user_df = load_user_data(username, jurisdiction)
     
     mask = (
         (user_df['term'] == term) &
@@ -407,12 +409,12 @@ def get_user_answer(username, term, category, subcategory, question, answer):
             return score
     return None
 
-def find_first_unanswered_category(username):
+def find_first_unanswered_category(username, jurisdiction, df):
     """Find the first category that hasn't been fully answered"""
-    user_df = load_user_data(username)
+    user_df = load_user_data(username, jurisdiction)
     
-    # Use the global term_category_pairs to ensure consistent ordering
-    # This matches the order used in the main application
+    # Get term_category_pairs for this jurisdiction
+    term_category_pairs = get_term_category_pairs(df)
     
     for idx, (term, category) in enumerate(term_category_pairs):
         # Get all subcategories for this term-category pair from base df
@@ -445,11 +447,13 @@ def find_first_unanswered_category(username):
     
     return len(term_category_pairs)  # All answered, return last index
 
-def restore_submitted_status(username):
+def restore_submitted_status(username, jurisdiction, df):
     """Restore submitted status for categories that have all answers in parquet file"""
-    user_df = load_user_data(username)
+    user_df = load_user_data(username, jurisdiction)
+    
+    # Get term_category_pairs for this jurisdiction
+    term_category_pairs = get_term_category_pairs(df)
     
-    # Use the global term_category_pairs to ensure consistent ordering
     submitted_pairs = set()
     for idx, (term, category) in enumerate(term_category_pairs):
         pair_key = f"{term}_{category}_{idx}"
@@ -585,22 +589,19 @@ def get_term_category_pairs(df):
     return [(term, category) for term, category in all_pairs 
             if category_has_subcategories(term, category, df)]
 
-# Create a list of unique (term, category) pairs for navigation
-term_category_pairs = get_term_category_pairs(df)
-total_pairs = len(term_category_pairs)
-
-# Cache for Term instances
+# Cache for Term instances (keyed by jurisdiction and term_name)
 term_cache = {}
 
-def get_term_instance(term_name):
-    """Get or create a Term instance"""
-    if term_name not in term_cache:
-        term_cache[term_name] = Term(term_name, df)
-    return term_cache[term_name]
+def get_term_instance(term_name, df):
+    """Get or create a Term instance for the given dataframe"""
+    cache_key = f"{id(df)}_{term_name}"  # Use df id to differentiate jurisdictions
+    if cache_key not in term_cache:
+        term_cache[cache_key] = Term(term_name, df)
+    return term_cache[cache_key]
 
-def get_category_for_pair(term_name, category_name):
+def get_category_for_pair(term_name, category_name, df):
     """Get Category instance for a term-category pair"""
-    term = get_term_instance(term_name)
+    term = get_term_instance(term_name, df)
     return term.get_category_by_name(category_name)
 
 # Initialize session state
@@ -608,6 +609,8 @@ if 'logged_in' not in st.session_state:
     st.session_state.logged_in = False
 if 'username' not in st.session_state:
     st.session_state.username = None
+if 'jurisdiction' not in st.session_state:
+    st.session_state.jurisdiction = None
 if 'current_index' not in st.session_state:
     st.session_state.current_index = 0
 if 'show_term_complete' not in st.session_state:
@@ -636,7 +639,11 @@ if 'has_unsaved_changes' not in st.session_state:
 # Login page
 if not st.session_state.logged_in:
     st.markdown("# Login")
-    st.markdown("Please enter your username and password to continue.")
+    st.markdown("Please select a jurisdiction and enter your username and password to continue.")
+    
+    # Jurisdiction selector
+    jurisdiction = st.selectbox("Jurisdiction", options=AVAILABLE_JURISDICTIONS, index=0 if st.session_state.jurisdiction is None else AVAILABLE_JURISDICTIONS.index(st.session_state.jurisdiction) if st.session_state.jurisdiction in AVAILABLE_JURISDICTIONS else 0)
+    st.session_state.jurisdiction = jurisdiction
     
     username = st.text_input("Username")
     password = st.text_input("Password", type="password")
@@ -644,44 +651,58 @@ if not st.session_state.logged_in:
     col1, col2 = st.columns(2)
     with col1:
         if st.button("Login", type="primary", use_container_width=True):
-            users = load_users()
-            
-            if username in users:
-                # Existing user - check password
-                if users[username]['password'] == hash_password(password):
-                    st.session_state.logged_in = True
-                    st.session_state.username = username
-                    # Restore submitted status for previously submitted categories
-                    st.session_state.submitted_pairs = restore_submitted_status(username)
-                    # Find first unanswered category and resume there
-                    resume_index = find_first_unanswered_category(username)
-                    st.session_state.current_index = resume_index
-                    st.rerun()
-                else:
-                    st.error("Incorrect password")
+            if not jurisdiction:
+                st.error("Please select a jurisdiction")
             else:
-                # Username not found - require registration
-                st.error("Username not found. Please register first using the 'Register New User' button.")
+                users = load_users(jurisdiction)
+                
+                if username in users:
+                    # Existing user - check password
+                    if users[username]['password'] == hash_password(password):
+                        st.session_state.logged_in = True
+                        st.session_state.username = username
+                        # Load grading template for this jurisdiction
+                        df = load_grading_template(jurisdiction)
+                        # Restore submitted status for previously submitted categories
+                        st.session_state.submitted_pairs = restore_submitted_status(username, jurisdiction, df)
+                        # Find first unanswered category and resume there
+                        resume_index = find_first_unanswered_category(username, jurisdiction, df)
+                        st.session_state.current_index = resume_index
+                        st.rerun()
+                    else:
+                        st.error("Incorrect password")
+                else:
+                    # Username not found - require registration
+                    st.error("Username not found. Please register first using the 'Register New User' button.")
     
     with col2:
         if st.button("Register New User", use_container_width=True):
-            users = load_users()
-            if username in users:
-                st.error("Username already exists")
-            elif username and password:
-                users[username] = {'password': hash_password(password)}
-                save_users(users)
-                # Auto-score all Unknown answers for the new user in the background
-                auto_score_all_unknown_answers_for_new_user(username)
-                st.success("User registered successfully! Please click Login.")
+            if not jurisdiction:
+                st.error("Please select a jurisdiction")
             else:
-                st.error("Please enter both username and password")
+                users = load_users(jurisdiction)
+                if username in users:
+                    st.error("Username already exists")
+                elif username and password:
+                    users[username] = {'password': hash_password(password)}
+                    save_users(users, jurisdiction)
+                    # Load grading template for this jurisdiction
+                    df = load_grading_template(jurisdiction)
+                    # Auto-score all Unknown answers for the new user in the background
+                    auto_score_all_unknown_answers_for_new_user(username, jurisdiction, df)
+                    st.success("User registered successfully! Please click Login.")
+                else:
+                    st.error("Please enter both username and password")
 
 # Main application (only shown if logged in)
 elif st.session_state.logged_in:
     username = st.session_state.username
+    jurisdiction = st.session_state.jurisdiction
     current_index = st.session_state.current_index
     
+    # Load grading template for the selected jurisdiction
+    df = load_grading_template(jurisdiction)
+    
     # Debug info (can be removed in production)
     with st.sidebar:
         with st.expander("Debug Info"):
@@ -689,8 +710,9 @@ elif st.session_state.logged_in:
             st.write(f"HF Token configured: {HF_TOKEN is not None}")
             st.write(f"HF API initialized: {hf_api is not None}")
             if username:
-                st.write(f"User parquet file: `users/{username}_answers.parquet`")
-                st.write(f"Users file: `users/users.json`")
+                st.write(f"Jurisdiction: `{jurisdiction}`")
+                st.write(f"User parquet file: `{jurisdiction}/users/{username}_answers.parquet`")
+                st.write(f"Users file: `{jurisdiction}/users/users.json`")
     
     # Check if we should show the annotation guide first
     if st.session_state.show_guide:
@@ -829,10 +851,14 @@ elif st.session_state.logged_in:
                 st.session_state.next_term = None
                 st.rerun()
     
+    # Get term_category_pairs for this jurisdiction
+    term_category_pairs = get_term_category_pairs(df)
+    total_pairs = len(term_category_pairs)
+    
     elif current_index < total_pairs:
         term_name, category_name = term_category_pairs[current_index]
-        category = get_category_for_pair(term_name, category_name)
-        term = get_term_instance(term_name)
+        category = get_category_for_pair(term_name, category_name, df)
+        term = get_term_instance(term_name, df)
         
         # Safety check: skip if category has no subcategories (shouldn't happen due to filtering, but just in case)
         if not category or len(category.subcategories) == 0:
@@ -858,7 +884,7 @@ elif st.session_state.logged_in:
         # Check visible subcategories
         for i, subcat in enumerate(category.subcategories):
             saved_score = get_user_answer(username, term_name, category_name, subcat.subcategory_name,
-                                         subcat.question, subcat.answer)
+                                         subcat.question, subcat.answer, jurisdiction)
             if saved_score is None:
                 category_fully_answered = False
                 break
@@ -867,7 +893,7 @@ elif st.session_state.logged_in:
         if category_fully_answered:
             unknown_answers = auto_score_unknown_answers(username, term_name, category_name, df)
             for subcategory, question, answer, score in unknown_answers:
-                saved_score = get_user_answer(username, term_name, category_name, subcategory, question, answer)
+                saved_score = get_user_answer(username, term_name, category_name, subcategory, question, answer, jurisdiction)
                 if saved_score is None:
                     category_fully_answered = False
                     break
@@ -890,7 +916,7 @@ elif st.session_state.logged_in:
                         st.session_state.original_selections[pair_key][radio_key] = st.session_state[radio_key]
                     else:
                         saved_score = get_user_answer(username, term_name, category_name, subcat.subcategory_name,
-                                                     subcat.question, subcat.answer)
+                                                     subcat.question, subcat.answer, jurisdiction)
                         if saved_score is not None:
                             score_to_option = {v: k for k, v in ASSESSMENT_TO_SCORE.items()}
                             if saved_score in score_to_option:
@@ -929,7 +955,7 @@ elif st.session_state.logged_in:
             # Get saved value to determine if we should set a default index
             # Don't pre-set session_state - only use index parameter to avoid conflicts
             saved_score = get_user_answer(username, term_name, category_name, subcat.subcategory_name,
-                                         subcat.question, subcat.answer)
+                                         subcat.question, subcat.answer, jurisdiction)
             default_index = None
             
             # Check if there's a saved value in parquet file
@@ -1014,7 +1040,7 @@ elif st.session_state.logged_in:
                             st.session_state.back_current_index = current_index
                             st.session_state.show_term_back_warning = True
                             st.session_state.back_current_term = term.formatted_name
-                            prev_term = get_term_instance(prev_term_name)
+                            prev_term = get_term_instance(prev_term_name, df)
                             st.session_state.back_previous_term = prev_term.formatted_name
                         else:
                             # Same term, just move back
@@ -1058,7 +1084,7 @@ elif st.session_state.logged_in:
                         
                         # Update all answers in a single commit
                         commit_message = f"Update answers for {username} - {term_name} - {category_name}"
-                        save_success = update_category_answers(username, term_name, category_name, answers_list, commit_message)
+                        save_success = update_category_answers(username, term_name, category_name, answers_list, jurisdiction, commit_message)
                         
                         if save_success:
                             # Save current selections as new originals
@@ -1090,7 +1116,7 @@ elif st.session_state.logged_in:
                         # Update all answers in a single commit
                         # Note: Unknown answers are already auto-scored during account creation
                         commit_message = f"Update answers for {username} - {term_name} - {category_name}"
-                        save_success = update_category_answers(username, term_name, category_name, answers_list, commit_message)
+                        save_success = update_category_answers(username, term_name, category_name, answers_list, jurisdiction, commit_message)
                         
                         if save_success:
                             # Mark as submitted and save original selections
@@ -1113,7 +1139,7 @@ elif st.session_state.logged_in:
                                     # Show intermediate page
                                     st.session_state.show_term_complete = True
                                     st.session_state.completed_term = term.formatted_name
-                                    next_term = get_term_instance(next_term_name)
+                                    next_term = get_term_instance(next_term_name, df)
                                     st.session_state.next_term = next_term.formatted_name
                                 else:
                                     # Same term, just move to next category
@@ -1148,7 +1174,7 @@ elif st.session_state.logged_in:
                             # Moving to a different term - show term switching page
                             st.session_state.show_term_complete = True
                             st.session_state.completed_term = term.formatted_name
-                            next_term = get_term_instance(next_term_name)
+                                next_term = get_term_instance(next_term_name, df)
                             st.session_state.next_term = next_term.formatted_name
                         else:
                             # Same term, just move forward
@@ -1170,6 +1196,7 @@ elif st.session_state.logged_in:
         if st.button("Logout"):
             st.session_state.logged_in = False
             st.session_state.username = None
+            st.session_state.jurisdiction = None
             st.session_state.current_index = 0
             st.session_state.show_guide = True
             st.session_state.submitted_pairs = set()