Coverage for src/lilbee/data/ingest/discovery.py: 100%

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, Mapping

10from enum import StrEnum

11from functools import cache

12from pathlib import Path

13from types import MappingProxyType

14from typing import NamedTuple

16from lilbee.core.config import active_config

17from lilbee.core.system import is_ignored_dir

18from lilbee.data.extract.code_chunker import is_code_file

19from lilbee.data.ingest.ignore import IgnoreRules

20from lilbee.data.types import IMAGE_CONTENT_TYPE, PDF_CONTENT_TYPE, ShardId

22log = logging.getLogger(__name__)

24_PDF_MIME = "application/pdf"

26# Base MIME subtypes that name a container of other files.

27_ARCHIVE_SUBTYPES = frozenset(

28 {

29 "7z-compressed",

30 "bzip",

31 "bzip2",

32 "compress",

33 "gtar",

34 "gzip",

35 "lzip",

36 "lzma",

37 "rar",

38 "rar-compressed",

39 "tar",

40 "xz",

41 "zip",

42 "zip-compressed",

43 "zstd",

44 }

45)

46# Prefixes that mark a subtype as unregistered or vendor-specific, not part of its name.

47_SUBTYPE_PREFIXES = ("x-", "vnd.", "prs.")

50class ExclusionReason(StrEnum):

51 """Why discovery refuses a file whose extension xberg could otherwise extract."""

53 VECTOR_GRAPHIC = "vector graphic, not a document"

54 NEEDS_TRANSCRIPTION = "audio or video, needs a transcription model lilbee does not run"

57# Refusals the MIME type cannot express: image/svg+xml is a drawing, not a scan.

58_DENIED_EXTENSIONS: dict[str, ExclusionReason] = {".svg": ExclusionReason.VECTOR_GRAPHIC}

59# Refusals by MIME type: xberg errors on these without a transcription config.

60_DENIED_MIME_PREFIXES: dict[str, ExclusionReason] = {

61 "audio/": ExclusionReason.NEEDS_TRANSCRIPTION,

62 "video/": ExclusionReason.NEEDS_TRANSCRIPTION,

63}

66def _content_type_for(ext: str, mime: str) -> str:

67 """content_type for a xberg format: PDFs and images grouped, others keyed by extension."""

68 if mime == _PDF_MIME:

69 return PDF_CONTENT_TYPE

70 if mime.startswith("image/"):

71 return IMAGE_CONTENT_TYPE

72 return ext.lstrip(".")

75def _normalized_ext(extension: str) -> str:

76 """A xberg format's extension as a lowercase suffix with its leading dot."""

77 ext = extension.lower()

78 return ext if ext.startswith(".") else f".{ext}"

81def _is_archive_mime(mime: str) -> bool:

82 """Whether the base subtype of *mime* names an archive.

84 ``application/epub+zip`` reads as ``epub``; ``application/x-tar`` reads as ``tar``.

85 """

86 subtype = mime.partition("/")[2].partition("+")[0].strip().lower()

87 for prefix in _SUBTYPE_PREFIXES:

88 subtype = subtype.removeprefix(prefix)

89 return subtype in _ARCHIVE_SUBTYPES

92def _denied_mime_reason(mime: str) -> ExclusionReason | None:

93 return next(

94 (reason for prefix, reason in _DENIED_MIME_PREFIXES.items() if mime.startswith(prefix)),

95 None,

96 )

99@cache

100def excluded_extension_reasons() -> Mapping[str, ExclusionReason]:

101 """Extension -> why discovery refuses it, for xberg formats lilbee will not ingest."""

102 from xberg import list_supported_formats

103

104 refused = dict(_DENIED_EXTENSIONS)

105 for fmt in list_supported_formats():

106 reason = _denied_mime_reason(fmt.mime_type)

107 if reason is not None:

108 refused[_normalized_ext(fmt.extension)] = reason

109 return MappingProxyType(refused)

110

111

112@cache

113def archive_content_types() -> frozenset[str]:

