import gradio as gr
from dotenv import load_dotenv
from llama_index import (
    ServiceContext,
    VectorStoreIndex,
    SimpleDirectoryReader,
    set_global_service_context,
)
from llama_index.llms import OpenAI
import openai
import constants

load_dotenv()
# openai.api_key = os.getenv("OPENAI_API_KEY")
openai.api_key = constants.APIKEY

service_context = ServiceContext.from_defaults(
    llm=OpenAI(model="gpt-3.5-turbo-16k", temperature=0)
)
set_global_service_context(service_context)

data = SimpleDirectoryReader(input_dir="./data/").load_data()

index = VectorStoreIndex.from_documents(data)
chat_engine = index.as_chat_engine(chat_mode="context")


def answer_question(question):
    response = chat_engine.stream_chat(question)
    response_text = ""
    for token in response.response_gen:
        print(token, end="")
        response_text += token
    return response_text


iface = gr.Interface(
    fn=answer_question,
    examples=[
        ["¿Quien eres?"],
        ["¿Cómo te llamas?"],
        ["¿Cómo puedes ayudarme?"],
        ["¿Qué puedo hacer si perdí mi contraseña?"],
        ["¿Qué puedo hacer si perdí mi tarjeta de cŕedito?"],
        ["Genera un reporte o lista de pasos de toda la info que necesito para abrir una cuenta"],
    ],
    inputs="text",
    outputs="text",
    title="Demo IA Conversacional",
    description="Pregúntame lo que quieras, te responderé con lo que sé."
)

iface.launch(share=True)