import streamlit as st
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# Chat input
if prompt := st.chat_input("Ask something..."):
    # Display user message
    st.chat_message("user").markdown(prompt)
    st.session_state.messages.append({"role": "user", "content": prompt})

    # Get bot response
    with st.chat_message("assistant"):
        with st.spinner("Thinking..."):
            response, st.session_state.conversation_history = generate_response(
                prompt, 
                st.session_state.conversation_history
            )
            st.markdown(response)
    
    # Store bot response
    st.session_state.messages.append({"role": "assistant", "content": response})