Cleanup session state handling

app.py

@@ -98,16 +98,38 @@ if "expander_toggle" not in st.session_state:
 if "evaluation_submitted" not in st.session_state:
     st.session_state.evaluation_submitted = False
 
+if "predict_button_clicked" not in st.session_state:
+    st.session_state.predict_button_clicked = False
+
+if "eval_selected" not in st.session_state:
+    st.session_state["eval_selected"] = False
+
+if "exp" not in st.session_state:
+    st.session_state.exp = None
+
+def set_prediction_view():
+    st.session_state.predict_button_clicked = True
+    st.session_state.evaluation_submitted = False    
+    hide_expander()
+
+def hide_expander():
+    st.session_state.expander_toggle = False
+
+def show_expander():
+    st.session_state.expander_toggle = True
+
 def submit_and_clear():
     try:
         # Call the explain endpoint as it also includes the prediction
         client.evaluate(
             deployment_id, request_log_id, prediction_log_id, st.session_state.evaluation_input
         )
-        # st.toast(":green-background[Feedback submitted successfully.]", icon="✅")
         st.session_state.eval_selected = False
         st.session_state.evaluation_submitted = True
         st.session_state.eval_selected = False
+        st.session_state.predict_button_clicked = False
+        st.session_state.exp = None
+        show_expander()
     except Exception as e:
         logging.error(e)
         st.error(
@@ -115,14 +137,6 @@ def submit_and_clear():
             + "Check whether you are using the right model URL and token for evaluations. "
             + "Contact Deeploy if the problem persists."
         )
-        # st.toast(f"Failed to submit feedback: {e}")
-
-
-def hide_expander():
-    st.session_state.expander_toggle = False
-
-def show_expander():
-    st.session_state.expander_toggle = True
 
 # Attributes
 st.subheader("Loan Application")
@@ -202,17 +216,13 @@ with st.expander("Application form", expanded=st.session_state.expander_toggle):
             ]
         ]
     }
-if "predict_button_clicked" not in st.session_state:
-    st.session_state.predict_button_clicked = False
 
-if deployment_token != "my-secret-token":
+# Show predict button if token is set
+if deployment_token != "my-secret-token" and st.session_state.exp is None:
     predict_button = st.button(
-        "Send loan application", key="predict_button", help="Click to get the AI prediction.", on_click=hide_expander
+        "Send loan application", key="predict_button", help="Click to get the AI prediction.", on_click=set_prediction_view
     )
     if predict_button:
-        st.session_state.predict_button_clicked = True
-        st.session_state.evaluation_submitted = False
-        # st.session_state.selected = "Loan Decision"
         with st.spinner("Loading prediction and explanation..."):
             # Call the explain endpoint as it also includes the prediction
             exp = client.explain(
@@ -220,7 +230,8 @@ if deployment_token != "my-secret-token":
             )
             st.session_state.exp = exp
 
-if st.session_state.predict_button_clicked:
+# Show prediction and explanation after predict button is clicked
+if st.session_state.predict_button_clicked and st.session_state.exp is not None:
     try:
         exp = st.session_state.exp
         # Read explanation to dataframe from json
@@ -272,85 +283,89 @@ if st.session_state.predict_button_clicked:
             f"{feat}: {neg_exp_df_t.loc[feat, 'Feature value']}"
             for feat in neg_feats
         ]
