import streamlit as st
import os
from langchain_huggingface import HuggingFaceEndpoint
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from langchain.chains import SequentialChain
from langchain import HuggingFaceHub

hf_token = os.getenv("HF_TKN")

model = HuggingFaceHub(
    repo_id='google/flan-t5-large',
    model_kwargs={"temperature":0.4, "max_length":180, 'max_new_tokens' : 240},
    huggingfacehub_api_token=hf_token
)

st.title('Resturant Name Generator')
def generate_rest_name_and_items(cuisine):
    template_name = PromptTemplate(input_variables=['cuisine'],
                              template = "I want to open a resturant for {cuisine} food. Suggest a fancy name for this."
                              )
    chain_1 = LLMChain(llm = model,
                  prompt = template_name,
                  output_key = "rest_name")

    template_items = PromptTemplate(input_variables=['rest_name'],
                               template = 'Suggest at least 5 menu items for {rest_name}. Return it as a comma separated string'
                               )
    chain_2 = LLMChain(llm = model,
                  prompt = template_items,
                  output_key = "menu_items")

    s_chain = SequentialChain(chains = [chain_1,chain_2],
                         input_variables = ['cuisine'],
                         output_variables = ['rest_name','menu_items'])

    response = s_chain({'cuisine':cuisine})
    return response
    
cuisine = st.sidebar.selectbox('Pick a cuisine',('Indian','Italian','Mexican','Arabic','American'))
if cuisine:
    response = generate_rest_name_and_items(cuisine)
    st.header(response['rest_name'])
    items = response['menu_items'].split(",")

    st.write('**Menu Items**')
    for item in items:
        st.write(">" ,item)