# import environment variables
from data.env_variables import AZURE_OPENAI_DEPLOYMENT_NAME, AZURE_OPENAI_MODEL_NAME, \
    AZURE_OPENAI_API_ENDPOINT, OPENAI_API_VERSION, AZURE_OPENAI_API_KEY, \
    HUGGINGFACE_API_TOKEN, LLAMA2_API_TOKEN, OPENAI_API_KEY, NVIDIANGC_API_KEY
from dotenv import load_dotenv

# import software general purpose libs
import os
import psutil
import logging as log

# import langchain debug mode
from langchain.globals import set_debug

# import langchain document loader
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.document_loaders import Docx2txtLoader

# import message handlers
from streamlit_chat import message
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler

# import embedding processing objects
from langchain.text_splitter import RecursiveCharacterTextSplitter

# import vector database
from langchain.vectorstores.chroma import Chroma

# import data retrieval chain
from langchain.chains import RetrievalQAWithSourcesChain

# import langchain models from huggingface
from langchain.llms.huggingface_hub import HuggingFaceHub
from langchain.llms.huggingface_pipeline import HuggingFacePipeline
from langchain_community.llms.huggingface_endpoint import HuggingFaceEndpoint

# import langchain models
from langchain.llms.gpt4all import GPT4All
from langchain.chat_models import AzureChatOpenAI
from langchain.chat_models import ChatOpenAI
from langchain_nvidia_ai_endpoints import ChatNVIDIA


from data.vectorize import get_embeddings_model
# import hugging face transformers lib - only for quantized models
# import transformers
# from transformers import BitsAndBytesConfig, AutoTokenizer, AutoModelForCausalLM, AutoConfig, pipeline

# import streamlit web framework
import streamlit as st


# start debugging
set_debug(True)

# start logging
log.basicConfig(filename="logs/app.log", level=log.DEBUG)


N_THREADS = psutil.cpu_count()


def load_vector_database():
    log.info("Initializing Vector DB")
    sentence_transformer_ef = get_embeddings_model()
    
    st.session_state.vectordb = Chroma(persist_directory="./documents_cache/qa_retrieval", embedding_function=sentence_transformer_ef)

def get_local_gpt4all_models():
    local_models = {}
    local_models["ggml-gpt4all-j-v1.3-groovy"] = "./model_cache/ggml-gpt4all-j-v1.3-groovy.bin"
    local_models["mistral-7b-openorca.Q4_0"] = "./model_cache/mistral-7b-openorca.Q4_0.gguf"
    # local_models["ggml-mpt-7b-instruct"] = "./model_cache/ggml-mpt-7b-instruct.bin"
    # local_models["ggml-gpt4all-l13b-snoozy"] = "./model_cache/ggml-gpt4all-l13b-snoozy.bin"
    # local_models["ggml-v3-13b-hermes-q5_1"] = "./model_cache/ggml-v3-13b-hermes-q5_1.bin"
    # local_models["ggml-vicuna-13b-1.1-q4_2"] = "./model_cache/ggml-vicuna-13b-1.1-q4_2.bin"
    
    return local_models


