import os
import gradio as gr
from typing import List
#from pydantic import BaseModel,Field
from langchain_core.pydantic_v1 import BaseModel, Field
#from langchain.chat_models import ChatOpenAI
from langchain_openai import ChatOpenAI
from langchain.prompts import HumanMessagePromptTemplate, ChatPromptTemplate
from langchain.output_parsers import PydanticOutputParser

# with open('../openai_api_key.txt') as f:
#     api_key = f.read()
    
# os.environ['OPENAI_API_KEY'] = api_key
    
# creating model instance
chat = ChatOpenAI()

# define pydantic model
class TravelPlanner(BaseModel):
    places: List[str] = Field(description='Python list containing list of places that can be visited. ')
    itenary: dict = Field(description="Python dictionary of itenary details")
    
# get format instructions
output_parser = PydanticOutputParser(pydantic_object=TravelPlanner)
format_instructions = output_parser.get_format_instructions()

def travel_planner(place : str) -> List:
    human_template = """
    I want to visit {place}. Suggest me a travel itenary. {format_instructions}"""
    human_message_prompt = HumanMessagePromptTemplate.from_template(human_template)


    chat_prompt  = ChatPromptTemplate.from_messages([human_message_prompt])
    prompt = chat_prompt.format_prompt(place = place, format_instructions = format_instructions)
    
    messages = prompt.messages
    response = chat(messages = messages)
    output = output_parser.parse(response.content)
    
    places, itenary = output.places, output.itenary
    return places, itenary

# building interface
with gr.Blocks() as demo:
    gr.HTML("<h1 align = 'center'> Travel Planner </h1>")
    
    inputs = [gr.Textbox(label = 'Enter the place that you want to go for vacation')]
    generate_btn = gr.Button(value = 'Generate')
    outputs = [gr.Text(label = 'List of places that you can visit'), gr.JSON(label = 'Itenary details')]
    generate_btn.click(fn = travel_planner,inputs=inputs, outputs=outputs) 
    
if __name__ == '__main__':
    demo.launch()