import streamlit as st
from langchain.prompts import PromptTemplate
import requests
import os

#from dotenv import load_dotenv
#load_dotenv()

#write a function to connect with LLAMA2
def GetinfoLLAMA(query):
    API_URL=os.getenv('HUGGING_API_KEY_LLAMA')
    API_KEY=os.getenv('HUGGING_API_KEY')
    headers = {"Authorization": API_KEY}
    template='''You are a commedian.Please tell joke on topic:{topic}'''

    prompt=PromptTemplate(input_variables=['topic'],
                        template=template)

    query=prompt.format(topic=query)

    payload={
	"inputs": query,
    }
    response=requests.post(API_URL,headers=headers,json=payload)
    response=response.json()
    #return response[0]["generated_text"]
    return str(response[0]["generated_text"]).replace(query,'')

st.set_page_config(page_title='Joke Generator',
                   page_icon=':circus_tent:',
                   layout='centered',
                   initial_sidebar_state='collapsed'
                   )

st.header('Joke Generator')

input_text=st.text_input("Enter the topic for joke: ")

submit=st.button('Generate')

if submit:
    st.write(GetinfoLLAMA(query=input_text))