import os
import gradio as gr

from langchain import PromptTemplate
from langchain.chains import LLMChain
from langchain.llms import OpenAI

openai_api_key = os.environ.get("OPENAI_API_KEY")

llm = OpenAI(temperature=0.9)

def generate_story(text):
    """Generate a story using the langchain library and OpenAI's GPT-3 model."""
    prompt = PromptTemplate(
        input_variables=["text"],
        template=""" 
         You are a fun and seasoned storyteller. Generate a short story for a 5 years old audience about {text}.
                 """
    )
    story = LLMChain(llm=llm, prompt=prompt)
    return story.run(text=text)

def app(text):
    story = generate_story(text)
    return story

with gr.Blocks() as demo:
    with gr.Column():
        text = gr.Textbox()
        submit_btn = gr.Button('Submit')
        story = gr.Textbox()

    submit_btn.click(fn=app, inputs=[text], outputs=[story])
        
demo.launch()