import openai 
import streamlit as st
import json

if 'exchanges' not in st.session_state:
    st.session_state.exchanges = []
    #st.session_state.exchanges = [{"role": "system", "content": "You are a chatbot."}]

def exchange_with_chatbot(prompt):
    st.session_state.exchanges.append({"role": "user", "content": prompt})
    # limit to 10 exchanges to save cost
    ex = st.session_state.exchanges if len(st.session_state.exchanges)<10 else st.session_state.exchanges[len(st.session_state.exchanges)-10:]    
    response = openai.ChatCompletion.create(
            model="gpt-3.5-turbo",
            messages= ex,
            max_tokens=1000,
            temperature=0.7
        )
    return response

def format_exchanges(exchanges):
    for i in range(len(exchanges)):
        if exchanges[i]["role"] == "user":
            icon, text, blank = st.columns([1,8,1])
        elif exchanges[i]["role"] == "assistant":
            blank, text, icon = st.columns([1,8,1])
        else:
            st.markdown("*" + exchanges[i]["role"] + ":* " + exchanges[i]["content"]) 
            continue
        
        with icon:
            st.image("icon_" + exchanges[i]["role"] + ".png", width=50)
        with text:
            st.markdown(exchanges[i]["content"])
        st.markdown("""---""")

openai.api_key = st.secrets["OPENAI_API_KEY"]

if openai.api_key:
    st.text_input("Prompt", placeholder="Ask me anything", key="prompt")
    
if st.session_state.prompt:
    try:
        response = exchange_with_chatbot(st.session_state.prompt)
    except Exception as e:
        st.error(e) 
        st.stop()
    json_object = json.loads(str(response))
    st.session_state.exchanges.append({"role": "assistant", "content": json_object['choices'][0]['message']['content']})
    format_exchanges(st.session_state.exchanges)