Issues with Streamlit Select Box Not Updating

Hello all,

I am currently testing a chatbot locally and want it to include a selectbox for users to provide feedback for the responses.
Every time I update the select box, it will automatically reset back to the original value. I have tried multiple things to fix this and cannot figure out what is going wrong. Below is the code I am using to render the chat bot.
Also note - I cannot use the streamlit feedback as I have to use version 1.32.0 to avoid a separate bug. I am happy to provide any additional information as needed.

def display_output(message) -> None:
    unique_key = str(uuid.uuid4())  # Generate a unique key for each widget
    if message["type"] == "SQL":
        with st.expander("SQL Query", expanded=False):
            st.code(message["content"], language="sql")
    elif message["role"] == "assistant":
        # Determine the index for the selectbox based on the current rating
        rating_index = message["rating"] - 1 if message["rating"] is not None else 0

        # Display the selectbox for rating
        rating = st.selectbox("Rate the Results out of 5", (1, 2, 3, 4, 5), index=rating_index, key=unique_key, on_change=update_function, args=(message,))

        # Update the message's rating based on the selected value
        message["rating"] = rating

        if message["type"] == "table":
            st.dataframe(data=message["content"], use_container_width=True,

update_function as of now is just a print statement.