Order of chat elements is messed up when using columns

[TomZ]

Hi there,

I’m using the chat elements and everything works fine, until all the chat elements into a column layout, then the order of the answers and the input box is messed up:

if __name__ == "__main__":

    st.set_page_config(layout="wide")
    col1, col2  = st.columns([0.7, 0.3], vertical_alignment="top")

    with col2:
        st.header("Prompts")
        
    with col1:
        if "messages" not in st.session_state:
            st.session_state["messages"] = [{"role": "assistant", "content": "Hello!"}]

        for msg in st.session_state.messages:
            st.chat_message(msg["role"]).write(msg["content"])

        if prompt := st.chat_input(placeholder="Ask me something..."):
            st.session_state.messages.append({"role": "user", "content": prompt})
            st.chat_message("user").write(prompt)
            response = rag_pipeline(prompt)

            st.session_state.messages.append({"role": "assistant", "content": response})
            st.chat_message("assistant").write(response)
            col2.markdown(prompt)

[edsaac]

Try defining a st.container before creating the st.chat_input element. Then, add the st.chat_messages or whatever other stuff to that container. Tweaking example 2 in the docs:

with col1:
    messages = st.container()
    if prompt := st.chat_input("Ask me something..."):
        messages.chat_message("user").write(prompt)
        response = rag_pipeline(prompt)
        messages.chat_message("assistant").write(response)
        ...