# Import packages
import openai
from llama_index import SimpleDirectoryReader, GPTListIndex, GPTSimpleVectorIndex, LLMPredictor, PromptHelper, ServiceContext
from langchain.chat_models import ChatOpenAI
import gradio as gr
import sys
import os
import PyPDF2

#os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY
# OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

'''
def construct_index(directory_path):
    max_input_size = 4096
    num_outputs = 512
    max_chunk_overlap = 20
    chunk_size_limit = 600
    prompt_helper = PromptHelper(max_input_size, num_outputs, max_chunk_overlap, chunk_size_limit=chunk_size_limit)
    llm_predictor = LLMPredictor(llm=ChatOpenAI(temperature=0.7, model_name="gpt-3.5-turbo", max_tokens=num_outputs))
    documents = SimpleDirectoryReader(directory_path).load_data()
    service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor, prompt_helper=prompt_helper)
    index = GPTSimpleVectorIndex.from_documents(documents, service_context=service_context)
    index.save_to_disk('index.json')
    return index
'''

def chatbot(input_text, openai_api_key):
    os.environ["OPENAI_API_KEY"] = openai_api_key
    index = GPTSimpleVectorIndex.load_from_disk('index.json')
    response = index.query(input_text, response_mode="compact")
    return response.response

# chat = gr.Interface(fn=chatbot,
#                     inputs=gr.components.Textbox(lines=7, label="Ask your question to ChatGPT"),
#                     outputs="text",
#                     title="Custom-trained AI Chatbot for employee tax assessment 2022")

# Documentation how to make gradio interfaces: https://gradio.app/quickstart/

with gr.Blocks() as chat:

    with gr.Column(elem_id="col-container"):
        gr.Markdown("""## Trained with custom data""",
                    elem_id="header")

    with gr.Column():
        gr.Markdown("Enter your OpenAI API Key.")
        openai_api_key = gr.Textbox(value='', placeholder="OpenAI API Key", type="password", label="Enter OpenAI API Key")
    
    text_input = gr.Textbox(lines=7, label="Enter your question")
    output = gr.Textbox(label="Response")
    greet_btn = gr.Button("Generate Response")
    greet_btn.click(fn=chatbot, inputs=[text_input, openai_api_key], outputs=output)

chat.launch()