from llama_index import Prompt, SimpleDirectoryReader, LLMPredictor, PromptHelper, StorageContext, ServiceContext, GPTVectorStoreIndex, load_index_from_storage
from llama_index.prompts import PromptTemplate
from langchain.chat_models import ChatOpenAI
import gradio as gr
import sys
import os

os.environ["OPENAI_API_KEY"]

def construct_index(directory_path):
    # setup parameters
    max_input_size = 4096
    num_outputs = 512
    max_chunk_overlap = 0.2
    chunk_size_limit = 600
    # create a prompt helper
    prompt_helper = PromptHelper(max_input_size, num_outputs, max_chunk_overlap, chunk_size_limit=chunk_size_limit)
    # initialize the predictor with a fine-tuned model
    llm_predictor = LLMPredictor(llm=ChatOpenAI(temperature=0.9, model_name="gpt-3.5-turbo", max_tokens=num_outputs))
    # load documents from the specified directory
    documents = SimpleDirectoryReader(directory_path).load_data()
    # construct the index
    index = GPTVectorStoreIndex(documents, llm_predictor=llm_predictor, prompt_helper=prompt_helper)
    # save the index to disk
    index.storage_context.persist(persist_dir="index.json")
    return index

def chatbot(input_text):
    # load the index from disk
    storage_context = StorageContext.from_defaults(persist_dir="index.json")
    # load index
    index = load_index_from_storage(storage_context)
    # define custom Prompt
    TEMPLATE_STR = (
        "Quiero que actues como un asistente personal de un cliente del Banco Galicia. Tu nombre es Gala. Me brindas información sobre mi resument de tarjeta de credito VISA. Si la respuesta no esta en el documento, respondeme de forma creativa que no lo sabes, pero que podes ayudarme con otra pregunta. Nunca te enojes y no contestes preguntas politicas o religiosas. \n"
        "---------------------\n"
        "Dado esto, por favor responde a la pregunta: {input_text}\n"
    )
    QA_TEMPLATE = PromptTemplate(TEMPLATE_STR)
    # query the index and get the response
    response = index.query(QA_TEMPLATE)
    return response.response

iface = gr.Interface(fn=chatbot,
                     inputs=gr.components.Textbox(lines=7, label="Ingresa tu pregunta"),
                     outputs="text",
                     title="Demo Galicia")

# construct the index
index = construct_index("docs")
iface.launch(share=True, debug=True)