'''
Create Vector Store from all documents in a folder, currently supports .pptx, .docx, .pdf files.

Created by Ric Zhou on 2021-03-27
'''

from langchain.document_loaders import (UnstructuredPowerPointLoader, UnstructuredWordDocumentLoader, PyPDFLoader, UnstructuredPDFLoader)
import glob
import langchain.text_splitter as text_splitter
from langchain.text_splitter import (RecursiveCharacterTextSplitter, CharacterTextSplitter)
from langchain.vectorstores import FAISS
from langchain.embeddings import OpenAIEmbeddings
from GlobalClasses import GlobalContext
from dotenv import load_dotenv
import os

load_dotenv()
GlobalContext()  # initialize global context

os.environ["OPENAI_API_TYPE"] = "azure"
os.environ["OPENAI_API_VERSION"] = "2022-12-01"
os.environ["OPENAI_API_BASE"] = GlobalContext.OPENAI_BASE
os.environ["OPENAI_API_KEY"] = GlobalContext.OPENAI_API_KEY

ENGLISH_CHUNK_SIZE = 1400
CHINESE_CHUNK_SIZE = 500

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=ENGLISH_CHUNK_SIZE, chunk_overlap=0)  # chunk_overlap=30

files = glob.glob(f"{GlobalContext.VECTOR_DB_PATH}/*.*")

all_docs = []
for p in files:
    if p.lower().endswith(".pptx"):
        loader = UnstructuredPowerPointLoader(p)
        docs = loader.load_and_split(text_splitter)
        print(p)
        print(len(docs))
        all_docs.extend(docs)
    elif p.lower().endswith(".docx"):
        loader = UnstructuredWordDocumentLoader(p)
        docs = loader.load_and_split(text_splitter)
        print(p)
        print(len(docs))
        all_docs.extend(docs)
    elif p.lower().endswith(".pdf"):
        loader = PyPDFLoader(p)
        docs = loader.load_and_split(text_splitter)
        print(p)
        print(len(docs))
        all_docs.extend(docs)


print(len(all_docs))

# vectorstore = FAISS.from_documents(all_docs, OpenAIEmbeddings(chunk_size=1, document_model_name="text-search-curie-doc-001", query_model_name="text-search-curie-query-001")) # text-search-curie-*-001 performance is worse than text-embedding-ada-002
vectorstore = FAISS.from_documents(all_docs, OpenAIEmbeddings(chunk_size=1))
#vectorstore = FAISS.from_documents(all_docs, OpenAIEmbeddings())

FAISS.save_local(vectorstore, GlobalContext.VECTOR_DB_PATH)

langchain.document_loaders.UnstructuredWordDocumentLoader langchain.text_splitter.RecursiveCharacterTextSplitter langchain.document_loaders.UnstructuredPowerPointLoader langchain.vectorstores.FAISS.save_local langchain.document_loaders.PyPDFLoader langchain.embeddings.OpenAIEmbeddings
[((604, 617), 'dotenv.load_dotenv', 'load_dotenv', ([], {}), '()\n', (615, 617), False, 'from dotenv import load_dotenv\n'), ((618, 633), 'GlobalClasses.GlobalContext', 'GlobalContext', ([], {}), '()\n', (631, 633), False, 'from GlobalClasses import GlobalContext\n'), ((939, 1017), 'langchain.text_splitter.RecursiveCharacterTextSplitter', 'RecursiveCharacterTextSplitter', ([], {'chunk_size': 'ENGLISH_CHUNK_SIZE', 'chunk_overlap': '(0)'}), '(chunk_size=ENGLISH_CHUNK_SIZE, chunk_overlap=0)\n', (969, 1017), False, 'from langchain.text_splitter import RecursiveCharacterTextSplitter, CharacterTextSplitter\n'), ((1052, 1100), 'glob.glob', 'glob.glob', (['f"""{GlobalContext.VECTOR_DB_PATH}/*.*"""'], {}), "(f'{GlobalContext.VECTOR_DB_PATH}/*.*')\n", (1061, 1100), False, 'import glob\n'), ((2166, 2225), 'langchain.vectorstores.FAISS.save_local', 'FAISS.save_local', (['vectorstore', 'GlobalContext.VECTOR_DB_PATH'], {}), '(vectorstore, GlobalContext.VECTOR_DB_PATH)\n', (2182, 2225), False, 'from langchain.vectorstores import FAISS\n'), ((2067, 2097), 'langchain.embeddings.OpenAIEmbeddings', 'OpenAIEmbeddings', ([], {'chunk_size': '(1)'}), '(chunk_size=1)\n', (2083, 2097), False, 'from langchain.embeddings import OpenAIEmbeddings\n'), ((1185, 1216), 'langchain.document_loaders.UnstructuredPowerPointLoader', 'UnstructuredPowerPointLoader', (['p'], {}), '(p)\n', (1213, 1216), False, 'from langchain.document_loaders import UnstructuredPowerPointLoader, UnstructuredWordDocumentLoader, PyPDFLoader, UnstructuredPDFLoader\n'), ((1396, 1429), 'langchain.document_loaders.UnstructuredWordDocumentLoader', 'UnstructuredWordDocumentLoader', (['p'], {}), '(p)\n', (1426, 1429), False, 'from langchain.document_loaders import UnstructuredPowerPointLoader, UnstructuredWordDocumentLoader, PyPDFLoader, UnstructuredPDFLoader\n'), ((1608, 1622), 'langchain.document_loaders.PyPDFLoader', 'PyPDFLoader', (['p'], {}), '(p)\n', (1619, 1622), False, 'from langchain.document_loaders import UnstructuredPowerPointLoader, UnstructuredWordDocumentLoader, PyPDFLoader, UnstructuredPDFLoader\n')]