Make one iteration of a loop each time a button is pressed

Hello,
I’m trying to make a simple app and I need to update and display a graph each time a button is pressed.
I’m trying to do this with a while statement but when the start button is pressed it goes out of the loop instead of reiterating.
Is this the most efficient way to make this? or is there a streamlit tool that may bbe more useful?
Thank you

The part of the code in question is the following

start_btn = st.button('Iniciar entrenamiento')
if start_btn:
    break_btn = st.button('Detener simulación')
    next_btn = st.button('Predecir siguiente mes')
    st.sidebar.subheader('Ajustar valores semanales de inversión deseados a lo largo del mes')
    inv_radio = st.sidebar.slider('Inversión en radio', 0,1000, 400, 50, key = 0)
    inv_diarios = st.sidebar.slider('Inversión en diarios', 0, 1000, 400, 50, key = 1)
    inv_tv = st.sidebar.slider('Inversión en Televisión', 0, 1000, 400, 50, key = 2)
        
    running = True
    display = False
    while running:
        if break_btn:
            running = False
        if next_btn:
            display = not display
            
        if display:
            part = df[0:4].copy()
            part['RD_GRP20'] = inv_radio
            part['DIARIOS_GRP'] = inv_diarios
            part['TV_GRPA20_40'] = inv_tv
            init_dt = build_model(init_dt, part)
            indexes = part.index
            df.drop(indexes[:-1], axis=0, inplace=True)
            display = not display
else:
    load_initial_model(init_dt)