114 """content_types of the containers whose members ingest as their own sources.

115

116 Found by the MIME type xberg reports, so ``application/epub+zip`` stays a book.

117 """

118 from xberg import list_supported_formats

119

120 return frozenset(

121 _content_type_for(_normalized_ext(fmt.extension), fmt.mime_type)

122 for fmt in list_supported_formats()

123 if _is_archive_mime(fmt.mime_type)

124 )

125

126

127def member_content_type(path: str, mime: str) -> str:

128 """content_type of an archive member: by MIME type, then extension, then MIME subtype."""

129 return _content_type_for(Path(path).suffix.lower(), mime) or mime.partition("/")[2]

130

131

132@cache

133def supported_extension_map() -> dict[str, str]:

134 """Extension -> content_type for every format lilbee ingests.

135

136 Built from ``xberg.list_supported_formats()`` so lilbee covers the full set

137 without a hand-maintained list, minus the formats ``excluded_extension_reasons``

138 refuses. Source-code files are routed separately (their extensions are absent

139 here), so ``classify_file`` falls through to the code path.

140 """

141 from xberg import list_supported_formats

142

143 excluded = excluded_extension_reasons()

144 out: dict[str, str] = {}

145 for fmt in list_supported_formats():

146 ext = _normalized_ext(fmt.extension)

147 if ext not in excluded:

148 out[ext] = _content_type_for(ext, fmt.mime_type)

149 return out

150

151

152# How often the discovery walk logs progress. The walk runs before the file count

153# is known (it is what produces the count), so it cannot show an ETA; it just

154# proves the run is alive. A large or NFS-backed tree can take minutes to walk,

155# during which the plan pass has not started and nothing else logs.

156_SCAN_LOG_INTERVAL_S = 10.0

157

158

159class _ScanProgress:

160 """Periodic progress for the pre-plan discovery walk.

161

162 Emitted at warning level, not info: the default LILBEE_LOG_LEVEL is WARNING,

163 so an info line would be filtered before any handler and a headless

164 ``lilbee sync`` would show nothing while the tree is walked. Interval-gated,

165 so a fast walk (the common case) stays silent -- the first line appears only

166 once the walk has run longer than the interval.

167 """

168

169 def __init__(self) -> None:

170 self._examined = 0

171 self._matched = 0

172 self._started = time.monotonic()

173 self._last = self._started

174

175 def tick(self, *, matched: bool) -> None:

176 self._examined += 1

177 if matched:

178 self._matched += 1

179 now = time.monotonic()

180 if now - self._last < _SCAN_LOG_INTERVAL_S:

181 return

182 self._last = now

183 elapsed = now - self._started

184 rate = self._examined / elapsed if elapsed > 0 else 0.0

185 log.warning(

186 "Scanning for files: examined %d, matched %d (%.0f files/s, %.0fs elapsed)",

187 self._examined,

188 self._matched,

189 rate,

190 elapsed,

191 )

192

193

194def file_hash(path: Path) -> str:

195 """Compute SHA-256 hex digest of a file."""

196 with open(path, "rb") as f:

197 return hashlib.file_digest(f, "sha256").hexdigest()

198

199

200def classify_file(path: Path) -> str | None:

201 """Classify a file by extension: a xberg content_type, "code", or None.

202

203 Ingestable xberg formats win; source code (not in xberg's set) routes to the

204 code chunker; a refused container and anything else is unsupported.

205 """

206 doc_type = supported_extension_map().get(path.suffix.lower())

207 if doc_type is not None:

208 return doc_type

209 if is_code_file(path):

210 return "code"

211 return None

212

213

214def resolve_source_path(filename: str) -> Path:

