import streamlit as st
from langchain.llms import OpenAI
from langchain.chains import PALChain
from langchain.chains.llm import LLMChain

# set page config
st.set_page_config(page_title="🦜🔗 This is a Program Aided Language model")
st.title("🦜🔗🧮 Program Aided Language Model: Helps with math problems")
st.markdown(
    """This example was adapt from Data Professor Github [repo](https://github.com/dataprofessor/langchain-quickstart/blob/master/streamlit_app.py)"""
)
st.markdown(
    """Paper: [Program-Aided Language Models for Program Synthesis](https://arxiv.org/pdf/2211.10435.pdf)"""
)
st.markdown(
    """Credit: [Sam Witteven](https://www.youtube.com/playlist?list=PL8motc6AQftk1Bs42EW45kwYbyJ4jOdiZ)"""
)

# sidebar for OpenAI API key
openai_api_key = st.sidebar.text_input("Enter your OpenAI API key", type="password")


def generate_response(input_text):
    """
    Generates response to input text using PALChain

    Parameters
    ----------
    input_text : str
      Input text to generate response for using PALChain

    Returns
    -------
    None

    Example
    -------
    """
    llm = OpenAI(temperature=0, openai_api_key=openai_api_key, max_tokens=512, skip_on_failure=True)
    pal_chain = PALChain.from_math_prompt(llm, verbose=True)
    st.write(pal_chain.run(input_text))


with st.form("my_form"):
    text = st.text_area(
        "Enter text:",
        "Leah had 32 chocolates and her sister had 42. If they ate 35, how many pieces do they have left in total?",
    )
    submitted = st.form_submit_button("Submit")
    if not openai_api_key.startswith("sk-"):
        st.warning("Please enter your OpenAI API key!", icon="⚠")
    if submitted and openai_api_key.startswith("sk-"):
        generate_response(text)