from langchain.document_loaders import UnstructuredFileLoader, TextLoader
from langchain.document_loaders import UnstructuredURLLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter, CharacterTextSplitter
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.chains.question_answering import load_qa_chain
from langchain import OpenAI
from langchain.prompts import PromptTemplate
from langchain.callbacks import get_openai_callback


#import openai
import os
from dotenv import load_dotenv, find_dotenv
_ = load_dotenv(find_dotenv()) # read local .env file
#openai.api_key  = os.environ['OPENAI_API_KEY']


# DEFINING A PERSIST DIRECTORY FOR CHROMADB SO THE VECTOR STORE WILL BE STORED LOCALLY
persist_directory = 'db'
embeddings = OpenAIEmbeddings()

# check to see if the db already exists
if not os.path.exists(persist_directory):
    # load book.txt encode and decode in ascii and write to file to avoid encoding errors
    with open("book.txt", "r",encoding="utf-8") as f:
        book = f.read().encode("ascii", "ignore").decode("ascii")
    with open("book_ascii.txt", "w") as f:
        f.write(book)

    print("Loading book_ascii.txt")
    # loader = UnstructuredFileLoader("book_ascii.txt")

    loader = TextLoader('book_ascii.txt')

    documents = loader.load()

    # print(docs[0].page_content[:])

    text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
    docs = text_splitter.split_documents(documents)
    # print(texts[0].page_content)


    
    print("embedding book_ascii.txt")
    db = Chroma.from_documents(docs, embeddings, persist_directory=persist_directory)
else:
    db = Chroma(persist_directory=persist_directory, embedding_function=embeddings)



# MAP_REDUCE REQUIRES pip install tiktoken
chain = load_qa_chain(OpenAI(temperature=0), chain_type="map_rerank", return_intermediate_steps=True)
while True:
    with get_openai_callback() as cb:
        query = input("Digite a query: ")
        docs = db.similarity_search(query)
        result = chain({"input_documents": docs, "question": query}, return_only_outputs=True)
        print(result['output_text'])
        print(result['intermediate_steps'])
        # print(docs)
        print("tokens used: ", cb.total_tokens)