-
-        if predictions[0]:
-            # Show prediction
-            st.subheader("Loan Decision: :green[Approve]", divider="green")
-            # Format subheader to green
-            st.markdown(
-                "<style>.css-1v3fvcr{color: green;}</style>", unsafe_allow_html=True
-            )
-
-            # If prediction is positive, first show positive features, then negative features
-            st.success(
-                "**Positive creditworthiness**. The most important characteristics in favor of loan approval are:  \n - "
-                + "  \n- ".join(pos_feats)
-            )
-            st.warning(
-                "However, the following features weight against the loan applicant:  \n - "
-                + "  \n- ".join(neg_feats)
-                + "  \n For more details, see full explanation of the credit assessment below.",
-            )
-        else:
-            st.subheader("Loan Decision: :red[Reject]", divider="red")
-            # If prediction is negative, first show negative features, then positive features
-            st.error(
-                "**Negative creditworthiness**. The most important characteristics in favor of loan rejection are:  \n - "
-                + "  \n - ".join(neg_feats)
-            )
-            st.warning(
-                "However, the following factors weigh in favor of the loan applicant:  \n - "
-                + "  \n - ".join(pos_feats)
-                + "  \n For more details, see full explanation of the credit assessment below.",
-            )
-        explanation_expander = st.expander("Show explanation")
-        with explanation_expander:
-            # Show explanation
-            col_pos, col_neg = st.columns(2)
-
-            with col_pos:
-                st.subheader("Factors :green[in favor] of loan approval")
-                # st.success("**Factors in favor of loan approval**")
-                st.dataframe(
-                    pos_exp_df_t,
-                    hide_index=True,
-                    width=600,
-                    column_config={
-                        "Weight": st.column_config.ProgressColumn(
-                            "Weight",
-                            width="small",
-                            format=" ",
-                            min_value=0,
-                            max_value=1,
-                        )
-                    },
+        if not st.session_state.evaluation_submitted:  
+            if predictions[0]:
+                # Show prediction
+                st.subheader("Loan Decision: :green[Approve]", divider="green")
+                # Format subheader to green
+                st.markdown(
+                    "<style>.css-1v3fvcr{color: green;}</style>", unsafe_allow_html=True
                 )
-
-            with col_neg:
-                st.subheader("Factors :red[against] loan approval")
-                # st.error("**Factors against loan approval**")
-                st.dataframe(
-                    neg_exp_df_t,
-                    hide_index=True,
-                    width=600,
-                    column_config={
-                        "Weight": st.column_config.ProgressColumn(
-                            "Weight",
-                            width="small",
-                            format=" ",
-                            min_value=0,
-                            max_value=1,
-                        )
-                    },
-                )
-
-        st.divider()
-
-        if not st.session_state.evaluation_submitted:            
+                col1, col2 = st.columns(2)
+                with col1:
+                    # If prediction is positive, first show positive features, then negative features
+                    st.success(
+                        "The most important characteristics in favor of loan approval are:  \n - "
+                        + "  \n- ".join(pos_feats)
+                    )
+                with col2:
+                    st.warning(
+                        "However, the following features weight against the loan applicant:  \n - "
+                        + "  \n- ".join(neg_feats)
+                        # + "  \n For more details, see full explanation of the credit assessment below.",
+                    )
+            else:
+                st.subheader("Loan Decision: :red[Reject]", divider="red")
+                col1, col2 = st.columns(2)
+                with col1:
+                    # If prediction is negative, first show negative features, then positive features
+                    st.error(
+                        "The most important characteristics in favor of loan rejection are:  \n - "
+                        + "  \n - ".join(neg_feats)
+                    )
+                with col2:
+                    st.warning(
+                        "However, the following factors weigh in favor of the loan applicant:  \n - "
+                        + "  \n - ".join(pos_feats)
+                        # + "  \n For more details, see full explanation of the credit assessment below.",
+                    )
+            explanation_expander = st.expander("Show explanation")
+            with explanation_expander:
+                # Show explanation
+                col_pos, col_neg = st.columns(2)
+
+                with col_pos:
+                    st.subheader("Factors :green[in favor] of loan approval")
+                    # st.success("**Factors in favor of loan approval**")
+                    st.dataframe(
+                        pos_exp_df_t,
+                        hide_index=True,
+                        width=600,
+                        column_config={
+                            "Weight": st.column_config.ProgressColumn(
+                                "Weight",
+                                width="small",
+                                format=" ",
+                                min_value=0,
+                                max_value=1,
+                            )
+                        },
+                    )
+
+                with col_neg:
+                    st.subheader("Factors :red[against] loan approval")
+                    # st.error("**Factors against loan approval**")
+                    st.dataframe(
+                        neg_exp_df_t,
+                        hide_index=True,
+                        width=600,
+                        column_config={
+                            "Weight": st.column_config.ProgressColumn(
+                                "Weight",
+                                width="small",
+                                format=" ",
+                                min_value=0,
+                                max_value=1,
+                            )
+                        },
+                    )
+
+            st.divider()
+          
             # Add prediction evaluation
             st.subheader("Evaluation: Do you agree with the loan assessment?")
-            st.info(
-                "AI model predictions always come with a certain level of uncertainty.  \nEvaluate the correctness of the prediction based on your expertise and experience."
+            st.write(
+                "AI model predictions always come with a certain level of uncertainty. Evaluate the correctness of the assessment based on your expertise and experience."
             )
             cols = st.columns(4)
             col_yes, col_no = cols[:2]
@@ -372,8 +387,7 @@ if st.session_state.predict_button_clicked:
                 )
                 ChangeButtonColour("No, I disagree", "white", "#DD360C")
                 # ChangeButtonColour("No, I disagree", "#DD360C", "#F0F0F0")
-            if "eval_selected" not in st.session_state:
-                st.session_state["eval_selected"] = False
+            
             if yes_button:
                 st.session_state.eval_selected = True
                 st.session_state.evaluation_input = {