components/chat_box.py

import streamlit as st


def chat_box(session_state, config):
    for message in session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])