# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_ingest.ipynb.

# %% auto 0
__all__ = ['chunk_size', 'chunk_overlap', 'DEFAULT_DB', 'MyElmLoader', 'ingest']

# %% ../nbs/01_ingest.ipynb 3
import os
import os.path
import glob
from typing import List
from dotenv import load_dotenv
from multiprocessing import Pool
from tqdm import tqdm

from langchain.document_loaders import (
    CSVLoader,
    EverNoteLoader,
    PyMuPDFLoader,
    TextLoader,
    UnstructuredEmailLoader,
    UnstructuredEPubLoader,
    UnstructuredHTMLLoader,
    UnstructuredMarkdownLoader,
    UnstructuredODTLoader,
    UnstructuredPowerPointLoader,
    UnstructuredWordDocumentLoader,
)

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Chroma
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.docstore.document import Document
import chromadb
from chromadb.config import Settings
chunk_size = 500
chunk_overlap = 50

# %% ../nbs/01_ingest.ipynb 4
class MyElmLoader(UnstructuredEmailLoader):
    """Wrapper to fallback to text/plain when default does not work"""

    def load(self) -> List[Document]:
        """Wrapper adding fallback for elm without html"""
        try:
            try:
                doc = UnstructuredEmailLoader.load(self)
            except ValueError as e:
                if 'text/html content not found in email' in str(e):
                    # Try plain text
                    self.unstructured_kwargs["content_source"]="text/plain"
                    doc = UnstructuredEmailLoader.load(self)
                else:
                    raise
        except Exception as e:
            # Add file_path to exception message
            raise type(e)(f"{self.file_path}: {e}") from e

        return doc

# %% ../nbs/01_ingest.ipynb 6
from typing import Any, Dict, Generator, List, Optional, Tuple, Union
from .utils import get_datadir
os.environ['TOKENIZERS_PARALLELISM'] = '0'
DEFAULT_DB = 'vectordb'
def ingest(source_directory:str, 
           embedding_model_name:str ='sentence-transformers/all-MiniLM-L6-v2',
           embedding_model_kwargs:dict ={'device': 'cpu'}
          ):
    """
    Ingests all documents in `source_folder` (previously-ingested documents are ignored)

    **Args**:
    
      - *source_directory*: path to folder containing document store
      - *embedding_model*: name of sentence-transformers model
      - *embedding_model_kwargs*: arguments to embedding model (e.g., `{device':'cpu'}`)
    
    **Returns**: `None`
    """
    if not os.path.exists(source_directory):
        raise ValueError('The source_directory does not exist.')
    persist_directory = os.path.join(get_datadir(), DEFAULT_DB)
    embeddings = HuggingFaceEmbeddings(model_name=embedding_model_name)
    chroma_settings = Settings(persist_directory=persist_directory,anonymized_telemetry=False)
    chroma_client = chromadb.PersistentClient(settings=chroma_settings , path=persist_directory)
    
    texts = None
    if does_vectorstore_exist(persist_directory, embeddings):
        # Update and store locally vectorstore
        print(f"Appending to existing vectorstore at {persist_directory}")
        db = Chroma(persist_directory=persist_directory, 
                    embedding_function=embeddings, 
                    client_settings=chroma_settings, client=chroma_client)
        collection = db.get()
        texts = process_documents(source_directory, 
                                  ignored_files=[metadata['source'] for metadata in collection['metadatas']])
        if texts:
            print(f"Creating embeddings. May take some minutes...")
            db.add_documents(texts)
    else:
        # Create and store locally vectorstore
        print("Creating new vectorstore")
        texts = process_documents(source_directory)
        if texts:
            print(f"Creating embeddings. May take some minutes...")
            db = Chroma.from_documents(texts, 
                                       embeddings, persist_directory=persist_directory, 
                                       client_settings=chroma_settings, client=chroma_client)
    if texts:
        db.persist()
        print(f"Ingestion complete! You can now query your documents using the prompt method")
    db = None
    return
