import gradio as gr
import os

from llama_index import (
    VectorStoreIndex,
    SimpleDirectoryReader,
    StorageContext,
    ServiceContext,
    load_index_from_storage,
)
from llama_index.llms import OpenAI
from llama_index.memory import ChatMemoryBuffer
from llama_index.prompts import ChatPromptTemplate, ChatMessage, MessageRole
from llama_index.embeddings import LangchainEmbedding
import tiktoken
from llama_index.text_splitter import SentenceSplitter

from langchain.embeddings import HuggingFaceEmbeddings

# criação do embeding LangChain
lc_embed_model = HuggingFaceEmbeddings(
    model_name="sentence-transformers/all-mpnet-base-v2"
)
# mapeamento do embeding LangChain para o embeding LlamaIndex
embed_model = LangchainEmbedding(lc_embed_model)

# max_tokens: o tamanho máximo da resposta a ser dada
llm = OpenAI(temperature=0.3, model='gpt-3.5-turbo', max_tokens=1024)

# quebra inteligênte das sentenças, combinando separadores, tokenizadores e chunks
text_splitter = SentenceSplitter(
  separator=" ", chunk_size=1000, chunk_overlap=200,
  paragraph_separator="  \n \n", secondary_chunking_regex="[^,.;。]+[,.;。]?",
  tokenizer=tiktoken.encoding_for_model("gpt-3.5-turbo").encode
)

# cria um serviço de contexto para configurar a criação do indice
service_context = ServiceContext.from_defaults(llm=llm, embed_model=embed_model, text_splitter=text_splitter)

# verifica se a pasta storage existe localmente
PERSIST_DIR = "./storage"
if not os.path.exists(PERSIST_DIR):
    # caso não exista lê os documentos da pasta e cria um índice
    documents = SimpleDirectoryReader("./data").load_data()

    # cria um indice utilizando um contexto de serviços
    index = VectorStoreIndex.from_documents(documents, service_context=service_context)

    # depois, armazena o índice na pasta
    index.storage_context.persist(persist_dir=PERSIST_DIR)
else:
    # caso a pasta exista, lê o índice existente
    storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
    index = load_index_from_storage(storage_context, service_context=service_context)

# define um prompt para o chat
message_templates = [
    ChatMessage(role=MessageRole.SYSTEM,
        content='''Você é um sistema especialista que sabe responder perguntas relacionadas ao contexto abaixo.
        ------------------
        {context_str}
        ------------------'''
    ),
    ChatMessage(role=MessageRole.SYSTEM,
        context='''Para responder leve em consideração as perguntas e respostas dadas anteriormente.'''
    ),
    ChatMessage(role=MessageRole.USER,
        content='''Utilizando o histórico de conversação e o contexto armazenado iteraja e ajude o usuário.'''
    ),
]
chat_template = ChatPromptTemplate(message_templates=message_templates)

# Definir uma memória de mensagens anteriores
memory = ChatMemoryBuffer.from_defaults(token_limit=3900)

chat_engine = index.as_chat_engine(
    chat_mode="condense_plus_context",
    memory=memory,
    context_prompt=chat_template,
)

# consulta o índice local
def slow_echo(message, history):
    response_gen = chat_engine.stream_chat(message)
    response = ""
    for token in response_gen.response_gen:
        response = response + token
        yield response

# cria a interface com o gradio
demo = gr.ChatInterface(slow_echo).queue()

if __name__ == "__main__":
    demo.launch(share=True)