import streamlit as st
from transformers import pipeline
pipe = pipeline("text-generation", model="infiagent/daagent_7b")


st.title("InfiAgent Code Interpreter Demo 🚀")
input_text = st.text_area("Write your prompt")
uploaded_files = st.file_uploader("Upload your files", accept_multiple_files=True)
button_pressed = st.button("Run code interpreter", use_container_width=True)

if 'chat_history' not in st.session_state:
    st.session_state.chat_history = []

if button_pressed and input_text != "":
    # Add user message to chat history
    st.session_state.chat_history.append({"role": "user", "message": input_text})
    response = pipe(uploaded_files)
    st.session_state.chat_history.append({"role": "assistant", "message": response})

for chat in st.session_state.chat_history:
    with st.chat_message(chat["role"]):
        st.write(chat["message"])