How to create multiple chat_input in same app

I figured out it was the st.rerun() which was preventing the display to output. However, I also realized, I should do on_click call back on submit form button. That can help in preventing earlier messages to again get refreshed on screen.

I have few questions around that I need help resolving. I will ask one by one so u can understand.

Here’s the small code I have for st.form submit button. I am doing a callback instead of doing if submit_button:

# Use columns to center the form
            col1, col2, col3 = st.columns([1,2,1])
            with col2:
                # Add details to st.form
                with st.form("correction_form"):
                    st.write("Please provide valid values for the following fields:")
                    st.session_state.corrected_fields = {}
                    # Iterating one by one for each field correction
                    for error in st.session_state.errors:
                        field = error['field']
                        formatted_field = field.replace('_', ' ').title()
                        placeholder = f"Enter a valid value for {formatted_field}"
                        if "Missing required field" not in error['message']:
                            placeholder += f" ({error['message']})"
                        corrected_value = st.text_input(
                            f"{formatted_field}:", 
                            key=f"correct_{field}",
                            placeholder=placeholder
                        )
                        if corrected_value:
                            st.session_state.corrected_fields[field] = corrected_value
                            print(f"within correct , corrected_values: {st.session_state.corrected_fields}")
                        
                    # User action to submit all corrections
                    st.form_submit_button("Submit Corrections",on_click = handle_form_submit_corrections)

The callback function is defined as below:


   # Callback function to handle submission of corrections
    def handle_form_submit_corrections():
        print(f"entered handle form")

         # Show thank you message before rerunning
        with st.chat_message("assistant", avatar="💁"):
            stream_message("Thank you for the corrections. I'm validating your input now.")
        
        print(f"corrected fields : {st.session_state.corrected_fields}")
        print(f"Session state corrected fields : {st.session_state.corrected_fields}")
        # Update merged_arguments with corrected values
        st.session_state.merged_arguments.update(st.session_state.corrected_fields)
        
        # Clear errors and change stage to 'validate'
        st.session_state.errors = []
        st.session_state.current_stage = "validate"
        
        
        # Rerun after the state changes
        st.rerun()

However, I am seeing it is not able to take the corrected_field dict as passed in callback to update my st.session_state.merged_arguments. Here is the output in terminal which shows that: Showing blank dict.

entered handle form
corrected fields : {}
Session state corrected fields : {}

How does the data get passed when we do callback on a button. I looked at docs and they said use st.session_state while sending variables data from form and I did, but it doesn’t seem to pass that.