Coverage for src/lilbee/data/ingest/discovery.py: 100%
210 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""File discovery, classification, hashing, and source-path resolution."""
3from __future__ import annotations
5import hashlib
6import logging
7import os
8import time
9from collections.abc import Iterable, Iterator, Mapping
10from enum import StrEnum
11from functools import cache
12from pathlib import Path
13from types import MappingProxyType
14from typing import NamedTuple
16from lilbee.core import settings
17from lilbee.core.config import active_config
18from lilbee.core.system import is_ignored_dir
19from lilbee.data.extract.code_chunker import is_code_file
20from lilbee.data.ingest.ignore import IgnoreRules
21from lilbee.data.types import CODE_CONTENT_TYPE, IMAGE_CONTENT_TYPE, PDF_CONTENT_TYPE, ShardId
23log = logging.getLogger(__name__)
25_PDF_MIME = "application/pdf"
27# Base MIME subtypes that name a container of other files.
28_ARCHIVE_SUBTYPES = frozenset(
29 {
30 "7z-compressed",
31 "bzip",
32 "bzip2",
33 "compress",
34 "gtar",
35 "gzip",
36 "lzip",
37 "lzma",
38 "rar",
39 "rar-compressed",
40 "tar",
41 "xz",
42 "zip",
43 "zip-compressed",
44 "zstd",
45 }
46)
47# Prefixes that mark a subtype as unregistered or vendor-specific, not part of its name.
48_SUBTYPE_PREFIXES = ("x-", "vnd.", "prs.")
51class ExclusionReason(StrEnum):
52 """Why discovery refuses a file whose extension xberg could otherwise extract."""
54 VECTOR_GRAPHIC = "vector graphic, not a document"
55 NEEDS_TRANSCRIPTION = "audio or video, needs a transcription model lilbee does not run"
58# Refusals the MIME type cannot express: image/svg+xml is a drawing, not a scan.
59_DENIED_EXTENSIONS: dict[str, ExclusionReason] = {".svg": ExclusionReason.VECTOR_GRAPHIC}
60# Refusals by MIME type: xberg errors on these without a transcription config.
61_DENIED_MIME_PREFIXES: dict[str, ExclusionReason] = {
62 "audio/": ExclusionReason.NEEDS_TRANSCRIPTION,
63 "video/": ExclusionReason.NEEDS_TRANSCRIPTION,
64}
67def _content_type_for(ext: str, mime: str) -> str:
68 """content_type for a xberg format: PDFs and images grouped, others keyed by extension."""
69 if mime == _PDF_MIME:
70 return PDF_CONTENT_TYPE
71 if mime.startswith("image/"):
72 return IMAGE_CONTENT_TYPE
73 return ext.lstrip(".")
76def _normalized_ext(extension: str) -> str:
77 """A xberg format's extension as a lowercase suffix with its leading dot."""
78 ext = extension.lower()
79 return ext if ext.startswith(".") else f".{ext}"
82def _is_archive_mime(mime: str) -> bool:
83 """Whether the base subtype of *mime* names an archive.
85 ``application/epub+zip`` reads as ``epub``; ``application/x-tar`` reads as ``tar``.
86 """
87 subtype = mime.partition("/")[2].partition("+")[0].strip().lower()
88 for prefix in _SUBTYPE_PREFIXES:
89 subtype = subtype.removeprefix(prefix)
90 return subtype in _ARCHIVE_SUBTYPES
93def _denied_mime_reason(mime: str) -> ExclusionReason | None:
94 return next(
95 (reason for prefix, reason in _DENIED_MIME_PREFIXES.items() if mime.startswith(prefix)),
96 None,
97 )
100@cache
101def excluded_extension_reasons() -> Mapping[str, ExclusionReason]:
102 """Extension -> why discovery refuses it, for xberg formats lilbee will not ingest."""
103 from xberg import list_supported_formats
105 refused = dict(_DENIED_EXTENSIONS)
106 for fmt in list_supported_formats():
107 reason = _denied_mime_reason(fmt.mime_type)
108 if reason is not None:
109 refused[_normalized_ext(fmt.extension)] = reason
110 return MappingProxyType(refused)
113@cache
114def archive_content_types() -> frozenset[str]:
115 """content_types of the containers whose members ingest as their own sources.
117 Found by the MIME type xberg reports, so ``application/epub+zip`` stays a book.
118 """
119 from xberg import list_supported_formats
121 return frozenset(
122 _content_type_for(_normalized_ext(fmt.extension), fmt.mime_type)
123 for fmt in list_supported_formats()
124 if _is_archive_mime(fmt.mime_type)
125 )
128@cache
129def supported_extension_map() -> dict[str, str]:
130 """Extension -> content_type for every format lilbee ingests.
132 Built from ``xberg.list_supported_formats()`` so lilbee covers the full set
133 without a hand-maintained list, minus the formats ``excluded_extension_reasons``
134 refuses. Source-code files are routed separately (their extensions are absent
135 here), so ``classify_file`` falls through to the code path.
136 """
137 from xberg import list_supported_formats
139 excluded = excluded_extension_reasons()
140 out: dict[str, str] = {}
141 for fmt in list_supported_formats():
142 ext = _normalized_ext(fmt.extension)
143 if ext not in excluded:
144 out[ext] = _content_type_for(ext, fmt.mime_type)
145 return out
148# How often the discovery walk logs progress. The walk runs before the file count
149# is known (it is what produces the count), so it cannot show an ETA; it just
150# proves the run is alive. A large or NFS-backed tree can take minutes to walk,
151# during which the plan pass has not started and nothing else logs.
152_SCAN_LOG_INTERVAL_S = 10.0
155class _ScanProgress:
156 """Periodic progress for the pre-plan discovery walk.
158 Emitted at warning level, not info: the default LILBEE_LOG_LEVEL is WARNING,
159 so an info line would be filtered before any handler and a headless
160 ``lilbee sync`` would show nothing while the tree is walked. Interval-gated,
161 so a fast walk (the common case) stays silent -- the first line appears only
162 once the walk has run longer than the interval.
163 """
165 def __init__(self) -> None:
166 self._examined = 0
167 self._matched = 0
168 self._started = time.monotonic()
169 self._last = self._started
170 self.unreadable: list[str] = []
172 def could_not_open(self, error: OSError) -> None:
173 """Record a folder the walk could not list, so the result can name it."""
174 self.unreadable.append(str(error.filename))
176 def tick(self, *, matched: bool) -> None:
177 self._examined += 1
178 if matched:
179 self._matched += 1
180 now = time.monotonic()
181 if now - self._last < _SCAN_LOG_INTERVAL_S:
182 return
183 self._last = now
184 elapsed = now - self._started
185 rate = self._examined / elapsed if elapsed > 0 else 0.0
186 log.warning(
187 "Scanning for files: examined %d, matched %d (%.0f files/s, %.0fs elapsed)",
188 self._examined,
189 self._matched,
190 rate,
191 elapsed,
192 )
195def file_hash(path: Path) -> str:
196 """Compute SHA-256 hex digest of a file."""
197 with open(path, "rb") as f:
198 return hashlib.file_digest(f, "sha256").hexdigest()
201def classify_file(path: Path) -> str | None:
202 """Classify a file by extension: a xberg content_type, "code", or None.
204 Ingestable xberg formats win; source code (not in xberg's set) routes to the
205 code chunker; a refused container and anything else is unsupported.
206 """
207 doc_type = supported_extension_map().get(path.suffix.lower())
208 if doc_type is not None:
209 return doc_type
210 if is_code_file(path):
211 return CODE_CONTENT_TYPE
212 return None
215def resolve_source_path(filename: str) -> Path:
216 """Map a stored source key back to the file it tracks on disk.
218 A key's first segment is a registered root label when ``add`` recorded that
219 root; the file then lives at ``linked_roots[label]/<rest>`` (or at the root
220 itself for a single-file root, where there is no rest). Every other key
221 belongs to a file lilbee owns under ``documents_dir`` and resolves there.
222 The path is returned whether or not it still exists: a source whose file was
223 moved or deleted keeps its index entry, and the dead path surfaces only when
224 something tries to open it.
226 A registered label owns its whole key namespace: if an owned ``documents_dir``
227 subtree of the same top-level name is created after the root is registered,
228 its files resolve to the root, not the owned copy. ``discover_files`` walks
229 the root after the owned tree and so keys the same file identically, keeping
230 resolution and discovery in agreement; ``add`` blocks the reverse collision
231 (registering a label that shadows an existing owned entry).
232 """
233 config = active_config()
234 settings.refresh_registry(config)
235 first, _, rest = filename.partition("/")
236 root = config.linked_roots.get(first)
237 if root is not None:
238 base = Path(root)
239 return base / rest if rest else base
240 return config.documents_dir / filename
243def resolve_source_root(filename: str) -> tuple[Path, Path] | None:
244 """The walked root and the resolved path for *filename*, or None if none walks it.
246 Pairs a source key with the base its patterns are written relative to, so the
247 index can be reconciled against ``.lilbeeignore`` without a second walk. A
248 single-file root is the file the user named, never a tree, so nothing walks
249 it and no pattern applies.
250 """
251 config = active_config()
252 settings.refresh_registry(config)
253 first, _, rest = filename.partition("/")
254 root = config.linked_roots.get(first)
255 if root is None:
256 return config.documents_dir, config.documents_dir / filename
257 if not rest:
258 return None
259 base = Path(root)
260 return base, base / rest
263def resolve_source_path_checked(filename: str) -> Path | None:
264 """Resolve *filename*, returning None if it escapes its owning root.
266 Guards a surface that resolves a caller-supplied source key (the HTTP
267 document-serving endpoint): a key with ``..`` that would climb out of
268 ``documents_dir`` or a registered root is rejected. Keys produced by
269 discovery never contain ``..``; this defends against a crafted request, not
270 stored data.
271 """
272 config = active_config()
273 with settings.fresh_registry_pinned(config):
274 resolved = resolve_source_path(filename).resolve(strict=False)
275 roots = [
276 config.documents_dir.resolve(),
277 *(Path(root).resolve() for root in config.linked_roots.values()),
278 ]
279 if any(resolved == root or root in resolved.parents for root in roots):
280 return resolved
281 return None
284class ScannedFile(NamedTuple):
285 """One file the corpus walk kept: its source key, its path, and why it is refused.
287 ``excluded`` is None for a file that will be ingested.
288 """
290 key: str
291 path: Path
292 excluded: ExclusionReason | None = None
295class CorpusScan(NamedTuple):
296 """One walk's outcome: the files to ingest, the refused ones, the folders it could not open."""
298 files: dict[str, Path]
299 excluded: dict[str, ExclusionReason]
300 unreadable: tuple[str, ...] = ()
303def _scan_entry(path: Path, key: str) -> ScannedFile | None:
304 """The walk's verdict for one file, or None when nothing here can be ingested."""
305 reason = excluded_extension_reasons().get(path.suffix.lower())
306 if reason is not None:
307 return ScannedFile(key, path, reason)
308 if classify_file(path) is None:
309 return None
310 return ScannedFile(key, path)
313def _walk_root(
314 base: Path,
315 label: str | None,
316 ignore_dirs: frozenset[str],
317 progress: _ScanProgress,
318 rules: IgnoreRules,
319 claimed: frozenset[Path] = frozenset(),
320) -> Iterator[ScannedFile]:
321 """Yield the files under *base* lilbee knows, keyed relative to it (prefixed by *label*).
323 A refused format is yielded with its reason; an unknown format is left out.
324 A path in *claimed* is another registered root: its own entry walks it, so
325 no file has two keys.
327 Symlinks are not followed (``followlinks=False``): each root is walked as the
328 real tree it names, so there is no traversal loop and no path can escape the
329 root it was registered under.
331 A directory ``.lilbeeignore`` excludes is pruned rather than filtered per
332 file, so an excluded tree costs nothing to skip and no pattern beneath it can
333 re-include a file -- git's rule, holding here because the walk never descends.
334 """
335 walk = os.walk(base, topdown=True, onerror=progress.could_not_open, followlinks=False)
336 for root, dirs, filenames in walk:
337 here = Path(root)
338 dirs[:] = [
339 d
340 for d in dirs
341 if not is_ignored_dir(d, ignore_dirs)
342 and not (claimed and here / d in claimed)
343 and not rules.excludes_entry(here / d, base=base, is_dir=True)
344 ]
345 for fname in filenames:
346 if fname.startswith("."):
347 continue
348 path = here / fname
349 if claimed and path in claimed:
350 continue
351 if rules.excludes_entry(path, base=base, is_dir=False):
352 progress.tick(matched=False)
353 continue
354 rel = path.relative_to(base).as_posix()
355 entry = _scan_entry(path, f"{label}/{rel}" if label else rel)
356 # tick per file visited, not per match: a skip-heavy tree still walks
357 # slowly and must still show a heartbeat.
358 progress.tick(matched=entry is not None and entry.excluded is None)
359 if entry is not None:
360 yield entry
363def walk_reaches(base: Path, target: Path, rules: IgnoreRules) -> bool:
364 """Whether the walk of *base* reaches *target*, a file or a directory below it.
366 Applies the checks ``_walk_root`` makes on the way down: a pruned directory
367 name, a dot-file, and what ``.lilbeeignore`` excludes. A path that does not
368 exist is read as a directory.
369 """
370 ignore_dirs = active_config().ignore_dirs
371 is_file = target.is_file()
372 parts = target.relative_to(base).parts
373 current = base
374 for index, part in enumerate(parts, start=1):
375 current = current / part
376 is_dir = not (is_file and index == len(parts))
377 hidden = is_ignored_dir(part, ignore_dirs) if is_dir else part.startswith(".")
378 if hidden or rules.excludes_entry(current, base=base, is_dir=is_dir):
379 return False
380 return True
383def _roots_below(base: Path, roots: Iterable[str]) -> frozenset[Path]:
384 """Each of *roots* that lies below *base*, as the path the walk of *base* meets it at."""
385 resolved = base.resolve()
386 return frozenset(
387 base / path.relative_to(resolved)
388 for path in (Path(root).resolve() for root in roots)
389 if path != resolved and path.is_relative_to(resolved)
390 )
393def _walk_corpus(
394 rules: IgnoreRules | None = None, progress: _ScanProgress | None = None
395) -> Iterator[ScannedFile]:
396 """Yield every file lilbee knows in the owned tree and in each registered root.
398 A single-file root is the file the user named at ``add`` time, so no ignore
399 pattern is consulted for it: naming a file is a stronger statement than a
400 pattern that would have swept it up. A root below another root is walked by
401 its own entry alone.
402 """
403 config = active_config()
404 progress = progress if progress is not None else _ScanProgress()
405 rules = rules if rules is not None else IgnoreRules.for_corpus(config.data_root)
406 targets = list(config.linked_roots.values())
407 if config.documents_dir.exists():
408 yield from _walk_root(config.documents_dir, None, config.ignore_dirs, progress, rules)
409 for label, root in config.linked_roots.items():
410 root_path = Path(root)
411 if root_path.is_dir():
412 claimed = _roots_below(root_path, targets)
413 yield from _walk_root(root_path, label, config.ignore_dirs, progress, rules, claimed)
414 elif root_path.is_file() and (entry := _scan_entry(root_path, label)) is not None:
415 yield entry
418def _as_scan(entries: Iterable[ScannedFile], progress: _ScanProgress) -> CorpusScan:
419 """Split walked entries into the files to ingest and the refused ones."""
420 files: dict[str, Path] = {}
421 excluded: dict[str, ExclusionReason] = {}
422 for entry in entries:
423 if entry.excluded is not None:
424 excluded[entry.key] = entry.excluded
425 else:
426 files[entry.key] = entry.path
427 return CorpusScan(files, excluded, tuple(progress.unreadable))
430def discover_corpus(shard: ShardId | None = None, rules: IgnoreRules | None = None) -> CorpusScan:
431 """Scan the owned documents dir and every registered root into a :class:`CorpusScan`.
433 A refused file (see ``excluded_extension_reasons``) lands in ``excluded``, not ``files``.
434 """
435 progress = _ScanProgress()
436 entries = _walk_corpus(rules, progress)
437 return _as_scan((e for e in entries if shard is None or shard.owns(e.key)), progress)
440def discover_dir(directory: Path) -> CorpusScan:
441 """Scan *directory* as ``add`` would walk it, with the ``.lilbeeignore`` rules applied."""
442 config = active_config()
443 rules = IgnoreRules.for_corpus(config.data_root)
444 ignore_dirs = config.ignore_dirs
445 progress = _ScanProgress()
446 return _as_scan(_walk_root(directory, None, ignore_dirs, progress, rules), progress)
449def discover_files(
450 shard: ShardId | None = None, rules: IgnoreRules | None = None
451) -> dict[str, Path]:
452 """Scan the owned documents dir and every registered root, return {key: path}.
454 Files lilbee owns under ``documents_dir`` (crawl and upload output) are keyed
455 by their path relative to it. Each root ``add`` registered is indexed where it
456 lives: a directory root contributes its files keyed under the root's label; a
457 single-file root contributes one entry keyed by the label alone. A root whose
458 path has since vanished contributes nothing this pass, and its already-indexed
459 sources are left in place (a dead path-link, not a removal).
461 A *shard* keeps only the keys that slice owns, so one worker of a multi-GPU
462 ingest holds the paths of its own slice and not the whole corpus.
464 A caller that also reconciles the index passes the *rules* it will reconcile
465 with, so the walk and that pass read one set of compiled patterns.
466 """
467 return discover_corpus(shard, rules).files
470def corpus_has_at_least(count: int) -> bool:
471 """Whether the corpus holds at least *count* ingestable files.
473 Stops at the threshold. The answer gates the multi-GPU ingest fan-out, and
474 walking a million-file tree to learn "yes, more than a few thousand" would
475 cost minutes before any work starts.
476 """
477 ingestable = (entry for entry in _walk_corpus() if entry.excluded is None)
478 return any(seen >= count for seen, _ in enumerate(ingestable, start=1))