import streamlit as st
from langchain.llms import HuggingFaceHub
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate

#Function to return the response 
def generate_answer(query):
    llm = HuggingFaceHub(
        repo_id = "google/flan-t5-xxl", 
        model_kwargs={"temperature": 0.7, "max_length": 64,"max_new_tokens":512}
    )
    
    template = """Question: {query}

                Answer: Let's think step by step.
                """
    
    PromptTemplate(template=template, input_variables=["query"])
    llm_chain = LLMChain(prompt=prompt, llm=llm)
    result = llm_chain.run(query)
    return result
    

#App UI starts here 
st.set_page_config(page_title = "LangChain Demo", page_icon = ":robot:")
st.header("LangChain Demo")


#Gets User Input 
def get_text():
    input_text = st.text_input("You: ", key="input")
    return input_text


user_input = get_text()
response = generate_answer(user_input)

submit = st.button("Generate")

#If the button clicked
if submit:
    st.subheader("Answer: ")
    st.write(response)