Coverage for src/lilbee/data/ingest/discovery.py: 100%
99 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +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 Iterator
10from functools import cache
11from pathlib import Path
13from lilbee.core.config import active_config
14from lilbee.core.system import is_ignored_dir
15from lilbee.data.extract.code_chunker import is_code_file
16from lilbee.data.types import IMAGE_CONTENT_TYPE, PDF_CONTENT_TYPE, ShardId
18log = logging.getLogger(__name__)
20_PDF_MIME = "application/pdf"
23def _content_type_for(ext: str, mime: str) -> str:
24 """content_type for a xberg format: PDFs and images grouped, others keyed by extension."""
25 if mime == _PDF_MIME:
26 return PDF_CONTENT_TYPE
27 if mime.startswith("image/"):
28 return IMAGE_CONTENT_TYPE
29 return ext.lstrip(".")
32@cache
33def supported_extension_map() -> dict[str, str]:
34 """Extension -> content_type for every format xberg can extract.
36 Built from ``xberg.list_supported_formats()`` so lilbee covers the full set
37 without a hand-maintained list. Source-code files are routed separately (their
38 extensions are absent here), so ``classify_file`` falls through to the code path.
39 """
40 from xberg import list_supported_formats
42 out: dict[str, str] = {}
43 for fmt in list_supported_formats():
44 ext = (fmt.extension if fmt.extension.startswith(".") else f".{fmt.extension}").lower()
45 out[ext] = _content_type_for(ext, fmt.mime_type)
46 return out
49# How often the discovery walk logs progress. The walk runs before the file count
50# is known (it is what produces the count), so it cannot show an ETA; it just
51# proves the run is alive. A large or NFS-backed tree can take minutes to walk,
52# during which the plan pass has not started and nothing else logs.
53_SCAN_LOG_INTERVAL_S = 10.0
56class _ScanProgress:
57 """Periodic progress for the pre-plan discovery walk.
59 Emitted at warning level, not info: the default LILBEE_LOG_LEVEL is WARNING,
60 so an info line would be filtered before any handler and a headless
61 ``lilbee sync`` would show nothing while the tree is walked. Interval-gated,
62 so a fast walk (the common case) stays silent -- the first line appears only
63 once the walk has run longer than the interval.
64 """
66 def __init__(self) -> None:
67 self._examined = 0
68 self._matched = 0
69 self._started = time.monotonic()
70 self._last = self._started
72 def tick(self, *, matched: bool) -> None:
73 self._examined += 1
74 if matched:
75 self._matched += 1
76 now = time.monotonic()
77 if now - self._last < _SCAN_LOG_INTERVAL_S:
78 return
79 self._last = now
80 elapsed = now - self._started
81 rate = self._examined / elapsed if elapsed > 0 else 0.0
82 log.warning(
83 "Scanning for files: examined %d, matched %d (%.0f files/s, %.0fs elapsed)",
84 self._examined,
85 self._matched,
86 rate,
87 elapsed,
88 )
91def file_hash(path: Path) -> str:
92 """Compute SHA-256 hex digest of a file."""
93 with open(path, "rb") as f:
94 return hashlib.file_digest(f, "sha256").hexdigest()
97def classify_file(path: Path) -> str | None:
98 """Classify a file by extension: a xberg content_type, "code", or None.
100 xberg-extractable formats win; source code (not in xberg's set) routes
101 to the code chunker; anything else is unsupported.
102 """
103 doc_type = supported_extension_map().get(path.suffix.lower())
104 if doc_type is not None:
105 return doc_type
106 if is_code_file(path):
107 return "code"
108 return None
111def resolve_source_path(filename: str) -> Path:
112 """Map a stored source key back to the file it tracks on disk.
114 A key's first segment is a registered root label when ``add`` recorded that
115 root; the file then lives at ``linked_roots[label]/<rest>`` (or at the root
116 itself for a single-file root, where there is no rest). Every other key
117 belongs to a file lilbee owns under ``documents_dir`` and resolves there.
118 The path is returned whether or not it still exists: a source whose file was
119 moved or deleted keeps its index entry, and the dead path surfaces only when
120 something tries to open it.
122 A registered label owns its whole key namespace: if an owned ``documents_dir``
123 subtree of the same top-level name is created after the root is registered,
124 its files resolve to the root, not the owned copy. ``discover_files`` walks
125 the root after the owned tree and so keys the same file identically, keeping
126 resolution and discovery in agreement; ``add`` blocks the reverse collision
127 (registering a label that shadows an existing owned entry).
128 """
129 config = active_config()
130 first, _, rest = filename.partition("/")
131 root = config.linked_roots.get(first)
132 if root is not None:
133 base = Path(root)
134 return base / rest if rest else base
135 return config.documents_dir / filename
138def resolve_source_path_checked(filename: str) -> Path | None:
139 """Resolve *filename*, returning None if it escapes its owning root.
141 Guards a surface that resolves a caller-supplied source key (the HTTP
142 document-serving endpoint): a key with ``..`` that would climb out of
143 ``documents_dir`` or a registered root is rejected. Keys produced by
144 discovery never contain ``..``; this defends against a crafted request, not
145 stored data.
146 """
147 config = active_config()
148 resolved = resolve_source_path(filename).resolve(strict=False)
149 roots = [
150 config.documents_dir.resolve(),
151 *(Path(root).resolve() for root in config.linked_roots.values()),
152 ]
153 if any(resolved == root or root in resolved.parents for root in roots):
154 return resolved
155 return None
158def _walk_root(
159 base: Path,
160 label: str | None,
161 ignore_dirs: frozenset[str],
162 progress: _ScanProgress,
163) -> Iterator[tuple[str, Path]]:
164 """Yield supported files under *base*, keyed relative to it (prefixed by *label*).
166 Symlinks are not followed (``followlinks=False``): each root is walked as the
167 real tree it names, so there is no traversal loop and no path can escape the
168 root it was registered under.
169 """
170 for root, dirs, filenames in os.walk(base, topdown=True, followlinks=False):
171 dirs[:] = [d for d in dirs if not is_ignored_dir(d, ignore_dirs)]
172 for fname in filenames:
173 if fname.startswith("."):
174 continue
175 path = Path(root) / fname
176 content_type = classify_file(path)
177 # tick per file visited, not per match: a skip-heavy tree still walks
178 # slowly and must still show a heartbeat.
179 progress.tick(matched=content_type is not None)
180 if content_type is None:
181 continue
182 rel = path.relative_to(base).as_posix()
183 yield f"{label}/{rel}" if label else rel, path
186def _walk_corpus() -> Iterator[tuple[str, Path]]:
187 """Yield every supported file in the owned tree and in each registered root."""
188 config = active_config()
189 progress = _ScanProgress()
190 if config.documents_dir.exists():
191 yield from _walk_root(config.documents_dir, None, config.ignore_dirs, progress)
192 for label, root in config.linked_roots.items():
193 root_path = Path(root)
194 if root_path.is_dir():
195 yield from _walk_root(root_path, label, config.ignore_dirs, progress)
196 elif root_path.is_file() and classify_file(root_path) is not None:
197 yield label, root_path
200def discover_files(shard: ShardId | None = None) -> dict[str, Path]:
201 """Scan the owned documents dir and every registered root, return {key: path}.
203 Files lilbee owns under ``documents_dir`` (crawl and upload output) are keyed
204 by their path relative to it. Each root ``add`` registered is indexed where it
205 lives: a directory root contributes its files keyed under the root's label; a
206 single-file root contributes one entry keyed by the label alone. A root whose
207 path has since vanished contributes nothing this pass, and its already-indexed
208 sources are left in place (a dead path-link, not a removal).
210 A *shard* keeps only the keys that slice owns, so one worker of a multi-GPU
211 ingest holds the paths of its own slice and not the whole corpus.
212 """
213 return {key: path for key, path in _walk_corpus() if shard is None or shard.owns(key)}
216def corpus_has_at_least(count: int) -> bool:
217 """Whether the corpus holds at least *count* supported files.
219 Stops at the threshold. The answer gates the multi-GPU ingest fan-out, and
220 walking a million-file tree to learn "yes, more than a few thousand" would
221 cost minutes before any work starts.
222 """
223 return any(seen >= count for seen, _ in enumerate(_walk_corpus(), start=1))