import gradio as gr

from datetime import date
from langchain.agents import tool
from langchain.agents.agent_toolkits import create_python_agent
from langchain.agents import load_tools, initialize_agent
from langchain.agents import AgentType
from langchain.tools.python.tool import PythonREPLTool
from langchain.python import PythonREPL
from langchain.chat_models import ChatOpenAI

config = {
    "max_tokens": 1000,
    "model": "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.")

    result = ""
    
    try:
        llm = ChatOpenAI(temperature=0, model="gpt-4")

        tools = load_tools(["llm-math","wikipedia"], llm=llm)

        agent= initialize_agent(
            tools + [time], 
            llm, 
            agent=AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION,
            handle_parsing_errors=True,
            verbose = True)
        
        result = agent(prompt)
    
        #content = completion.choices[0].message.content
    except Exception as e:
        err_msg = e

        raise gr.Error(e)

    return result

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)],
                    outputs = [gr.Textbox(label = "Completion", lines = 1)],
                    title = "Generative AI - LLM & Agent",
                    description = description,)

demo.launch()