fixing displaying wrong number of total messages

pages/convosim.py

@@ -94,7 +94,6 @@ memoryA = st.session_state[list(memories.keys())[0]]
 llm_chain, stopper = get_chain(issue, language, source, memoryA, temperature, texter_name=st.session_state["texter_name"])
 
 st.title("💬 Simulator") 
-st.session_state['total_messages'] = len(memoryA.chat_memory.messages)
 for msg in memoryA.buffer_as_messages:
     role = "user" if type(msg) == HumanMessage else "assistant"
     st.chat_message(role).write(msg.content)
@@ -105,6 +104,7 @@ def sent_request_llm(llm_chain, prompt):
     responses = custom_chain_predict(llm_chain, prompt, stopper)
     for response in responses:
         st.chat_message("assistant").write(response)
+    st.session_state['total_messages'] = len(memoryA.chat_memory.messages)
     transcript = memoryA.load_memory_variables({})[memoryA.memory_key]
     update_convo(st.session_state["db_client"], st.session_state["convo_id"], transcript)
 
@@ -143,7 +143,7 @@ def push_latest_confirmation_db():
     else:
         bp_push2db({x['label']:x['score'] for x in st.session_state['bp_prediction']})
 
-disable_chat = (st.session_state['total_messages'] > MAX_MSG_COUNT - 4) and not st.session_state['scored']
+disable_chat = (st.session_state['total_messages'] > MAX_MSG_COUNT - 2) and not st.session_state['scored']
 if prompt := st.chat_input(disabled=disable_chat): #account for next interaction
     st.session_state['last_message_ts'] = time.time()
     if 'convo_id' not in st.session_state:
@@ -206,7 +206,6 @@ with st.sidebar:
             update_convo(st.session_state["db_client"], st.session_state["convo_id"], transcript)
             st.switch_page("pages/training_adherence.py")
 
-st.session_state['total_messages'] = len(memoryA.chat_memory.messages)
 if st.session_state['total_messages'] >= MAX_MSG_COUNT:
     st.toast(f"Total of {MAX_MSG_COUNT} Messages reached. Conversation Ended", icon=":material/verified:")
 elif st.session_state['total_messages'] >= WARN_MSG_COUT: