import gradio as gr
import os

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

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

AGENT_OFF = False
AGENT_ON  = True

@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, agent_option):
    if (openai_api_key == ""):
        raise gr.Error("OpenAI API Key is required.")
    if (prompt == ""):
        raise gr.Error("Prompt is required.")
    if (agent_option is None):
        raise gr.Error("Use Agent is required.")

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

        if (agent_option == AGENT_OFF):
            agent = initialize_agent([],
                                     llm,
                                     agent = AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION,
                                     handle_parsing_errors = True,
                                     verbose = True)
        else:
            tools = load_tools(["openweathermap-api"], llm)
            
            agent = initialize_agent(tools + [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, value = "sk-"),
                              gr.Textbox(label = "Prompt", lines = 1, value = "What is today's date?"),
                              gr.Radio([AGENT_OFF, AGENT_ON], label = "Use Agent", value = AGENT_OFF)],
                    outputs = [gr.Textbox(label = "Completion", lines = 1)],
                    title = "Generative AI - LLM & Agent",
                    description = description,
                    examples = [["sk-", "What is today's date?", AGENT_ON],
                                ["sk-", "What is the weather like in Irvine, CA?", AGENT_ON]],
                                cache_examples = False)

demo.launch()