def get_llm_instance(model_interface: str):
    if model_interface == "azure":
        llm_instance = AzureChatOpenAI(
            deployment_name=AZURE_OPENAI_DEPLOYMENT_NAME,
            model_name=AZURE_OPENAI_MODEL_NAME,
            azure_endpoint=AZURE_OPENAI_API_ENDPOINT,
            openai_api_version=OPENAI_API_VERSION,
            openai_api_key=AZURE_OPENAI_API_KEY,
            openai_api_type="azure"
        )
    elif model_interface == "openai":
        llm_instance = ChatOpenAI(
            temperature=0.1,
            openai_api_key=""
        )
    elif model_interface == "gpt4all":
        local_models = get_local_gpt4all_models()
        callbacks = [StreamingStdOutCallbackHandler()]

        llm_instance = GPT4All(
            # model=local_models["mistral-7b-openorca.Q4_0"],
            model="model_cache/zephyr-7b-beta.Q3_K_S.gguf",
            # allow_download=True,
            callbacks=callbacks,
            verbose=True,
            # device="gpu",
            device="nvidia",
            # n_threads=16,
            # n_threads=N_THREADS,
        )
    elif model_interface == "huggingface-falcon":
        llm_instance = HuggingFaceHub(
            verbose=True,
            task="text-generation",
            repo_id="tiiuae/falcon-40b-instruct"
        )
    elif model_interface == "huggingface-mistral-7b":
        llm_instance = HuggingFacePipeline.from_model_id(
            # model_id="mistralai/Mistral-7B-Instruct-v0.1",
            model_id="Open-Orca/Mistral-7B-OpenOrca",
            task="text-generation",
            pipeline_kwargs={"max_new_tokens": 10},
            device=0
        )
    elif model_interface == "huggingface-endpoint-zephyr-7b":
        endpoint_url = "https://api-inference.huggingface.co/models/HuggingFaceH4/zephyr-7b-beta"
        headers = {"Authorization": "Bearer "}
        llm_instance = HuggingFaceEndpoint(
            endpoint_url=endpoint_url,
            task="text-generation",
            huggingfacehub_api_token=HUGGINGFACE_API_TOKEN
        )
    elif model_interface == "zephyr-7b-beta":
        llm_instance = HuggingFacePipeline.from_model_id(
            model_id="HuggingFaceH4/zephyr-7b-beta",
            task="text-generation",
            # pipeline_kwargs={"max_new_tokens": 10},
            device=0
        )
    elif model_interface == "huggingface-api-llama2":
        llm_instance = HuggingFacePipeline.from_model_id(
            model_id="meta-llama/Llama-2-7b-chat-hf",
            task="text-generation",
            device="cuda",
            pipeline_kwargs={
                "token": LLAMA2_API_TOKEN
            }
        )
    elif model_interface == "nvidia-mixtral":
        callbacks = [StreamingStdOutCallbackHandler()]
        llm_instance = ChatNVIDIA(
            model="mixtral_8x7b",
            nvidia_api_key=NVIDIANGC_API_KEY,
            callbacks=callbacks,
            temperature=0.2,
            top_p=0.7,
            max_tokens=1024,
            seed=42
        )
    
    return llm_instance


def initialize_conversation_chain():
    vectordb = st.session_state.vectordb
    callbacks = [StreamingStdOutCallbackHandler()]
    local_models = get_local_gpt4all_models()
    
    retriever_instance = vectordb.as_retriever(search_kwargs={'k':4})

    # llm_instance = get_llm_instance("huggingface-endpoint-zephyr-7b")
    llm_instance = get_llm_instance("nvidia-mixtral")
    # llm_instance = get_llm_instance("gpt4all")

    log.info("Inicializando")
    st.session_state.qa_chain = RetrievalQAWithSourcesChain.from_chain_type(
        llm=llm_instance,
        chain_type="stuff",
        retriever=retriever_instance
    )


def handle_user_input(user_question, response_container):
    if user_question is None:
        return
    
    qa_chain:RetrievalQAWithSourcesChain = st.session_state.qa_chain
    
    response_container.empty()
    # Handle user Queries
    with response_container.container():
        with st.spinner("Gerando resposta..."):
            log.info(f"Gerando resposta para consulta do cliente: {user_question}")
            user_question += " (responda resumidamente em pt-br)"
            response = qa_chain({"question":user_question}, return_only_outputs=True)
            
            # st.write(response)
            st.write(response["answer"])
            
            with st.expander(label="Sources", expanded=False):
                for source in response["sources"]:
                    st.write(source)


def process_new_uploads(uploaded_files):
    vectordb:Chroma = st.session_state.vectordb
    for doc in uploaded_files:
        log.info(f"Processa arquivo: {doc.name}")
        
        with open(os.path.join("tmp_documents", doc.name), "wb") as f:
            f.write(doc.getbuffer())
        
        extension = doc.name.split(".")[-1]
        filepath = f"./tmp_documents/{doc.name}"

        if extension == "pdf":
            loader = PyPDFLoader(file_path=filepath)
        elif extension == "docx" or extension == "doc":
            loader = Docx2txtLoader(file_path=filepath)
        
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=640, chunk_overlap=128)
        
        log.info("Particiona texto")
        text_chunks = text_splitter.split_documents(loader.load())
        # log.info("Chunks: %s", text_chunks)
        
        log.info("Processa embeddings e adiciona documento ao Vector DB")
        vectordb.add_documents(documents=text_chunks)
        vectordb.persist()

        os.remove(f"./tmp_documents/{doc.name}")
        log.info(f"Arquivo processado com sucesso: {doc.name}")
    

