Coverage for src/lilbee/app/ingest.py: 100%
182 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""Register external source roots, and remove indexed documents durably."""
3from __future__ import annotations
5import fnmatch
6import logging
7from collections.abc import Generator, Iterable
8from contextlib import contextmanager
9from dataclasses import dataclass, field
10from pathlib import Path
12from lilbee.app.services import get_services
13from lilbee.core import settings
14from lilbee.core.config import active_config
15from lilbee.data.store.types import RemoveResult
18@dataclass
19class RegisterResult:
20 """Result of registering source roots into the knowledge base."""
22 registered: list[str] = field(default_factory=list) # labels newly registered
23 skipped: list[str] = field(default_factory=list)
24 tracked: list[str] = field(default_factory=list)
25 """Named sources the knowledge base already tracks, so nothing was registered.
27 Either the path already lives under ``documents_dir`` or this exact source is
28 already registered under that label. Split from ``skipped`` because there is
29 nothing wrong to report and ``--force`` would change nothing: the sync that
30 follows covers them. ``skipped`` is only what could not be registered -- a
31 label held by a different source, or an overlap that would double-index.
32 """
35def _resolve_label(
36 base: str, roots: dict[str, str], docs_resolved: Path, *, force: bool
37) -> str | None:
38 """Choose the source-key label for a new root, or None when the name is taken.
40 An owned ``documents_dir`` top-level entry of the same name always wins and is
41 never shadowed, even under ``force`` -- a label that shadows it would make
42 resolve_source_path disagree with how discovery keyed the owned file. Reuses
43 the label when a root of that name was registered before but its path has since
44 vanished (the source moved: re-register in place, no ``--force`` needed) or
45 when ``force`` overwrites a live registered root of the same name.
46 """
47 if (docs_resolved / base).exists():
48 return None # an owned entry holds this name; never shadow it
49 existing = roots.get(base)
50 if existing is not None and not Path(existing).exists():
51 return base # dangling root; the source moved, re-point it to the new path
52 if force:
53 return base
54 if base in roots:
55 return None
56 return base
59def _overlaps_existing(src: Path, docs_resolved: Path, roots: dict[str, str]) -> bool:
60 """Whether *src* overlaps ``documents_dir`` or a live registered root.
62 Two roots covering the same tree would walk the same file twice and index it
63 under two keys (double-index). The caller already rejects *src* inside
64 ``documents_dir``; this rejects *src* being an ANCESTOR of it, and *src*
65 nesting under or over any live registered root. A vanished root cannot
66 double-index, so it is ignored.
67 """
68 if docs_resolved.is_relative_to(src):
69 return True
70 for target in roots.values():
71 root = Path(target)
72 if not root.exists():
73 continue
74 root = root.resolve()
75 if src.is_relative_to(root) or root.is_relative_to(src):
76 return True
77 return False
80def source_label_taken(name: str, target: Path | None = None) -> bool:
81 """Whether registering *target* under *name* would collide with a different source.
83 The confirm-before-overwrite affordance in the TUI reads this. The label
84 rule is delegated to :func:`_resolve_label` (the authority register_sources
85 itself applies) rather than mirrored, and a *target* whose exact path is
86 already registered under *name* is not a collision: re-adding the same
87 source is idempotent, matching register_sources' by-target no-op.
88 """
89 config = active_config()
90 roots = dict(config.linked_roots)
91 if target is not None and roots.get(name) == str(target.resolve()):
92 return False
93 return _resolve_label(name, roots, config.documents_dir.resolve(), force=False) is None
96def register_sources(paths: list[Path], *, force: bool = False) -> RegisterResult:
97 """Register each path as a root lilbee indexes where it already lives.
99 A prepared corpus is already on local disk, so ``add`` records where it is
100 rather than copying or linking it: discovery walks the registered root and
101 keys its files under the root's label (its basename). A path already inside
102 ``documents_dir`` is left to the owned-files walk; a path already registered
103 under the same target is a no-op; a label already taken by a different live
104 root or an owned entry is skipped unless ``force``. The registry is persisted
105 so later processes index the same roots.
106 """
107 config = active_config()
108 documents_dir = config.documents_dir
109 documents_dir.mkdir(parents=True, exist_ok=True)
110 docs_resolved = documents_dir.resolve()
111 result = RegisterResult()
112 if not paths:
113 return result
115 def _mutate(persisted: dict[str, str] | None) -> tuple[dict[str, str], RegisterResult]:
116 # Read the registry from config.toml INSIDE the lock (not the possibly
117 # stale in-memory copy) so two processes registering roots concurrently
118 # cannot lose each other's entry.
119 roots = dict(persisted or {})
120 by_target = {target: label for label, target in roots.items()}
121 for p in paths:
122 src = p.resolve()
123 if src == docs_resolved or docs_resolved in src.parents:
124 result.tracked.append(p.name) # already owned by the knowledge base
125 continue
126 already = by_target.get(str(src))
127 if already is not None:
128 result.tracked.append(already) # this exact source is already registered
129 continue
130 if _overlaps_existing(src, docs_resolved, roots):
131 result.skipped.append(p.name) # nests under/over another root; would double-index
132 continue
133 label = _resolve_label(src.name, roots, docs_resolved, force=force)
134 if label is None:
135 result.skipped.append(src.name) # name taken; --force to overwrite
136 continue
137 roots[label] = str(src)
138 by_target[str(src)] = label
139 result.registered.append(label)
140 config.linked_roots = roots # refresh the in-process view (picks up merges)
141 return roots, result
143 result = settings.mutate_value(config.data_root, "linked_roots", _mutate)
144 _unmark_sources_under(paths)
145 return result
148def _unmark_sources_under(paths: list[Path]) -> None:
149 """Drop the skip markers and reasons for every source *paths* covers.
151 A marker exists to stop *discovery* from resurrecting a source the user
152 removed, or from re-paying the extract cost on a file that yielded nothing.
153 Naming the path outranks it: ``add`` is the user asking for that source
154 back, so the marker goes and the sync that follows ingests the file again.
155 Without this a removal would be permanent, undoable only by
156 ``retry-skipped`` or ``rebuild``, neither of which the user has any reason
157 to reach for after typing the path they want.
159 Runs after the registry update so a root this call registered resolves.
160 """
161 from lilbee.data.ingest.skip_marker import (
162 load_skip_markers,
163 load_skip_reasons,
164 write_skip_markers,
165 write_skip_reasons,
166 )
168 config = active_config()
169 markers = load_skip_markers(config.data_root)
170 covered = _markers_covering(markers, paths)
171 if not covered:
172 return
173 write_skip_markers(config.data_root, {k: v for k, v in markers.items() if k not in covered})
174 reasons = load_skip_reasons(config.data_root)
175 write_skip_reasons(config.data_root, {k: v for k, v in reasons.items() if k not in covered})
178def _markers_covering(markers: dict[str, str], paths: list[Path]) -> set[str]:
179 """Marker keys whose file is one of *paths* or lives beneath one.
181 Each key is resolved back to the file it tracks -- the same mapping
182 discovery keyed it by -- so an owned ``documents_dir`` entry, a file under a
183 registered root, and a single-file root are all matched by the one rule
184 instead of three shape-specific ones.
185 """
186 from lilbee.data.ingest.discovery import resolve_source_path
188 named = [p.resolve() for p in paths]
189 covered = set()
190 for name in markers:
191 tracked = resolve_source_path(name).resolve(strict=False)
192 if any(tracked == path or path in tracked.parents for path in named):
193 covered.add(name)
194 return covered
197_REMOVED_SKIP_REASON = "removed via remove (re-add the source or run retry-skipped to restore)"
199_GLOB_CHARS = frozenset("*?[")
202def _is_glob(name: str) -> bool:
203 """Whether *name* should be matched as a glob rather than a literal source."""
204 return any(char in _GLOB_CHARS for char in name)
207def folder_members(name: str, known: Iterable[str]) -> list[str]:
208 """Indexed sources under folder *name*, matched on whole path segments.
210 ``myrepo`` covers ``myrepo/a.py`` but never ``myrepo-2/x``. Empty when *name*
211 is not a parent directory of any known source.
212 """
213 prefix = name.rstrip("/") + "/"
214 return [source for source in known if source.startswith(prefix)]
217def expand_remove_targets(names: list[str], known: list[str] | None = None) -> list[str]:
218 """Expand folder names and glob patterns to the indexed sources they cover.
220 An exact source name is kept. A folder name (a parent directory of indexed
221 sources) expands to every source beneath it. A glob (a name containing
222 ``* ? [``) expands to every source it fnmatches. A name matching none of
223 these is kept unchanged so the caller reports it not-found. Order and
224 de-duplication are preserved. *known* (the indexed source filenames) is read
225 from the store when not supplied; a caller that already has it passes it to
226 avoid a second read.
227 """
228 if known is None:
229 known = [s["filename"] for s in get_services().store.get_sources()]
230 known_set = set(known)
231 expanded: list[str] = []
232 seen: set[str] = set()
234 def _add(candidate: str) -> None:
235 if candidate not in seen:
236 seen.add(candidate)
237 expanded.append(candidate)
239 for name in names:
240 if name in known_set:
241 _add(name)
242 continue
243 if _is_glob(name):
244 matches = [source for source in known if fnmatch.fnmatchcase(source, name)]
245 else:
246 matches = folder_members(name, known)
247 if matches:
248 for match in matches:
249 _add(match)
250 else:
251 _add(name) # not-found; reported by the store
252 return expanded
255def unregister_roots(names: Iterable[str]) -> list[str]:
256 """Un-register any top-level source root named in *names*. Returns removed labels.
258 ``add`` registers a source root; removing it by its label drops the registry
259 entry so discovery stops finding its files, which then need no skip marker.
260 The source bytes on disk are never touched. Nested names (``corpus/a.txt``)
261 are not roots and are left alone.
262 """
263 config = active_config()
264 labels = list(names)
265 removed: list[str] = []
266 if not labels:
267 return removed
269 def _mutate(persisted: dict[str, str] | None) -> tuple[dict[str, str], list[str]]:
270 roots = dict(persisted or {})
271 for name in labels:
272 label = name.strip("/")
273 if "/" in label or label not in roots:
274 continue
275 del roots[label]
276 removed.append(label)
277 config.linked_roots = roots # refresh the in-process view
278 return roots, removed
280 return settings.mutate_value(config.data_root, "linked_roots", _mutate)
283log = logging.getLogger(__name__)
286def remove_documents_durably(names: list[str], targets: list[str] | None = None) -> RemoveResult:
287 """Remove documents from the index (folders and globs expand) and make it stick.
289 Never deletes source bytes. A folder or glob argument expands to every
290 indexed source it covers. Each removed source gets a skip-marker keyed on its
291 current hash so the next sync treats it as unchanged-and-skipped instead of
292 re-ingesting it. Removing a top-level registered root instead un-registers it
293 (discovery then can't re-find its files, so no markers are needed for them).
294 Editing the source (new hash), ``retry-skipped``, or ``rebuild`` restores it.
295 *targets* (the expanded names) is computed when not supplied; a caller that
296 already expanded for a confirmation prompt passes it to avoid re-expanding.
297 """
298 from lilbee.data.ingest.discovery import file_hash, resolve_source_path
299 from lilbee.data.ingest.skip_marker import (
300 load_skip_markers,
301 load_skip_reasons,
302 write_skip_markers,
303 write_skip_reasons,
304 )
306 config = active_config()
307 if targets is None:
308 targets = expand_remove_targets(names)
309 result = get_services().store.remove_documents(targets)
310 if not result.removed:
311 return result
312 unregistered = unregister_roots(names)
313 markers = load_skip_markers(config.data_root)
314 reasons = load_skip_reasons(config.data_root)
315 for name in result.removed:
316 if any(name == root or name.startswith(root + "/") for root in unregistered):
317 continue # the root is gone; discovery won't resurrect these
318 path = resolve_source_path(name)
319 # Imported sources have no file on disk; sync never re-ingests them, so a
320 # marker is only needed for a real file that would otherwise be re-found.
321 if path.exists():
322 markers[name] = file_hash(path)
323 reasons[name] = _REMOVED_SKIP_REASON
324 write_skip_markers(config.data_root, markers)
325 write_skip_reasons(config.data_root, reasons)
326 _forget_removed_from_wiki_index(list(result.removed))
327 return result
330def _forget_removed_from_wiki_index(removed: list[str]) -> None:
331 """Drop removed documents from the wiki's browse index.
333 Their skip markers keep them out of later syncs, so no refresh would ever
334 revisit their entries and the tree would keep offering pages the library
335 can no longer support. Best effort: the removal itself already succeeded.
336 """
337 if not active_config().wiki or not removed:
338 return
339 from lilbee.wiki.stubs import drop_sources_from_index
341 try:
342 drop_sources_from_index(set(removed))
343 except Exception:
344 log.warning("Failed to drop removed documents from the wiki index", exc_info=True)
347@contextmanager
348def temporary_ocr_config(
349 enable_ocr: bool | None = None,
350 ocr_timeout: float | None = None,
351) -> Generator[None, None, None]:
352 """Override OCR config for the duration of the block, per request.
354 Backed by a ContextVar rather than a global ``cfg`` mutation, so concurrent
355 ingests on the shared HTTP daemon do not clobber one another's OCR settings.
356 """
357 from lilbee.data.extract.document import ocr_override
359 with ocr_override(enable_ocr, ocr_timeout):
360 yield