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

import os
os.environ['OPENAI_API_KEY'] = os.getenv("API_KEY")

llm = OpenAI(temperature=0.7)

def generate_restaurant_name_and_items(travel_place):
    # Chain 1: Restaurant Name
    prompt_template_name = PromptTemplate(
        input_variables=['travel_place'],
        template="I want to travel to {travel_place}.  Suggest some places to visi with relevant emoji's for each places and description"
    )

    name_chain = LLMChain(llm=llm, prompt=prompt_template_name, output_key="places")

    # Chain 2: Menu Items
    prompt_template_items = PromptTemplate(
        input_variables=['travel_place'],
        template="""Suggest some must try food items at {travel_place}. Return it as a comma separate items with relevant emoji's for food items"""
    )

    food_items_chain = LLMChain(llm=llm, prompt=prompt_template_items, output_key="food")

    # Chain 3: Things to do
    prompt_template_toDo = PromptTemplate(
        input_variables=['travel_place'],
        template="""Suggest Top 10 things to do in {travel_place}. Return with estimated cost and with relevant emoji's for each."""
    )

    todo_items_chain = LLMChain(llm=llm, prompt=prompt_template_toDo, output_key="toDo")

    chain = SequentialChain(
        chains=[name_chain, food_items_chain, todo_items_chain],
        input_variables=['travel_place'],
        output_variables=['places', 'food', 'toDo']
    )

    response = chain({'travel_place': travel_place})


    return response

if __name__ == "__main__":
    print(generate_restaurant_name_and_items("Italian"))