def main():
    load_dotenv()
    st.set_page_config(page_title="Converse com seus documentos", page_icon=":books:")
    
    st.header("Converse com seus documentos :books:")
    
    if "vectordb" not in st.session_state:
        with st.spinner("Inicializando Vector DB..."):
            load_vector_database()
    
    if "qa_chain" not in st.session_state:
        with st.spinner("Inicializando AI Model..."):
            initialize_conversation_chain()
    
    user_question = st.text_input("Faça sua pergunta aqui")
    
    response_container = st.empty()
    
    if user_question:
        handle_user_input(user_question, response_container)
        user_question = None

    with st.sidebar:
        st.subheader("Seus documentos")
        uploaded_files = st.file_uploader(
            "Insira seu arquivo aqui (.pdf, .docx) e clique em 'Processar'",
            accept_multiple_files=True
        )
        if st.button("Processar"):
            with st.spinner("Processando..."):
                process_new_uploads(uploaded_files)


if __name__ == "__main__":
    main()

langchain_community.document_loaders.PyPDFLoader langchain_community.document_loaders.Docx2txtLoader langchain.llms.huggingface_pipeline.HuggingFacePipeline.from_model_id langchain.vectorstores.chroma.Chroma langchain.text_splitter.RecursiveCharacterTextSplitter langchain.callbacks.streaming_stdout.StreamingStdOutCallbackHandler langchain.llms.gpt4all.GPT4All langchain.chat_models.ChatOpenAI langchain.chains.RetrievalQAWithSourcesChain.from_chain_type langchain.chat_models.AzureChatOpenAI langchain.globals.set_debug langchain_community.llms.huggingface_endpoint.HuggingFaceEndpoint langchain.llms.huggingface_hub.HuggingFaceHub langchain_nvidia_ai_endpoints.ChatNVIDIA
[((1820, 1835), 'langchain.globals.set_debug', 'set_debug', (['(True)'], {}), '(True)\n', (1829, 1835), False, 'from langchain.globals import set_debug\n'), ((1853, 1910), 'logging.basicConfig', 'log.basicConfig', ([], {'filename': '"""logs/app.log"""', 'level': 'log.DEBUG'}), "(filename='logs/app.log', level=log.DEBUG)\n", (1868, 1910), True, 'import logging as log\n'), ((1925, 1943), 'psutil.cpu_count', 'psutil.cpu_count', ([], {}), '()\n', (1941, 1943), False, 'import psutil\n'), ((1978, 2012), 'logging.info', 'log.info', (['"""Initializing Vector DB"""'], {}), "('Initializing Vector DB')\n", (1986, 2012), True, 'import logging as log\n'), ((2043, 2065), 'data.vectorize.get_embeddings_model', 'get_embeddings_model', ([], {}), '()\n', (2063, 2065), False, 'from data.vectorize import get_embeddings_model\n'), ((2103, 2209), 'langchain.vectorstores.chroma.Chroma', 'Chroma', ([], {'persist_directory': '"""./documents_cache/qa_retrieval"""', 'embedding_function': 'sentence_transformer_ef'}), "(persist_directory='./documents_cache/qa_retrieval',\n    embedding_function=sentence_transformer_ef)\n", (2109, 2209), False, 'from langchain.vectorstores.chroma import Chroma\n'), ((6313, 6338), 'logging.info', 'log.info', (['"""Inicializando"""'], {}), "('Inicializando')\n", (6321, 6338), True, 'import logging as log\n'), ((6371, 6487), 'langchain.chains.RetrievalQAWithSourcesChain.from_chain_type', 'RetrievalQAWithSourcesChain.from_chain_type', ([], {'llm': 'llm_instance', 'chain_type': '"""stuff"""', 'retriever': 'retriever_instance'}), "(llm=llm_instance, chain_type=\n    'stuff', retriever=retriever_instance)\n", (6414, 6487), False, 'from langchain.chains import RetrievalQAWithSourcesChain\n'), ((8476, 8489), 'dotenv.load_dotenv', 'load_dotenv', ([], {}), '()\n', (8487, 8489), False, 'from dotenv import load_dotenv\n'), ((8494, 8581), 'streamlit.set_page_config', 'st.set_page_config', ([], {'page_title': '"""Converse com seus documentos"""', 'page_icon': '""":books:"""'}), "(page_title='Converse com seus documentos', page_icon=\n    ':books:')\n", (8512, 8581), True, 'import streamlit as st\n'), ((8586, 8635), 'streamlit.header', 'st.header', (['"""Converse com seus documentos :books:"""'], {}), "('Converse com seus documentos :books:')\n", (8595, 8635), True, 'import streamlit as st\n'), ((8945, 8984), 'streamlit.text_input', 'st.text_input', (['"""Faça sua pergunta aqui"""'], {}), "('Faça sua pergunta aqui')\n", (8958, 8984), True, 'import streamlit as st\n'), ((9015, 9025), 'streamlit.empty', 'st.empty', ([], {}), '()\n', (9023, 9025), True, 'import streamlit as st\n'), ((2949, 3203), 'langchain.chat_models.AzureChatOpenAI', 'AzureChatOpenAI', ([], {'deployment_name': 'AZURE_OPENAI_DEPLOYMENT_NAME', 'model_name': 'AZURE_OPENAI_MODEL_NAME', 'azure_endpoint': 'AZURE_OPENAI_API_ENDPOINT', 'openai_api_version': 'OPENAI_API_VERSION', 'openai_api_key': 'AZURE_OPENAI_API_KEY', 'openai_api_type': '"""azure"""'}), "(deployment_name=AZURE_OPENAI_DEPLOYMENT_NAME, model_name=\n    AZURE_OPENAI_MODEL_NAME, azure_endpoint=AZURE_OPENAI_API_ENDPOINT,\n    openai_api_version=OPENAI_API_VERSION, openai_api_key=\n    AZURE_OPENAI_API_KEY, openai_api_type='azure')\n", (2964, 3203), False, 'from langchain.chat_models import AzureChatOpenAI\n'), ((5977, 6009), 'langchain.callbacks.streaming_stdout.StreamingStdOutCallbackHandler', 'StreamingStdOutCallbackHandler', ([], {}), '()\n', (6007, 6009), False, 'from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler\n'), ((7460, 7501), 'logging.info', 'log.info', (['f"""Processa arquivo: {doc.name}"""'], {}), "(f'Processa arquivo: {doc.name}')\n", (7468, 7501), True, 'import logging as log\n'), ((7951, 8016), 'langchain.text_splitter.RecursiveCharacterTextSplitter', 'RecursiveCharacterTextSplitter', ([], {'chunk_size': '(640)', 'chunk_overlap': '(128)'}), '(chunk_size=640, chunk_overlap=128)\n', (7981, 8016), False, 'from langchain.text_splitter import RecursiveCharacterTextSplitter\n'), ((8034, 8062), 'logging.info', 'log.info', (['"""Particiona texto"""'], {}), "('Particiona texto')\n", (8042, 8062), True, 'import logging as log\n'), ((8193, 8258), 'logging.info', 'log.info', (['"""Processa embeddings e adiciona documento ao Vector DB"""'], {}), "('Processa embeddings e adiciona documento ao Vector DB')\n", (8201, 8258), True, 'import logging as log\n'), ((8349, 8389), 'os.remove', 'os.remove', (['f"""./tmp_documents/{doc.name}"""'], {}), "(f'./tmp_documents/{doc.name}')\n", (8358, 8389), False, 'import os\n'), ((8398, 8453), 'logging.info', 'log.info', (['f"""Arquivo processado com sucesso: {doc.name}"""'], {}), "(f'Arquivo processado com sucesso: {doc.name}')\n", (8406, 8453), True, 'import logging as log\n'), ((9173, 9204), 'streamlit.subheader', 'st.subheader', (['"""Seus documentos"""'], {}), "('Seus documentos')\n", (9185, 9204), True, 'import streamlit as st\n'), ((9230, 9348), 'streamlit.file_uploader', 'st.file_uploader', (['"""Insira seu arquivo aqui (.pdf, .docx) e clique em \'Processar\'"""'], {'accept_multiple_files': '(True)'}), '(\n    "Insira seu arquivo aqui (.pdf, .docx) e clique em \'Processar\'",\n    accept_multiple_files=True)\n', (9246, 9348), True, 'import streamlit as st\n'), ((9385, 9407), 'streamlit.button', 'st.button', (['"""Processar"""'], {}), "('Processar')\n", (9394, 9407), True, 'import streamlit as st\n'), ((3333, 3379), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'temperature': '(0.1)', 'openai_api_key': '""""""'}), "(temperature=0.1, openai_api_key='')\n", (3343, 3379), False, 'from langchain.chat_models import ChatOpenAI\n'), ((6808, 6841), 'streamlit.spinner', 'st.spinner', (['"""Gerando resposta..."""'], {}), "('Gerando resposta...')\n", (6818, 6841), True, 'import streamlit as st\n'), ((6855, 6926), 'logging.info', 'log.info', (['f"""Gerando resposta para consulta do cliente: {user_question}"""'], {}), "(f'Gerando resposta para consulta do cliente: {user_question}')\n", (6863, 6926), True, 'import logging as log\n'), ((7137, 7165), 'streamlit.write', 'st.write', (["response['answer']"], {}), "(response['answer'])\n", (7145, 7165), True, 'import streamlit as st\n'), ((7774, 7805), 'langchain_community.document_loaders.PyPDFLoader', 'PyPDFLoader', ([], {'file_path': 'filepath'}), '(file_path=filepath)\n', (7785, 7805), False, 'from langchain_community.document_loaders import PyPDFLoader\n'), ((8697, 8737), 'streamlit.spinner', 'st.spinner', (['"""Inicializando Vector DB..."""'], {}), "('Inicializando Vector DB...')\n", (8707, 8737), True, 'import streamlit as st\n'), ((8835, 8874), 'streamlit.spinner', 'st.spinner', (['"""Inicializando AI Model..."""'], {}), "('Inicializando AI Model...')\n", (8845, 8874), True, 'import streamlit as st\n'), ((3582, 3693), 'langchain.llms.gpt4all.GPT4All', 'GPT4All', ([], {'model': '"""model_cache/zephyr-7b-beta.Q3_K_S.gguf"""', 'callbacks': 'callbacks', 'verbose': '(True)', 'device': '"""nvidia"""'}), "(model='model_cache/zephyr-7b-beta.Q3_K_S.gguf', callbacks=callbacks,\n    verbose=True, device='nvidia')\n", (3589, 3693), False, 'from langchain.llms.gpt4all import GPT4All\n'), ((7196, 7240), 'streamlit.expander', 'st.expander', ([], {'label': '"""Sources"""', 'expanded': '(False)'}), "(label='Sources', expanded=False)\n", (7207, 7240), True, 'import streamlit as st\n'), ((7529, 7568), 'os.path.join', 'os.path.join', (['"""tmp_documents"""', 'doc.name'], {}), "('tmp_documents', doc.name)\n", (7541, 7568), False, 'import os\n'), ((7883, 7917), 'langchain_community.document_loaders.Docx2txtLoader', 'Docx2txtLoader', ([], {'file_path': 'filepath'}), '(file_path=filepath)\n', (7897, 7917), False, 'from langchain_community.document_loaders import Docx2txtLoader\n'), ((9426, 9454), 'streamlit.spinner', 'st.spinner', (['"""Processando..."""'], {}), "('Processando...')\n", (9436, 9454), True, 'import streamlit as st\n'), ((3524, 3556), 'langchain.callbacks.streaming_stdout.StreamingStdOutCallbackHandler', 'StreamingStdOutCallbackHandler', ([], {}), '()\n', (3554, 3556), False, 'from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler\n'), ((4010, 4105), 'langchain.llms.huggingface_hub.HuggingFaceHub', 'HuggingFaceHub', ([], {'verbose': '(True)', 'task': '"""text-generation"""', 'repo_id': '"""tiiuae/falcon-40b-instruct"""'}), "(verbose=True, task='text-generation', repo_id=\n    'tiiuae/falcon-40b-instruct')\n", (4024, 4105), False, 'from langchain.llms.huggingface_hub import HuggingFaceHub\n'), ((7313, 7329), 'streamlit.write', 'st.write', (['source'], {}), '(source)\n', (7321, 7329), True, 'import streamlit as st\n'), ((4224, 4377), 'langchain.llms.huggingface_pipeline.HuggingFacePipeline.from_model_id', 'HuggingFacePipeline.from_model_id', ([], {'model_id': '"""Open-Orca/Mistral-7B-OpenOrca"""', 'task': '"""text-generation"""', 'pipeline_kwargs': "{'max_new_tokens': 10}", 'device': '(0)'}), "(model_id='Open-Orca/Mistral-7B-OpenOrca',\n    task='text-generation', pipeline_kwargs={'max_new_tokens': 10}, device=0)\n", (4257, 4377), False, 'from langchain.llms.huggingface_pipeline import HuggingFacePipeline\n'), ((4723, 4845), 'langchain_community.llms.huggingface_endpoint.HuggingFaceEndpoint', 'HuggingFaceEndpoint', ([], {'endpoint_url': 'endpoint_url', 'task': '"""text-generation"""', 'huggingfacehub_api_token': 'HUGGINGFACE_API_TOKEN'}), "(endpoint_url=endpoint_url, task='text-generation',\n    huggingfacehub_api_token=HUGGINGFACE_API_TOKEN)\n", (4742, 4845), False, 'from langchain_community.llms.huggingface_endpoint import HuggingFaceEndpoint\n'), ((4957, 5069), 'langchain.llms.huggingface_pipeline.HuggingFacePipeline.from_model_id', 'HuggingFacePipeline.from_model_id', ([], {'model_id': '"""HuggingFaceH4/zephyr-7b-beta"""', 'task': '"""text-generation"""', 'device': '(0)'}), "(model_id='HuggingFaceH4/zephyr-7b-beta',\n    task='text-generation', device=0)\n", (4990, 5069), False, 'from langchain.llms.huggingface_pipeline import HuggingFacePipeline\n'), ((5243, 5410), 'langchain.llms.huggingface_pipeline.HuggingFacePipeline.from_model_id', 'HuggingFacePipeline.from_model_id', ([], {'model_id': '"""meta-llama/Llama-2-7b-chat-hf"""', 'task': '"""text-generation"""', 'device': '"""cuda"""', 'pipeline_kwargs': "{'token': LLAMA2_API_TOKEN}"}), "(model_id='meta-llama/Llama-2-7b-chat-hf',\n    task='text-generation', device='cuda', pipeline_kwargs={'token':\n    LLAMA2_API_TOKEN})\n", (5276, 5410), False, 'from langchain.llms.huggingface_pipeline import HuggingFacePipeline\n'), ((5615, 5760), 'langchain_nvidia_ai_endpoints.ChatNVIDIA', 'ChatNVIDIA', ([], {'model': '"""mixtral_8x7b"""', 'nvidia_api_key': 'NVIDIANGC_API_KEY', 'callbacks': 'callbacks', 'temperature': '(0.2)', 'top_p': '(0.7)', 'max_tokens': '(1024)', 'seed': '(42)'}), "(model='mixtral_8x7b', nvidia_api_key=NVIDIANGC_API_KEY,\n    callbacks=callbacks, temperature=0.2, top_p=0.7, max_tokens=1024, seed=42)\n", (5625, 5760), False, 'from langchain_nvidia_ai_endpoints import ChatNVIDIA\n'), ((5558, 5590), 'langchain.callbacks.streaming_stdout.StreamingStdOutCallbackHandler', 'StreamingStdOutCallbackHandler', ([], {}), '()\n', (5588, 5590), False, 'from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler\n')]