215 """Map a stored source key back to the file it tracks on disk.

216

217 A key's first segment is a registered root label when ``add`` recorded that

218 root; the file then lives at ``linked_roots[label]/<rest>`` (or at the root

219 itself for a single-file root, where there is no rest). Every other key

220 belongs to a file lilbee owns under ``documents_dir`` and resolves there.

221 The path is returned whether or not it still exists: a source whose file was

222 moved or deleted keeps its index entry, and the dead path surfaces only when

223 something tries to open it.

224

225 A registered label owns its whole key namespace: if an owned ``documents_dir``

226 subtree of the same top-level name is created after the root is registered,

227 its files resolve to the root, not the owned copy. ``discover_files`` walks

228 the root after the owned tree and so keys the same file identically, keeping

229 resolution and discovery in agreement; ``add`` blocks the reverse collision

230 (registering a label that shadows an existing owned entry).

231 """

232 config = active_config()

233 first, _, rest = filename.partition("/")

234 root = config.linked_roots.get(first)

235 if root is not None:

236 base = Path(root)

237 return base / rest if rest else base

238 return config.documents_dir / filename

239

240

241def resolve_source_root(filename: str) -> tuple[Path, Path] | None:

242 """The walked root and the resolved path for *filename*, or None if none walks it.

243

244 Pairs a source key with the base its patterns are written relative to, so the

245 index can be reconciled against ``.lilbeeignore`` without a second walk. A

246 single-file root is the file the user named, never a tree, so nothing walks

247 it and no pattern applies.

248 """

249 config = active_config()

250 first, _, rest = filename.partition("/")

251 root = config.linked_roots.get(first)

252 if root is None:

253 return config.documents_dir, config.documents_dir / filename

254 if not rest:

255 return None

256 base = Path(root)

257 return base, base / rest

258

259

260def resolve_source_path_checked(filename: str) -> Path | None:

261 """Resolve *filename*, returning None if it escapes its owning root.

262

263 Guards a surface that resolves a caller-supplied source key (the HTTP

264 document-serving endpoint): a key with ``..`` that would climb out of

265 ``documents_dir`` or a registered root is rejected. Keys produced by

266 discovery never contain ``..``; this defends against a crafted request, not

267 stored data.

268 """

269 config = active_config()

270 resolved = resolve_source_path(filename).resolve(strict=False)

271 roots = [

272 config.documents_dir.resolve(),

273 *(Path(root).resolve() for root in config.linked_roots.values()),

274 ]

275 if any(resolved == root or root in resolved.parents for root in roots):

276 return resolved

277 return None

278

279

280class ScannedFile(NamedTuple):

281 """One file the corpus walk kept: its source key, its path, and why it is refused.

282

283 ``excluded`` is None for a file that will be ingested.

284 """

285

286 key: str

287 path: Path

288 excluded: ExclusionReason | None = None

289

290

291class CorpusScan(NamedTuple):

292 """One walk's outcome: the files to ingest, and the refused ones keyed to their reason."""

293

294 files: dict[str, Path]

295 excluded: dict[str, ExclusionReason]

296

297

298def _scan_entry(path: Path, key: str) -> ScannedFile | None:

299 """The walk's verdict for one file, or None when nothing here can be ingested."""

300 reason = excluded_extension_reasons().get(path.suffix.lower())

301 if reason is not None:

302 return ScannedFile(key, path, reason)

303 if classify_file(path) is None:

304 return None

305 return ScannedFile(key, path)

306

307

308def _walk_root(

309 base: Path,

310 label: str | None,

311 ignore_dirs: frozenset[str],

312 progress: _ScanProgress,

313 rules: IgnoreRules,

314) -> Iterator[ScannedFile]:

315 """Yield the files under *base* lilbee knows, keyed relative to it (prefixed by *label*).

316

317 A refused format is yielded with its reason; an unknown format is left out.

318

319 Symlinks are not followed (``followlinks=False``): each root is walked as the

320 real tree it names, so there is no traversal loop and no path can escape the

321 root it was registered under.

322

323 A directory ``.lilbeeignore`` excludes is pruned rather than filtered per

324 file, so an excluded tree costs nothing to skip and no pattern beneath it can

325 re-include a file -- git's rule, holding here because the walk never descends.

326 """

