hacky but works :D

app.py

@@ -32,10 +32,8 @@ def initialization():
         st.session_state.start_new_survey = False
     if 'ratings' not in st.session_state:
         st.session_state.ratings = {}
-    if 'previous_gemini_ratings' not in st.session_state:
-        st.session_state.previous_gemini_ratings = {}
-    if 'previous_llama_ratings' not in st.session_state:
-        st.session_state.previous_llama_ratings = {}
+    if 'previous_ratings' not in st.session_state:
+        st.session_state.previous_ratings = {}
 
 
 def exit_screen():

views/continue_survey.py

@@ -36,7 +36,6 @@ def continue_survey_screen(data):
                 # Set survey_continued flag to True only when there's saved progress
                 st.session_state.survey_continued = True
                 st.rerun()
-                # questions_screen(data)
 
             else:
                 st.warning("No previous progress found. Starting a new survey.")

views/nav_buttons.py

@@ -42,30 +42,29 @@ def navigation_buttons(data, response: Response):
     col1, col2, col3 = st.columns([1, 1, 2])
 
     with col1:  # Back button #TODO fix: only gets ratings for the session, not from previous session
-        if st.button("Back"):
+        if st.button("Back", disabled=st.session_state.current_index == 0):
             if current_index > 0:
+                st.session_state.previous_ratings[
+                    data.iloc[st.session_state.current_index]['config_id']] = response.model_ratings
                 st.session_state.current_index -= 1
-                previous_ratings = st.session_state.ratings.get(st.session_state.current_index, {})
-                # st.session_state.previous_ratings = previous_ratings
-                # st.session_state.previous_gemini_ratings = previous_ratings.get("gemini", {})
-                # st.session_state.previous_llama_ratings = previous_ratings.get("llama", {})
                 st.rerun()
             else:
                 st.warning("You are at the beginning of the survey, can't go back.")
             # st.rerun()
 
-    with col2:  # Next button
-        if st.button("Next"):
+    with col2:  # Next button TODO might be buggy
+        if st.button("Next", disabled=st.session_state.current_index == len(data) - 1):
             all_ratings = flatten_ratings(response)
             if any(rating == 0 for rating in all_ratings):
                 st.warning("Please provide all ratings before proceeding.")
             else:
                 if current_index < len(data) - 1:
+                    st.session_state.previous_ratings[data.iloc[st.session_state.current_index]['config_id']] = response.model_ratings
                     st.session_state.current_index += 1
-                    # st.session_state.previous_ratings = {}
                     st.rerun()
                 else:
-                    submit_feedback(current_index)
+                    if st.button("Finish"):
+                        submit_feedback(current_index)
 
     with col3:  # Save & Resume Later button
         if st.button("Exit & Resume Later"):

views/questions_screen.py

@@ -48,27 +48,58 @@ def display_ratings_row(model_name, config, current_index):
 
 def render_query_ratings(model_name, query_label, config, query_key, current_index, has_persona_alignment=False):
     """Helper function to render ratings for a given query."""
-    # Get stored ratings if they exist
-    if current_index < st.session_state.current_index:
-        previous_ratings = st.session_state.get("previous_ratings", {}).get(model_name, None)
+
+    previous_ratings = {}
+    # If the user is coming from a next button press, then previous rating will not exist
+    if current_index < st.session_state.current_index and len(st.session_state.responses) > current_index:
+        if st.session_state.previous_ratings:
+            previous_ratings = st.session_state.previous_ratings.get(config["config_id"], {})
+            if "gemini" == model_name:
+                previous_ratings = previous_ratings.get("gemini", None)
+            else:
+                previous_ratings = previous_ratings.get("llama", None)
+    # This means there were no previous responses, i.e first time opening the page, or just clicking continue
+    elif len(st.session_state.responses) <= current_index:
+        previous_ratings = {}
+    # User has already entered some response in the page they are in
     else:
-        previous_ratings = None  # Ensure new questions start fresh
+        # get the saved ratings from session state for this question
+        response_from_session = st.session_state.responses[current_index]
+        if "gemini" == model_name:
+            try:
+                previous_ratings = response_from_session.model_ratings.get("gemini", {})
+            except AttributeError:
+                previous_ratings = response_from_session["model_ratings"].get("gemini", {})
+        else:
+            try:
+                previous_ratings = response_from_session.model_ratings.get("llama", {})
+            except AttributeError:
+                previous_ratings = response_from_session["model_ratings"].get("llama", {})
+
     stored_query_ratings = {}
+
     if previous_ratings:
         if "query_v" in query_key:
-            stored_query_ratings = previous_ratings.query_v_ratings
+            try:
+                stored_query_ratings = previous_ratings.query_v_ratings
+            except AttributeError:
+                stored_query_ratings = previous_ratings["query_v_ratings"]
         elif "query_p0" in query_key:
-            stored_query_ratings = previous_ratings.query_p0_ratings
+            try:
+                stored_query_ratings = previous_ratings.query_p0_ratings
+            except AttributeError:
+                stored_query_ratings = previous_ratings["query_p0_ratings"]
         elif "query_p1" in query_key:
-            stored_query_ratings = previous_ratings.query_p1_ratings
+            try:
+                stored_query_ratings = previous_ratings.query_p1_ratings
+            except AttributeError:
+                stored_query_ratings = previous_ratings["query_p1_ratings"]
         else:
             stored_query_ratings = {}
-    # Extract individual stored values, or default to None
-    # stored_query_ratings = previous_ratings.get(f"query_{query_key}_ratings", {})
 
-    stored_relevance = stored_query_ratings.get("relevance", 0)
-    stored_clarity = stored_query_ratings.get("clarity", 0)
-    stored_persona_alignment = stored_query_ratings.get("persona_alignment", 0) if has_persona_alignment else None
+    stored_relevance = stored_query_ratings.get("relevance", 0) if stored_query_ratings else 0
+    stored_clarity = stored_query_ratings.get("clarity", 0) if stored_query_ratings else 0
+    stored_persona_alignment = stored_query_ratings.get("persona_alignment", 0) if has_persona_alignment and stored_query_ratings else 0
 
     if model_name == "gemini":
         bg_color = "#e0f7fa"
@@ -163,7 +194,11 @@ def questions_screen(data):
             comment=comment,
             timestamp=datetime.now().isoformat()
         )
-        st.session_state.ratings[current_index] = response.model_ratings
+        print(response)
+        try:
+            st.session_state.ratings[current_index] = response["model_ratings"]
+        except TypeError:
+            st.session_state.ratings[current_index] = response.model_ratings
         if len(st.session_state.responses) > current_index:
             st.session_state.responses[current_index] = response
         else: