Coverage for src/local_deep_research/document_loaders/loader_registry.py: 91%
122 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1"""
2Document loader registry - maps file extensions to LangChain loaders.
4This module provides a centralized registry for document loaders that can be used
5by both collection uploads (bytes) and local search (file paths).
6"""
8import importlib.util
9import shutil
10from pathlib import Path
11from typing import Optional
13from langchain_community.document_loaders import (
14 MHTMLLoader,
15 NotebookLoader,
16 PyPDFLoader,
17 UnstructuredEmailLoader,
18 UnstructuredExcelLoader,
19 UnstructuredHTMLLoader,
20 UnstructuredMarkdownLoader,
21 UnstructuredPowerPointLoader,
22 UnstructuredWordDocumentLoader,
23 UnstructuredXMLLoader,
24)
25from langchain_core.document_loaders import BaseLoader
26from loguru import logger
28# Import loaders that require external system tools (pandoc, tesseract, etc.)
29# These may fail at runtime if the tools aren't installed
30try:
31 from langchain_community.document_loaders import UnstructuredODTLoader
33 HAS_ODT_LOADER = True
34except ImportError:
35 HAS_ODT_LOADER = False
36 logger.debug("UnstructuredODTLoader not available - ODT support disabled")
38try:
39 from langchain_community.document_loaders import UnstructuredEPubLoader
41 HAS_EPUB_LOADER = True
42except ImportError:
43 HAS_EPUB_LOADER = False
44 logger.debug("UnstructuredEPubLoader not available - EPUB support disabled")
46try:
47 from langchain_community.document_loaders import UnstructuredRTFLoader
49 HAS_RTF_LOADER = True
50except ImportError:
51 HAS_RTF_LOADER = False
52 logger.debug("UnstructuredRTFLoader not available - RTF support disabled")
54try:
55 from langchain_community.document_loaders import UnstructuredRSTLoader
57 HAS_RST_LOADER = True
58except ImportError:
59 HAS_RST_LOADER = False
60 logger.debug("UnstructuredRSTLoader not available - RST support disabled")
62try:
63 from langchain_community.document_loaders import UnstructuredOrgModeLoader
65 HAS_ORG_LOADER = True
66except ImportError:
67 HAS_ORG_LOADER = False
68 logger.debug(
69 "UnstructuredOrgModeLoader not available - Org support disabled"
70 )
72try:
73 from langchain_community.document_loaders import UnstructuredImageLoader
75 HAS_IMAGE_LOADER = True
76except ImportError:
77 HAS_IMAGE_LOADER = False
78 logger.debug(
79 "UnstructuredImageLoader not available - Image/OCR support disabled"
80 )
82# Import our custom loaders
83from .json_loader import SimpleJSONLoader
84from .text_encoding import EncodingAwareCSVLoader, EncodingAwareTextLoader
85from .evernote_loader import EverNoteLoader
86from .toml_loader import TOMLLoader
87from .xls_loader import XLSLoader
88from .yaml_loader import YAMLLoader
91def _module_available(module_name: str) -> bool:
92 """Return True if *module_name* can be imported in this environment.
94 Uses ``importlib.util.find_spec`` so the (potentially heavy) module is
95 located but not actually imported.
96 """
97 try:
98 return importlib.util.find_spec(module_name) is not None
99 except (ImportError, ValueError, ModuleNotFoundError):
100 # A parent package may itself be missing, which surfaces as
101 # ModuleNotFoundError/ImportError from find_spec.
102 return False
105# Whether a loader *class* can be imported is NOT sufficient to know whether a
106# format actually works: the ``unstructured`` loaders import their real parser
107# (and its third-party dependency) lazily inside ``.load()``. So a format like
108# ODT can be advertised as supported, accept an upload, and only then fail with
109# ``ModuleNotFoundError: No module named 'docx'`` (issue #4414).
110#
111# These flags probe the actual runtime dependency so the registry only
112# advertises formats that will really extract text. Missing deps mean the
113# format is omitted and the upload path returns a clear "Unsupported format"
114# instead of a swallowed extraction failure.
115HAS_DOCX_DEP = _module_available("docx") # python-docx: .docx
116HAS_PPTX_DEP = _module_available("pptx") # python-pptx: .pptx
117# Legacy OLE binary office formats (.doc/.ppt) are NOT read by python-docx/
118# python-pptx; unstructured converts them to the modern format by shelling out
119# to LibreOffice (``soffice``). Without that binary they fail at runtime with
120# "soffice command was not found", so only advertise them when it is present.
121HAS_LIBREOFFICE = bool(
122 shutil.which("soffice") or shutil.which("libreoffice")
123) # .doc/.ppt conversion
124# unstructured.partition.xlsx imports msoffcrypto at module load, so openpyxl
125# alone is not enough to extract a modern .xlsx spreadsheet.
126HAS_XLSX_DEP = _module_available("openpyxl") and _module_available(
127 "msoffcrypto"
128) # openpyxl + msoffcrypto: .xlsx
129# Legacy .xls is read by our own XLSLoader (pandas + xlrd), which avoids the
130# msoffcrypto pre-check that crashes on some .xls files in the unstructured path.
131HAS_XLS_DEP = _module_available("xlrd") # xlrd: .xls
132HAS_PANDOC_DEP = _module_available(
133 "pypandoc"
134) # pandoc bridge: epub/rtf/rst/org/odt
135HAS_OCR_DEP = _module_available("pytesseract") or _module_available(
136 "unstructured.pytesseract"
137) # tesseract OCR for image formats (also needs the tesseract binary)
140# Extension to loader mapping
141# Each entry contains:
142# - loader_class: The LangChain loader class
143# - loader_kwargs: Optional kwargs to pass to the loader
144# - requires_path: Whether the loader requires a file path (vs bytes)
145LOADER_REGISTRY: dict = {
146 # PDF
147 ".pdf": {
148 "loader_class": PyPDFLoader,
149 "loader_kwargs": {},
150 },
151 # Plain text, read in the encoding it was written in
152 ".txt": {
153 "loader_class": EncodingAwareTextLoader,
154 "loader_kwargs": {},
155 },
156 # Markdown
157 ".md": {
158 "loader_class": UnstructuredMarkdownLoader,
159 "loader_kwargs": {},
160 },
161 ".markdown": {
162 "loader_class": UnstructuredMarkdownLoader,
163 "loader_kwargs": {},
164 },
165 # CSV (Word/Excel are registered below, gated on their runtime deps)
166 ".csv": {
167 "loader_class": EncodingAwareCSVLoader,
168 "loader_kwargs": {},
169 },
170 # HTML
171 ".html": {
172 "loader_class": UnstructuredHTMLLoader,
173 "loader_kwargs": {},
174 },
175 ".htm": {
176 "loader_class": UnstructuredHTMLLoader,
177 "loader_kwargs": {},
178 },
179}
181# Word documents - .docx needs python-docx; legacy .doc additionally needs
182# LibreOffice to convert the OLE binary to .docx first.
183if HAS_DOCX_DEP:
184 LOADER_REGISTRY[".docx"] = {
185 "loader_class": UnstructuredWordDocumentLoader,
186 "loader_kwargs": {},
187 }
188 if HAS_LIBREOFFICE: 188 ↛ 189line 188 didn't jump to line 189 because the condition on line 188 was never true
189 LOADER_REGISTRY[".doc"] = {
190 "loader_class": UnstructuredWordDocumentLoader,
191 "loader_kwargs": {},
192 }
194# Modern spreadsheets - require openpyxl + msoffcrypto (unstructured path)
195if HAS_XLSX_DEP: 195 ↛ 202line 195 didn't jump to line 202 because the condition on line 195 was always true
196 LOADER_REGISTRY[".xlsx"] = {
197 "loader_class": UnstructuredExcelLoader,
198 "loader_kwargs": {},
199 }
201# Legacy spreadsheets - read directly via our XLSLoader (pandas + xlrd)
202if HAS_XLS_DEP: 202 ↛ 213line 202 didn't jump to line 213 because the condition on line 202 was always true
203 LOADER_REGISTRY[".xls"] = {
204 "loader_class": XLSLoader,
205 "loader_kwargs": {},
206 }
208# ODT - uploads convert odt -> HTML through the bounded pandoc wrapper
209# (utilities/pandoc_conversion), exactly like EPUB/RTF/RST/Org below, so
210# ODT is gated on the same pandoc bridge. python-docx is no longer on
211# the ODT path (unstructured's partition_odt, which needed it, is not
212# used for uploads).
213if HAS_ODT_LOADER and HAS_PANDOC_DEP:
214 LOADER_REGISTRY[".odt"] = {
215 "loader_class": UnstructuredODTLoader,
216 "loader_kwargs": {},
217 }
219# PowerPoint presentations - .pptx needs python-pptx; legacy .ppt additionally
220# needs LibreOffice to convert the OLE binary to .pptx first.
221if HAS_PPTX_DEP:
222 LOADER_REGISTRY[".pptx"] = {
223 "loader_class": UnstructuredPowerPointLoader,
224 "loader_kwargs": {},
225 }
226 if HAS_LIBREOFFICE: 226 ↛ 227line 226 didn't jump to line 227 because the condition on line 226 was never true
227 LOADER_REGISTRY[".ppt"] = {
228 "loader_class": UnstructuredPowerPointLoader,
229 "loader_kwargs": {},
230 }
232# EPUB (ebooks, technical manuals) - requires pandoc
233if HAS_EPUB_LOADER and HAS_PANDOC_DEP:
234 LOADER_REGISTRY[".epub"] = {
235 "loader_class": UnstructuredEPubLoader,
236 "loader_kwargs": {},
237 }
239# RTF (Rich Text Format) - requires pandoc
240if HAS_RTF_LOADER and HAS_PANDOC_DEP:
241 LOADER_REGISTRY[".rtf"] = {
242 "loader_class": UnstructuredRTFLoader,
243 "loader_kwargs": {},
244 }
246# XML (important for USPTO patent data)
247LOADER_REGISTRY[".xml"] = {
248 "loader_class": UnstructuredXMLLoader,
249 "loader_kwargs": {},
250}
252# RST (reStructuredText) - requires pandoc
253if HAS_RST_LOADER and HAS_PANDOC_DEP:
254 LOADER_REGISTRY[".rst"] = {
255 "loader_class": UnstructuredRSTLoader,
256 "loader_kwargs": {},
257 }
259# Org-mode files - requires pandoc
260if HAS_ORG_LOADER and HAS_PANDOC_DEP:
261 LOADER_REGISTRY[".org"] = {
262 "loader_class": UnstructuredOrgModeLoader,
263 "loader_kwargs": {},
264 }
266# Email files. Attachments are NOT partitioned: unstructured's partition_email
267# (and partition_msg, which UnstructuredEmailLoader dispatches to when the
268# content sniffs as an Outlook .msg) default to process_attachments=True and
269# hand every attachment to unstructured.partition.auto.partition. That path
270# bypasses load_from_bytes and its guards: .doc/.ppt attachments reach soffice
271# with no timeout, .rtf/.rst/.org/.odt/.epub reach pypandoc with no timeout or
272# heap cap, and .docx/.pptx/.xlsx reach their parsers without the zip-container
273# guard. Only the message body is indexed.
274LOADER_REGISTRY[".eml"] = {
275 "loader_class": UnstructuredEmailLoader,
276 "loader_kwargs": {"process_attachments": False},
277}
279# TSV (Tab-Separated Values) - the CSV loader with a tab delimiter
280LOADER_REGISTRY[".tsv"] = {
281 "loader_class": EncodingAwareCSVLoader,
282 "loader_kwargs": {"csv_args": {"delimiter": "\t"}},
283}
285# JSON loader using our custom SimpleJSONLoader (no jq dependency)
286LOADER_REGISTRY[".json"] = {
287 "loader_class": SimpleJSONLoader,
288 "loader_kwargs": {},
289}
291# YAML loader using our custom YAMLLoader
292LOADER_REGISTRY[".yaml"] = {
293 "loader_class": YAMLLoader,
294 "loader_kwargs": {},
295}
296LOADER_REGISTRY[".yml"] = {
297 "loader_class": YAMLLoader,
298 "loader_kwargs": {},
299}
301# Images with OCR support - require pytesseract (and the tesseract binary)
302if HAS_IMAGE_LOADER and HAS_OCR_DEP: 302 ↛ 303line 302 didn't jump to line 303 because the condition on line 302 was never true
303 for ext in [".png", ".jpg", ".jpeg", ".tiff", ".tif", ".bmp", ".heic"]:
304 LOADER_REGISTRY[ext] = {
305 "loader_class": UnstructuredImageLoader,
306 "loader_kwargs": {},
307 }
309# Jupyter notebooks (research code, data analysis)
310LOADER_REGISTRY[".ipynb"] = {
311 "loader_class": NotebookLoader,
312 "loader_kwargs": {
313 "include_outputs": True,
314 "remove_newline": True,
315 },
316}
318# Evernote exports
319LOADER_REGISTRY[".enex"] = {
320 "loader_class": EverNoteLoader,
321 "loader_kwargs": {"load_single_document": False},
322}
324# TOML config files
325LOADER_REGISTRY[".toml"] = {
326 "loader_class": TOMLLoader,
327 "loader_kwargs": {},
328}
330# MHTML web archives (saved web pages)
331LOADER_REGISTRY[".mhtml"] = {
332 "loader_class": MHTMLLoader,
333 "loader_kwargs": {},
334}
335LOADER_REGISTRY[".mht"] = {
336 "loader_class": MHTMLLoader,
337 "loader_kwargs": {},
338}
341def get_supported_extensions() -> list[str]:
342 """Get list of all supported file extensions."""
343 return list(LOADER_REGISTRY.keys())
346def is_extension_supported(extension: str) -> bool:
347 """Check if a file extension is supported."""
348 ext = (
349 extension.lower()
350 if extension.startswith(".")
351 else f".{extension.lower()}"
352 )
353 return ext in LOADER_REGISTRY
356def get_loader_for_path(file_path: str | Path) -> Optional[BaseLoader]:
357 """
358 Get an appropriate document loader for a file based on its extension.
360 Args:
361 file_path: Path to the file to load
363 Returns:
364 A LangChain BaseLoader instance, or None if the extension is not supported
366 Warning:
367 This returns the raw registered loader and bypasses every bound
368 that ``bytes_loader.load_from_bytes`` applies: the zip-container
369 decompression guard, the bounded soffice conversion for
370 ``.doc``/``.ppt``, and the bounded pandoc conversion for
371 ``.rtf``/``.epub``/``.rst``/``.org``/``.odt`` (whose unstructured
372 loaders run pandoc with no timeout or memory cap). Never use it
373 on untrusted files; route those through ``load_from_bytes``.
374 """
375 file_path = Path(file_path)
376 extension = file_path.suffix.lower()
378 loader_info = get_loader_class_for_extension(extension)
379 if loader_info is None:
380 logger.warning(
381 f"Unsupported file extension: {extension} for {file_path}"
382 )
383 return None
385 loader_class, loader_kwargs = loader_info
387 try:
388 return loader_class(str(file_path), **loader_kwargs) # type: ignore[no-any-return]
389 except Exception:
390 logger.exception(f"Error creating loader for {file_path}")
391 return None
394def get_loader_class_for_extension(
395 extension: str,
396) -> Optional[tuple[type, dict]]:
397 """
398 Get the loader class and kwargs for an extension.
400 Args:
401 extension: File extension (with or without leading dot)
403 Returns:
404 Tuple of (loader_class, loader_kwargs) or None if not supported
405 """
406 ext = (
407 extension.lower()
408 if extension.startswith(".")
409 else f".{extension.lower()}"
410 )
412 if ext not in LOADER_REGISTRY:
413 return None
415 entry = LOADER_REGISTRY[ext]
416 return entry["loader_class"], entry.get("loader_kwargs", {})