import streamlit as st
import openai

def main():
    st.title("GPT-3.5 챗봇")

    # User API key input
    api_key = st.text_input("발급받은 OpenAI API Key를 여기에 입력:")
    if api_key:
        openai.api_key = api_key

        init_msgs = [{"role": "system", "content": "You are a helpful assistant."}]
        if st.button("대화 기록 삭제"):
            st.session_state['messages'] = init_msgs
        # Initialize Session State for messages
        if 'messages' not in st.session_state:
            st.session_state['messages'] = init_msgs

        user_input = st.text_area("프롬프트:", height=200)
        submit_button = st.button("작성 완료")
        conversation = st.empty()
        
        if submit_button:
            if user_input:
                st.session_state.messages.append({"role": "user", "content": user_input})
                response = openai.ChatCompletion.create(
                    model="gpt-3.5-turbo",
                    messages=st.session_state.messages
                )
                new_msg = {"role": "assistant",
                           "content": response['choices'][0]['message']['content']}
                st.session_state.messages.append(new_msg)
                for message in st.session_state.messages:
                    if message["role"] == "user":
                        conversation.write(f"User: {message['content']}")
                    else:
                        conversation.write(f"Assistant: {message['content']}")
            else:
                st.write("유효한 질문을 입력하세요.")
    else:
        st.warning("유효한 OpenAI API Key를 입력하세요.")

if __name__ == "__main__":
    main()