I am trying to implement a chat with Streamlit, where I need to keep a counter of words being exchanged, essentially starting from this.

So I initialize a session variable at the beginning of the main() function as follows:

st.session_state['words'] = 0

After the question has been asked and response has been provided, i.e. at the end of the st.chat_input block, after the spinner block, I calculate the count of words in the question and response (wd_cnt). I increment the count and write it:

st.session_state.words += wd_cnt
st.write(st.session_state.words)

The problem is: the counter is not incremented. The write statement above always outputs the value of wd_cnt, as if st.session_state.words was reset every time.

What am I missing?