Update app.py

app.py

@@ -1,24 +1,22 @@
 import json
 import os
+import numpy as np
+import pandas as pd
+
 from dataclasses import dataclass, field
 from typing import List, Optional, Dict
 from PIL import Image
 
 import streamlit as st
+import streamlit.components.v1 as components
+from streamlit.components.v1 import html as st_html
 from markdown import markdown
 st.set_page_config(layout="wide", page_title="Annotation of Simulated Patients", initial_sidebar_state="collapsed")
-
 import re, textwrap, html as py_html
 from pathlib import Path
-from streamlit.components.v1 import html as st_html
-
-import numpy as np
-import pandas as pd
 from fsspec.implementations.local import LocalFileSystem
 from huggingface_hub import HfFileSystem
 
-import streamlit.components.v1 as components
-
 @dataclass
 class Field:
     type: str
@@ -119,7 +117,7 @@ Please indicate, in the box below, that you are at least 18 years old, have read
 - *I agree to participate in this research and I want to continue with the survey.*
 '''
 
-fields: List[Field] = [
+fields0: List[Field] = [
     Field(name="patient", type="input_col", title=" "),
     Field(type="expander", title="**Session Transcription:** *(expand)*", children=[
         Field(name="dialogue_name", type="input_col", title=""),
@@ -140,7 +138,8 @@ fields: List[Field] = [
         Field(name="rupture_marker", type="rupture_markers",
           title="Select rupture markers noted in the session, include line numbers where rupture is found.", mandatory=False),
     ]),
-
+]
+fields1: List[Field] = [
     Field(type="container", title="#### True-To-Patient-Prompt Features", children=[
         Field(type="expander", title="**Patient Role Description:** *(expand)*", children=[
         Field(name="role_name", type="input_col", title=""),
@@ -198,6 +197,9 @@ fields: List[Field] = [
         Field(name="other_comments", type="text", title="Please provide any additional details or information:", mandatory=False),
     ]),
 ]
+
+STEPS: List[List[Field]] = [fields0, fields1]
+
 url_conditional_fields = [
     Field(name="skip", type="skip_checkbox",
           title="*I am uncomfortable annotating this text and  voluntarily skip this instance*", mandatory=False)
@@ -214,6 +216,8 @@ INPUT_FIELD_DEFAULT_VALUES = {'slider': 0,
 SHOW_HELP_ICON = False
 SHOW_VALIDATION_ERROR_MESSAGE = True 
 
+
+
 ########################################################################################
 if filesystem == 'hf':
     HF_TOKEN = os.environ.get("HF_TOKEN_WRITE")
@@ -535,6 +539,7 @@ def validate_current_page(fields: List[Field], index: int) -> bool:
 
 # Function to navigate rows
 def navigate(index_change):
+    st.session_state.step = 0
     st.session_state.current_index += index_change
     # only works consistently if done before rerun
     js = '''
@@ -640,8 +645,6 @@ def show_field(f: Field, index: int, data_collected):
             if data_collected else INPUT_FIELD_DEFAULT_VALUES[f.type]
         )
         value = _ensure_key(key, default_val)
-        #value = st.session_state[key] if key in st.session_state else \
-        #    (data_collected[f.name] if data_collected else INPUT_FIELD_DEFAULT_VALUES[f.type])
         if not SHOW_HELP_ICON:
             f.title = f'**{f.title}**\n\n{f.help}' if f.help else f.title
 
@@ -765,14 +768,33 @@ def show_fields(fields: List[Field]):
     # mark that the page has been rendered at least once (if you still use this elsewhere)
     st.session_state.form_displayed = st.session_state.current_index
 
-def prep_and_save_data(index, skip_sample):
+def iter_all_input_fields():
+    """Yield all Field objects which are real input widgets, across all steps."""
+    def walk(nodes):
+        for f in nodes:
+            if f.children:
+                walk(f.children)
+            elif f.type in INPUT_FIELD_DEFAULT_VALUES:
+                yield f
+    for step_fields in STEPS:
+        yield from walk(step_fields)
+
+    
+def prep_and_save_data(index, skip_sample, completed: bool):
     payload = {
         'user_id': st.session_state.user_id,
         'index': st.session_state.current_index,
         **(st.session_state.data.iloc[index][COLS_TO_SAVE].to_dict() if 0 <= index < len(st.session_state.data) else {}),
         **{k: st.session_state[k + str(index)] for k in st.session_state.data_inputs_keys},
-        'skip': skip_sample
+        'skip': skip_sample,
+        'completed': completed,
     }
+    
+    for f in iter_all_input_fields():
+        key = f.name + str(index)
+        val = st.session_state.get(key, INPUT_FIELD_DEFAULT_VALUES[f.type])
+        payload[f.name] = val
+        
     # normalize rupture markers -> always write 8 slots
     count_key = f"rupture_count_{index}"
     count = st.session_state.get(count_key, 1)
@@ -868,10 +890,14 @@ if 'current_index' not in st.session_state:
         st.session_state.current_index = start_index+1
     st.session_state.form_displayed = -2
 
+if 'step' not in st.session_state:
+    st.session_state.step = 0
+
 if get_param_from_url('show_extra_fields'):
-    fields += url_conditional_fields
+    fields1 += url_conditional_fields
 else:
-    fields += url_conditional_fields
+    fields1 += url_conditional_fields
+
 
 def add_validated_submit(fields, message):
     st.session_state.form_displayed = st.session_state.current_index
@@ -907,17 +933,29 @@ elif st.session_state.current_index == -1:
             add_validated_submit([st.session_state.user_id], "Please enter a valid user ID")
 
 elif st.session_state.current_index < len(st.session_state.data):
-    show_fields(fields)
+    step = st.session_state.step
+    total_steps = len(STEPS)
+
+    show_fields(STEPS[step])
 
     # Action buttons
     c1, c2, c3 = st.columns([1,1,6])
     with c1:
-        if st.button("Submit"):
-            if validate_current_page(fields, st.session_state.current_index):
+        label = "Next" if step < total_steps - 1 else "Submit & next session"
+        if st.button(label):
+            if validate_current_page(STEPS[step], st.session_state.current_index):
+                is_last_page = (step == total_steps - 1)
+
                 with st.spinner("saving"):
-                    prep_and_save_data(st.session_state.current_index, ('skip' in st.session_state and st.session_state['skip']))
-                st.success("Saved!")
-                navigate(1)
+                    prep_and_save_data(st.session_state.current_index, 
+                                       ('skip' in st.session_state and st.session_state['skip']),
+                                       completed=is_last_page)
+                if is_last_page:
+                    st.success("Saved!")    
+                    navigate(1)
+                else:
+                    st.session_state.step += 1
+                    st.rerun()
 
 elif st.session_state.current_index == len(st.session_state.data):
     st.write(f"**Thank you for taking part in this study!** \n ")
@@ -927,10 +965,16 @@ elif st.session_state.current_index == len(st.session_state.data):
 # Navigation buttons
 if 0 < st.session_state.current_index < len(st.session_state.data):
     if st.button("Previous"):
-        navigate(-1)
+        if step > 0:
+            st.session_state.step -= 1
+            st.rerun()
+        else:
+            st.session_state.current_index -= 1
+            st.session_state.step = total_steps - 1
+            st.rerun()
 
 if 0 <= st.session_state.current_index < len(st.session_state.data):
-    st.write(f"Page {st.session_state.current_index + 1} out of {len(st.session_state.data)}")
+    st.write(f"Session {st.session_state.current_index + 1} out of {len(st.session_state.data)}")
 
 st.markdown(
     """<style>