codekingpro/portable-devtools

1import concurrent2import logging3import random4from pathlib import Path5from typing import Any, Callable, Iterator, List, Optional, Sequence, Tuple, Type, Union6 7from langchain_core.documents import Document8 9from langchain_community.document_loaders.base import BaseLoader10from langchain_community.document_loaders.csv_loader import CSVLoader11from langchain_community.document_loaders.html_bs import BSHTMLLoader12from langchain_community.document_loaders.text import TextLoader13from langchain_community.document_loaders.unstructured import UnstructuredFileLoader14 15FILE_LOADER_TYPE = Union[16    Type[UnstructuredFileLoader], Type[TextLoader], Type[BSHTMLLoader], Type[CSVLoader]17]18logger = logging.getLogger(__name__)19 20 21def _is_visible(p: Path) -> bool:22    parts = p.parts23    for _p in parts:24        if _p.startswith("."):25            return False26    return True27 28 29class DirectoryLoader(BaseLoader):30    """Load from a directory."""31 32    def __init__(33        self,34        path: str,35        glob: Union[List[str], Tuple[str], str] = "**/[!.]*",36        silent_errors: bool = False,37        load_hidden: bool = False,38        loader_cls: FILE_LOADER_TYPE = UnstructuredFileLoader,39        loader_kwargs: Union[dict, None] = None,40        recursive: bool = False,41        show_progress: bool = False,42        use_multithreading: bool = False,43        max_concurrency: int = 4,44        *,45        exclude: Union[Sequence[str], str] = (),46        sample_size: int = 0,47        randomize_sample: bool = False,48        sample_seed: Union[int, None] = None,49    ):50        """Initialize with a path to directory and how to glob over it.51 52        Args:53            path: Path to directory.54            glob: A glob pattern or list of glob patterns to use to find files.55                Defaults to "**/[!.]*" (all files except hidden).56            exclude: A pattern or list of patterns to exclude from results.57                Use glob syntax.58            silent_errors: Whether to silently ignore errors. Defaults to False.59            load_hidden: Whether to load hidden files. Defaults to False.60            loader_cls: Loader class to use for loading files.61              Defaults to UnstructuredFileLoader.62            loader_kwargs: Keyword arguments to pass to loader_cls. Defaults to None.63            recursive: Whether to recursively search for files. Defaults to False.64            show_progress: Whether to show a progress bar. Defaults to False.65            use_multithreading: Whether to use multithreading. Defaults to False.66            max_concurrency: The maximum number of threads to use. Defaults to 4.67            sample_size: The maximum number of files you would like to load from the68                directory.69            randomize_sample: Shuffle the files to get a random sample.70            sample_seed: set the seed of the random shuffle for reproducibility.71 72        Examples:73 74            .. code-block:: python75                from langchain_community.document_loaders import DirectoryLoader76 77                # Load all non-hidden files in a directory.78                loader = DirectoryLoader("/path/to/directory")79 80                # Load all text files in a directory without recursion.81                loader = DirectoryLoader("/path/to/directory", glob="*.txt")82 83                # Recursively load all text files in a directory.84                loader = DirectoryLoader(85                    "/path/to/directory", glob="*.txt", recursive=True86                )87 88                # Load all files in a directory, except for py files.89                loader = DirectoryLoader("/path/to/directory", exclude="*.py")90 91                # Load all files in a directory, except for py or pyc files.92                loader = DirectoryLoader(93                    "/path/to/directory", exclude=["*.py", "*.pyc"]94                )95        """96        if loader_kwargs is None:97            loader_kwargs = {}98        if isinstance(exclude, str):99            exclude = (exclude,)100        self.path = path101        self.glob = glob102        self.exclude = exclude103        self.load_hidden = load_hidden104        self.loader_cls = loader_cls105        self.loader_kwargs = loader_kwargs106        self.silent_errors = silent_errors107        self.recursive = recursive108        self.show_progress = show_progress109        self.use_multithreading = use_multithreading110        self.max_concurrency = max_concurrency111        self.sample_size = sample_size112        self.randomize_sample = randomize_sample113        self.sample_seed = sample_seed114 115    def load(self) -> List[Document]:116        """Load documents."""117        return list(self.lazy_load())118 119    def lazy_load(self) -> Iterator[Document]:120        """Load documents lazily."""121        p = Path(self.path)122        if not p.exists():123            raise FileNotFoundError(f"Directory not found: '{self.path}'")124        if not p.is_dir():125            raise ValueError(f"Expected directory, got file: '{self.path}'")126 127        # glob multiple patterns if a list is provided, e.g., multiple file extensions128        if isinstance(self.glob, (list, tuple)):129            paths = []130            for pattern in self.glob:131                paths.extend(132                    list(p.rglob(pattern) if self.recursive else p.glob(pattern))133                )134        elif isinstance(self.glob, str):135            paths = list(p.rglob(self.glob) if self.recursive else p.glob(self.glob))136        else:137            raise TypeError(138                f"Expected glob to be str or sequence of str, but got {type(self.glob)}"139            )140 141        items = [142            path143            for path in paths144            if not (self.exclude and any(path.match(glob) for glob in self.exclude))145            and path.is_file()146        ]147 148        if self.sample_size > 0:149            if self.randomize_sample:150                randomizer = random.Random(151                    self.sample_seed if self.sample_seed else None152                )153                randomizer.shuffle(items)154            items = items[: min(len(items), self.sample_size)]155 156        pbar = None157        if self.show_progress:158            try:159                from tqdm import tqdm160 161                pbar = tqdm(total=len(items))162            except ImportError as e:163                logger.warning(164                    "To log the progress of DirectoryLoader you need to install tqdm, "165                    "`pip install tqdm`"166                )167                if self.silent_errors:168                    logger.warning(e)169                else:170                    raise ImportError(171                        "To log the progress of DirectoryLoader "172                        "you need to install tqdm, "173                        "`pip install tqdm`"174                    )175 176        if self.use_multithreading:177            futures = []178            with concurrent.futures.ThreadPoolExecutor(179                max_workers=self.max_concurrency180            ) as executor:181                for i in items:182                    futures.append(183                        executor.submit(184                            self._lazy_load_file_to_non_generator(self._lazy_load_file),185                            i,186                            p,187                            pbar,188                        )189                    )190                for future in concurrent.futures.as_completed(futures):191                    for item in future.result():192                        yield item193        else:194            for i in items:195                yield from self._lazy_load_file(i, p, pbar)196 197        if pbar:198            pbar.close()199 200    def _lazy_load_file_to_non_generator(self, func: Callable) -> Callable:201        def non_generator(item: Path, path: Path, pbar: Optional[Any]) -> List:202            return [x for x in func(item, path, pbar)]203 204        return non_generator205 206    def _lazy_load_file(207        self, item: Path, path: Path, pbar: Optional[Any]208    ) -> Iterator[Document]:209        """Load a file.210 211        Args:212            item: File path.213            path: Directory path.214            pbar: Progress bar. Defaults to None.215 216        """217        if item.is_file():218            if _is_visible(item.relative_to(path)) or self.load_hidden:219                try:220                    logger.debug(f"Processing file: {str(item)}")221                    loader = self.loader_cls(str(item), **self.loader_kwargs)222                    try:223                        for subdoc in loader.lazy_load():224                            yield subdoc225                    except NotImplementedError:226                        for subdoc in loader.load():227                            yield subdoc228                except Exception as e:229                    if self.silent_errors:230                        logger.warning(f"Error loading file {str(item)}: {e}")231                    else:232                        logger.error(f"Error loading file {str(item)}")233                        raise e234                finally:235                    if pbar:236                        pbar.update(1)237