import openai 
import streamlit as st
import json

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

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(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 (i%2) == 0:
            icon, text, blank = st.columns([1,8,2])
        else:
            blank, text, icon = st.columns([2,8,1])
        with icon:
            st.image("icon_" + exchanges[i]["role"] + ".png", width=50)
        with text:
            st.markdown(exchanges[i]["content"])
        st.divider()

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:
    response = exchange_with_chatbot(st.session_state.prompt)
    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)