327 for root, dirs, filenames in os.walk(base, topdown=True, followlinks=False):

328 here = Path(root)

329 dirs[:] = [

330 d

331 for d in dirs

332 if not is_ignored_dir(d, ignore_dirs)

333 and not rules.excludes_entry(here / d, base=base, is_dir=True)

334 ]

335 for fname in filenames:

336 if fname.startswith("."):

337 continue

338 path = here / fname

339 if rules.excludes_entry(path, base=base, is_dir=False):

340 progress.tick(matched=False)

341 continue

342 rel = path.relative_to(base).as_posix()

343 entry = _scan_entry(path, f"{label}/{rel}" if label else rel)

344 # tick per file visited, not per match: a skip-heavy tree still walks

345 # slowly and must still show a heartbeat.

346 progress.tick(matched=entry is not None and entry.excluded is None)

347 if entry is not None:

348 yield entry

349

350

351def _walk_corpus(rules: IgnoreRules | None = None) -> Iterator[ScannedFile]:

352 """Yield every file lilbee knows in the owned tree and in each registered root.

353

354 A single-file root is the file the user named at ``add`` time, so no ignore

355 pattern is consulted for it: naming a file is a stronger statement than a

356 pattern that would have swept it up.

357 """

358 config = active_config()

359 progress = _ScanProgress()

360 rules = rules if rules is not None else IgnoreRules.for_corpus()

361 if config.documents_dir.exists():

362 yield from _walk_root(config.documents_dir, None, config.ignore_dirs, progress, rules)

363 for label, root in config.linked_roots.items():

364 root_path = Path(root)

365 if root_path.is_dir():

366 yield from _walk_root(root_path, label, config.ignore_dirs, progress, rules)

367 elif root_path.is_file() and (entry := _scan_entry(root_path, label)) is not None:

368 yield entry

369

370

371def discover_corpus(shard: ShardId | None = None, rules: IgnoreRules | None = None) -> CorpusScan:

372 """Scan the owned documents dir and every registered root into a :class:`CorpusScan`.

373

374 A refused file (see ``excluded_extension_reasons``) lands in ``excluded``, not ``files``.

375 """

376 files: dict[str, Path] = {}

377 excluded: dict[str, ExclusionReason] = {}

378 for entry in _walk_corpus(rules):

379 if shard is not None and not shard.owns(entry.key):

380 continue

381 if entry.excluded is not None:

382 excluded[entry.key] = entry.excluded

383 else:

384 files[entry.key] = entry.path

385 return CorpusScan(files, excluded)

386

387

388def discover_files(

389 shard: ShardId | None = None, rules: IgnoreRules | None = None

390) -> dict[str, Path]:

391 """Scan the owned documents dir and every registered root, return {key: path}.

392

393 Files lilbee owns under ``documents_dir`` (crawl and upload output) are keyed

394 by their path relative to it. Each root ``add`` registered is indexed where it

395 lives: a directory root contributes its files keyed under the root's label; a

396 single-file root contributes one entry keyed by the label alone. A root whose

397 path has since vanished contributes nothing this pass, and its already-indexed

398 sources are left in place (a dead path-link, not a removal).

399

400 A *shard* keeps only the keys that slice owns, so one worker of a multi-GPU

401 ingest holds the paths of its own slice and not the whole corpus.

402

403 A caller that also reconciles the index passes the *rules* it will reconcile

404 with, so the walk and that pass read one set of compiled patterns.

405 """

406 return discover_corpus(shard, rules).files

407

408

409def corpus_has_at_least(count: int) -> bool:

410 """Whether the corpus holds at least *count* ingestable files.

411

412 Stops at the threshold. The answer gates the multi-GPU ingest fan-out, and

413 walking a million-file tree to learn "yes, more than a few thousand" would

414 cost minutes before any work starts.

415 """

416 ingestable = (entry for entry in _walk_corpus() if entry.excluded is None)

417 return any(seen >= count for seen, _ in enumerate(ingestable, start=1))