Customizing StreamlitSurvey Submit Button

Hello, I am currently using a package called StreamlitSurvey to create a fairly complex survey.

I am having issues with scope. Is there a way to create the submit button (and therefore submitted()) after the if-statements? Or is there a better way to access variables within the if-statement? I’m pretty new to this, so anything helps!

Here’s some example code:

def submitted():
    st.success("Your responses have been recorded. Thank you!")

survey = ss.StreamlitSurvey("survey")
pages = survey.pages(2, on_submit=lambda: submitted())

with pages:
    if pages.current == 0:
        test = survey.checkbox("N/A", key="vb_0")

I was able to solve this with st.session_state!