Coverage for src/lilbee/data/ingest/pipeline.py: 100%
1003 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""Top-level sync orchestration: discovery, dispatch, batching, post-sync hooks."""
3from __future__ import annotations
5import asyncio
6import contextlib
7import functools
8import logging
9import os
10import threading
11import time
12from collections import deque
13from collections.abc import (
14 AsyncGenerator,
15 Callable,
16 Collection,
17 Coroutine,
18 Iterable,
19 Iterator,
20 Mapping,
21)
22from concurrent.futures import Future, ThreadPoolExecutor, as_completed
23from dataclasses import dataclass, field
24from itertools import count
25from pathlib import Path
26from typing import Any, ParamSpec, TypeVar, cast
28from rich.progress import (
29 BarColumn,
30 MofNCompleteColumn,
31 Progress,
32 SpinnerColumn,
33 TimeElapsedColumn,
34)
36from lilbee.app.services import get_services
37from lilbee.core import settings
38from lilbee.core.config import Config, active_config
39from lilbee.core.web_names import is_web_source
40from lilbee.data.extract.chunk import ChunkLimitError
41from lilbee.data.extract.document import (
42 extract_batching,
43 ingest_archive,
44 ingest_document,
45 ingest_markdown,
46 ocr_backend,
47 warn_if_table_model_ignored,
48)
49from lilbee.data.extract.trace import configure_from_env as configure_trace_from_env
50from lilbee.data.ingest.adaptive import (
51 AdaptiveController,
52 ResizableGate,
53 enumerate_fleet_devices,
54 make_signal_sampler,
55 profile_for,
56 resolve_mode,
57)
58from lilbee.data.ingest.code import ingest_code_sync
59from lilbee.data.ingest.discovery import (
60 ExclusionReason,
61 archive_content_types,
62 classify_file,
63 discover_corpus,
64 discover_files,
65 file_hash,
66 resolve_source_path,
67 resolve_source_root,
68)
69from lilbee.data.ingest.errors import error_reason
70from lilbee.data.ingest.fanout import (
71 WORKER_LOG_NAME,
72 ShardDone,
73 ShardOptions,
74 ShardSpec,
75 aggregate_results,
76 plan_fanout,
77 run_workers,
78)
79from lilbee.data.ingest.ignore import IgnoreRules
80from lilbee.data.offload import (
81 embed_inflight_target,
82 ingest_thread_future,
83 max_workers,
84 owns_ingest_pool,
85 to_executor,
86 to_ingest_thread,
87)
88from lilbee.data.store import (
89 SOURCE_STAT_UNKNOWN,
90 ChunkWrite,
91 ConceptRecords,
92 HeldOut,
93 HeldRecord,
94 IndexMismatch,
95 PageTextRecord,
96 SourceMeta,
97 SourceMove,
98 SourceRecord,
99 SourceStat,
100 SourceStatBackfill,
101 SourceType,
102 Store,
103 source_stat,
104)
105from lilbee.data.title import derive_title
106from lilbee.data.types import (
107 CODE_CONTENT_TYPE,
108 ChunkRecord,
109 DocumentRecords,
110 FileChangePlan,
111 FileToProcess,
112 MemberRecords,
113 OcrReport,
114 ShardId,
115 SkippedSource,
116 SyncResult,
117 _IngestResult,
118)
119from lilbee.providers.base import ProviderError
120from lilbee.runtime.asyncio_loop import is_executor_shutdown
121from lilbee.runtime.cancellation import CancelSignal, TaskCancelledError
122from lilbee.runtime.console import PlainConsole
123from lilbee.runtime.cpu import available_cpu_count, cpu_quota
124from lilbee.runtime.lock import (
125 LockTimeoutError,
126 SyncTurnOwner,
127 sync_running,
128 sync_turn,
129 sync_wait_message,
130)
131from lilbee.runtime.progress import (
132 BatchProgressEvent,
133 BatchStatus,
134 DetailedProgressCallback,
135 EmbedEvent,
136 EventType,
137 ExtractEvent,
138 FileDoneEvent,
139 FileStartEvent,
140 OcrBackendUsed,
141 OcrStartEvent,
142 ProgressEvent,
143 SyncDoneEvent,
144 SyncWaitEvent,
145 noop_callback,
146)
147from lilbee.runtime.progress.columns import literal_text_column
149log = logging.getLogger(__name__)
152def _max_concurrent() -> int:
153 """Files allowed in their compute phase at once.
155 ``cpu_quota()`` (cpu_count // 2) keeps worker storms from starving the TUI's asyncio
156 main thread, and is the cap for text/code ingest. Vision OCR is different: every file
157 in compute holds a continuous-batching slot on a vision server, so an OCR run is bounded
158 by the vision slot capacity. Once the fleet is up that capacity is the servers' real
159 fitted ``--parallel`` slots (a memory-constrained card fits fewer than requested); until
160 then it is estimated from ``replicas x per-server pages``. Sizing to real capacity keeps
161 the OCR queue shallow instead of piling pages behind the dispatcher with their deadlines
162 ticking, and still keeps every slot fed.
163 """
164 from lilbee.providers.fleet.replicas import gpu_device_count, resolve_replica_count
165 from lilbee.providers.roles import WorkerRole
167 config = active_config()
168 if ocr_backend() is OcrBackendUsed.VISION:
169 fitted = get_services().provider.vision_slot_capacity()
170 if fitted is not None:
171 return fitted
172 replicas = resolve_replica_count(WorkerRole.VISION, gpu_device_count())
173 return max(1, replicas * config.vision_ocr_concurrency)
174 if config.ingest_max_inflight > 0:
175 return config.ingest_max_inflight # explicit override
176 # Auto: keep every embed replica fed. The CPU-bound quota alone leaves a
177 # many-core multi-GPU box starved (~4 files/card), so scale admission with
178 # the detected fleet size -- no manual cap needed.
179 return max(cpu_quota(), embed_inflight_target())
182async def _rebuild_concept_clusters(
183 added: Collection[str] = (),
184 updated: Collection[str] = (),
185 removed: Collection[str] = (),
186 *,
187 full: bool = False,
188) -> None:
189 """Re-run Leiden clustering after sync. No-op if disabled.
191 Removals force a full pass: their rows are already gone, so the touched
192 communities are unrecoverable, and so does a rebuild, which re-ingests every
193 source. Adds and updates recluster incrementally.
194 """
195 if not active_config().concept_graph:
196 return
197 from lilbee.retrieval.concepts import concepts_available
199 if not concepts_available():
200 return
201 try:
202 cg = get_services().concepts
203 if not cg.get_graph():
204 return
205 if removed or full:
206 await to_ingest_thread(cg.rebuild_clusters)
207 else:
208 await to_ingest_thread(cg.rebuild_clusters, set(added), set(updated))
209 except Exception:
210 log.warning("Concept cluster rebuild failed", exc_info=True)
213async def build_entity_records(records: list[ChunkRecord], source_name: str) -> list[dict] | None:
214 """Extract typed entities for ingested chunks. None when the mode is off.
216 Gated twice: the ``entity_extraction`` config flag, and a schema already
217 induced into the index; absent either, syncs cost nothing. Extraction
218 failures degrade to no rows for the file, mirroring concept extraction.
219 """
220 config = active_config()
221 if not config.entity_extraction or not records:
222 return None
223 from lilbee.retrieval.entities import ExtractorKind, extract_entities, load_schema
225 schema = load_schema(get_services().store)
226 if schema is None:
227 return None
228 nlp = None
229 if any(t.kind is ExtractorKind.SPACY for t in schema.types):
230 from lilbee.retrieval.concepts import concepts_available
231 from lilbee.retrieval.concepts.nlp import load_spacy_pipeline
233 if concepts_available():
234 try:
235 nlp = load_spacy_pipeline()
236 except ImportError:
237 log.warning("spaCy model unavailable; spacy-kind entity types skipped")
238 provider = None
239 if any(t.kind is ExtractorKind.LLM for t in schema.types):
240 provider = get_services().provider
241 try:
242 return await to_ingest_thread(
243 extract_entities,
244 cast("list[Mapping[str, Any]]", records),
245 schema,
246 provider=provider,
247 nlp=nlp,
248 )
249 except Exception:
250 log.warning("Entity extraction failed for %s", source_name, exc_info=True)
251 return None
254async def build_concept_records(
255 records: list[ChunkRecord], source_name: str
256) -> ConceptRecords | None:
257 """Extract concepts for ingested chunks and build their table rows. None if disabled.
259 Pure record building, no store access: the rows are buffered on the file's
260 ingest result and written once per flush (see :func:`_flush_concept_records`),
261 so a large sync pays one concept-table write per flush, not per file.
262 """
263 if not active_config().concept_graph or not records:
264 return None
265 from lilbee.retrieval.concepts import concepts_available
267 if not concepts_available():
268 return None
269 try:
270 cg = get_services().concepts
271 texts = [r["chunk"] for r in records]
272 concept_lists = await to_ingest_thread(cg.extract_concepts_batch, texts)
273 chunk_ids = [(source_name, r["chunk_index"]) for r in records]
274 return await to_ingest_thread(cg.build_concept_records, chunk_ids, concept_lists)
275 except Exception:
276 log.warning("Concept extraction failed for %s", source_name, exc_info=True)
277 return None
280async def produce_records(
281 path: Path,
282 source_name: str,
283 content_type: str,
284 *,
285 quiet: bool = False,
286 on_progress: DetailedProgressCallback = noop_callback,
287 page_texts_out: list[PageTextRecord] | None = None,
288 cancel: CancelSignal | None = None,
289) -> DocumentRecords:
290 """Extract, chunk, and embed a single file into its records, metadata and OCR report.
292 The LanceDB write is deferred: records are returned to the caller and written
293 in a batched flush (see :func:`_flush_writes`), so bulk ingest pays one
294 write-lock acquisition per batch instead of one per file. The per-page text
295 dataset rows land in ``page_texts_out`` and are written by the same flush.
296 The returned metadata (extraction-provided when available, stem-derived title
297 otherwise) stamps every record's ``title`` and updates the source row. The OCR
298 report is ``None`` for code and markdown, which never reach OCR.
299 """
300 records: list[ChunkRecord]
301 ocr: OcrReport | None = None
302 page_texts: list[PageTextRecord] = page_texts_out if page_texts_out is not None else []
303 if content_type == CODE_CONTENT_TYPE:
304 records = await to_ingest_thread(ingest_code_sync, path, source_name, on_progress)
305 meta = SourceMeta(title=derive_title(source_name))
306 elif path.suffix.lower() == ".md":
307 records, meta = await ingest_markdown(
308 path, source_name, on_progress, page_texts_out=page_texts
309 )
310 else:
311 records, meta, ocr = await ingest_document(
312 path,
313 source_name,
314 content_type,
315 quiet=quiet,
316 on_progress=on_progress,
317 page_texts_out=page_texts,
318 cancel=cancel,
319 )
321 for record in records:
322 # NULL (not "") for an absent title, so chunk rows match the migration
323 # and the _sources table, which both persist absence as NULL.
324 record["title"] = meta.title or None
325 return DocumentRecords(records, meta, ocr)
328def _disk_stat(path: Path) -> SourceStat | None:
329 """Current size/mtime of *path* stamped with now, or None when it cannot be stat'd."""
330 try:
331 st = path.stat()
332 except OSError:
333 return None
334 return SourceStat(st.st_size, st.st_mtime_ns, time.time_ns())
337def _stat_unchanged(stored: SourceStat, current: SourceStat) -> bool:
338 """Whether the stored stat proves the file unchanged without hashing it.
340 Git-style racily-clean guard: a matching (size, mtime) only counts when the
341 mtime is strictly older than the time the stat was recorded; a same-size
342 edit landing in the same mtime tick is otherwise missed forever.
343 """
344 if (stored.size_bytes, stored.mtime_ns) != (current.size_bytes, current.mtime_ns):
345 return False
346 # Unknown capture or mtime >= capture hashes anyway; clock skew can only widen
347 # hashing, never widen skipping past the pre-existing same-tick window.
348 return stored.captured_ns != SOURCE_STAT_UNKNOWN and current.mtime_ns < stored.captured_ns
351@dataclass(frozen=True)
352class _FileChangeVerdict:
353 """One file's sync verdict: process it, hold it out on its record, or unchanged.
355 A held file is not in the index, so it is never counted as unchanged.
356 """
358 to_process: FileToProcess | None = None
359 backfill: SourceStatBackfill | None = None
360 is_update: bool = False
361 held: bool = False
364def _page_titles() -> dict[str, str]:
365 """The title of each crawled page, read once for a whole plan pass."""
366 return get_services().store.crawl_pages.titles()
369def _title_is_stale(name: str, record: SourceRecord, titles: Mapping[str, str]) -> bool:
370 """Whether the crawl page table gives the crawled page *name* a title its source row lacks."""
371 if not is_web_source(name):
372 return False
373 title = titles.get(name)
374 return title is not None and record.get("title") != title
377def _classify_file_change(
378 name: str,
379 path: Path,
380 record: SourceRecord | None,
381 held_hashes: dict[str, str],
382 titles: Mapping[str, str],
383) -> _FileChangeVerdict:
384 """Decide one file's verdict: stat-unchanged, hash-unchanged, held out, or process."""
385 content_type = classify_file(path)
386 if content_type is None:
387 raise ValueError(f"Unsupported file slipped through discovery: {name}")
388 stored_stat = source_stat(record) if record is not None else None
389 current_stat = _disk_stat(path)
390 stale = record is not None and _title_is_stale(name, record, titles)
391 if (
392 record is not None
393 and not stale
394 and stored_stat is not None
395 and current_stat is not None
396 and _stat_unchanged(stored_stat, current_stat)
397 ):
398 return _FileChangeVerdict()
399 old_hash = record["file_hash"] if record is not None else None
400 current_hash = file_hash(path)
401 if old_hash == current_hash and not stale:
402 # Content verified unchanged; persist the stat pair so the next
403 # sync skips the hash entirely.
404 backfill = (
405 SourceStatBackfill(record, current_stat)
406 if record is not None and current_stat is not None
407 else None
408 )
409 return _FileChangeVerdict(backfill=backfill)
410 if held_hashes.get(name) == current_hash:
411 # Failed last sync at this exact hash; skip the retry.
412 return _FileChangeVerdict(held=True)
413 # needs_cleanup=True unconditionally: delete_by_source is idempotent,
414 # and this closes the race where a prior ingest wrote chunks but died
415 # before upsert_source, leaving orphaned chunks that would duplicate.
416 return _FileChangeVerdict(
417 to_process=FileToProcess(
418 name, path, content_type, current_hash, needs_cleanup=True, stat=current_stat
419 ),
420 is_update=old_hash is not None,
421 )
424def _plan_workers() -> int:
425 """Worker count for the parallel planning pass: config override, else auto.
427 ``config.ingest_workers`` (also set per run by ``add --max-cpus``) wins when
428 positive; otherwise size to the container-aware CPU budget so a big corpus
429 hashes on every core the pod actually has, not the host's vCPU count.
430 """
431 configured = active_config().ingest_workers
432 return configured if configured > 0 else available_cpu_count()
435# How often the plan pass logs progress. The pass can run for tens of minutes on
436# a multi-million-file corpus while the Rich bar renders nothing without a TTY
437# and stdout is block-buffered when piped; a periodic line (which logging flushes
438# per record) keeps a headless run observable instead of looking hung.
439_PLAN_LOG_INTERVAL_S = 10.0
442class _PlanProgress:
443 """Periodic progress for the plan/hash pass, with rate and ETA.
445 Emitted at warning level, not info: the default LILBEE_LOG_LEVEL is WARNING,
446 so an info line would be filtered before any handler and a headless
447 ``lilbee sync`` would show nothing during the plan pass and still look hung.
448 """
450 def __init__(self, total: int) -> None:
451 self._total = total
452 self._done = 0
453 self._started = time.monotonic()
454 self._last = self._started
456 def tick(self) -> None:
457 self._done += 1
458 now = time.monotonic()
459 if now - self._last < _PLAN_LOG_INTERVAL_S:
460 return
461 self._last = now
462 elapsed = now - self._started
463 rate = self._done / elapsed if elapsed > 0 else 0.0
464 remaining = (self._total - self._done) / rate if rate > 0 else 0.0
465 log.warning(
466 "Planning: examined %d/%d files (%.0f%%, %.0f files/s, ~%.0fs left)",
467 self._done,
468 self._total,
469 100.0 * self._done / self._total,
470 rate,
471 remaining,
472 )
475class _StreamStop:
476 """Stop signal for a streamed plan: the caller's cancel, or the stream closing.
478 Closing the stream has to reach the plan batch in flight, not just the next one.
479 Build-vs-buy: the hashers run in a thread pool, so the flag must be a
480 thread-visible ``threading.Event``; ``anyio.CancelScope`` is async-only.
481 """
483 def __init__(self, cancel: CancelSignal | None) -> None:
484 self._cancel = cancel
485 self._closed = threading.Event()
487 def close(self) -> None:
488 self._closed.set()
490 def is_set(self) -> bool:
491 return self._closed.is_set() or (self._cancel is not None and self._cancel.is_set())
494def _classify_pooled(
495 pool: ThreadPoolExecutor,
496 items: list[tuple[str, Path]],
497 classify: Callable[[str, Path], _FileChangeVerdict],
498 cancel: CancelSignal | None,
499 progress: _PlanProgress,
500) -> dict[str, _FileChangeVerdict]:
501 """Fan *items* across *pool*, returning name -> verdict for what completed."""
502 verdicts: dict[str, _FileChangeVerdict] = {}
503 futures: dict[Future[_FileChangeVerdict], str] = {}
504 for name, path in items:
505 if cancel and cancel.is_set():
506 break
507 futures[pool.submit(classify, name, path)] = name
508 for future in as_completed(futures):
509 if cancel and cancel.is_set():
510 # Drop queued-but-unstarted work; running tasks drain on their own.
511 for pending in futures:
512 pending.cancel()
513 break
514 verdicts[futures[future]] = future.result()
515 progress.tick()
516 return verdicts
519def _classify_changes(
520 items: list[tuple[str, Path]],
521 existing_sources: dict[str, SourceRecord],
522 held_hashes: dict[str, str],
523 titles: Mapping[str, str],
524 cancel: CancelSignal | None,
525 *,
526 progress: _PlanProgress | None = None,
527 pool: ThreadPoolExecutor | None = None,
528) -> dict[str, _FileChangeVerdict]:
529 """Classify each file (stat + hash) by name, fanning across a thread pool.
531 Independent per file and side-effect-free, so pooled classification matches a
532 serial pass; ``hashlib`` releases the GIL during digest, giving real speedup
533 on a large corpus. Returns name -> verdict. A set ``cancel`` stops promptly:
534 submission halts and queued-but-unstarted work is cancelled, so a mid-pass
535 cancel over a huge corpus does not hash every remaining file. A *pool* and
536 *progress* passed in are shared across the batches of a streamed plan, so the
537 ETA covers the whole corpus and the workers are spun up once.
538 """
540 def _classify(name: str, path: Path) -> _FileChangeVerdict:
541 return _classify_file_change(name, path, existing_sources.get(name), held_hashes, titles)
543 total = len(items)
544 progress = progress or _PlanProgress(total)
545 if pool is not None:
546 return _classify_pooled(pool, items, _classify, cancel, progress)
547 workers = _plan_workers()
548 if workers <= 1 or total <= 1:
549 verdicts: dict[str, _FileChangeVerdict] = {}
550 for name, path in items:
551 if cancel and cancel.is_set():
552 break
553 verdicts[name] = _classify(name, path)
554 progress.tick()
555 return verdicts
556 with ThreadPoolExecutor(max_workers=workers, thread_name_prefix="lilbee-plan") as owned:
557 return _classify_pooled(owned, items, _classify, cancel, progress)
560def _plan_file_changes(
561 disk_files: dict[str, Path],
562 existing_sources: dict[str, SourceRecord],
563 cancel: CancelSignal | None,
564 held_hashes: dict[str, str] | None = None,
565) -> FileChangePlan:
566 """Diff every disk file against the store in one pass (see :func:`_plan_items`)."""
567 return _plan_items(
568 sorted(disk_files.items()), existing_sources, cancel, held_hashes or {}, _page_titles()
569 )
572def _plan_items(
573 items: list[tuple[str, Path]],
574 existing_sources: dict[str, SourceRecord],
575 cancel: CancelSignal | None,
576 held_hashes: dict[str, str],
577 titles: Mapping[str, str],
578 *,
579 progress: _PlanProgress | None = None,
580 pool: ThreadPoolExecutor | None = None,
581) -> FileChangePlan:
582 """Diff *items* (sorted name -> path) against the store, hashing only drifted files.
584 A tracked file whose stored (size, mtime) matches the disk stat, and whose
585 mtime predates the stat capture (see :func:`_stat_unchanged`), is unchanged
586 without reading its bytes; everything else is SHA-256 hashed. A file whose
587 current hash matches its record in ``held_hashes`` (set by a prior failed
588 attempt) is held out rather than retried every sync, and is reported as held
589 out rather than counted unchanged. Edit the file or run
590 ``/sync --retry-skipped`` to release the record and try again.
592 Classification fans across a thread pool (see :func:`_classify_changes`); the
593 plan is assembled from the results in the original sorted order, so a partial
594 or reordered completion never yields a wrong or reordered plan -- only a
595 shorter one when cancelled mid-pass.
596 """
597 verdicts = _classify_changes(
598 items, existing_sources, held_hashes, titles, cancel, progress=progress, pool=pool
599 )
601 files_to_process: list[FileToProcess] = []
602 added: dict[str, None] = {}
603 updated: dict[str, None] = {}
604 stat_backfills: list[SourceStatBackfill] = []
605 held_out: list[str] = []
606 unchanged = 0
607 # Assemble in the original sorted order so a partial (cancelled) or reordered
608 # completion never changes the plan the serial pass would have produced.
609 for name, _path in items:
610 verdict = verdicts.get(name)
611 if verdict is None:
612 continue # cancelled before this file was classified
613 if verdict.to_process is None:
614 if verdict.held:
615 held_out.append(name)
616 else:
617 unchanged += 1
618 if verdict.backfill is not None:
619 stat_backfills.append(verdict.backfill)
620 continue
621 files_to_process.append(verdict.to_process)
622 if verdict.is_update:
623 updated[name] = None
624 else:
625 added[name] = None
626 return FileChangePlan(files_to_process, added, updated, unchanged, stat_backfills, held_out)
629class _MovePool:
630 """Absent sources indexed by content hash, offered to the new files of that content.
632 Built once per sync and read across the streamed plan's batches. The store
633 gives each moved file one candidate under its write lock, in name order, so
634 a file takes the old key a single-pass plan pairs it with however the corpus
635 is sharded and whichever writer asks first.
636 """
638 def __init__(
639 self,
640 absent: list[str],
641 existing_sources: dict[str, SourceRecord],
642 failed: Mapping[str, str] | None = None,
643 ) -> None:
644 by_hash: dict[str, list[str]] = {}
645 for name in absent:
646 record = existing_sources.get(name)
647 if record is not None:
648 by_hash.setdefault(record["file_hash"], []).append(name)
649 # An absent failed file is offered at the hash it failed at.
650 for name, digest in (failed or {}).items():
651 by_hash.setdefault(digest, []).append(name)
652 # One tuple for each content: every new file of that content is offered the same one.
653 self._by_hash = {digest: tuple(sorted(names)) for digest, names in by_hash.items()}
655 def load(self, hashes: Iterable[str]) -> None:
656 """Make the absent sources with these content hashes known; all are from the start."""
658 def candidates(self, file_hash: str) -> tuple[str, ...]:
659 """The absent sources with this content hash, in name order."""
660 return self._by_hash.get(file_hash, ())
663class _IndexMovePool(_MovePool):
664 """Absent sources of every slice, found by content hash in the index.
666 A fan-out worker holds the source rows of its own slice, and the old name of
667 a moved file can belong to any slice. The lookup reads the source table as
668 the sync found it: a row written since belongs to a file that is on disk.
669 """
671 def __init__(self, store: Store, is_absent: Callable[[str], bool]) -> None:
672 super().__init__([], {})
673 self._store = store
674 self._version = store.sources_version()
675 self._is_absent = is_absent
676 self._loaded: set[str] = set()
678 def load(self, hashes: Iterable[str]) -> None:
679 """Look up the absent sources with these content hashes, once for each hash."""
680 wanted = sorted(set(hashes) - self._loaded)
681 self._loaded.update(wanted)
682 if self._version is None or not wanted:
683 return
684 found = self._store.sources_by_hash(wanted, version=self._version)
685 for digest, names in found.items():
686 absent = tuple(sorted(name for name in names if self._is_absent(name)))
687 if absent:
688 self._by_hash[digest] = absent
691def _moved_failures(store: Store, relocated: list[str]) -> set[str]:
692 """The *relocated* files a failure holds out: they moved with their record and stay out."""
693 held = store.held_records()
694 return {name for name in relocated if name in held and held[name].kind is HeldOut.FAILED}
697def _detect_moves(
698 files_to_process: list[FileToProcess],
699 added: dict[str, None],
700 pool: _MovePool,
701) -> list[SourceMove]:
702 """Offer each brand-new file the absent sources of the same content hash.
704 Only additions (files with a new name) can be moves; an update keeps its name.
705 When several absent sources share a hash, the store pairs them with the new
706 files one-to-one in name order, so a duplicated file that moved takes exactly
707 one old key and any leftovers stay indexed under their old key.
708 """
709 moves: list[SourceMove] = []
710 for entry in files_to_process:
711 if entry.name not in added:
712 continue
713 if candidates := pool.candidates(entry.file_hash):
714 moves.append(SourceMove(candidates, entry.name, entry.stat))
715 return moves
718def _apply_moves(
719 moved_new: Collection[str],
720 files_to_process: list[FileToProcess],
721 added: dict[str, None],
722) -> tuple[list[FileToProcess], list[str]]:
723 """Fold the files a relocation re-keyed out of the add set.
725 Drops each moved file from the ingest list and the added set: its chunks were
726 re-keyed onto the new source name, not rebuilt. Returns the trimmed
727 ``(files_to_process, relocated)``.
728 """
729 for name in moved_new:
730 added.pop(name, None)
731 remaining = [e for e in files_to_process if e.name not in moved_new]
732 return remaining, sorted(moved_new)
735def _in_archive(name: str, on_disk: Callable[[str], bool]) -> bool:
736 """Whether source *name* is a member of an archive that *on_disk* finds.
738 A member's key is its archive's key, a slash, and its path in the archive.
739 """
740 cut = name.find("/")
741 while cut != -1:
742 if on_disk(name[:cut]):
743 return True
744 cut = name.find("/", cut + 1)
745 return False
748def _absent_sources(sources: list[SourceRecord], disk_files: dict[str, Path]) -> list[str]:
749 """Document sources whose backing file is not on disk this sync.
751 A vanished file is never removed on its own (its chunks stay searchable, a
752 dead path-link discovered at open time); this set exists only to pair a
753 reappeared identical file to its old key in move detection. Imported sources
754 have no backing file, so they are excluded. An archive member has no file
755 of its own and carries its archive's hash: it is absent only when its
756 archive is.
757 """
758 return [
759 s["filename"]
760 for s in sources
761 if s["filename"] not in disk_files
762 and s["source_type"] != SourceType.IMPORTED
763 and not _in_archive(s["filename"], disk_files.__contains__)
764 ]
767def _worker_move_pool(
768 store: Store,
769 shard: ShardId,
770 disk_files: dict[str, Path],
771 gone: Collection[str],
772 rules: IgnoreRules,
773) -> _IndexMovePool:
774 """The move pool of one fan-out worker, over the absent sources of every slice.
776 The worker walked its own slice, so for a source of another slice it asks the
777 disk: the source is absent when no file is at its path or a pattern excludes it.
778 An archive member is absent only when its archive is.
779 """
781 def _on_disk(name: str) -> bool:
782 if shard.owns(name):
783 return name in disk_files
784 resolved = resolve_source_root(name)
785 if resolved is None:
786 return resolve_source_path(name).is_file()
787 base, path = resolved
788 return path.is_file() and not rules.excludes_path(path, base=base)
790 def _is_absent(name: str) -> bool:
791 return name not in gone and not _on_disk(name) and not _in_archive(name, _on_disk)
793 return _IndexMovePool(store, _is_absent)
796def _move_pool(
797 store: Store,
798 shard: ShardId | None,
799 sources: list[SourceRecord],
800 disk_files: dict[str, Path],
801 gone: Collection[str],
802 rules: IgnoreRules,
803) -> _MovePool:
804 """The absent sources a brand-new file of this sync can be paired with as a move.
806 A failed file moves with its folder, so it stays held out at its new path.
807 """
808 if shard is not None:
809 return _worker_move_pool(store, shard, disk_files, gone, rules)
810 absent = [name for name in _absent_sources(sources, disk_files) if name not in gone]
811 failed = {
812 name: record.file_hash
813 for name, record in store.held_records().items()
814 if record.kind is HeldOut.FAILED
815 and name not in disk_files
816 and name not in gone
817 and not _in_archive(name, disk_files.__contains__)
818 }
819 return _MovePool(absent, {source["filename"]: source for source in sources}, failed)
822# A plan batch is only done when its slowest hash is, so the first one is small (work
823# reaches the fleet within a second) and later ones amortize that barrier.
824_PLAN_SHARD_MIN_FILES = 256
825_PLAN_SHARD_MAX_FILES = 8192
828def _plan_batch_bounds(total: int) -> Iterator[tuple[int, int]]:
829 """(start, stop) slices covering *total* files, doubling up to the cap.
831 Build-vs-buy: ``itertools.batched`` is the stock slicer but is 3.12+ (floor is
832 3.11) and fixed-size, so it cannot ramp the batch size.
833 """
834 start = 0
835 size = _PLAN_SHARD_MIN_FILES
836 while start < total:
837 stop = min(start + size, total)
838 yield start, stop
839 start = stop
840 size = min(size * 2, _PLAN_SHARD_MAX_FILES)
843@dataclass
844class _StreamedPlan:
845 """Bookkeeping a streamed plan accumulates across its batches.
847 ``added`` and ``updated`` are the dicts the ingest pass mutates as files land.
848 """
850 added: dict[str, None] = field(default_factory=dict)
851 updated: dict[str, None] = field(default_factory=dict)
852 # Processed files' content hashes, for the held-out records written after the run.
853 pending_hashes: dict[str, str] = field(default_factory=dict)
854 relocated: list[str] = field(default_factory=list)
855 unchanged: int = 0
856 # Files a held-out record held out of this run, in plan order (ordered set).
857 held_out: dict[str, None] = field(default_factory=dict)
858 planned: int = 0
859 # Files this pass's slice holds, from the discovery walk. Fixed before the
860 # first batch is planned, so it is what progress is measured against: the
861 # plan's own running totals grow as batches land and cannot say how much of
862 # the corpus is left. Summed across a fan-out's workers it is the corpus.
863 corpus_total: int = 0
865 @property
866 def resolved(self) -> int:
867 """Files the plan disposed of without ingest, as it disposes of them.
869 Unchanged files, held-out held-out files and repointed moves are done
870 as far as the corpus is concerned, and nothing downstream ingests them,
871 so an incremental sync would otherwise show a handful of changed files
872 against the whole corpus. Files a cancelled plan never classified are
873 absent from every count and so are not claimed as done.
874 """
875 return self.unchanged + len(self.held_out) + len(self.relocated)
878async def _absorb_plan_batch(
879 plan: FileChangePlan, state: _StreamedPlan, moves: _MovePool
880) -> list[FileToProcess]:
881 """Fold one batch's plan into *state* and return the files it queues for ingest.
883 Relocations and stat backfills are written per batch, so they contend with the
884 batch flushes and take the same one-shot lock retry.
885 """
886 store = get_services().store
887 entries = plan.files_to_process
888 state.unchanged += plan.unchanged
889 state.held_out.update(dict.fromkeys(plan.held_out))
890 state.added.update(plan.added)
891 state.updated.update(plan.updated)
893 if new_hashes := [e.file_hash for e in entries if e.name in plan.added]:
894 await to_ingest_thread(moves.load, new_hashes)
895 detected = _detect_moves(entries, plan.added, moves)
896 if detected:
897 taken = await to_ingest_thread(_relocate_in_holds, store, detected)
898 # A file whose candidates other writers all took is an add.
899 entries, relocated = _apply_moves(taken.keys(), entries, plan.added)
900 for name in relocated:
901 state.added.pop(name, None)
902 failures = await to_ingest_thread(_moved_failures, store, relocated)
903 state.held_out.update(dict.fromkeys(failures))
904 state.relocated.extend(name for name in relocated if name not in failures)
905 if plan.stat_backfills:
906 backfills = plan.stat_backfills
907 await to_ingest_thread(
908 _retry_after_lock_timeout, lambda: store.update_source_stats(backfills)
909 )
911 state.pending_hashes.update((entry.name, entry.file_hash) for entry in entries)
912 state.planned += len(entries)
913 return entries
916async def _plan_batches(
917 disk_files: dict[str, Path],
918 existing_sources: dict[str, SourceRecord],
919 held_hashes: dict[str, str],
920 moves: _MovePool,
921 state: _StreamedPlan,
922 cancel: CancelSignal | None,
923) -> AsyncGenerator[list[FileToProcess]]:
924 """Plan the corpus batch by batch, yielding each batch's files to ingest.
926 Shards are contiguous slices of one sorted item list consumed in order, so
927 this delivers the single-pass plan of :func:`_plan_items` in pieces. The next
928 batch is planned while the current one ingests, on a dedicated thread: the
929 shared ingest pool is saturated by extraction and would stall the stream it
930 feeds. Empty batches are not yielded, so the first yield means there is work.
931 """
932 items = sorted(disk_files.items())
933 titles = await to_ingest_thread(_page_titles)
934 progress = _PlanProgress(len(items))
935 stop = _StreamStop(cancel)
936 workers = _plan_workers()
937 hashers = ThreadPoolExecutor(max_workers=workers, thread_name_prefix="lilbee-plan")
938 driver = ThreadPoolExecutor(max_workers=1, thread_name_prefix="lilbee-plan-driver")
940 def _plan(lo: int, hi: int) -> FileChangePlan:
941 return _plan_items(
942 items[lo:hi],
943 existing_sources,
944 stop,
945 held_hashes,
946 titles,
947 progress=progress,
948 pool=hashers if workers > 1 else None,
949 )
951 bounds = list(_plan_batch_bounds(len(items)))
952 try:
953 ahead = asyncio.ensure_future(to_executor(driver, _plan, *bounds[0])) if bounds else None
954 for index, _bound in enumerate(bounds):
955 if ahead is None:
956 break
957 plan = await ahead
958 ahead = (
959 asyncio.ensure_future(to_executor(driver, _plan, *bounds[index + 1]))
960 if index + 1 < len(bounds) and not stop.is_set()
961 else None
962 )
963 entries = await _absorb_plan_batch(plan, state, moves)
964 if entries:
965 yield entries
966 finally:
967 stop.close()
968 if ahead is not None:
969 ahead.cancel()
970 # Retrieve the outcome: a prefetch that had already failed would
971 # otherwise surface as an unretrieved task exception.
972 with contextlib.suppress(Exception, asyncio.CancelledError):
973 await ahead
974 # No wait: the running batch notices the stop and exits on its own, and a
975 # generator being closed must not block the event loop on a hash in flight.
976 hashers.shutdown(wait=False)
977 driver.shutdown(wait=False)
980def detect_pending() -> int:
981 """Count files in documents/ that are out of sync with the store.
983 Cheap operation: filesystem walk + stat-gated SHA-256 hashing + a single
984 sources-table read. No embedding, no writes. Returns the count of files that
985 would be ingested (added + updated), which is what the TaskBar hint surfaces.
986 A vanished file is not pending work: sync leaves it indexed, so it is not
987 counted. Reuses ``_plan_file_changes`` so the diff logic stays single-sourced.
988 Honors the held-out records: a file that failed last time at this hash does
989 not show up as pending. Blocking: callers on the event loop run it via
990 ``asyncio.to_thread``.
991 """
992 disk_files = discover_files()
993 if not disk_files:
994 return 0
995 store = get_services().store
996 existing_sources = {s["filename"]: s for s in store.get_sources()}
997 held_hashes = _held_hashes(store)
998 plan = _plan_file_changes(disk_files, existing_sources, cancel=None, held_hashes=held_hashes)
999 return len(plan.files_to_process)
1002# Refused files named in the log line before it falls back to a count.
1003_EXCLUDED_LOG_SAMPLE = 5
1004# The reason held out with a failed file whose failure named none.
1005_NO_REASON = "ingestion produced nothing"
1006# The text a rebuild embeds once to prove the embedder answers before it drops an index.
1007_REBUILD_PROBE = "lilbee"
1010def _log_excluded(excluded: dict[str, ExclusionReason]) -> None:
1011 """Warn once per exclusion reason, naming at most ``_EXCLUDED_LOG_SAMPLE`` files."""
1012 by_reason: dict[ExclusionReason, list[str]] = {}
1013 for name, why in excluded.items():
1014 by_reason.setdefault(why, []).append(name)
1015 for why, names in by_reason.items():
1016 shown = sorted(names)[:_EXCLUDED_LOG_SAMPLE]
1017 more = f" and {len(names) - len(shown)} more" if len(names) > len(shown) else ""
1018 log.warning("Skipped %d file(s), %s: %s%s", len(names), why.value, ", ".join(shown), more)
1021def _describe_failures(store: Store, names: Iterable[str]) -> list[SkippedSource]:
1022 """Each of *names* a failure holds out, with its reason, in order."""
1023 held = store.held_records()
1024 return [
1025 SkippedSource(filename=name, reason=record.reason)
1026 for name in names
1027 if (record := held.get(name)) is not None and record.kind is HeldOut.FAILED
1028 ]
1031def _held_hashes(store: Store) -> dict[str, str]:
1032 """The content hash each held-out path is held out at, by path."""
1033 return {path: record.file_hash for path, record in store.held_records().items()}
1036async def _corpus_root(store: Store, shard: ShardId | None, *, retry_skipped: bool) -> Path:
1037 """The data root whose ``.lilbeeignore`` this sync uses, after a retry released the failures.
1039 Only a sync that is not a worker releases them, so a fan-out releases them once.
1040 """
1041 if shard is not None:
1042 return shard.corpus_root
1043 if retry_skipped:
1044 await to_ingest_thread(store.release_failed)
1045 return active_config().data_root
1048def _persist_failures(
1049 store: Store, pending_hashes: dict[str, str], reasons: dict[str, str], failed: Iterable[str]
1050) -> None:
1051 """Hold each of *failed* out of later syncs at the hash this sync planned it at.
1053 A file that ingests cleanly needs no step: its new source row holds nothing out.
1054 A write that cannot take the lock is logged, and those files are tried again next sync.
1055 """
1056 held = {
1057 name: HeldRecord(HeldOut.FAILED, fhash, reasons.get(name, _NO_REASON))
1058 for name in failed
1059 if (fhash := pending_hashes.get(name))
1060 }
1061 try:
1062 _retry_after_lock_timeout(lambda: store.hold_failed(held))
1063 except LockTimeoutError as error:
1064 log.error("Could not hold %d failed file(s) out of the next sync: %s", len(held), error)
1067def _report_index_mismatch(store: Store) -> IndexMismatch | None:
1068 """Name an index built with another embedder than cfg; the sync leaves it as it is.
1070 An unchanged corpus never reaches the write gate that refuses such an index,
1071 so without this a sync finishes green over an index search refuses. Wiping
1072 it is the caller's decision, never the sync's.
1073 """
1074 mismatch = store.index_mismatch()
1075 if mismatch is None:
1076 return None
1077 log.warning("Sync left the index as it is: %s", mismatch)
1078 return mismatch.describe()
1081def _prepare_forced_rebuild(store: Any) -> bool:
1082 """Ready the index for a forced rebuild; True when the rebuild runs in place (blocking).
1084 A store with chunks and no embedder mismatch, with or without a stamp, is rebuilt in
1085 place and a file that fails keeps its chunks. An empty store, or one built by another
1086 embedder, is dropped after one probe request. Memories are re-embedded either way.
1087 Run off the event loop.
1088 """
1089 in_place = store.index_mismatch() is None and (
1090 store.get_meta() is not None or store.count_chunks() > 0
1091 )
1092 if not in_place:
1093 _require_embedding_model()
1094 embedder = get_services().embedder
1095 try:
1096 refreshed = 0
1097 if embedder.embedding_available():
1098 refreshed = store.rebuild_memory_embeddings(lambda texts: embedder.embed_batch(texts))
1099 if not (in_place or refreshed):
1100 embedder.embed_batch([_REBUILD_PROBE])
1101 except Exception as error:
1102 raise RuntimeError(
1103 "Rebuild stopped and the index is unchanged: "
1104 f"the embedding model did not answer ({error})."
1105 ) from error
1106 if not in_place:
1107 store.drop_all()
1108 return in_place
1111def _reconcile_missing(
1112 disk_files: dict[str, Path],
1113 sources: list[SourceRecord],
1114 failed: Iterable[str],
1115 skipped: Iterable[str],
1116 held: Iterable[str],
1117) -> list[str]:
1118 """On-disk document files absent from the store that no mechanism accounts for.
1120 A file discovery found and classified that ended up in neither the sources table
1121 nor any of the accounting sets was dropped with no signal -- the silent
1122 data-loss case (a scanned PDF that never made it into the index yet reported no
1123 error). Everything legitimately not indexed is excluded: a failed extraction is in
1124 ``failed``, a zero-text file this run attempted is in ``skipped``, a file this run
1125 held out on its record is in ``held``, and an unsupported type was never
1126 returned by discovery in the first place. ``held`` is the run's held-out set, not
1127 every record: a stale record (file edited since it failed) must not hide a drop.
1128 """
1129 accounted = {s["filename"] for s in sources} | set(failed) | set(skipped) | set(held)
1130 return sorted(name for name in disk_files if name not in accounted)
1133def _missing_from_index(
1134 store: Store,
1135 shard: ShardId | None,
1136 disk_files: dict[str, Path],
1137 failed: Iterable[str],
1138 skipped: Iterable[str],
1139 held: Iterable[str],
1140) -> list[str]:
1141 """The files of this sync's scope that :func:`_reconcile_missing` names.
1143 A worker asks the index for names alone, in one pass over the name column:
1144 it does not read the source rows of its slice a second time.
1145 """
1146 if shard is None:
1147 return _reconcile_missing(disk_files, store.get_sources(), failed, skipped, held)
1148 waived = {*failed, *skipped, *held}
1149 return sorted(store.missing_sources(name for name in disk_files if name not in waived))
1152def _ignored_sources(sources: list[SourceRecord], rules: IgnoreRules) -> list[str]:
1153 """Indexed sources a ``.lilbeeignore`` now excludes.
1155 Asked of the patterns rather than of discovery, which cannot tell an excluded
1156 file from a lost one -- both simply stop being yielded, and the two need
1157 opposite handling. Reading every source rather than only the undiscovered
1158 ones keeps the answer independent of which slice of the corpus this pass
1159 walked. Imported sources have no file to match a pattern against.
1160 """
1161 ignored = []
1162 for source in sources:
1163 if source["source_type"] == SourceType.IMPORTED:
1164 continue
1165 resolved = resolve_source_root(source["filename"])
1166 if resolved is not None and rules.excludes_path(resolved[1], base=resolved[0]):
1167 ignored.append(source["filename"])
1168 return ignored
1171def _forget_ignored(sources: list[SourceRecord], rules: IgnoreRules) -> list[str]:
1172 """Drop sources the patterns now exclude from the index. Returns what went.
1174 A corpus-wide pass: it reads every source, so one worker of a fan-out must
1175 not run it. No held-out record is written either, because the ignore file is
1176 itself the durable statement -- a record could only outlive it and hold a
1177 source out after its pattern was deleted. Deleting the pattern brings the
1178 source back on the next sync.
1179 """
1180 names = _ignored_sources(sources, rules)
1181 if not names:
1182 return []
1183 removed = list(get_services().store.remove_documents(names).removed)
1184 _forget_in_wiki(removed)
1185 return removed
1188def _forget_refused(
1189 excluded: Mapping[str, ExclusionReason],
1190 existing: Mapping[str, SourceRecord],
1191 *,
1192 forget_in_wiki: bool,
1193) -> list[str]:
1194 """Drop indexed sources that discovery now refuses. Returns what went."""
1195 names = [name for name in excluded if name in existing]
1196 if not names:
1197 return []
1198 removed = list(get_services().store.remove_documents(names).removed)
1199 if forget_in_wiki:
1200 _forget_in_wiki(removed)
1201 return removed
1204def _finish_in_place_rebuild(store: Store, rules: IgnoreRules, any_failed: bool) -> list[str]:
1205 """Drop the sources of vanished files; stamp the store when every file was rebuilt."""
1206 if not any_failed:
1207 store.restamp_meta()
1208 return _forget_vanished(rules)
1211def _forget_vanished(rules: IgnoreRules) -> list[str]:
1212 """Drop the document sources whose file is gone, as a rebuild that dropped the index would.
1214 A corpus-wide pass, run by the parent after the files were ingested again.
1215 """
1216 store = get_services().store
1217 names = _absent_sources(store.get_sources(), discover_corpus(None, rules).files)
1218 if not names:
1219 return []
1220 removed = list(store.remove_documents(names).removed)
1221 _forget_in_wiki(removed)
1222 return removed
1225def _forget_in_wiki(removed: list[str]) -> None:
1226 """Drop removed sources from the wiki's browse index."""
1227 # circular: lilbee.app.ingest imports this module for sync
1228 from lilbee.app.ingest import forget_removed_from_wiki_index
1230 forget_removed_from_wiki_index(removed)
1233def _show_current_paths_in_wiki() -> None:
1234 """Render again each wiki page that shows a path its cited source no longer has."""
1235 # circular: lilbee.app.ingest imports this module for sync
1236 from lilbee.app.ingest import show_current_paths_in_wiki
1238 show_current_paths_in_wiki()
1241def _forget_missing_in_wiki() -> None:
1242 """Drop from the wiki's browse index what the workers of a fan-out that stopped removed."""
1243 # circular: lilbee.app.ingest imports this module for sync
1244 from lilbee.app.ingest import forget_missing_from_wiki_index
1246 forget_missing_from_wiki_index()
1249def _sources_in_scope(store: Store, shard: ShardId | None) -> list[SourceRecord]:
1250 """The source rows this sync plans against: all of them, or a worker's own slice."""
1251 if shard is None:
1252 return store.get_sources()
1253 return store.get_sources_where(shard.owns)
1256def _require_embedding_model() -> None:
1257 """Refuse ingest without an embedding model.
1259 Ingest has no degraded mode: without one, every chunk fails to embed after
1260 the run has already paid the parse and OCR cost. Search and chat fall back
1261 to keyword via embedding_available() and carry on.
1262 """
1263 if not get_services().embedder.validate_model():
1264 ref = active_config().embedding_model
1265 detail = (
1266 f"{ref!r} is not available: pull it, or set a different embedding_model."
1267 if ref
1268 else "None is configured: pull one and set embedding_model."
1269 )
1270 raise RuntimeError(f"Ingest needs an embedding model. {detail}")
1273async def _run_post_ingest_passes(
1274 store: Any,
1275 *,
1276 indexed_anything: bool,
1277 cluster_added: Collection[str],
1278 cluster_updated: Collection[str],
1279 cluster_removed: Collection[str],
1280 touched: set[str],
1281 cancel: CancelSignal | None,
1282 recluster_all: bool = False,
1283) -> None:
1284 """Index maintenance, concept clusters, the wiki hook, and the entity lifecycle.
1286 The index pass only runs when this sync indexed something. The concept
1287 clusters are recomputed only when a source was added, updated or removed: a
1288 relocation changes no co-occurrence. Every sync renders again each wiki
1289 page that shows a path its cited source no longer has, so the next sync
1290 repairs a page a stopped one left. The wiki hook and the entity lifecycle
1291 run every sync (a cheap no-op when off or already current): the wiki index
1292 takes what any sync wrote and no refresh read, and turning the entity
1293 setting on takes effect without a separate operation.
1294 """
1295 if indexed_anything:
1296 store.ensure_fts_index()
1297 store.ensure_scalar_indexes()
1298 store.ensure_vector_index()
1299 store.optimize_sources()
1300 if cluster_added or cluster_updated or cluster_removed:
1301 await _rebuild_concept_clusters(
1302 cluster_added, cluster_updated, cluster_removed, full=recluster_all
1303 )
1304 await to_ingest_thread(_show_current_paths_in_wiki)
1305 await _update_wiki(touched, active_config())
1307 from lilbee.retrieval.entities.lifecycle import ensure_entities
1309 await to_ingest_thread(ensure_entities, cancel)
1312async def _update_wiki(named: set[str], config: Config) -> None:
1313 """Refresh the wiki index, and regenerate pages when auto-update is on.
1315 The index refresh runs whenever the wiki is enabled, because it spends no
1316 LLM call and it is what lets the browse tree list a page the moment its
1317 document lands. It takes *named* and every source the store holds that no
1318 refresh has read, whichever sync wrote it. Regeneration is the expensive
1319 half and stays behind wiki_auto_update.
1321 Best effort: the ingest itself already succeeded and `lilbee wiki update`
1322 re-runs the regeneration, so a wiki failure must not skip the post-ingest
1323 entity pass or the reconciliation guard.
1324 """
1325 if not config.wiki:
1326 return
1327 # circular: lilbee.wiki imports lilbee.data.ingest.file_hash, so the
1328 # post-ingest hook stays function-local at this boundary.
1329 from lilbee.wiki.ingest import incremental_update
1330 from lilbee.wiki.stubs import refresh_owed_sources
1332 changed = named
1333 try:
1334 changed = named | await to_ingest_thread(
1335 refresh_owed_sources, get_services().store, config, named
1336 )
1337 except Exception:
1338 log.warning("Wiki index refresh failed after sync", exc_info=True)
1340 if not config.wiki_auto_update or not changed:
1341 return
1342 try:
1343 await incremental_update(changed, config)
1344 except Exception:
1345 log.warning("Wiki auto-update failed after sync", exc_info=True)
1348def _worker_failure_message(failures: list[ShardDone], specs: list[ShardSpec]) -> str:
1349 """What a failed fan-out reports; the files its workers finished are in the index."""
1350 detail = "; ".join(
1351 f"worker {failure.index} ({specs[failure.index].config.data_root / WORKER_LOG_NAME}): "
1352 f"{failure.error}"
1353 for failure in failures
1354 )
1355 return (
1356 f"{len(failures)} ingest worker(s) failed: {detail}. "
1357 "The files the workers finished are in the index: "
1358 "re-running sync continues from where they stopped."
1359 )
1362async def _sync_across_workers(
1363 specs: list[ShardSpec],
1364 store: Store,
1365 *,
1366 prune_ignored: bool = False,
1367 rebuild: bool = False,
1368 options: ShardOptions,
1369 quiet: bool,
1370 on_progress: DetailedProgressCallback,
1371 cancel: CancelSignal | None,
1372) -> SyncResult:
1373 """Ingest on one worker per GPU, each writing the one index, then run the corpus-wide passes."""
1374 try:
1375 verdicts = await run_workers(
1376 specs, options=options, quiet=quiet, on_progress=on_progress, cancel=cancel
1377 )
1378 if cancel is not None and cancel.is_set():
1379 raise asyncio.CancelledError
1380 if failures := [verdict for verdict in verdicts if verdict.error is not None]:
1381 raise RuntimeError(_worker_failure_message(failures, specs))
1382 except BaseException:
1383 await to_ingest_thread(_forget_missing_in_wiki)
1384 raise
1385 result = aggregate_results(verdicts)
1386 if rebuild:
1387 rules = IgnoreRules.for_corpus(active_config().data_root)
1388 vanished = await to_ingest_thread(
1389 _finish_in_place_rebuild, store, rules, bool(result.failed)
1390 )
1391 result.removed = [*vanished, *result.removed]
1392 touched = set(result.added) | set(result.updated) | set(result.relocated)
1393 # No worker tells the wiki index what it removed, and a pass with no extractor drops nothing.
1394 await to_ingest_thread(_forget_in_wiki, result.removed)
1395 # No worker sees the whole corpus, so each one leaves this pass to the parent.
1396 if prune_ignored:
1397 ignored = await to_ingest_thread(
1398 _forget_ignored, store.get_sources(), IgnoreRules.for_corpus(active_config().data_root)
1399 )
1400 result.removed = [*ignored, *result.removed]
1401 await _run_post_ingest_passes(
1402 store,
1403 indexed_anything=bool(touched),
1404 cluster_added=result.added,
1405 cluster_updated=result.updated,
1406 cluster_removed=result.removed,
1407 touched=touched,
1408 cancel=cancel,
1409 recluster_all=rebuild,
1410 )
1411 on_progress(
1412 EventType.SYNC_DONE,
1413 SyncDoneEvent(
1414 added=len(result.added),
1415 updated=len(result.updated),
1416 removed=len(result.removed),
1417 failed=len(result.failed),
1418 skipped=len(result.skipped),
1419 relocated=len(result.relocated),
1420 ),
1421 )
1422 return result
1425_SyncParams = ParamSpec("_SyncParams")
1426# What a sync that names no starter is called in the message of a sync that waits for it.
1427DEFAULT_SYNC_STARTER = "lilbee"
1430def _marks_sync_running(
1431 run: Callable[_SyncParams, Coroutine[Any, Any, SyncResult]],
1432) -> Callable[_SyncParams, Coroutine[Any, Any, SyncResult]]:
1433 """Hold the data root's sync mark for the whole run, so a reset refuses meanwhile.
1435 Under the mark, the sync also takes the data root's sync turn before it plans, so
1436 syncs run one after another and none syncs the files of another that runs. A
1437 fan-out worker runs under its parent's turn, and a data root that cannot hold the
1438 mark cannot hold the turn either; neither takes it.
1439 """
1441 @functools.wraps(run)
1442 async def _marked(*args: _SyncParams.args, **kwargs: _SyncParams.kwargs) -> SyncResult:
1443 data_root = active_config().data_root
1444 # ``cancel``, ``shard``, ``on_progress`` and ``started_by`` are keyword-only on a
1445 # sync, so they are never in *args*.
1446 cancel = cast("CancelSignal | None", kwargs.get("cancel"))
1447 worker = kwargs.get("shard") is not None
1448 # The services come before the mark.
1449 store = (await asyncio.to_thread(get_services)).store
1451 async def _run() -> SyncResult:
1452 # The sync walks and names sources by the registry config.toml holds at its
1453 # start; a worker keeps its parent's.
1454 if not worker:
1455 await asyncio.to_thread(settings.overlay_persisted_roots, active_config())
1456 with store.registry_pinned():
1457 return await run(*args, **kwargs)
1459 async with sync_running(data_root, cancel) as marked:
1460 if not marked or worker:
1461 return await _run()
1462 async with _sync_turn_for(data_root, cancel, kwargs):
1463 return await _run()
1465 return _marked
1468def _sync_turn_for(
1469 data_root: Path, cancel: CancelSignal | None, kwargs: dict[str, Any]
1470) -> contextlib.AbstractAsyncContextManager[None]:
1471 """The sync turn of *data_root* for the sync called with *kwargs*; a wait is reported."""
1472 on_progress = cast("DetailedProgressCallback", kwargs.get("on_progress", noop_callback))
1473 started_by = cast("str", kwargs.get("started_by", DEFAULT_SYNC_STARTER))
1475 def _report_wait(owner: SyncTurnOwner | None) -> None:
1476 message = sync_wait_message(owner)
1477 log.info(message)
1478 on_progress(
1479 EventType.SYNC_WAIT,
1480 SyncWaitEvent(
1481 message=message,
1482 holder=owner.started_by if owner else None,
1483 pid=owner.pid if owner else None,
1484 ),
1485 )
1487 return sync_turn(data_root, cancel, started_by=started_by, on_wait=_report_wait)
1490@_marks_sync_running
1491@owns_ingest_pool
1492async def sync(
1493 force_rebuild: bool = False,
1494 quiet: bool = False,
1495 *,
1496 on_progress: DetailedProgressCallback = noop_callback,
1497 cancel: CancelSignal | None = None,
1498 retry_skipped: bool = False,
1499 prune_ignored: bool = False,
1500 shard: ShardId | None = None,
1501 started_by: str = DEFAULT_SYNC_STARTER,
1502) -> SyncResult:
1503 """Sync documents/ with the vector store.
1504 Returns a SyncResult with the added/updated/removed/unchanged/failed/skipped lists.
1505 When *quiet* is True, the Rich progress bar is suppressed (for JSON output).
1506 When *cancel* is set, planning and processing stop between files without
1507 data loss (completed work is flushed) and CancelledError is raised.
1508 When *retry_skipped* is set, the records of failed files are released so this
1509 sync attempts those files again; *force_rebuild* clears removals too.
1510 When *prune_ignored* is set, sources a ``.lilbeeignore`` now excludes are
1511 dropped from the index. Off by default: the patterns govern what sync takes
1512 in, and removing what a past sync already indexed is the caller's decision.
1513 A *shard* runs this sync as one worker of a multi-GPU fan-out: it writes its
1514 slice of the corpus to the index and leaves the corpus-wide passes to the parent.
1515 *started_by* names what started this sync, for the message a sync that waits for
1516 it shows; ``_marks_sync_running`` reads it.
1517 """
1518 config = active_config()
1519 _store = get_services().store
1521 in_place = force_rebuild and (
1522 shard is not None or await to_ingest_thread(_prepare_forced_rebuild, _store)
1523 )
1525 config.documents_dir.mkdir(parents=True, exist_ok=True)
1526 index_mismatch = _report_index_mismatch(_store)
1527 corpus_root = await _corpus_root(_store, shard, retry_skipped=retry_skipped)
1528 if shard is None and (specs := plan_fanout()):
1529 merged = await _sync_across_workers(
1530 specs,
1531 _store,
1532 prune_ignored=prune_ignored,
1533 options=ShardOptions(parent_pid=os.getpid(), rebuild=force_rebuild),
1534 rebuild=in_place,
1535 quiet=quiet,
1536 on_progress=on_progress,
1537 cancel=cancel,
1538 )
1539 return merged.model_copy(update={"index_mismatch": index_mismatch})
1541 rules = IgnoreRules.for_corpus(corpus_root)
1542 scan = discover_corpus(shard, rules)
1543 disk_files = scan.files
1544 sources = _sources_in_scope(_store, shard)
1545 # A rebuild plans every file as new: nothing is unchanged, held out or moved, so each
1546 # file is chunked and embedded again and replaces its own chunks when its write lands.
1547 existing_sources = {} if force_rebuild else {s["filename"]: s for s in sources}
1548 held_hashes = {} if force_rebuild else await to_ingest_thread(_held_hashes, _store)
1550 failed: dict[str, None] = {}
1551 # Refused formats start the run skipped; they get no held-out record (no planned hash).
1552 skipped: dict[str, OcrReport | None] = dict.fromkeys(scan.excluded)
1553 # filename → why it was skipped/failed (for reporting)
1554 reasons: dict[str, str] = {name: why.value for name, why in scan.excluded.items()}
1555 retry_next: set[str] = set()
1556 _log_excluded(scan.excluded)
1558 # Opt-in, and corpus-wide: a worker reads one slice of the sources, so it
1559 # leaves this pass to the parent.
1560 ignored = _forget_ignored(sources, rules) if prune_ignored and shard is None else []
1561 # A worker removes only keys from its own slice; the parent reconciles the wiki index.
1562 refused = _forget_refused(scan.excluded, existing_sources, forget_in_wiki=shard is None)
1564 # Sources whose backing file is not on disk this pass. A vanished file is NOT
1565 # removed: it stays indexed and searchable, a dead path-link the user
1566 # discovers only when they try to open it, and the set pairs a reappeared
1567 # identical file to its old key below. What was just removed leaves the set,
1568 # where it could otherwise capture a real move.
1569 gone = set(ignored) | set(refused)
1570 moves = (
1571 _MovePool([], {})
1572 if force_rebuild
1573 else _move_pool(_store, shard, sources, disk_files, gone, rules)
1574 )
1576 # The planning pass stats (and where needed hashes) every file on disk, batch by
1577 # batch off the event loop and overlapped with ingest. A brand-new file whose
1578 # content hash matches an absent source is folded in per batch as a move, not an
1579 # add: repointed in place so its chunks and embeddings are reused, not rebuilt.
1580 state = _StreamedPlan(corpus_total=len(disk_files))
1581 added, updated, pending_hashes = state.added, state.updated, state.pending_hashes
1582 plan_batches = _plan_batches(disk_files, existing_sources, held_hashes, moves, state, cancel)
1584 # Snapshot the cumulative truncation counter so the delta over this sync can
1585 # surface "N chunks truncated" instead of being lost in per-chunk debug logs.
1586 truncated_before = get_services().embedder.truncated_total
1588 # Ingest files (with optional progress bar). Only non-empty batches are yielded,
1589 # so the first one arriving is what proves there is work to do.
1590 try:
1591 first = await anext(plan_batches, None)
1592 if first is not None:
1593 # Hold the embed fleet resident for the whole batch: an unevenly loaded
1594 # replica must not idle-unload and reload cold mid-run (which snowballs
1595 # into a fleet collapse). The ContextVar propagates into the ingest
1596 # thread pool, where the fleet actually spawns on the first embed.
1597 from lilbee.providers.fleet.ingest_warmth import keep_fleet_warm
1599 with keep_fleet_warm():
1600 _require_embedding_model()
1601 await ingest_stream(
1602 _chain_plan_batches(first, plan_batches),
1603 added,
1604 updated,
1605 failed,
1606 skipped,
1607 plan=state,
1608 quiet=quiet,
1609 on_progress=on_progress,
1610 cancel=cancel,
1611 retry_next=retry_next,
1612 reasons=reasons,
1613 )
1614 if cancel is not None and cancel.is_set():
1615 # The stream stops feeding on cancel, so ingest can drain its
1616 # admitted files and return without raising. A cancelled run must
1617 # not go on to write held-out records or reconcile an unplanned corpus.
1618 raise asyncio.CancelledError
1619 finally:
1620 # Idempotent, and the only close when the stream is never consumed.
1621 await plan_batches.aclose()
1622 relocated = state.relocated
1624 # A flush or engine failure is not a verdict on the file: it is planned again next sync.
1625 held = [name for name in (*failed, *skipped) if name not in retry_next]
1626 await to_ingest_thread(_persist_failures, _store, pending_hashes, reasons, held)
1628 vanished = (
1629 await to_ingest_thread(_finish_in_place_rebuild, _store, rules, bool(failed))
1630 if in_place and shard is None
1631 else []
1632 )
1633 if shard is None:
1634 # The passes are corpus-wide, so a fan-out runs them once, in the parent.
1635 await _run_post_ingest_passes(
1636 _store,
1637 indexed_anything=bool(state.planned or relocated),
1638 cluster_added=added,
1639 cluster_updated=updated,
1640 cluster_removed=[*ignored, *refused, *vanished],
1641 touched=set(added) | set(updated) | set(relocated),
1642 cancel=cancel,
1643 recluster_all=in_place,
1644 )
1646 # Reconciliation guard against silent data loss: any on-disk document file that
1647 # ended up in neither the index nor an accounting set was dropped without a
1648 # signal. Surface it loudly instead of letting a whole dataset vanish quietly.
1649 if missing := _missing_from_index(_store, shard, disk_files, failed, skipped, state.held_out):
1650 log.warning(
1651 "Sync reconciliation: %d document file(s) on disk are absent from the index "
1652 "with no failure reported (possible silent drop): %s",
1653 len(missing),
1654 ", ".join(missing[:20]),
1655 )
1657 result = SyncResult(
1658 added=list(added),
1659 updated=list(updated),
1660 removed=ignored + refused + vanished,
1661 unchanged=state.unchanged,
1662 relocated=relocated,
1663 failed=list(failed),
1664 skipped=list(skipped),
1665 skipped_ocr={name: ocr for name, ocr in skipped.items() if ocr is not None},
1666 held_out=_describe_failures(_store, state.held_out),
1667 truncated=get_services().embedder.truncated_total - truncated_before,
1668 unreadable_folders=list(scan.unreadable),
1669 index_mismatch=index_mismatch,
1670 )
1671 on_progress(
1672 EventType.SYNC_DONE,
1673 SyncDoneEvent(
1674 added=len(result.added),
1675 updated=len(result.updated),
1676 removed=len(result.removed),
1677 failed=len(result.failed),
1678 skipped=len(result.skipped),
1679 relocated=len(result.relocated),
1680 ),
1681 )
1682 return result
1685def _phase_progress_callback(
1686 progress: Progress, ptask: Any, chain: DetailedProgressCallback
1687) -> DetailedProgressCallback:
1688 """Wrap *chain*, updating the bar's description on per-page / per-chunk events.
1690 EXTRACT (page i/N, named for the OCR backend that ran), OCR_START (Tesseract
1691 running on a file) and EMBED (chunk i/N) events would otherwise leave the bar
1692 frozen between file completions; surfacing them on the spinner description
1693 keeps a single large file's row visibly moving. All events still forward to
1694 *chain* so the caller's own callback (TUI / JSON) is unaffected.
1695 """
1697 def _callback(event_type: EventType, data: ProgressEvent) -> None:
1698 if event_type is EventType.EXTRACT and isinstance(data, ExtractEvent):
1699 progress.update(
1700 ptask, description=f"{data.step} {data.file} (page {data.page}/{data.total_pages})"
1701 )
1702 elif event_type is EventType.OCR_START and isinstance(data, OcrStartEvent):
1703 progress.update(ptask, description=data.status_text)
1704 elif event_type is EventType.EMBED and isinstance(data, EmbedEvent):
1705 progress.update(
1706 ptask, description=f"Embedding {data.file} ({data.chunk}/{data.total_chunks})"
1707 )
1708 chain(event_type, data)
1710 return _callback
1713def _ingest_progress_bar(quiet: bool) -> Progress:
1714 """The transient bar an in-process ingest reports on, disabled when *quiet*."""
1715 return Progress(
1716 SpinnerColumn(),
1717 literal_text_column("{task.description}"),
1718 BarColumn(),
1719 MofNCompleteColumn(),
1720 TimeElapsedColumn(),
1721 console=PlainConsole(),
1722 transient=True,
1723 disable=quiet,
1724 )
1727# In-flight task cap, as a multiple of _max_concurrent(): enough queued tasks to
1728# keep every compute slot fed, without materializing one task object per file.
1729_TASK_WINDOW_MULTIPLIER = 2
1732def _build_admission(
1733 baseline: int, pages_done: list[int]
1734) -> tuple[asyncio.Semaphore | ResizableGate, int, asyncio.Task[None] | None]:
1735 """The batch's admission control, plus its task-window size and controller task.
1737 Static mode (the default) returns a fixed semaphore and no controller. Adaptive
1738 mode, when a GPU fleet is present to feed, returns a resizable gate and a running
1739 :class:`AdaptiveController` that tunes it toward this box's throughput knee; with
1740 no fleet it falls back to the static path so a GPU-less host is never affected.
1741 """
1742 profile = profile_for(resolve_mode())
1743 devices = enumerate_fleet_devices() if profile is not None else []
1744 if profile is None or not devices:
1745 return asyncio.Semaphore(baseline), baseline * _TASK_WINDOW_MULTIPLIER, None
1746 permit_max = max_workers()
1747 gate = ResizableGate(min(baseline, permit_max))
1748 controller = AdaptiveController(
1749 gate,
1750 profile,
1751 make_signal_sampler(devices),
1752 lambda: pages_done[0],
1753 permit_min=1,
1754 permit_max=permit_max,
1755 )
1756 task = asyncio.ensure_future(controller.run())
1757 # warning, not info: the default LILBEE_LOG_LEVEL is WARNING, so the
1758 # auto-chosen concurrency would otherwise never surface on a headless sync.
1759 log.warning(
1760 "Adaptive ingest concurrency (%s): start %d, max %d", profile.name, gate.limit, permit_max
1761 )
1762 return gate, permit_max * _TASK_WINDOW_MULTIPLIER, task
1765def _failed_result(
1766 exc: Exception,
1767 entry: FileToProcess,
1768 *,
1769 pages_done: list[int],
1770 on_progress: DetailedProgressCallback,
1771 cancel: CancelSignal | None,
1772) -> _IngestResult:
1773 """A file's failure as a result, or cancellation when the run is stopping.
1775 During shutdown, worker pools raise RuntimeError from submit(). Those are
1776 cancellation, not ingest failures: the cancel flag is the source of truth,
1777 and the executor's shutdown message covers the race where cancel was set
1778 after the submit.
1779 """
1780 if (cancel and cancel.is_set()) or is_executor_shutdown(exc):
1781 raise asyncio.CancelledError from exc
1782 # Suppress TaskCancelledError on the FILE_DONE notice: the user already
1783 # cancelled, and re-raising here would strand sibling tasks awaiting in
1784 # _collect_results.
1785 with contextlib.suppress(TaskCancelledError):
1786 on_progress(EventType.FILE_DONE, FileDoneEvent(file=entry.name, status="error", chunks=0))
1787 pages_done[0] += 1 # cleared the gate (as a failure); still a throughput tick
1788 return _IngestResult(entry.name, entry.path, 0, error=exc)
1791async def _archive_result(
1792 entry: FileToProcess,
1793 on_progress: DetailedProgressCallback,
1794 pages_done: list[int],
1795 cancel: CancelSignal | None,
1796) -> _IngestResult:
1797 """Ingest an archive: its members become sources, the archive row keeps the disk stat."""
1798 members = await ingest_archive(
1799 entry.path, entry.name, entry.content_type, on_progress=on_progress, cancel=cancel
1800 )
1801 concept_batches = [await build_concept_records(m.records, m.name) for m in members]
1802 entity_rows = [
1803 row for m in members for row in (await build_entity_records(m.records, m.name) or [])
1804 ]
1805 chunk_total = sum(len(m.records) for m in members)
1806 on_progress(
1807 EventType.FILE_DONE, FileDoneEvent(file=entry.name, status="ok", chunks=chunk_total)
1808 )
1809 pages_done[0] += max(1, sum(len(m.page_texts) for m in members))
1810 found = [batch for batch in concept_batches if batch is not None]
1811 return _IngestResult(
1812 entry.name,
1813 entry.path,
1814 chunk_total,
1815 error=None,
1816 file_hash=entry.file_hash,
1817 records=[],
1818 needs_cleanup=entry.needs_cleanup,
1819 page_texts=[],
1820 stat=entry.stat,
1821 concept_records=ConceptRecords.merged(found) if found else None,
1822 entity_rows=entity_rows or None,
1823 meta=SourceMeta(title=derive_title(entry.name)),
1824 members=members,
1825 )
1828def _over_limit_result(
1829 exc: ChunkLimitError,
1830 entry: FileToProcess,
1831 *,
1832 pages_done: list[int],
1833 on_progress: DetailedProgressCallback,
1834) -> _IngestResult:
1835 """A file over the per-file chunk limit, recorded as skipped and held out at its hash."""
1836 log.warning("Skipped %s: %s", entry.name, exc)
1837 with contextlib.suppress(TaskCancelledError):
1838 on_progress(EventType.FILE_DONE, FileDoneEvent(file=entry.name, status="skipped", chunks=0))
1839 pages_done[0] += 1
1840 return _IngestResult(
1841 entry.name,
1842 entry.path,
1843 0,
1844 error=None,
1845 file_hash=entry.file_hash,
1846 skip_reason=str(exc),
1847 needs_cleanup=entry.needs_cleanup,
1848 )
1851async def _stream_tasks(
1852 plan_batches: AsyncGenerator[list[FileToProcess]],
1853 make_task: Callable[[FileToProcess, int], Coroutine[Any, Any, _IngestResult]],
1854) -> AsyncGenerator[list[Coroutine[Any, Any, _IngestResult]]]:
1855 """Each plan batch's per-file coroutines, in plan order."""
1856 index = count(1)
1857 async for plan_batch in plan_batches:
1858 yield [make_task(entry, next(index)) for entry in plan_batch]
1861async def _chain_plan_batches(
1862 first: list[FileToProcess], rest: AsyncGenerator[list[FileToProcess]]
1863) -> AsyncGenerator[list[FileToProcess]]:
1864 """Yield an already-pulled plan batch, then the remainder of its stream."""
1865 yield first
1866 async for plan_batch in rest:
1867 yield plan_batch
1870async def ingest_stream(
1871 plan_batches: AsyncGenerator[list[FileToProcess]],
1872 added: dict[str, None],
1873 updated: dict[str, None],
1874 failed: dict[str, None],
1875 skipped: dict[str, OcrReport | None],
1876 *,
1877 plan: _StreamedPlan | None = None,
1878 quiet: bool = False,
1879 on_progress: DetailedProgressCallback = noop_callback,
1880 cancel: CancelSignal | None = None,
1881 retry_next: set[str] | None = None,
1882 reasons: dict[str, str] | None = None,
1883) -> None:
1884 """Ingest a stream of planned file batches, optionally showing a Rich progress bar.
1886 Files are admitted as their batch is planned, so ingest starts on the first
1887 batch instead of waiting for the whole corpus to be diffed. Old chunks are
1888 deleted in the same transaction as the new write, so the two are atomic per
1889 file. When *cancel* is set, pending files raise CancelledError before starting.
1891 *plan* is the bookkeeping the batches were planned into; it carries the corpus
1892 the run is measured against. Without it progress is reported with no total,
1893 since a bare stream of batches does not say what corpus it came from.
1894 """
1895 # Honor LILBEE_INGEST_TRACE once per batch: it raises the trace loggers above
1896 # the default WARNING so per-file extraction lines actually surface.
1897 configure_trace_from_env()
1898 warn_if_table_model_ignored()
1899 # Throughput is measured in OCR pages, not documents: a document's cost scales
1900 # with its page count (a 500-page scan is 500x a memo), so pages are the unbiased
1901 # unit of GPU-feeding work for the adaptive controller to hill-climb on.
1902 pages_done = [0]
1903 # Sized off the files with no source row yet, not the planned count: the plan
1904 # streams in batches and its total is unknown until the stream drains, but the
1905 # pool has to be decided before the first batch is dispatched. Unindexed files
1906 # are the one part of the plan that is known without diffing, exact for a first
1907 # ingest or a rebuild and near zero for an incremental sync, so a small sync
1908 # over a large corpus stays in-process. Undercounting only keeps a run
1909 # in-process, which is the safe direction.
1910 admission, window, controller_task = _build_admission(_max_concurrent(), pages_done)
1911 progress = _ingest_progress_bar(quiet)
1912 # The corpus the discovery walk found, known before the first batch is planned.
1913 corpus_total = plan.corpus_total if plan is not None else 0
1914 ptask = progress.add_task("Ingesting documents...", total=corpus_total or None)
1915 # Rebound once, so every event the run emits, per file or per batch, passes the
1916 # bar: a single large file's OCR and embed phases move it between completions.
1917 on_progress = _phase_progress_callback(progress, ptask, on_progress)
1919 async def _process_one(entry: FileToProcess, file_index: int) -> _IngestResult:
1920 name = entry.name
1921 async with admission:
1922 if cancel and cancel.is_set():
1923 raise asyncio.CancelledError
1925 try:
1926 on_progress(
1927 EventType.FILE_START,
1928 FileStartEvent(file=name, total_files=feed.planned, current_file=file_index),
1929 )
1930 except TaskCancelledError as exc:
1931 # FILE_START itself can raise the cooperative cancel signal;
1932 # normalize so _collect_results can drain siblings cleanly.
1933 raise asyncio.CancelledError from exc
1934 try:
1935 # The source's old chunks are deleted in the same locked
1936 # transaction as the new write (see _flush_writes), so cleanup is
1937 # carried on the result rather than run eagerly here.
1938 if entry.content_type in archive_content_types():
1939 return await _archive_result(entry, on_progress, pages_done, cancel)
1940 page_texts: list[PageTextRecord] = []
1941 records, meta, ocr = await produce_records(
1942 entry.path,
1943 name,
1944 entry.content_type,
1945 quiet=quiet,
1946 on_progress=on_progress,
1947 page_texts_out=page_texts,
1948 cancel=cancel,
1949 )
1950 concept_records = await build_concept_records(records, name)
1951 entity_rows = await build_entity_records(records, name)
1952 on_progress(
1953 EventType.FILE_DONE,
1954 FileDoneEvent(file=name, status="ok", chunks=len(records)),
1955 )
1956 pages_done[0] += max(1, len(page_texts)) # OCR pages cleared: the throughput signal
1957 return _IngestResult(
1958 name,
1959 entry.path,
1960 len(records),
1961 error=None,
1962 file_hash=entry.file_hash,
1963 records=records,
1964 needs_cleanup=entry.needs_cleanup,
1965 page_texts=page_texts,
1966 stat=entry.stat,
1967 concept_records=concept_records,
1968 entity_rows=entity_rows,
1969 meta=meta,
1970 ocr=ocr,
1971 )
1972 except ChunkLimitError as exc:
1973 return _over_limit_result(
1974 exc, entry, pages_done=pages_done, on_progress=on_progress
1975 )
1976 except (asyncio.CancelledError, TaskCancelledError) as exc:
1977 # TaskCancelledError is the TUI's cooperative cancel signal raised
1978 # by reporter.check_cancelled() inside on_progress; treat it as
1979 # asyncio cancellation so _collect_results can drain siblings
1980 # cleanly instead of orphaning their pending exceptions.
1981 raise asyncio.CancelledError from exc
1982 except Exception as exc:
1983 return _failed_result(
1984 exc, entry, pages_done=pages_done, on_progress=on_progress, cancel=cancel
1985 )
1987 feed = _ResultFeed(_stream_tasks(plan_batches, _process_one), plan)
1988 try:
1989 # extract_batching coalesces extractions into xberg batch calls when the
1990 # toggle is on (off by default); the per-file collect contract is unchanged.
1991 with progress:
1992 async with extract_batching():
1993 await _collect_results(
1994 feed,
1995 added,
1996 updated,
1997 failed,
1998 skipped,
1999 window=window,
2000 on_progress=on_progress,
2001 progress=progress,
2002 ptask=ptask,
2003 retry_next=retry_next,
2004 reasons=reasons,
2005 cancel=cancel,
2006 )
2007 finally:
2008 # Stop the adaptive controller (if any) before returning: its background
2009 # loop must not outlive the batch it was tuning.
2010 if controller_task is not None:
2011 controller_task.cancel()
2012 with contextlib.suppress(asyncio.CancelledError):
2013 await controller_task
2016# Accumulate roughly this many chunks across documents before one batched
2017# LanceDB write. Bounds buffered-vector memory while amortizing the write lock
2018# and per-transaction overhead over many documents instead of one write per file.
2019_WRITE_FLUSH_CHUNKS = 2000
2022class _ResultFeed:
2023 """Pull-based source of per-file ingest coroutines over a streamed plan.
2025 ``take(wait=False)`` hands back an already-planned file without blocking, so
2026 the collector waits on the planner only when it has nothing left to run.
2027 Build-vs-buy: an ``asyncio.Queue`` is the stock bounded channel, but it would
2028 need a separate producer task to pump the plan-batch generator into it and a
2029 sentinel to close it; pulling ``anext`` on demand keeps the plan stream the
2030 single driver and needs neither. ``planned`` is the file count seen so far:
2031 the run's total once the stream is drained.
2032 """
2034 def __init__(
2035 self,
2036 plan_batches: AsyncGenerator[list[Coroutine[Any, Any, _IngestResult]]],
2037 plan: _StreamedPlan | None = None,
2038 ) -> None:
2039 self._plan_batches = plan_batches
2040 self._buffer: deque[Coroutine[Any, Any, _IngestResult]] = deque()
2041 self._pull: asyncio.Task[list[Coroutine[Any, Any, _IngestResult]] | None] | None = None
2042 self._drained = False
2043 self._plan = plan if plan is not None else _StreamedPlan()
2044 self.planned = 0
2046 @property
2047 def corpus_total(self) -> int:
2048 """Files the run's slice holds, or 0 when the caller declared no corpus."""
2049 return self._plan.corpus_total
2051 @property
2052 def resolved(self) -> int:
2053 """Files already disposed of by the plan, which never reach this feed."""
2054 return self._plan.resolved
2056 def pull(self) -> asyncio.Task[list[Coroutine[Any, Any, _IngestResult]] | None] | None:
2057 """The in-flight plan-batch prefetch, so a waiting collector wakes when it lands."""
2058 return self._pull
2060 async def take(self, *, wait: bool) -> Coroutine[Any, Any, _IngestResult] | None:
2061 """The next planned file, or None when the stream is drained (or, with
2062 *wait* False, when the next batch is not planned yet)."""
2063 while not self._buffer:
2064 if self._drained:
2065 return None
2066 if self._pull is None:
2067 self._pull = asyncio.ensure_future(anext(self._plan_batches, None))
2068 if not wait and not self._pull.done():
2069 return None
2070 plan_batch = await self._pull
2071 self._pull = None
2072 if plan_batch is None:
2073 self._drained = True
2074 return None
2075 self._buffer.extend(plan_batch)
2076 self.planned += len(plan_batch)
2077 return self._buffer.popleft()
2079 async def aclose(self) -> None:
2080 """Close the plan stream and discard files that were never started."""
2081 if self._pull is not None:
2082 self._pull.cancel()
2083 with contextlib.suppress(asyncio.CancelledError):
2084 # A batch that landed before the cancel took effect still owns
2085 # coroutines; recover it so they are closed rather than leaked.
2086 plan_batch = await self._pull
2087 if plan_batch:
2088 self._buffer.extend(plan_batch)
2089 self._pull = None
2090 for coro in self._buffer:
2091 coro.close()
2092 self._buffer.clear()
2093 await self._plan_batches.aclose()
2096async def _refill_window(
2097 in_flight: set[asyncio.Task[_IngestResult]],
2098 feed: _ResultFeed,
2099 window: int,
2100) -> None:
2101 """Top up the in-flight task set from *feed*, capped at *window* tasks.
2103 Waits on the planner only when nothing is running, so a slow batch never
2104 stalls files that are already planned.
2105 """
2106 while len(in_flight) < window:
2107 coro = await feed.take(wait=not in_flight)
2108 if coro is None:
2109 return
2110 in_flight.add(asyncio.ensure_future(coro))
2113async def _next_completions(
2114 in_flight: set[asyncio.Task[_IngestResult]], prefetch: asyncio.Future[Any] | None
2115) -> tuple[Iterable[asyncio.Future[Any]], set[asyncio.Task[_IngestResult]]]:
2116 """Wait for the next file to finish, returning (completed, still running).
2118 The feed's plan-batch prefetch waits alongside the running files, so work is
2119 admitted as soon as it is planned rather than on the next file completion,
2120 and it is filtered out of the still-running set here since it is not a file.
2121 """
2122 waiting: set[asyncio.Future[Any]] = set(in_flight)
2123 if prefetch is not None:
2124 waiting.add(prefetch)
2125 done, still_running = await asyncio.wait(waiting, return_when=asyncio.FIRST_COMPLETED)
2126 # Explicit loop, like _cancel_in_flight: Nuitka miscompiled the comprehension
2127 # form of this task-set filtering.
2128 remaining: set[asyncio.Task[_IngestResult]] = set()
2129 for task in still_running:
2130 if task is not prefetch:
2131 remaining.add(cast("asyncio.Task[_IngestResult]", task))
2132 return done, remaining
2135async def _collect_results(
2136 feed: _ResultFeed,
2137 added: dict[str, None],
2138 updated: dict[str, None],
2139 failed: dict[str, None],
2140 skipped: dict[str, OcrReport | None],
2141 *,
2142 window: int,
2143 on_progress: DetailedProgressCallback = noop_callback,
2144 progress: Progress | None = None,
2145 ptask: Any = None,
2146 retry_next: set[str] | None = None,
2147 reasons: dict[str, str] | None = None,
2148 cancel: CancelSignal | None = None,
2149) -> None:
2150 """Run *feed* through a bounded task window, batching writes and progress.
2152 At most *window* tasks exist at once: results are consumed as they complete
2153 and the window is refilled from the feed, so memory stays flat however many
2154 files a sync covers. Successful files are buffered and flushed to LanceDB in
2155 batches (one locked transaction per batch) rather than one write per file.
2156 The buffer is flushed on the way out too -- even on cancel -- so
2157 completed-but-unwritten work is persisted. On exception (typically
2158 asyncio.CancelledError from a user cancel), cancel every in-flight sibling
2159 and await them with ``return_exceptions=True`` so their pending
2160 CancelledErrors don't surface as "Task exception was never retrieved".
2161 """
2162 buffer: list[_IngestResult] = []
2163 buffered_chunks = 0
2164 completed_count = 0
2165 to_purge: list[str] = []
2166 in_flight: set[asyncio.Task[_IngestResult]] = set()
2167 try:
2168 await _refill_window(in_flight, feed, window)
2169 while in_flight:
2170 prefetch = feed.pull()
2171 done, in_flight = await _next_completions(in_flight, prefetch)
2172 saw_cancel = False
2173 for fut in done:
2174 if fut is prefetch:
2175 continue # a planned batch landing, not a file result
2176 try:
2177 result = fut.result()
2178 except asyncio.CancelledError:
2179 # A user cancel completes several futures together. Flag it but
2180 # keep draining `done` so a sibling that genuinely finished in
2181 # the same batch is still buffered and flushed (the
2182 # cancel-persists contract), then propagate after the loop. A
2183 # non-cancel exception still propagates immediately, as before,
2184 # so a genuine ingest bug surfaces and cancels the siblings.
2185 saw_cancel = True
2186 continue
2187 completed_count += 1
2188 status = _classify_result(result, added, updated, failed, skipped, reasons)
2189 if retry_next is not None and isinstance(result.error, ProviderError):
2190 # The model engine failed, not the file: a later sync tries it again.
2191 retry_next.add(result.name)
2192 if status is BatchStatus.INGESTED:
2193 buffered_chunks = await _buffer_and_maybe_flush(
2194 result,
2195 buffer,
2196 buffered_chunks,
2197 added,
2198 updated,
2199 failed,
2200 skipped,
2201 retry_next,
2202 cancel,
2203 )
2204 elif status is BatchStatus.SKIPPED and result.needs_cleanup:
2205 # Zero-text result is never buffered; collect it for the
2206 # purge pass (see _purge_emptied_sources).
2207 to_purge.append(result.name)
2208 _report_file_progress(
2209 result,
2210 status,
2211 feed.resolved + completed_count,
2212 feed.corpus_total,
2213 on_progress,
2214 progress,
2215 ptask,
2216 )
2217 if saw_cancel:
2218 # Completed siblings in this batch are now buffered; propagate the
2219 # cancel so the finally flushes them and cancels still-running work.
2220 raise asyncio.CancelledError
2221 await _refill_window(in_flight, feed, window)
2222 finally:
2223 # The inner finally guarantees the sibling cancel even if the flush
2224 # itself raises (e.g. a cancellation landing on the to_thread await).
2225 try:
2226 await _flush_to_end(buffer, added, updated, failed, skipped, retry_next, cancel)
2227 await to_ingest_thread(_purge_emptied_sources, to_purge)
2228 finally:
2229 try:
2230 await _cancel_in_flight(in_flight)
2231 finally:
2232 # Closing the feed stops the planner behind it, so a cancelled
2233 # sync does not keep hashing the rest of the corpus.
2234 await feed.aclose()
2237async def _cancel_in_flight(in_flight: set[asyncio.Task[_IngestResult]]) -> None:
2238 """Cancel still-running tasks and await them so their CancelledErrors are retrieved."""
2239 # Explicit loop: Nuitka miscompiled the comprehension form of this cleanup.
2240 still_pending = []
2241 for t in in_flight:
2242 if not t.done():
2243 still_pending.append(t)
2244 for task in still_pending:
2245 task.cancel()
2246 if still_pending:
2247 await asyncio.gather(*still_pending, return_exceptions=True)
2250async def _buffer_and_maybe_flush(
2251 result: _IngestResult,
2252 buffer: list[_IngestResult],
2253 buffered_chunks: int,
2254 added: dict[str, None],
2255 updated: dict[str, None],
2256 failed: dict[str, None],
2257 skipped: dict[str, OcrReport | None],
2258 retry_next: set[str] | None,
2259 cancel: CancelSignal | None = None,
2260) -> int:
2261 """Buffer one ingested file, flushing at the chunk threshold; returns the new count."""
2262 buffer.append(result)
2263 # Zero-chunk files count one unit so the buffer stays bounded.
2264 buffered_chunks += max(result.chunk_count, 1)
2265 if buffered_chunks >= _WRITE_FLUSH_CHUNKS:
2266 await _flush_to_end(buffer, added, updated, failed, skipped, retry_next, cancel)
2267 buffered_chunks = 0
2268 return buffered_chunks
2271async def _flush_to_end(
2272 buffer: list[_IngestResult],
2273 added: dict[str, None],
2274 updated: dict[str, None],
2275 failed: dict[str, None],
2276 skipped: dict[str, OcrReport | None],
2277 retry_next: set[str] | None,
2278 cancel: CancelSignal | None = None,
2279) -> None:
2280 """Run :func:`_flush_writes` on the ingest pool; a cancel waits for it, then propagates.
2282 The write thread owns *buffer* until it returns, so a flush the cancel abandoned
2283 would run beside the next flush of the same buffer. The flush is a plain future:
2284 a cancel of every task on the loop reaches this caller and never the write.
2285 A write that fails under a cancel, of this caller or a set *cancel*, leaves as
2286 the cause of the cancel; one that escaped :func:`_flush_writes` is logged here.
2287 """
2288 flush = ingest_thread_future(_flush_writes, buffer, added, updated, failed, skipped, retry_next)
2289 cancelled = False
2290 while not flush.done():
2291 try:
2292 await asyncio.wait([flush])
2293 except asyncio.CancelledError:
2294 cancelled = True
2295 escaped = flush.exception()
2296 error = flush.result() if escaped is None else escaped
2297 if cancelled or (error is not None and cancel is not None and cancel.is_set()):
2298 if escaped is not None:
2299 log.warning("The write of a cancelled sync failed", exc_info=escaped)
2300 raise asyncio.CancelledError from error
2301 if escaped is not None:
2302 raise escaped
2305def _report_file_progress(
2306 result: _IngestResult,
2307 status: BatchStatus,
2308 completed_count: int,
2309 total: int,
2310 on_progress: DetailedProgressCallback,
2311 progress: Progress | None,
2312 ptask: Any,
2313) -> None:
2314 """Advance the Rich bar (when present) and emit one BATCH_PROGRESS event.
2316 *completed_count* counts every file the pass has disposed of and *total* is
2317 the corpus the discovery walk found, so the pair answers how much of the
2318 corpus is done rather than how much of the plan so far is.
2319 """
2320 if progress is not None and ptask is not None:
2321 desc = f"Ingested {result.name}" if result.error is None else f"Failed {result.name}"
2322 # Set, not advanced: files the plan resolved without ingest produce no
2323 # result of their own and would otherwise never reach the bar.
2324 progress.update(ptask, description=desc, completed=completed_count)
2325 with contextlib.suppress(TaskCancelledError):
2326 on_progress(
2327 EventType.BATCH_PROGRESS,
2328 BatchProgressEvent(
2329 file=result.name,
2330 status=status,
2331 current=completed_count,
2332 total=total,
2333 ),
2334 )
2337def _classify_result(
2338 result: _IngestResult,
2339 added: dict[str, None],
2340 updated: dict[str, None],
2341 failed: dict[str, None],
2342 skipped: dict[str, OcrReport | None],
2343 reasons: dict[str, str] | None = None,
2344) -> BatchStatus:
2345 """Record a completed file's outcome and return its batch status.
2347 Failures, refusals and zero-chunk files are tracked here; a successful file is
2348 reported as ``INGESTED`` and its chunks are persisted by the batched flush, so
2349 it stays in ``added`` / ``updated`` until then. When *reasons* is given, the
2350 human-readable cause is recorded there (filename → reason) for reporting.
2351 """
2352 if result.skip_reason is not None:
2353 added.pop(result.name, None)
2354 updated.pop(result.name, None)
2355 skipped[result.name] = None
2356 if reasons is not None:
2357 reasons[result.name] = result.skip_reason
2358 return BatchStatus.SKIPPED
2359 if result.error is not None:
2360 # A traceback here would bleed into the TUI chat pane; the full trace stays at DEBUG.
2361 log.warning("Failed to ingest %s: %s", result.name, result.error)
2362 log.debug("Traceback for failed ingest of %s", result.name, exc_info=result.error)
2363 added.pop(result.name, None)
2364 updated.pop(result.name, None)
2365 failed[result.name] = None
2366 if reasons is not None:
2367 reasons[result.name] = error_reason(result.error)
2368 return BatchStatus.FAILED
2369 if result.chunk_count == 0:
2370 # No searchable chunks: never report it as added/updated. With no page
2371 # texts either, nothing is persisted and the file retries next sync. With
2372 # page texts, it stays INGESTED so its pages persist (export/recon) and it
2373 # stops replanning, but it is reported as skipped since search can't see it.
2374 added.pop(result.name, None)
2375 updated.pop(result.name, None)
2376 skipped[result.name] = result.ocr
2377 if reasons is not None:
2378 if result.page_texts:
2379 reasons[result.name] = "stored page text only (0 searchable chunks)"
2380 elif result.ocr is not None and result.ocr.backend is OcrBackendUsed.NONE:
2381 reasons[result.name] = "no text extracted, OCR is off (0 chunks)"
2382 else:
2383 reasons[result.name] = "no text extracted (0 chunks)"
2384 return BatchStatus.SKIPPED if not result.page_texts else BatchStatus.INGESTED
2385 return BatchStatus.INGESTED
2388# Back off briefly before the single flush retry: the usual contender is a
2389# search-triggered FTS optimize holding the store lock past its 30s timeout.
2390_FLUSH_RETRY_DELAY_SECONDS = 2.0
2392# Twice the poll of the store's file lock: a writer that waits looks at least
2393# once between two holds of a re-key.
2394_RELOCATE_PAUSE_SECONDS = 0.1
2396_Written = TypeVar("_Written")
2399def _retry_after_lock_timeout(write: Callable[[], _Written]) -> _Written:
2400 """Run one store write, retrying once after a lock timeout."""
2401 try:
2402 return write()
2403 except LockTimeoutError:
2404 log.warning(
2405 "Store write lock busy; retrying batch flush in %.0fs", _FLUSH_RETRY_DELAY_SECONDS
2406 )
2407 time.sleep(_FLUSH_RETRY_DELAY_SECONDS)
2408 return write()
2411def _relocate_in_holds(store: Store, moves: list[SourceMove]) -> dict[str, str]:
2412 """Re-key *moves* one hold of the write lock at a time; the old name each new name took.
2414 Each hold takes the one-shot lock retry, so a timeout keeps what earlier
2415 holds re-keyed. The pause between two holds lets a writer that waits take
2416 the lock.
2417 """
2418 taken: dict[str, str] = {}
2419 settled = 0
2420 while settled < len(moves):
2421 if settled:
2422 time.sleep(_RELOCATE_PAUSE_SECONDS)
2423 hold = functools.partial(store.relocate_sources, moves[settled:])
2424 step = _retry_after_lock_timeout(hold)
2425 taken.update(step.taken)
2426 settled += step.settled
2427 return taken
2430def _flush_batch(buffer: list[_IngestResult]) -> None:
2431 """Persist one flush unit in a single locked ``write_chunks_batch`` transaction.
2433 Page texts travel inside each :class:`ChunkWrite` so the store writes them
2434 after the cleanup delete (which clears the source's old page-text rows) and
2435 before the source row: a page-text failure leaves the row stale and the file
2436 replans next sync instead of losing its pages forever behind the stat
2437 short-circuit. The write retries once on a lock timeout. A file that is
2438 gone when the write takes the lock is left out, its concepts and entities
2439 with it: a crawl that retired it after this sync read it removed its source.
2440 """
2441 store = get_services().store
2442 items: list[ChunkWrite] = []
2443 stale: list[str] = []
2444 for r in buffer:
2445 digest = r.file_hash or file_hash(r.path)
2446 items.append(
2447 ChunkWrite(
2448 source=r.name,
2449 file_hash=digest,
2450 records=cast(list[dict], r.records or []),
2451 needs_cleanup=r.needs_cleanup,
2452 stat=r.stat,
2453 page_texts=cast(list[dict], r.page_texts or []),
2454 meta=r.meta,
2455 )
2456 )
2457 if r.members is None:
2458 continue
2459 items.extend(_member_write(m, digest, r.stat) for m in r.members)
2460 current = {m.name for m in r.members}
2461 stale.extend(n for n in store.member_sources(r.name) if n not in current)
2462 if stale:
2463 store.remove_documents(stale)
2464 paths = {r.name: r.path for r in buffer}
2465 paths.update((m.name, r.path) for r in buffer for m in r.members or [])
2466 gone: set[str] = set()
2468 def present(name: str) -> bool:
2469 if paths[name].is_file():
2470 return True
2471 gone.add(name)
2472 return False
2474 _retry_after_lock_timeout(lambda: store.write_chunks_batch(items, present))
2475 if gone:
2476 log.info("Not indexed, the file is gone: %s", ", ".join(sorted(gone)))
2477 kept = [r for r in buffer if r.name not in gone]
2478 _flush_concept_records(kept)
2479 _flush_entity_rows(kept)
2482def _member_write(member: MemberRecords, digest: str, stat: SourceStat | None) -> ChunkWrite:
2483 """A member's write item: its own source row, keyed to the archive's hash and stat."""
2484 return ChunkWrite(
2485 source=member.name,
2486 file_hash=digest,
2487 records=cast(list[dict], member.records),
2488 needs_cleanup=True,
2489 stat=stat,
2490 page_texts=cast(list[dict], member.page_texts),
2491 meta=member.meta,
2492 )
2495def _flush_concept_records(buffer: list[_IngestResult]) -> None:
2496 """Write the flush unit's buffered concept rows in one batched pass.
2498 Runs after the chunk write so a failed flush (files moved to ``failed``
2499 and replanned) never lands concept rows for unwritten chunks. A concept
2500 write failure is logged and never fails the files, matching the
2501 per-file extraction failure semantics.
2502 """
2503 batches = [r.concept_records for r in buffer if r.concept_records is not None]
2504 if not batches:
2505 return
2506 try:
2507 get_services().concepts.write_concept_records(ConceptRecords.merged(batches))
2508 except Exception:
2509 log.warning("Concept indexing failed for %d-file batch", len(batches), exc_info=True)
2512def _flush_entity_rows(buffer: list[_IngestResult]) -> None:
2513 """Write the flush unit's buffered entity rows in one batched pass.
2515 Runs after the chunk write, which also performed the per-source deletes,
2516 so replacement never leaves a source's stale entity rows behind. A write
2517 failure is logged and never fails the files, matching concept semantics.
2518 """
2519 rows = [row for r in buffer if r.entity_rows for row in r.entity_rows]
2520 if not rows:
2521 return
2522 try:
2523 get_services().store.add_entities(rows)
2524 except Exception:
2525 log.warning("Entity indexing failed for a %d-row batch", len(rows), exc_info=True)
2528def _purge_emptied_sources(names: list[str]) -> None:
2529 """Remove the prior index entry for files that now extract to nothing.
2531 An already-indexed file edited to yield zero chunks and zero page texts is
2532 classified SKIPPED and never buffered, so the batched cleanup delete never
2533 runs and its old chunks and source row would linger in search results. Full
2534 removal here keeps the index consistent; ``remove_documents`` is a no-op for
2535 never-indexed (brand-new empty) files, so unindexed inputs cost nothing.
2536 """
2537 if not names:
2538 return
2539 get_services().store.remove_documents(names)
2542def _flush_writes(
2543 buffer: list[_IngestResult],
2544 added: dict[str, None],
2545 updated: dict[str, None],
2546 failed: dict[str, None],
2547 skipped: dict[str, OcrReport | None],
2548 retry_next: set[str] | None = None,
2549) -> Exception | None:
2550 """Flush the buffered documents to the store; track and return a write failure.
2552 Each buffered file's page texts, chunks, cleanup delete, and source upsert
2553 are written by :func:`_flush_batch`. If that fails, every file in the batch
2554 is moved to ``failed`` since its source row did not land, and recorded in
2555 *retry_next* so the caller replans them next sync instead of holding out
2556 them; the exception never escapes, so the caller's sibling-cancel and
2557 held-out path always runs. The buffer is cleared either way.
2558 """
2559 if not buffer:
2560 return None
2561 try:
2562 _flush_batch(buffer)
2563 except Exception as exc:
2564 for r in buffer:
2565 log.warning("Failed to write %s: %s", r.name, exc)
2566 added.pop(r.name, None)
2567 updated.pop(r.name, None)
2568 # A page-text-only file was pre-marked skipped at classification; on a
2569 # flush failure it belongs in failed only, never both.
2570 skipped.pop(r.name, None)
2571 failed[r.name] = None
2572 if retry_next is not None:
2573 retry_next.add(r.name)
2574 return exc
2575 finally:
2576 buffer.clear()
2577 return None