import gradio as gr

from datetime import date
from langchain.agents import AgentType, initialize_agent, tool
from langchain.chat_models import ChatOpenAI

config = {
    "max_tokens": 1000,
    "model_name": "gpt-4",
    "temperature": 0,
}

@tool
def time(text: str) -> str:
    """Returns todays date, use this for any \
    questions related to knowing todays date. \
    The input should always be an empty string, \
    and this function will always return todays \
    date - any date mathmatics should occur \
    outside this function."""
    return str(date.today())

def invoke(openai_api_key, prompt):
    if (openai_api_key == ""):
        raise gr.Error("OpenAI API Key is required.")
    if (prompt == ""):
        raise gr.Error("Prompt is required.")

    output = ""
    
    try:
        llm = ChatOpenAI(model_name = config["model_name"],
                         openai_api_key = openai_api_key, 
                         temperature = config["temperature"])

        agent = initialize_agent([time], 
                                 llm, 
                                 agent = AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION,
                                 handle_parsing_errors = True,
                                 verbose = True)
        
        completion = agent(prompt)
    
        output = completion.output
    except Exception as e:
        err_msg = e

        raise gr.Error(e)

    return output

description = """<a href='https://www.gradio.app/'>Gradio</a> UI using the <a href='https://openai.com/'>OpenAI</a> API 
                 with <a href='https://openai.com/research/gpt-4'>gpt-4</a> model."""

gr.close_all()

demo = gr.Interface(fn = invoke, 
                    inputs = [gr.Textbox(label = "OpenAI API Key", type = "password", lines = 1),
                              gr.Textbox(label = "Prompt", lines = 1, value = "What is today's date?")],
                    outputs = [gr.Textbox(label = "Completion", lines = 1)],
                    title = "Generative AI - LLM & Agent",
                    description = description)

demo.launch()