diff --git a/CHANGELOG.md b/CHANGELOG.md index 077e103..fbb67ab 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Structured logging for ingest, query, and retriever stages, including timing and empty-ingest warnings (closes #41). - `RecursiveChunker.from_config()` for creating a chunker from pipeline configuration (closes #25). +- `DirectoryLoader` for loading supported documents from a directory with per-extension loader dispatch, including recursive and non-recursive traversal, `on_error` handling, and a skipped-file count (closes #34). ### Fixed - `TextFileLoader` and `MarkdownLoader` now convert Unicode decoding and unknown encoding errors to `LoaderError`, with configurable text decoding error handling (closes #24). diff --git a/ragframework/document/__init__.py b/ragframework/document/__init__.py index d09b1bb..d5b6d96 100644 --- a/ragframework/document/__init__.py +++ b/ragframework/document/__init__.py @@ -8,7 +8,7 @@ ) from .html import HTMLLoader -from .loaders import DocxLoader, MarkdownLoader, PDFLoader, TextFileLoader +from .loaders import DirectoryLoader, DocxLoader, MarkdownLoader, PDFLoader, TextFileLoader from .tabular import CSVLoader, JSONLLoader __all__ = [ @@ -19,6 +19,7 @@ "CSVLoader", "JSONLLoader", "HTMLLoader", + "DirectoryLoader", "FixedSizeChunker", "RecursiveChunker", "SentenceChunker", diff --git a/ragframework/document/loaders.py b/ragframework/document/loaders.py index 3447e71..e3c882d 100644 --- a/ragframework/document/loaders.py +++ b/ragframework/document/loaders.py @@ -8,6 +8,7 @@ from __future__ import annotations import hashlib +import importlib.util from pathlib import Path from ragframework.base import Document, DocumentLoader @@ -214,3 +215,78 @@ def load(self, source: str) -> list[Document]: ) ) return documents + + +class DirectoryLoader(DocumentLoader): + """Load documents from a directory using loaders selected by file extension.""" + + def __init__( + self, + loaders: dict[str, DocumentLoader] | None = None, + glob: str = "**/*", + recursive: bool = True, + on_error: str = "raise", + ) -> None: + if on_error not in {"skip", "raise"}: + raise ValueError("on_error must be either 'skip' or 'raise'") + + if loaders is None: + loaders = { + ".txt": TextFileLoader(), + ".md": MarkdownLoader(), + ".markdown": MarkdownLoader(), + } + if importlib.util.find_spec("pypdf") is not None: + loaders[".pdf"] = PDFLoader() + + self.loaders = loaders + self.glob = glob + self.recursive = recursive + self.on_error = on_error + self.skipped_count = 0 + + + def load(self, source: str) -> list[Document]: + root = Path(source) + + if not root.exists(): + raise LoaderError(f"Directory not found: {source}") + + if not root.is_dir(): + raise LoaderError(f"Not a directory: {source}") + + paths = root.rglob(self.glob) if self.recursive else root.glob(self.glob) + + documents: list[Document] = [] + + for path in sorted(paths): + if not path.is_file(): + continue + + if not self.recursive and path.parent != root: + continue + + extension = path.suffix.lower() + loader = self.loaders.get(extension) + + if loader is None: + self.skipped_count += 1 + continue + + try: + loaded_documents = loader.load(str(path)) + except Exception as exc: + if self.on_error == "skip": + continue + + raise LoaderError( + f"Failed to load {path}: {exc}" + ) from exc + + relative_path = path.relative_to(root).as_posix() + + for document in loaded_documents: + document.metadata["relative_path"] = relative_path + documents.append(document) + + return documents diff --git a/tests/test_document/test_loaders.py b/tests/test_document/test_loaders.py index 2a3ac80..4f61f75 100644 --- a/tests/test_document/test_loaders.py +++ b/tests/test_document/test_loaders.py @@ -3,10 +3,13 @@ import builtins import sys import types +from pathlib import Path import pytest +from ragframework.base import Document, DocumentLoader from ragframework.document.loaders import ( + DirectoryLoader, DocxLoader, MarkdownLoader, PDFLoader, @@ -262,3 +265,215 @@ def __init__(self, source): assert len(docs_whole) == 1 assert "page1" in docs_whole[0].content assert "page2" in docs_whole[0].content + +def test_directory_loader_loads_matching_files(tmp_path): + docs_dir = tmp_path / "docs" + docs_dir.mkdir() + + (docs_dir / "hello.txt").write_text( + "Hello from text.", + encoding="utf-8", + ) + + (docs_dir / "readme.md").write_text( + "# README", + encoding="utf-8", + ) + + (docs_dir / "ignored.py").write_text( + "print('ignored')", + encoding="utf-8", + ) + + loader = DirectoryLoader( + loaders={ + ".txt": TextFileLoader(), + ".md": MarkdownLoader(), + } + ) + + docs = loader.load(str(docs_dir)) + + assert len(docs) == 2 + assert docs[0].metadata["relative_path"] == "hello.txt" + assert docs[1].metadata["relative_path"] == "readme.md" + +def test_directory_loader_loads_nested_files_in_sorted_order(tmp_path): + docs_dir = tmp_path / "docs" + nested_dir = docs_dir / "nested" + nested_dir.mkdir(parents=True) + + (docs_dir / "z.txt").write_text("Z", encoding="utf-8") + (docs_dir / "a.txt").write_text("A", encoding="utf-8") + (nested_dir / "b.txt").write_text("B", encoding="utf-8") + + loader = DirectoryLoader( + loaders={ + ".txt": TextFileLoader(), + } + ) + + docs = loader.load(str(docs_dir)) + + assert len(docs) == 3 + + assert [ + doc.metadata["relative_path"] + for doc in docs + ] == [ + "a.txt", + "nested/b.txt", + "z.txt", + ] + +def test_directory_loader_skips_loader_errors(tmp_path): + docs_dir = tmp_path / "docs" + docs_dir.mkdir() + + good_file = docs_dir / "good.txt" + bad_file = docs_dir / "bad.txt" + + good_file.write_text("Good", encoding="utf-8") + bad_file.write_text("Bad", encoding="utf-8") + + class FailingLoader(DocumentLoader): + def load(self, source: str) -> list[Document]: + raise LoaderError(f"Cannot load {source}") + + loader = DirectoryLoader( + loaders={".txt": FailingLoader()}, + on_error="skip", + ) + + docs = loader.load(str(docs_dir)) + + assert docs == [] + +def test_directory_loader_raises_loader_error(tmp_path): + docs_dir = tmp_path / "docs" + docs_dir.mkdir() + + bad_file = docs_dir / "bad.txt" + bad_file.write_text("Bad", encoding="utf-8") + + class FailingLoader(DocumentLoader): + def load(self, source: str) -> list[Document]: + raise LoaderError(f"Cannot load {source}") + + loader = DirectoryLoader( + loaders={".txt": FailingLoader()}, + on_error="raise", + ) + + with pytest.raises(LoaderError, match="bad.txt"): + loader.load(str(docs_dir)) + +def test_directory_loader_has_default_loaders(tmp_path): + docs_dir = tmp_path / "docs" + docs_dir.mkdir() + + (docs_dir / "notes.txt").write_text( + "Some notes", + encoding="utf-8", + ) + + (docs_dir / "readme.md").write_text( + "# README", + encoding="utf-8", + ) + + (docs_dir / "guide.markdown").write_text( + "# Guide", + encoding="utf-8", + ) + + loader = DirectoryLoader() + + docs = loader.load(str(docs_dir)) + + assert len(docs) == 3 + assert { + doc.metadata["relative_path"] + for doc in docs + } == { + "notes.txt", + "readme.md", + "guide.markdown", + } + +def test_directory_loader_rejects_invalid_on_error(tmp_path): + with pytest.raises(ValueError, match="on_error"): + DirectoryLoader(on_error="banana") + +def test_directory_loader_handles_unreadable_file(tmp_path, monkeypatch): + docs_dir = tmp_path / "docs" + docs_dir.mkdir() + + good_file = docs_dir / "good.txt" + bad_file = docs_dir / "bad.txt" + + good_file.write_text("Good", encoding="utf-8") + bad_file.write_text("Bad", encoding="utf-8") + + original_read_text = Path.read_text + + def fake_read_text(self, *args, **kwargs): + if self == bad_file: + raise OSError("Permission denied") + return original_read_text(self, *args, **kwargs) + + monkeypatch.setattr(Path, "read_text", fake_read_text) + + skip_loader = DirectoryLoader(on_error="skip") + docs = skip_loader.load(str(docs_dir)) + + assert len(docs) == 1 + assert docs[0].metadata["relative_path"] == "good.txt" + + raise_loader = DirectoryLoader(on_error="raise") + + with pytest.raises(LoaderError, match="bad.txt"): + raise_loader.load(str(docs_dir)) + +def test_directory_loader_non_recursive_ignores_nested_files(tmp_path): + docs_dir = tmp_path / "docs" + nested_dir = docs_dir / "nested" + nested_dir.mkdir(parents=True) + + (docs_dir / "top.txt").write_text("Top", encoding="utf-8") + (nested_dir / "deep.txt").write_text("Deep", encoding="utf-8") + + loader = DirectoryLoader( + loaders={".txt": TextFileLoader()}, + recursive=False, + ) + + docs = loader.load(str(docs_dir)) + + assert [doc.metadata["relative_path"] for doc in docs] == ["top.txt"] + +def test_directory_loader_default_loaders_skip_pdf_without_pypdf(monkeypatch): + monkeypatch.setattr( + "ragframework.document.loaders.importlib.util.find_spec", + lambda name: None, + ) + + loader = DirectoryLoader() + + assert ".pdf" not in loader.loaders + assert ".txt" in loader.loaders + +def test_directory_loader_counts_skipped_unsupported_files(tmp_path): + docs_dir = tmp_path / "docs" + docs_dir.mkdir() + + (docs_dir / "keep.txt").write_text("Keep", encoding="utf-8") + (docs_dir / "ignored.py").write_text("print('x')", encoding="utf-8") + (docs_dir / "ignored.csv").write_text("a,b", encoding="utf-8") + + loader = DirectoryLoader(loaders={".txt": TextFileLoader()}) + + docs = loader.load(str(docs_dir)) + + assert len(docs) == 1 + assert loader.skipped_count == 2