Coverage for src/lilbee/data/ingest/pipeline.py: 100%
816 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-12 00:44 +0000
1"""Top-level sync orchestration: discovery, dispatch, batching, post-sync hooks."""
3from __future__ import annotations
5import asyncio
6import contextlib
7import logging
8import os
9import threading
10import time
11from collections import deque
12from collections.abc import AsyncGenerator, Callable, Coroutine, Iterable, Iterator, Mapping
13from concurrent.futures import Future, ThreadPoolExecutor, as_completed
14from dataclasses import dataclass, field
15from itertools import count
16from pathlib import Path
17from typing import Any, cast
19from rich.progress import (
20 BarColumn,
21 MofNCompleteColumn,
22 Progress,
23 SpinnerColumn,
24 TextColumn,
25 TimeElapsedColumn,
26)
28from lilbee.app.services import get_services
29from lilbee.core.config import Config, active_config
30from lilbee.data.extract.chunk import ChunkLimitError
31from lilbee.data.extract.document import (
32 extract_batching,
33 ingest_archive,
34 ingest_document,
35 ingest_markdown,
36 warn_if_table_model_ignored,
37)
38from lilbee.data.extract.trace import configure_from_env as configure_trace_from_env
39from lilbee.data.ingest.adaptive import (
40 AdaptiveController,
41 ResizableGate,
42 enumerate_fleet_devices,
43 make_signal_sampler,
44 profile_for,
45 resolve_mode,
46)
47from lilbee.data.ingest.code import ingest_code_sync
48from lilbee.data.ingest.discovery import (
49 ExclusionReason,
50 archive_content_types,
51 classify_file,
52 discover_corpus,
53 discover_files,
54 file_hash,
55 resolve_source_root,
56)
57from lilbee.data.ingest.errors import error_reason
58from lilbee.data.ingest.fanout import (
59 WORKER_LOG_NAME,
60 ShardDone,
61 ShardOptions,
62 ShardSpec,
63 aggregate_results,
64 plan_fanout,
65 run_workers,
66)
67from lilbee.data.ingest.ignore import IgnoreRules
68from lilbee.data.ingest.skip_marker import (
69 clear_skip_markers,
70 describe_skips,
71 load_skip_markers,
72 load_skip_reasons,
73 write_skip_markers,
74 write_skip_reasons,
75)
76from lilbee.data.offload import (
77 embed_inflight_target,
78 max_workers,
79 to_executor,
80 to_ingest_thread,
81)
82from lilbee.data.store import (
83 SOURCE_STAT_UNKNOWN,
84 ChunkWrite,
85 ConceptRecords,
86 IndexMismatch,
87 PageTextRecord,
88 SourceMeta,
89 SourceRecord,
90 SourceStat,
91 SourceStatBackfill,
92 SourceType,
93 Store,
94 source_stat,
95)
96from lilbee.data.title import derive_title
97from lilbee.data.types import (
98 ChunkRecord,
99 FileChangePlan,
100 FileToProcess,
101 MemberRecords,
102 ShardId,
103 SyncResult,
104 _IngestResult,
105)
106from lilbee.runtime.asyncio_loop import is_executor_shutdown
107from lilbee.runtime.cancellation import CancelSignal, TaskCancelledError
108from lilbee.runtime.cpu import available_cpu_count, cpu_quota
109from lilbee.runtime.lock import LockTimeoutError
110from lilbee.runtime.progress import (
111 BatchProgressEvent,
112 BatchStatus,
113 DetailedProgressCallback,
114 EmbedEvent,
115 EventType,
116 ExtractEvent,
117 FileDoneEvent,
118 FileStartEvent,
119 ProgressEvent,
120 SyncDoneEvent,
121 noop_callback,
122)
124log = logging.getLogger(__name__)
127def _max_concurrent() -> int:
128 """Files allowed in their compute phase at once.
130 ``cpu_quota()`` (cpu_count // 2) keeps worker storms from starving the TUI's asyncio
131 main thread, and is the cap for text/code ingest. Vision OCR is different: every file
132 in compute holds a continuous-batching slot on a vision server, so an OCR run is bounded
133 by the vision slot capacity. Once the fleet is up that capacity is the servers' real
134 fitted ``--parallel`` slots (a memory-constrained card fits fewer than requested); until
135 then it is estimated from ``replicas x per-server pages``. Sizing to real capacity keeps
136 the OCR queue shallow instead of piling pages behind the dispatcher with their deadlines
137 ticking, and still keeps every slot fed.
138 """
139 from lilbee.providers.fleet.replicas import gpu_device_count, resolve_replica_count
140 from lilbee.providers.roles import WorkerRole
142 config = active_config()
143 if config.vision_model:
144 fitted = get_services().provider.vision_slot_capacity()
145 if fitted is not None:
146 return fitted
147 replicas = resolve_replica_count(WorkerRole.VISION, gpu_device_count())
148 return max(1, replicas * config.vision_ocr_concurrency)
149 if config.ingest_max_inflight > 0:
150 return config.ingest_max_inflight # explicit override
151 # Auto: keep every embed replica fed. The CPU-bound quota alone leaves a
152 # many-core multi-GPU box starved (~4 files/card), so scale admission with
153 # the detected fleet size -- no manual cap needed.
154 return max(cpu_quota(), embed_inflight_target())
157async def _rebuild_concept_clusters() -> None:
158 """Re-run Leiden clustering after sync. No-op if disabled."""
159 if not active_config().concept_graph:
160 return
161 from lilbee.retrieval.concepts import concepts_available
163 if not concepts_available():
164 return
165 try:
166 cg = get_services().concepts
167 if not cg.get_graph():
168 return
169 await to_ingest_thread(cg.rebuild_clusters)
170 except Exception:
171 log.warning("Concept cluster rebuild failed", exc_info=True)
174async def build_entity_records(records: list[ChunkRecord], source_name: str) -> list[dict] | None:
175 """Extract typed entities for ingested chunks. None when the mode is off.
177 Gated twice: the ``entity_extraction`` config flag, and a schema already
178 induced into the index; absent either, syncs cost nothing. Extraction
179 failures degrade to no rows for the file, mirroring concept extraction.
180 """
181 config = active_config()
182 if not config.entity_extraction or not records:
183 return None
184 from lilbee.retrieval.entities import ExtractorKind, extract_entities, load_schema
186 schema = load_schema(get_services().store)
187 if schema is None:
188 return None
189 nlp = None
190 if any(t.kind is ExtractorKind.SPACY for t in schema.types):
191 from lilbee.retrieval.concepts import concepts_available
192 from lilbee.retrieval.concepts.nlp import load_spacy_pipeline
194 if concepts_available():
195 try:
196 nlp = load_spacy_pipeline()
197 except ImportError:
198 log.warning("spaCy model unavailable; spacy-kind entity types skipped")
199 provider = None
200 if any(t.kind is ExtractorKind.LLM for t in schema.types):
201 provider = get_services().provider
202 try:
203 return await to_ingest_thread(
204 extract_entities,
205 cast("list[Mapping[str, Any]]", records),
206 schema,
207 provider=provider,
208 nlp=nlp,
209 )
210 except Exception:
211 log.warning("Entity extraction failed for %s", source_name, exc_info=True)
212 return None
215async def build_concept_records(
216 records: list[ChunkRecord], source_name: str
217) -> ConceptRecords | None:
218 """Extract concepts for ingested chunks and build their table rows. None if disabled.
220 Pure record building, no store access: the rows are buffered on the file's
221 ingest result and written once per flush (see :func:`_flush_concept_records`),
222 so a large sync pays one concept-table write per flush, not per file.
223 """
224 if not active_config().concept_graph or not records:
225 return None
226 from lilbee.retrieval.concepts import concepts_available
228 if not concepts_available():
229 return None
230 try:
231 cg = get_services().concepts
232 texts = [r["chunk"] for r in records]
233 concept_lists = await to_ingest_thread(cg.extract_concepts_batch, texts)
234 chunk_ids = [(source_name, r["chunk_index"]) for r in records]
235 return await to_ingest_thread(cg.build_concept_records, chunk_ids, concept_lists)
236 except Exception:
237 log.warning("Concept extraction failed for %s", source_name, exc_info=True)
238 return None
241async def produce_records(
242 path: Path,
243 source_name: str,
244 content_type: str,
245 *,
246 quiet: bool = False,
247 on_progress: DetailedProgressCallback = noop_callback,
248 page_texts_out: list[PageTextRecord] | None = None,
249) -> tuple[list[ChunkRecord], SourceMeta]:
250 """Extract, chunk, and embed a single file into (records, source metadata).
252 The LanceDB write is deferred: records are returned to the caller and written
253 in a batched flush (see :func:`_flush_writes`), so bulk ingest pays one
254 write-lock acquisition per batch instead of one per file. The per-page text
255 dataset rows land in ``page_texts_out`` and are written by the same flush.
256 The returned metadata (extraction-provided when available, stem-derived title
257 otherwise) stamps every record's ``title`` and updates the source row.
258 """
259 records: list[ChunkRecord]
260 page_texts: list[PageTextRecord] = page_texts_out if page_texts_out is not None else []
261 if content_type == "code":
262 records = await to_ingest_thread(ingest_code_sync, path, source_name, on_progress)
263 meta = SourceMeta(title=derive_title(source_name))
264 elif path.suffix.lower() == ".md":
265 records, meta = await ingest_markdown(
266 path, source_name, on_progress, page_texts_out=page_texts
267 )
268 else:
269 records, meta = await ingest_document(
270 path,
271 source_name,
272 content_type,
273 quiet=quiet,
274 on_progress=on_progress,
275 page_texts_out=page_texts,
276 )
278 for record in records:
279 # NULL (not "") for an absent title, so chunk rows match the migration
280 # and the _sources table, which both persist absence as NULL.
281 record["title"] = meta.title or None
282 return records, meta
285def _disk_stat(path: Path) -> SourceStat | None:
286 """Current size/mtime of *path* stamped with now, or None when it cannot be stat'd."""
287 try:
288 st = path.stat()
289 except OSError:
290 return None
291 return SourceStat(st.st_size, st.st_mtime_ns, time.time_ns())
294def _stat_unchanged(stored: SourceStat, current: SourceStat) -> bool:
295 """Whether the stored stat proves the file unchanged without hashing it.
297 Git-style racily-clean guard: a matching (size, mtime) only counts when the
298 mtime is strictly older than the time the stat was recorded; a same-size
299 edit landing in the same mtime tick is otherwise missed forever.
300 """
301 if (stored.size_bytes, stored.mtime_ns) != (current.size_bytes, current.mtime_ns):
302 return False
303 # Unknown capture or mtime >= capture hashes anyway; clock skew can only widen
304 # hashing, never widen skipping past the pre-existing same-tick window.
305 return stored.captured_ns != SOURCE_STAT_UNKNOWN and current.mtime_ns < stored.captured_ns
308@dataclass(frozen=True)
309class _FileChangeVerdict:
310 """One file's sync verdict: process it, hold it out on its skip marker, or unchanged.
312 A held file is not in the index, so it is never counted as unchanged.
313 """
315 to_process: FileToProcess | None = None
316 backfill: SourceStatBackfill | None = None
317 is_update: bool = False
318 held: bool = False
321def _classify_file_change(
322 name: str,
323 path: Path,
324 record: SourceRecord | None,
325 skip_markers: dict[str, str],
326) -> _FileChangeVerdict:
327 """Decide one file's verdict: stat-unchanged, hash-unchanged, skip-marked, or process."""
328 content_type = classify_file(path)
329 if content_type is None:
330 raise ValueError(f"Unsupported file slipped through discovery: {name}")
331 stored_stat = source_stat(record) if record is not None else None
332 current_stat = _disk_stat(path)
333 if (
334 record is not None
335 and stored_stat is not None
336 and current_stat is not None
337 and _stat_unchanged(stored_stat, current_stat)
338 ):
339 return _FileChangeVerdict()
340 old_hash = record["file_hash"] if record is not None else None
341 current_hash = file_hash(path)
342 if old_hash == current_hash:
343 # Content verified unchanged; persist the stat pair so the next
344 # sync skips the hash entirely.
345 backfill = (
346 SourceStatBackfill(record, current_stat)
347 if record is not None and current_stat is not None
348 else None
349 )
350 return _FileChangeVerdict(backfill=backfill)
351 if skip_markers.get(name) == current_hash:
352 # Failed last sync at this exact hash; skip the retry.
353 return _FileChangeVerdict(held=True)
354 # needs_cleanup=True unconditionally: delete_by_source is idempotent,
355 # and this closes the race where a prior ingest wrote chunks but died
356 # before upsert_source, leaving orphaned chunks that would duplicate.
357 return _FileChangeVerdict(
358 to_process=FileToProcess(
359 name, path, content_type, current_hash, needs_cleanup=True, stat=current_stat
360 ),
361 is_update=old_hash is not None,
362 )
365def _plan_workers() -> int:
366 """Worker count for the parallel planning pass: config override, else auto.
368 ``config.ingest_workers`` (also set per run by ``add --max-cpus``) wins when
369 positive; otherwise size to the container-aware CPU budget so a big corpus
370 hashes on every core the pod actually has, not the host's vCPU count.
371 """
372 configured = active_config().ingest_workers
373 return configured if configured > 0 else available_cpu_count()
376# How often the plan pass logs progress. The pass can run for tens of minutes on
377# a multi-million-file corpus while the Rich bar renders nothing without a TTY
378# and stdout is block-buffered when piped; a periodic line (which logging flushes
379# per record) keeps a headless run observable instead of looking hung.
380_PLAN_LOG_INTERVAL_S = 10.0
383class _PlanProgress:
384 """Periodic progress for the plan/hash pass, with rate and ETA.
386 Emitted at warning level, not info: the default LILBEE_LOG_LEVEL is WARNING,
387 so an info line would be filtered before any handler and a headless
388 ``lilbee sync`` would show nothing during the plan pass and still look hung.
389 """
391 def __init__(self, total: int) -> None:
392 self._total = total
393 self._done = 0
394 self._started = time.monotonic()
395 self._last = self._started
397 def tick(self) -> None:
398 self._done += 1
399 now = time.monotonic()
400 if now - self._last < _PLAN_LOG_INTERVAL_S:
401 return
402 self._last = now
403 elapsed = now - self._started
404 rate = self._done / elapsed if elapsed > 0 else 0.0
405 remaining = (self._total - self._done) / rate if rate > 0 else 0.0
406 log.warning(
407 "Planning: examined %d/%d files (%.0f%%, %.0f files/s, ~%.0fs left)",
408 self._done,
409 self._total,
410 100.0 * self._done / self._total,
411 rate,
412 remaining,
413 )
416class _StreamStop:
417 """Stop signal for a streamed plan: the caller's cancel, or the stream closing.
419 Closing the stream has to reach the plan batch in flight, not just the next one.
420 Build-vs-buy: the hashers run in a thread pool, so the flag must be a
421 thread-visible ``threading.Event``; ``anyio.CancelScope`` is async-only.
422 """
424 def __init__(self, cancel: CancelSignal | None) -> None:
425 self._cancel = cancel
426 self._closed = threading.Event()
428 def close(self) -> None:
429 self._closed.set()
431 def is_set(self) -> bool:
432 return self._closed.is_set() or (self._cancel is not None and self._cancel.is_set())
435def _classify_pooled(
436 pool: ThreadPoolExecutor,
437 items: list[tuple[str, Path]],
438 classify: Callable[[str, Path], _FileChangeVerdict],
439 cancel: CancelSignal | None,
440 progress: _PlanProgress,
441) -> dict[str, _FileChangeVerdict]:
442 """Fan *items* across *pool*, returning name -> verdict for what completed."""
443 verdicts: dict[str, _FileChangeVerdict] = {}
444 futures: dict[Future[_FileChangeVerdict], str] = {}
445 for name, path in items:
446 if cancel and cancel.is_set():
447 break
448 futures[pool.submit(classify, name, path)] = name
449 for future in as_completed(futures):
450 if cancel and cancel.is_set():
451 # Drop queued-but-unstarted work; running tasks drain on their own.
452 for pending in futures:
453 pending.cancel()
454 break
455 verdicts[futures[future]] = future.result()
456 progress.tick()
457 return verdicts
460def _classify_changes(
461 items: list[tuple[str, Path]],
462 existing_sources: dict[str, SourceRecord],
463 skip_markers: dict[str, str],
464 cancel: CancelSignal | None,
465 *,
466 progress: _PlanProgress | None = None,
467 pool: ThreadPoolExecutor | None = None,
468) -> dict[str, _FileChangeVerdict]:
469 """Classify each file (stat + hash) by name, fanning across a thread pool.
471 Independent per file and side-effect-free, so pooled classification matches a
472 serial pass; ``hashlib`` releases the GIL during digest, giving real speedup
473 on a large corpus. Returns name -> verdict. A set ``cancel`` stops promptly:
474 submission halts and queued-but-unstarted work is cancelled, so a mid-pass
475 cancel over a huge corpus does not hash every remaining file. A *pool* and
476 *progress* passed in are shared across the batches of a streamed plan, so the
477 ETA covers the whole corpus and the workers are spun up once.
478 """
480 def _classify(name: str, path: Path) -> _FileChangeVerdict:
481 return _classify_file_change(name, path, existing_sources.get(name), skip_markers)
483 total = len(items)
484 progress = progress or _PlanProgress(total)
485 if pool is not None:
486 return _classify_pooled(pool, items, _classify, cancel, progress)
487 workers = _plan_workers()
488 if workers <= 1 or total <= 1:
489 verdicts: dict[str, _FileChangeVerdict] = {}
490 for name, path in items:
491 if cancel and cancel.is_set():
492 break
493 verdicts[name] = _classify(name, path)
494 progress.tick()
495 return verdicts
496 with ThreadPoolExecutor(max_workers=workers, thread_name_prefix="lilbee-plan") as owned:
497 return _classify_pooled(owned, items, _classify, cancel, progress)
500def _plan_file_changes(
501 disk_files: dict[str, Path],
502 existing_sources: dict[str, SourceRecord],
503 cancel: CancelSignal | None,
504 skip_markers: dict[str, str] | None = None,
505) -> FileChangePlan:
506 """Diff every disk file against the store in one pass (see :func:`_plan_items`)."""
507 return _plan_items(sorted(disk_files.items()), existing_sources, cancel, skip_markers or {})
510def _plan_items(
511 items: list[tuple[str, Path]],
512 existing_sources: dict[str, SourceRecord],
513 cancel: CancelSignal | None,
514 skip_markers: dict[str, str],
515 *,
516 progress: _PlanProgress | None = None,
517 pool: ThreadPoolExecutor | None = None,
518) -> FileChangePlan:
519 """Diff *items* (sorted name -> path) against the store, hashing only drifted files.
521 A tracked file whose stored (size, mtime) matches the disk stat, and whose
522 mtime predates the stat capture (see :func:`_stat_unchanged`), is unchanged
523 without reading its bytes; everything else is SHA-256 hashed. A file whose
524 current hash matches a marker in ``skip_markers`` (set by a prior failed
525 attempt) is held out rather than retried every sync, and is reported as held
526 out rather than counted unchanged. Edit the file or run
527 ``/sync --force-rebuild`` to clear the marker and try again.
529 Classification fans across a thread pool (see :func:`_classify_changes`); the
530 plan is assembled from the results in the original sorted order, so a partial
531 or reordered completion never yields a wrong or reordered plan -- only a
532 shorter one when cancelled mid-pass.
533 """
534 verdicts = _classify_changes(
535 items, existing_sources, skip_markers, cancel, progress=progress, pool=pool
536 )
538 files_to_process: list[FileToProcess] = []
539 added: dict[str, None] = {}
540 updated: dict[str, None] = {}
541 stat_backfills: list[SourceStatBackfill] = []
542 held_out: list[str] = []
543 unchanged = 0
544 # Assemble in the original sorted order so a partial (cancelled) or reordered
545 # completion never changes the plan the serial pass would have produced.
546 for name, _path in items:
547 verdict = verdicts.get(name)
548 if verdict is None:
549 continue # cancelled before this file was classified
550 if verdict.to_process is None:
551 if verdict.held:
552 held_out.append(name)
553 else:
554 unchanged += 1
555 if verdict.backfill is not None:
556 stat_backfills.append(verdict.backfill)
557 continue
558 files_to_process.append(verdict.to_process)
559 if verdict.is_update:
560 updated[name] = None
561 else:
562 added[name] = None
563 return FileChangePlan(files_to_process, added, updated, unchanged, stat_backfills, held_out)
566@dataclass(frozen=True)
567class _Move:
568 """One relocated source: its old key, its new key, and the new file's stat."""
570 old: str
571 new: str
572 stat: SourceStat | None
575class _MovePool:
576 """Absent sources indexed by content hash, consumed as moves are paired.
578 Built once per sync and drained across the streamed plan's batches, so a file
579 that moved matches exactly the one old key it would have matched in a
580 single-pass plan however the corpus is sharded.
581 """
583 def __init__(self, absent: list[str], existing_sources: dict[str, SourceRecord]) -> None:
584 by_hash: dict[str, list[str]] = {}
585 for name in absent:
586 record = existing_sources.get(name)
587 if record is not None:
588 by_hash.setdefault(record["file_hash"], []).append(name)
589 for candidates in by_hash.values():
590 candidates.sort()
591 self._by_hash = by_hash
593 def take(self, file_hash: str) -> str | None:
594 """The next absent source with this content hash, or None."""
595 matches = self._by_hash.get(file_hash)
596 return matches.pop(0) if matches else None
599def _detect_moves(
600 files_to_process: list[FileToProcess],
601 added: dict[str, None],
602 pool: _MovePool,
603) -> list[_Move]:
604 """Pair brand-new files with absent sources of the same content hash.
606 Only additions (files with a new name) can be moves; an update keeps its name.
607 When several absent sources share a hash, pairing is deterministic (sorted)
608 and one-to-one, so a duplicated file that moved matches exactly one old key
609 and any leftovers stay indexed under their old key.
610 """
611 moves: list[_Move] = []
612 for entry in files_to_process:
613 if entry.name not in added:
614 continue
615 old = pool.take(entry.file_hash)
616 if old is not None:
617 moves.append(_Move(old, entry.name, entry.stat))
618 return moves
621def _apply_moves(
622 moves: list[_Move],
623 files_to_process: list[FileToProcess],
624 added: dict[str, None],
625) -> tuple[list[FileToProcess], list[str]]:
626 """Fold detected moves out of the add set after they were relocated.
628 Drops each moved file from the ingest list and the added set: its chunks were
629 re-keyed onto the new source name, not rebuilt. Returns the trimmed
630 ``(files_to_process, relocated)``.
631 """
632 moved_new = {m.new for m in moves}
633 for name in moved_new:
634 added.pop(name, None)
635 remaining = [e for e in files_to_process if e.name not in moved_new]
636 return remaining, sorted(moved_new)
639def _absent_sources(sources: list[SourceRecord], disk_files: dict[str, Path]) -> list[str]:
640 """Document sources whose backing file is not on disk this sync.
642 A vanished file is never removed on its own (its chunks stay searchable, a
643 dead path-link discovered at open time); this set exists only to pair a
644 reappeared identical file to its old key in move detection. Imported sources
645 have no backing file, so they are excluded.
646 """
647 return [
648 s["filename"]
649 for s in sources
650 if s["filename"] not in disk_files and s["source_type"] != SourceType.IMPORTED
651 ]
654# A plan batch is only done when its slowest hash is, so the first one is small (work
655# reaches the fleet within a second) and later ones amortize that barrier.
656_PLAN_SHARD_MIN_FILES = 256
657_PLAN_SHARD_MAX_FILES = 8192
660def _plan_batch_bounds(total: int) -> Iterator[tuple[int, int]]:
661 """(start, stop) slices covering *total* files, doubling up to the cap.
663 Build-vs-buy: ``itertools.batched`` is the stock slicer but is 3.12+ (floor is
664 3.11) and fixed-size, so it cannot ramp the batch size.
665 """
666 start = 0
667 size = _PLAN_SHARD_MIN_FILES
668 while start < total:
669 stop = min(start + size, total)
670 yield start, stop
671 start = stop
672 size = min(size * 2, _PLAN_SHARD_MAX_FILES)
675@dataclass
676class _StreamedPlan:
677 """Bookkeeping a streamed plan accumulates across its batches.
679 ``added`` and ``updated`` are the dicts the ingest pass mutates as files land.
680 """
682 added: dict[str, None] = field(default_factory=dict)
683 updated: dict[str, None] = field(default_factory=dict)
684 # Processed files' content hashes, for the skip markers written after the run.
685 pending_hashes: dict[str, str] = field(default_factory=dict)
686 relocated: list[str] = field(default_factory=list)
687 # Old keys of relocated sources. The wiki index is keyed by source name, so
688 # without these a move leaves the old name in it forever: its mentions
689 # double-count and its dead chunk refs occupy the per-subject cap.
690 relocated_from: list[str] = field(default_factory=list)
691 unchanged: int = 0
692 # Files a skip marker held out of this run, in plan order (ordered set).
693 held_out: dict[str, None] = field(default_factory=dict)
694 planned: int = 0
695 # Files this pass's slice holds, from the discovery walk. Fixed before the
696 # first batch is planned, so it is what progress is measured against: the
697 # plan's own running totals grow as batches land and cannot say how much of
698 # the corpus is left. Summed across a fan-out's workers it is the corpus.
699 corpus_total: int = 0
701 @property
702 def resolved(self) -> int:
703 """Files the plan disposed of without ingest, as it disposes of them.
705 Unchanged files, skip-marker held-out files and repointed moves are done
706 as far as the corpus is concerned, and nothing downstream ingests them,
707 so an incremental sync would otherwise show a handful of changed files
708 against the whole corpus. Files a cancelled plan never classified are
709 absent from every count and so are not claimed as done.
710 """
711 return self.unchanged + len(self.held_out) + len(self.relocated)
714async def _absorb_plan_batch(
715 plan: FileChangePlan, state: _StreamedPlan, moves: _MovePool
716) -> list[FileToProcess]:
717 """Fold one batch's plan into *state* and return the files it queues for ingest.
719 Relocations and stat backfills are written per batch, so they contend with the
720 batch flushes and take the same one-shot lock retry.
721 """
722 store = get_services().store
723 entries = plan.files_to_process
724 state.unchanged += plan.unchanged
725 state.held_out.update(dict.fromkeys(plan.held_out))
726 state.added.update(plan.added)
727 state.updated.update(plan.updated)
729 detected = _detect_moves(entries, plan.added, moves)
730 if detected:
731 relocations = [(m.old, m.new, m.stat) for m in detected]
732 await to_ingest_thread(
733 _retry_after_lock_timeout, lambda: store.relocate_sources(relocations)
734 )
735 entries, relocated = _apply_moves(detected, entries, plan.added)
736 state.relocated_from.extend(m.old for m in detected)
737 for name in relocated:
738 state.added.pop(name, None)
739 state.relocated.extend(relocated)
740 if plan.stat_backfills:
741 backfills = plan.stat_backfills
742 await to_ingest_thread(
743 _retry_after_lock_timeout, lambda: store.update_source_stats(backfills)
744 )
746 state.pending_hashes.update((entry.name, entry.file_hash) for entry in entries)
747 state.planned += len(entries)
748 return entries
751async def _plan_batches(
752 disk_files: dict[str, Path],
753 existing_sources: dict[str, SourceRecord],
754 skip_markers: dict[str, str],
755 absent: list[str],
756 state: _StreamedPlan,
757 cancel: CancelSignal | None,
758) -> AsyncGenerator[list[FileToProcess]]:
759 """Plan the corpus batch by batch, yielding each batch's files to ingest.
761 Shards are contiguous slices of one sorted item list consumed in order, so
762 this delivers the single-pass plan of :func:`_plan_items` in pieces. The next
763 batch is planned while the current one ingests, on a dedicated thread: the
764 shared ingest pool is saturated by extraction and would stall the stream it
765 feeds. Empty batches are not yielded, so the first yield means there is work.
766 """
767 items = sorted(disk_files.items())
768 moves = _MovePool(absent, existing_sources)
769 progress = _PlanProgress(len(items))
770 stop = _StreamStop(cancel)
771 workers = _plan_workers()
772 hashers = ThreadPoolExecutor(max_workers=workers, thread_name_prefix="lilbee-plan")
773 driver = ThreadPoolExecutor(max_workers=1, thread_name_prefix="lilbee-plan-driver")
775 def _plan(lo: int, hi: int) -> FileChangePlan:
776 return _plan_items(
777 items[lo:hi],
778 existing_sources,
779 stop,
780 skip_markers,
781 progress=progress,
782 pool=hashers if workers > 1 else None,
783 )
785 bounds = list(_plan_batch_bounds(len(items)))
786 try:
787 ahead = asyncio.ensure_future(to_executor(driver, _plan, *bounds[0])) if bounds else None
788 for index, _bound in enumerate(bounds):
789 if ahead is None:
790 break
791 plan = await ahead
792 ahead = (
793 asyncio.ensure_future(to_executor(driver, _plan, *bounds[index + 1]))
794 if index + 1 < len(bounds) and not stop.is_set()
795 else None
796 )
797 entries = await _absorb_plan_batch(plan, state, moves)
798 if entries:
799 yield entries
800 finally:
801 stop.close()
802 if ahead is not None:
803 ahead.cancel()
804 # Retrieve the outcome: a prefetch that had already failed would
805 # otherwise surface as an unretrieved task exception.
806 with contextlib.suppress(Exception, asyncio.CancelledError):
807 await ahead
808 # No wait: the running batch notices the stop and exits on its own, and a
809 # generator being closed must not block the event loop on a hash in flight.
810 hashers.shutdown(wait=False)
811 driver.shutdown(wait=False)
814def detect_pending() -> int:
815 """Count files in documents/ that are out of sync with the store.
817 Cheap operation: filesystem walk + stat-gated SHA-256 hashing + a single
818 sources-table read. No embedding, no writes. Returns the count of files that
819 would be ingested (added + updated), which is what the TaskBar hint surfaces.
820 A vanished file is not pending work: sync leaves it indexed, so it is not
821 counted. Reuses ``_plan_file_changes`` so the diff logic stays single-sourced.
822 Honors skip markers: a file that failed last time at this hash does
823 not show up as pending. Blocking: callers on the event loop run it via
824 ``asyncio.to_thread``.
825 """
826 config = active_config()
827 disk_files = discover_files()
828 if not disk_files:
829 return 0
830 existing_sources = {s["filename"]: s for s in get_services().store.get_sources()}
831 skip_markers = load_skip_markers(config.data_root)
832 plan = _plan_file_changes(disk_files, existing_sources, cancel=None, skip_markers=skip_markers)
833 return len(plan.files_to_process)
836# Refused files named in the log line before it falls back to a count.
837_EXCLUDED_LOG_SAMPLE = 5
840def _log_excluded(excluded: dict[str, ExclusionReason]) -> None:
841 """Warn once per exclusion reason, naming at most ``_EXCLUDED_LOG_SAMPLE`` files."""
842 by_reason: dict[ExclusionReason, list[str]] = {}
843 for name, why in excluded.items():
844 by_reason.setdefault(why, []).append(name)
845 for why, names in by_reason.items():
846 shown = sorted(names)[:_EXCLUDED_LOG_SAMPLE]
847 more = f" and {len(names) - len(shown)} more" if len(names) > len(shown) else ""
848 log.warning("Skipped %d file(s), %s: %s%s", len(names), why.value, ", ".join(shown), more)
851def _load_sync_skip_markers(*, clear_first: bool) -> dict[str, str]:
852 """Read the skip-marker file, optionally clearing it first.
854 Entries are kept whether or not this pass discovered the file. A marker is
855 what holds a removed or unextractable file out of the next sync, so a pass
856 that cannot see the file -- its root is unmounted or moved, or this worker
857 owns only a shard of the corpus -- must not erase the record and re-offer
858 the file the moment it comes back. A marker is dropped when the file
859 ingests cleanly, or by ``retry-skipped`` / ``rebuild``.
860 """
861 data_root = active_config().data_root
862 if clear_first:
863 # Clearing the markers makes the diff re-include the skipped files.
864 clear_skip_markers(data_root)
865 return load_skip_markers(data_root)
868def _persist_skip_markers(
869 markers: dict[str, str],
870 pending_hashes: dict[str, str],
871 *,
872 succeeded: Iterable[str],
873 failed: Iterable[str],
874) -> None:
875 """Mark files that produced no chunks so the next sync skips them, clear the
876 markers for files that ingested cleanly, then write the file back."""
877 for name in succeeded:
878 markers.pop(name, None)
879 for name in failed:
880 fhash = pending_hashes.get(name)
881 if fhash:
882 markers[name] = fhash
883 write_skip_markers(active_config().data_root, markers)
886def _persist_skip_reasons(markers: dict[str, str], reasons: dict[str, str]) -> None:
887 """Write the reasons sidecar so it explains exactly the markers that survive.
889 Merged onto what is already recorded, not replaced: the reasons for files
890 this sync never touched (a removal, an earlier failure) explain markers that
891 are still in force, and dropping them would leave the user with a file held
892 out of every sync and nothing saying why. Reasons whose marker is gone are
893 dropped in the same step, so the two sidecars cannot drift apart.
894 """
895 data_root = active_config().data_root
896 merged = load_skip_reasons(data_root) | reasons
897 write_skip_reasons(data_root, {name: why for name, why in merged.items() if name in markers})
900def _report_index_mismatch(store: Store) -> IndexMismatch | None:
901 """Name an index built with another embedder than cfg; the sync leaves it as it is.
903 An unchanged corpus never reaches the write gate that refuses such an index,
904 so without this a sync finishes green over an index search refuses. Wiping
905 it is the caller's decision, never the sync's.
906 """
907 mismatch = store.index_mismatch()
908 if mismatch is None:
909 return None
910 log.warning("Sync left the index as it is: %s", mismatch)
911 return mismatch.describe()
914def _force_rebuild_store(store: Any) -> None:
915 """Drop the store and re-embed the preserved memories table (blocking).
917 Run off the event loop by ``sync``. ``drop_all`` keeps the memories table, so
918 its vectors are refreshed under the (possibly changed) embedding model;
919 a no-op when empty or no embedder.
920 """
921 store.drop_all()
922 embedder = get_services().embedder
923 if embedder.embedding_available():
924 store.rebuild_memory_embeddings(lambda texts: embedder.embed_batch(texts))
927def _reconcile_missing(
928 disk_files: dict[str, Path],
929 sources: list[SourceRecord],
930 failed: Iterable[str],
931 skipped: Iterable[str],
932 held: Iterable[str],
933) -> list[str]:
934 """On-disk document files absent from the store that no mechanism accounts for.
936 A file discovery found and classified that ended up in neither the sources table
937 nor any of the accounting sets was dropped with no signal -- the silent
938 data-loss case (a scanned PDF that never made it into the index yet reported no
939 error). Everything legitimately not indexed is excluded: a failed extraction is in
940 ``failed``, a zero-text file this run attempted is in ``skipped``, a file this run
941 held out on its skip marker is in ``held``, and an unsupported type was never
942 returned by discovery in the first place. ``held`` is the run's held-out set, not
943 the marker file: a stale marker (file edited since it failed) must not hide a drop.
944 """
945 accounted = {s["filename"] for s in sources} | set(failed) | set(skipped) | set(held)
946 return sorted(name for name in disk_files if name not in accounted)
949def _ignored_sources(sources: list[SourceRecord], rules: IgnoreRules) -> list[str]:
950 """Indexed sources a ``.lilbeeignore`` now excludes.
952 Asked of the patterns rather than of discovery, which cannot tell an excluded
953 file from a lost one -- both simply stop being yielded, and the two need
954 opposite handling. Reading every source rather than only the undiscovered
955 ones keeps the answer independent of which slice of the corpus this pass
956 walked. Imported sources have no file to match a pattern against.
957 """
958 ignored = []
959 for source in sources:
960 if source["source_type"] == SourceType.IMPORTED:
961 continue
962 resolved = resolve_source_root(source["filename"])
963 if resolved is not None and rules.excludes_path(resolved[1], base=resolved[0]):
964 ignored.append(source["filename"])
965 return ignored
968def _forget_ignored(sources: list[SourceRecord], rules: IgnoreRules) -> list[str]:
969 """Drop sources the patterns now exclude from the index. Returns what went.
971 A corpus-wide pass: it reads every source, so one worker of a fan-out must
972 not run it. No skip marker is written either, because the ignore file is
973 itself the durable statement -- a marker could only outlive it and hold a
974 source out after its pattern was deleted. Deleting the pattern brings the
975 source back on the next sync.
976 """
977 names = _ignored_sources(sources, rules)
978 if not names:
979 return []
980 from lilbee.app.ingest import forget_removed_from_wiki_index
982 removed = list(get_services().store.remove_documents(names).removed)
983 forget_removed_from_wiki_index(removed)
984 return removed
987def _forget_refused(
988 excluded: Mapping[str, ExclusionReason], existing: Mapping[str, SourceRecord]
989) -> list[str]:
990 """Drop indexed sources that discovery now refuses. Returns what went."""
991 names = [name for name in excluded if name in existing]
992 if not names:
993 return []
994 from lilbee.app.ingest import forget_removed_from_wiki_index
996 removed = list(get_services().store.remove_documents(names).removed)
997 forget_removed_from_wiki_index(removed)
998 return removed
1001def _require_embedding_model() -> None:
1002 """Refuse ingest without an embedding model.
1004 Ingest has no degraded mode: without one, every chunk fails to embed after
1005 the run has already paid the parse and OCR cost. Search and chat fall back
1006 to keyword via embedding_available() and carry on.
1007 """
1008 if not get_services().embedder.validate_model():
1009 ref = active_config().embedding_model
1010 detail = (
1011 f"{ref!r} is not available: pull it, or set a different embedding_model."
1012 if ref
1013 else "None is configured: pull one and set embedding_model."
1014 )
1015 raise RuntimeError(f"Ingest needs an embedding model. {detail}")
1018async def _run_post_ingest_passes(
1019 store: Any,
1020 *,
1021 indexed_anything: bool,
1022 touched: set[str],
1023 cancel: CancelSignal | None,
1024) -> None:
1025 """Index maintenance, concept clusters, the wiki hook, and the entity lifecycle.
1027 The index and wiki passes only run when this sync indexed something. The
1028 entity lifecycle runs every sync (a cheap no-op when off or already current)
1029 so turning the setting on takes effect without a separate operation.
1030 """
1031 if indexed_anything:
1032 store.ensure_fts_index()
1033 store.ensure_scalar_indexes()
1034 store.ensure_vector_index()
1035 store.optimize_sources()
1036 await _rebuild_concept_clusters()
1037 await _update_wiki(touched, active_config())
1039 from lilbee.retrieval.entities.lifecycle import ensure_entities
1041 await to_ingest_thread(ensure_entities, cancel)
1044async def _update_wiki(changed_sources: set[str], config: Config) -> None:
1045 """Refresh the wiki index, and regenerate pages when auto-update is on.
1047 The index refresh runs whenever the wiki is enabled, because it spends no
1048 LLM call and it is what lets the browse tree list a page the moment its
1049 document lands. Regeneration is the expensive half and stays behind
1050 wiki_auto_update.
1052 Best effort: the ingest itself already succeeded and `lilbee wiki update`
1053 re-runs the regeneration, so a wiki failure must not skip the post-ingest
1054 entity pass or the reconciliation guard.
1055 """
1056 if not config.wiki:
1057 return
1058 # circular: lilbee.wiki imports lilbee.data.ingest.file_hash, so the
1059 # post-ingest hook stays function-local at this boundary.
1060 from lilbee.wiki.ingest import incremental_update
1061 from lilbee.wiki.stubs import refresh_stub_index
1063 try:
1064 await to_ingest_thread(
1065 refresh_stub_index, get_services().store, config, sources=changed_sources
1066 )
1067 except Exception:
1068 log.warning("Wiki index refresh failed after sync", exc_info=True)
1070 if not config.wiki_auto_update:
1071 return
1072 try:
1073 await incremental_update(changed_sources, config)
1074 except Exception:
1075 log.warning("Wiki auto-update failed after sync", exc_info=True)
1078def _worker_failure_message(failures: list[ShardDone], specs: list[ShardSpec]) -> str:
1079 """What a failed fan-out reports; its workers' shards are kept for the re-run."""
1080 detail = "; ".join(
1081 f"worker {failure.index} ({specs[failure.index].config.data_root / WORKER_LOG_NAME}): "
1082 f"{failure.error}"
1083 for failure in failures
1084 )
1085 return (
1086 f"{len(failures)} ingest worker(s) failed, so the index was not updated: {detail}. "
1087 "Their work is kept: re-running sync continues from where they stopped."
1088 )
1091def _merge_worker_shards(store: Store, specs: list[ShardSpec], touched: set[str]) -> None:
1092 """Fold every worker's shard into this index.
1094 A store with no chunks of its own takes the shards whole; one that already
1095 holds a corpus takes only the sources this run touched, so a re-sync replaces
1096 those rows instead of appending a second copy of everything.
1097 """
1098 from lilbee.data.store.shard_merge import merge_shards
1100 scope = touched if store.has_chunks() else None
1101 merge_shards(store, [spec.config.lancedb_dir for spec in specs], sources=scope)
1104async def _sync_across_workers(
1105 specs: list[ShardSpec],
1106 store: Store,
1107 *,
1108 prune_ignored: bool = False,
1109 options: ShardOptions,
1110 quiet: bool,
1111 on_progress: DetailedProgressCallback,
1112 cancel: CancelSignal | None,
1113) -> SyncResult:
1114 """Ingest on one worker per GPU, then fold their shards into the one index."""
1115 verdicts = await run_workers(
1116 specs, options=options, quiet=quiet, on_progress=on_progress, cancel=cancel
1117 )
1118 if cancel is not None and cancel.is_set():
1119 raise asyncio.CancelledError
1120 if failures := [verdict for verdict in verdicts if verdict.error is not None]:
1121 raise RuntimeError(_worker_failure_message(failures, specs))
1122 result = aggregate_results(verdicts)
1123 touched = set(result.added) | set(result.updated) | set(result.relocated)
1124 await to_ingest_thread(_merge_worker_shards, store, specs, touched)
1125 # No worker sees the whole corpus, so each one leaves this pass to the parent.
1126 if prune_ignored:
1127 result.removed = await to_ingest_thread(
1128 _forget_ignored, store.get_sources(), IgnoreRules.for_corpus()
1129 )
1130 await _run_post_ingest_passes(
1131 store, indexed_anything=bool(touched), touched=touched, cancel=cancel
1132 )
1133 on_progress(
1134 EventType.DONE,
1135 SyncDoneEvent(
1136 added=len(result.added),
1137 updated=len(result.updated),
1138 removed=len(result.removed),
1139 failed=len(result.failed),
1140 skipped=len(result.skipped),
1141 relocated=len(result.relocated),
1142 ),
1143 )
1144 return result
1147async def sync(
1148 force_rebuild: bool = False,
1149 quiet: bool = False,
1150 *,
1151 on_progress: DetailedProgressCallback = noop_callback,
1152 cancel: CancelSignal | None = None,
1153 retry_skipped: bool = False,
1154 prune_ignored: bool = False,
1155 shard: ShardId | None = None,
1156) -> SyncResult:
1157 """Sync documents/ with the vector store.
1158 Returns a SyncResult with the added/updated/removed/unchanged/failed/skipped lists.
1159 When *quiet* is True, the Rich progress bar is suppressed (for JSON output).
1160 When *cancel* is set mid-run, planning and processing stop between files
1161 without data loss (completed work is flushed) and CancelledError is raised;
1162 a cancel already set on entry returns an empty result instead.
1163 When *retry_skipped* (or *force_rebuild*) is set, the failed-file skip
1164 markers are cleared so this sync attempts every file.
1165 When *prune_ignored* is set, sources a ``.lilbeeignore`` now excludes are
1166 dropped from the index. Off by default: the patterns govern what sync takes
1167 in, and removing what a past sync already indexed is the caller's decision.
1168 A *shard* runs this sync as one worker of a multi-GPU fan-out: it sees only
1169 its slice of the corpus and leaves the corpus-wide passes to the parent.
1170 """
1171 config = active_config()
1172 _store = get_services().store
1174 if force_rebuild:
1175 # drop_all + memory re-embedding are heavy blocking store work; run them
1176 # off the event loop so a rebuild doesn't stall other admitted requests.
1177 await to_ingest_thread(_force_rebuild_store, _store)
1179 config.documents_dir.mkdir(parents=True, exist_ok=True)
1180 index_mismatch = _report_index_mismatch(_store)
1182 if shard is None and (specs := plan_fanout()):
1183 merged = await _sync_across_workers(
1184 specs,
1185 _store,
1186 prune_ignored=prune_ignored,
1187 options=ShardOptions(
1188 parent_pid=os.getpid(),
1189 force_rebuild=force_rebuild,
1190 retry_skipped=retry_skipped,
1191 ),
1192 quiet=quiet,
1193 on_progress=on_progress,
1194 cancel=cancel,
1195 )
1196 return merged.model_copy(update={"index_mismatch": index_mismatch})
1198 rules = IgnoreRules.for_corpus()
1199 scan = discover_corpus(shard, rules)
1200 disk_files = scan.files
1201 sources = _store.get_sources()
1202 existing_sources = {s["filename"]: s for s in sources}
1203 skip_markers = _load_sync_skip_markers(clear_first=force_rebuild or retry_skipped)
1205 failed: dict[str, None] = {}
1206 # Refused formats start the run skipped; they get no skip marker (no planned hash).
1207 skipped: dict[str, None] = dict.fromkeys(scan.excluded)
1208 # filename → why it was skipped/failed (for reporting)
1209 reasons: dict[str, str] = {name: why.value for name, why in scan.excluded.items()}
1210 flush_failed: set[str] = set()
1211 _log_excluded(scan.excluded)
1213 # Opt-in, and corpus-wide: a worker sees one slice but the whole sources
1214 # table, so it leaves this pass to the parent rather than racing its siblings.
1215 ignored = _forget_ignored(sources, rules) if prune_ignored and shard is None else []
1216 # Shard-safe: a worker removes only keys from its own slice.
1217 refused = _forget_refused(scan.excluded, existing_sources)
1219 # Sources whose backing file is not on disk this pass. A vanished file is NOT
1220 # removed: it stays indexed and searchable, a dead path-link the user
1221 # discovers only when they try to open it, and the set pairs a reappeared
1222 # identical file to its old key below. What was just removed leaves the set,
1223 # where it could otherwise capture a real move.
1224 gone = set(ignored) | set(refused)
1225 absent = [name for name in _absent_sources(sources, disk_files) if name not in gone]
1227 # The planning pass stats (and where needed hashes) every file on disk, batch by
1228 # batch off the event loop and overlapped with ingest. A brand-new file whose
1229 # content hash matches an absent source is folded in per batch as a move, not an
1230 # add: repointed in place so its chunks and embeddings are reused, not rebuilt.
1231 state = _StreamedPlan(corpus_total=len(disk_files))
1232 added, updated, pending_hashes = state.added, state.updated, state.pending_hashes
1233 plan_batches = _plan_batches(disk_files, existing_sources, skip_markers, absent, state, cancel)
1235 # Snapshot the cumulative truncation counter so the delta over this sync can
1236 # surface "N chunks truncated" instead of being lost in per-chunk debug logs.
1237 truncated_before = get_services().embedder.truncated_total
1239 # Ingest files (with optional progress bar). Only non-empty batches are yielded,
1240 # so the first one arriving is what proves there is work to do.
1241 try:
1242 first = await anext(plan_batches, None)
1243 if first is not None:
1244 # Hold the embed fleet resident for the whole batch: an unevenly loaded
1245 # replica must not idle-unload and reload cold mid-run (which snowballs
1246 # into a fleet collapse). The ContextVar propagates into the ingest
1247 # thread pool, where the fleet actually spawns on the first embed.
1248 from lilbee.providers.fleet.ingest_warmth import keep_fleet_warm
1250 with keep_fleet_warm():
1251 _require_embedding_model()
1252 await ingest_stream(
1253 _chain_plan_batches(first, plan_batches),
1254 added,
1255 updated,
1256 failed,
1257 skipped,
1258 plan=state,
1259 quiet=quiet,
1260 on_progress=on_progress,
1261 cancel=cancel,
1262 flush_failed=flush_failed,
1263 reasons=reasons,
1264 )
1265 if cancel is not None and cancel.is_set():
1266 # The stream stops feeding on cancel, so ingest can drain its
1267 # admitted files and return without raising. A cancelled run must
1268 # not go on to write skip markers or reconcile an unplanned corpus.
1269 raise asyncio.CancelledError
1270 finally:
1271 # Idempotent, and the only close when the stream is never consumed.
1272 await plan_batches.aclose()
1273 relocated = state.relocated
1275 # A flush failure is a transient store-side problem, not a verdict on the
1276 # file: leaving it unmarked re-plans it next sync instead of skipping it.
1277 marker_failed = [name for name in (*failed, *skipped) if name not in flush_failed]
1278 _persist_skip_markers(
1279 skip_markers, pending_hashes, succeeded=[*added, *updated], failed=marker_failed
1280 )
1281 # Record why each file this run skip-marked was marked (informational; the
1282 # hash markers above drive the resume logic). Only marker_failed files, so a
1283 # transient flush failure doesn't leave a stale reason behind.
1284 _persist_skip_reasons(skip_markers, {n: reasons[n] for n in marker_failed if n in reasons})
1286 if shard is None:
1287 # A worker's shard is merged before the indexes are built, so the passes
1288 # run once corpus-wide in the parent instead of once per shard.
1289 await _run_post_ingest_passes(
1290 _store,
1291 indexed_anything=bool(state.planned or relocated),
1292 # The old names of relocated sources ride along so the wiki index
1293 # subtracts them in the same pass that merges their new ones.
1294 touched=set(added) | set(updated) | set(relocated) | set(state.relocated_from),
1295 cancel=cancel,
1296 )
1298 # Reconciliation guard against silent data loss: any on-disk document file that
1299 # ended up in neither the index nor an accounting set was dropped without a
1300 # signal. Surface it loudly instead of letting a whole dataset vanish quietly.
1301 if missing := _reconcile_missing(
1302 disk_files, _store.get_sources(), failed, skipped, state.held_out
1303 ):
1304 log.warning(
1305 "Sync reconciliation: %d document file(s) on disk are absent from the index "
1306 "with no failure reported (possible silent drop): %s",
1307 len(missing),
1308 ", ".join(missing[:20]),
1309 )
1311 result = SyncResult(
1312 added=list(added),
1313 updated=list(updated),
1314 removed=ignored + refused,
1315 unchanged=state.unchanged,
1316 relocated=relocated,
1317 failed=list(failed),
1318 skipped=list(skipped),
1319 held_out=describe_skips(config.data_root, state.held_out),
1320 truncated=get_services().embedder.truncated_total - truncated_before,
1321 index_mismatch=index_mismatch,
1322 )
1323 on_progress(
1324 EventType.DONE,
1325 SyncDoneEvent(
1326 added=len(result.added),
1327 updated=len(result.updated),
1328 removed=len(result.removed),
1329 failed=len(result.failed),
1330 skipped=len(result.skipped),
1331 relocated=len(result.relocated),
1332 ),
1333 )
1334 return result
1337def _phase_progress_callback(
1338 progress: Progress, ptask: Any, chain: DetailedProgressCallback
1339) -> DetailedProgressCallback:
1340 """Wrap *chain*, updating the bar's description on per-page / per-chunk events.
1342 EXTRACT (vision OCR page i/N) and EMBED (chunk i/N) events would otherwise
1343 leave the bar frozen between file completions; surfacing them on the spinner
1344 description keeps a single large file's row visibly moving. All events still
1345 forward to *chain* so the caller's own callback (TUI / JSON) is unaffected.
1346 """
1348 def _callback(event_type: EventType, data: ProgressEvent) -> None:
1349 if event_type is EventType.EXTRACT and isinstance(data, ExtractEvent):
1350 progress.update(
1351 ptask, description=f"OCR {data.file} (page {data.page}/{data.total_pages})"
1352 )
1353 elif event_type is EventType.EMBED and isinstance(data, EmbedEvent):
1354 progress.update(
1355 ptask, description=f"Embedding {data.file} ({data.chunk}/{data.total_chunks})"
1356 )
1357 chain(event_type, data)
1359 return _callback
1362# In-flight task cap, as a multiple of _max_concurrent(): enough queued tasks to
1363# keep every compute slot fed, without materializing one task object per file.
1364_TASK_WINDOW_MULTIPLIER = 2
1367def _build_admission(
1368 baseline: int, pages_done: list[int]
1369) -> tuple[asyncio.Semaphore | ResizableGate, int, asyncio.Task[None] | None]:
1370 """The batch's admission control, plus its task-window size and controller task.
1372 Static mode (the default) returns a fixed semaphore and no controller. Adaptive
1373 mode, when a GPU fleet is present to feed, returns a resizable gate and a running
1374 :class:`AdaptiveController` that tunes it toward this box's throughput knee; with
1375 no fleet it falls back to the static path so a GPU-less host is never affected.
1376 """
1377 profile = profile_for(resolve_mode())
1378 devices = enumerate_fleet_devices() if profile is not None else []
1379 if profile is None or not devices:
1380 return asyncio.Semaphore(baseline), baseline * _TASK_WINDOW_MULTIPLIER, None
1381 permit_max = max_workers()
1382 gate = ResizableGate(min(baseline, permit_max))
1383 controller = AdaptiveController(
1384 gate,
1385 profile,
1386 make_signal_sampler(devices),
1387 lambda: pages_done[0],
1388 permit_min=1,
1389 permit_max=permit_max,
1390 )
1391 task = asyncio.ensure_future(controller.run())
1392 # warning, not info: the default LILBEE_LOG_LEVEL is WARNING, so the
1393 # auto-chosen concurrency would otherwise never surface on a headless sync.
1394 log.warning(
1395 "Adaptive ingest concurrency (%s): start %d, max %d", profile.name, gate.limit, permit_max
1396 )
1397 return gate, permit_max * _TASK_WINDOW_MULTIPLIER, task
1400def _failed_result(
1401 exc: Exception,
1402 entry: FileToProcess,
1403 *,
1404 pages_done: list[int],
1405 on_progress: DetailedProgressCallback,
1406 cancel: CancelSignal | None,
1407) -> _IngestResult:
1408 """A file's failure as a result, or cancellation when the run is stopping.
1410 During shutdown, worker pools raise RuntimeError from submit(). Those are
1411 cancellation, not ingest failures: the cancel flag is the source of truth,
1412 and the executor's shutdown message covers the race where cancel was set
1413 after the submit.
1414 """
1415 if (cancel and cancel.is_set()) or is_executor_shutdown(exc):
1416 raise asyncio.CancelledError from exc
1417 # Suppress TaskCancelledError on the FILE_DONE notice: the user already
1418 # cancelled, and re-raising here would strand sibling tasks awaiting in
1419 # _collect_results.
1420 with contextlib.suppress(TaskCancelledError):
1421 on_progress(EventType.FILE_DONE, FileDoneEvent(file=entry.name, status="error", chunks=0))
1422 pages_done[0] += 1 # cleared the gate (as a failure); still a throughput tick
1423 return _IngestResult(entry.name, entry.path, 0, error=exc)
1426async def _archive_result(
1427 entry: FileToProcess, on_progress: DetailedProgressCallback, pages_done: list[int]
1428) -> _IngestResult:
1429 """Ingest an archive: its members become sources, the archive row keeps the disk stat."""
1430 members = await ingest_archive(
1431 entry.path, entry.name, entry.content_type, on_progress=on_progress
1432 )
1433 concept_batches = [await build_concept_records(m.records, m.name) for m in members]
1434 entity_rows = [
1435 row for m in members for row in (await build_entity_records(m.records, m.name) or [])
1436 ]
1437 chunk_total = sum(len(m.records) for m in members)
1438 on_progress(
1439 EventType.FILE_DONE, FileDoneEvent(file=entry.name, status="ok", chunks=chunk_total)
1440 )
1441 pages_done[0] += max(1, sum(len(m.page_texts) for m in members))
1442 found = [batch for batch in concept_batches if batch is not None]
1443 return _IngestResult(
1444 entry.name,
1445 entry.path,
1446 chunk_total,
1447 error=None,
1448 file_hash=entry.file_hash,
1449 records=[],
1450 needs_cleanup=entry.needs_cleanup,
1451 page_texts=[],
1452 stat=entry.stat,
1453 concept_records=ConceptRecords.merged(found) if found else None,
1454 entity_rows=entity_rows or None,
1455 meta=SourceMeta(title=derive_title(entry.name)),
1456 members=members,
1457 )
1460def _over_limit_result(
1461 exc: ChunkLimitError,
1462 entry: FileToProcess,
1463 *,
1464 pages_done: list[int],
1465 on_progress: DetailedProgressCallback,
1466) -> _IngestResult:
1467 """A file over the per-file chunk limit, recorded as skipped and skip-marked at its hash."""
1468 log.warning("Skipped %s: %s", entry.name, exc)
1469 with contextlib.suppress(TaskCancelledError):
1470 on_progress(EventType.FILE_DONE, FileDoneEvent(file=entry.name, status="skipped", chunks=0))
1471 pages_done[0] += 1
1472 return _IngestResult(
1473 entry.name,
1474 entry.path,
1475 0,
1476 error=None,
1477 file_hash=entry.file_hash,
1478 skip_reason=str(exc),
1479 needs_cleanup=entry.needs_cleanup,
1480 )
1483async def _stream_tasks(
1484 plan_batches: AsyncGenerator[list[FileToProcess]],
1485 make_task: Callable[[FileToProcess, int], Coroutine[Any, Any, _IngestResult]],
1486) -> AsyncGenerator[list[Coroutine[Any, Any, _IngestResult]]]:
1487 """Each plan batch's per-file coroutines, in plan order."""
1488 index = count(1)
1489 async for plan_batch in plan_batches:
1490 yield [make_task(entry, next(index)) for entry in plan_batch]
1493async def _chain_plan_batches(
1494 first: list[FileToProcess], rest: AsyncGenerator[list[FileToProcess]]
1495) -> AsyncGenerator[list[FileToProcess]]:
1496 """Yield an already-pulled plan batch, then the remainder of its stream."""
1497 yield first
1498 async for plan_batch in rest:
1499 yield plan_batch
1502async def ingest_stream(
1503 plan_batches: AsyncGenerator[list[FileToProcess]],
1504 added: dict[str, None],
1505 updated: dict[str, None],
1506 failed: dict[str, None],
1507 skipped: dict[str, None],
1508 *,
1509 plan: _StreamedPlan | None = None,
1510 quiet: bool = False,
1511 on_progress: DetailedProgressCallback = noop_callback,
1512 cancel: CancelSignal | None = None,
1513 flush_failed: set[str] | None = None,
1514 reasons: dict[str, str] | None = None,
1515) -> None:
1516 """Ingest a stream of planned file batches, optionally showing a Rich progress bar.
1518 Files are admitted as their batch is planned, so ingest starts on the first
1519 batch instead of waiting for the whole corpus to be diffed. Old chunks are
1520 deleted in the same transaction as the new write, so the two are atomic per
1521 file. When *cancel* is set, pending files raise CancelledError before starting.
1523 *plan* is the bookkeeping the batches were planned into; it carries the corpus
1524 the run is measured against. Without it progress is reported with no total,
1525 since a bare stream of batches does not say what corpus it came from.
1526 """
1527 # Honor LILBEE_INGEST_TRACE once per batch: it raises the trace loggers above
1528 # the default WARNING so per-file extraction lines actually surface.
1529 configure_trace_from_env()
1530 warn_if_table_model_ignored()
1531 # Throughput is measured in OCR pages, not documents: a document's cost scales
1532 # with its page count (a 500-page scan is 500x a memo), so pages are the unbiased
1533 # unit of GPU-feeding work for the adaptive controller to hill-climb on.
1534 pages_done = [0]
1535 # Sized off the files with no source row yet, not the planned count: the plan
1536 # streams in batches and its total is unknown until the stream drains, but the
1537 # pool has to be decided before the first batch is dispatched. Unindexed files
1538 # are the one part of the plan that is known without diffing, exact for a first
1539 # ingest or a rebuild and near zero for an incremental sync, so a small sync
1540 # over a large corpus stays in-process. Undercounting only keeps a run
1541 # in-process, which is the safe direction.
1542 admission, window, controller_task = _build_admission(_max_concurrent(), pages_done)
1544 async def _process_one(entry: FileToProcess, file_index: int) -> _IngestResult:
1545 name = entry.name
1546 async with admission:
1547 if cancel and cancel.is_set():
1548 raise asyncio.CancelledError
1550 try:
1551 on_progress(
1552 EventType.FILE_START,
1553 FileStartEvent(file=name, total_files=feed.planned, current_file=file_index),
1554 )
1555 except TaskCancelledError as exc:
1556 # FILE_START itself can raise the cooperative cancel signal;
1557 # normalize so _collect_results can drain siblings cleanly.
1558 raise asyncio.CancelledError from exc
1559 try:
1560 # The source's old chunks are deleted in the same locked
1561 # transaction as the new write (see _flush_writes), so cleanup is
1562 # carried on the result rather than run eagerly here.
1563 if entry.content_type in archive_content_types():
1564 return await _archive_result(entry, on_progress, pages_done)
1565 page_texts: list[PageTextRecord] = []
1566 records, meta = await produce_records(
1567 entry.path,
1568 name,
1569 entry.content_type,
1570 quiet=quiet,
1571 on_progress=on_progress,
1572 page_texts_out=page_texts,
1573 )
1574 concept_records = await build_concept_records(records, name)
1575 entity_rows = await build_entity_records(records, name)
1576 on_progress(
1577 EventType.FILE_DONE,
1578 FileDoneEvent(file=name, status="ok", chunks=len(records)),
1579 )
1580 pages_done[0] += max(1, len(page_texts)) # OCR pages cleared: the throughput signal
1581 return _IngestResult(
1582 name,
1583 entry.path,
1584 len(records),
1585 error=None,
1586 file_hash=entry.file_hash,
1587 records=records,
1588 needs_cleanup=entry.needs_cleanup,
1589 page_texts=page_texts,
1590 stat=entry.stat,
1591 concept_records=concept_records,
1592 entity_rows=entity_rows,
1593 meta=meta,
1594 )
1595 except ChunkLimitError as exc:
1596 return _over_limit_result(
1597 exc, entry, pages_done=pages_done, on_progress=on_progress
1598 )
1599 except (asyncio.CancelledError, TaskCancelledError) as exc:
1600 # TaskCancelledError is the TUI's cooperative cancel signal raised
1601 # by reporter.check_cancelled() inside on_progress; treat it as
1602 # asyncio cancellation so _collect_results can drain siblings
1603 # cleanly instead of orphaning their pending exceptions.
1604 raise asyncio.CancelledError from exc
1605 except Exception as exc:
1606 return _failed_result(
1607 exc, entry, pages_done=pages_done, on_progress=on_progress, cancel=cancel
1608 )
1610 feed = _ResultFeed(_stream_tasks(plan_batches, _process_one), plan)
1611 collect = _collect_results if quiet else _collect_under_bar
1612 try:
1613 # extract_batching coalesces extractions into xberg batch calls when the
1614 # toggle is on (off by default); the per-file collect contract is unchanged.
1615 async with extract_batching():
1616 await collect(
1617 feed,
1618 added,
1619 updated,
1620 failed,
1621 skipped,
1622 window=window,
1623 on_progress=on_progress,
1624 flush_failed=flush_failed,
1625 reasons=reasons,
1626 )
1627 finally:
1628 # Stop the adaptive controller (if any) before returning: its background
1629 # loop must not outlive the batch it was tuning.
1630 if controller_task is not None:
1631 controller_task.cancel()
1632 with contextlib.suppress(asyncio.CancelledError):
1633 await controller_task
1636# Accumulate roughly this many chunks across documents before one batched
1637# LanceDB write. Bounds buffered-vector memory while amortizing the write lock
1638# and per-transaction overhead over many documents instead of one write per file.
1639_WRITE_FLUSH_CHUNKS = 2000
1642class _ResultFeed:
1643 """Pull-based source of per-file ingest coroutines over a streamed plan.
1645 ``take(wait=False)`` hands back an already-planned file without blocking, so
1646 the collector waits on the planner only when it has nothing left to run.
1647 Build-vs-buy: an ``asyncio.Queue`` is the stock bounded channel, but it would
1648 need a separate producer task to pump the plan-batch generator into it and a
1649 sentinel to close it; pulling ``anext`` on demand keeps the plan stream the
1650 single driver and needs neither. ``planned`` is the file count seen so far:
1651 the run's total once the stream is drained.
1652 """
1654 def __init__(
1655 self,
1656 plan_batches: AsyncGenerator[list[Coroutine[Any, Any, _IngestResult]]],
1657 plan: _StreamedPlan | None = None,
1658 ) -> None:
1659 self._plan_batches = plan_batches
1660 self._buffer: deque[Coroutine[Any, Any, _IngestResult]] = deque()
1661 self._pull: asyncio.Task[list[Coroutine[Any, Any, _IngestResult]] | None] | None = None
1662 self._drained = False
1663 self._plan = plan if plan is not None else _StreamedPlan()
1664 self.planned = 0
1666 @property
1667 def corpus_total(self) -> int:
1668 """Files the run's slice holds, or 0 when the caller declared no corpus."""
1669 return self._plan.corpus_total
1671 @property
1672 def resolved(self) -> int:
1673 """Files already disposed of by the plan, which never reach this feed."""
1674 return self._plan.resolved
1676 def pull(self) -> asyncio.Task[list[Coroutine[Any, Any, _IngestResult]] | None] | None:
1677 """The in-flight plan-batch prefetch, so a waiting collector wakes when it lands."""
1678 return self._pull
1680 async def take(self, *, wait: bool) -> Coroutine[Any, Any, _IngestResult] | None:
1681 """The next planned file, or None when the stream is drained (or, with
1682 *wait* False, when the next batch is not planned yet)."""
1683 while not self._buffer:
1684 if self._drained:
1685 return None
1686 if self._pull is None:
1687 self._pull = asyncio.ensure_future(anext(self._plan_batches, None))
1688 if not wait and not self._pull.done():
1689 return None
1690 plan_batch = await self._pull
1691 self._pull = None
1692 if plan_batch is None:
1693 self._drained = True
1694 return None
1695 self._buffer.extend(plan_batch)
1696 self.planned += len(plan_batch)
1697 return self._buffer.popleft()
1699 async def aclose(self) -> None:
1700 """Close the plan stream and discard files that were never started."""
1701 if self._pull is not None:
1702 self._pull.cancel()
1703 with contextlib.suppress(asyncio.CancelledError):
1704 # A batch that landed before the cancel took effect still owns
1705 # coroutines; recover it so they are closed rather than leaked.
1706 plan_batch = await self._pull
1707 if plan_batch:
1708 self._buffer.extend(plan_batch)
1709 self._pull = None
1710 for coro in self._buffer:
1711 coro.close()
1712 self._buffer.clear()
1713 await self._plan_batches.aclose()
1716async def _refill_window(
1717 in_flight: set[asyncio.Task[_IngestResult]],
1718 feed: _ResultFeed,
1719 window: int,
1720) -> None:
1721 """Top up the in-flight task set from *feed*, capped at *window* tasks.
1723 Waits on the planner only when nothing is running, so a slow batch never
1724 stalls files that are already planned.
1725 """
1726 while len(in_flight) < window:
1727 coro = await feed.take(wait=not in_flight)
1728 if coro is None:
1729 return
1730 in_flight.add(asyncio.ensure_future(coro))
1733async def _next_completions(
1734 in_flight: set[asyncio.Task[_IngestResult]], prefetch: asyncio.Future[Any] | None
1735) -> tuple[Iterable[asyncio.Future[Any]], set[asyncio.Task[_IngestResult]]]:
1736 """Wait for the next file to finish, returning (completed, still running).
1738 The feed's plan-batch prefetch waits alongside the running files, so work is
1739 admitted as soon as it is planned rather than on the next file completion,
1740 and it is filtered out of the still-running set here since it is not a file.
1741 """
1742 waiting: set[asyncio.Future[Any]] = set(in_flight)
1743 if prefetch is not None:
1744 waiting.add(prefetch)
1745 done, still_running = await asyncio.wait(waiting, return_when=asyncio.FIRST_COMPLETED)
1746 # Explicit loop, like _cancel_in_flight: Nuitka miscompiled the comprehension
1747 # form of this task-set filtering.
1748 remaining: set[asyncio.Task[_IngestResult]] = set()
1749 for task in still_running:
1750 if task is not prefetch:
1751 remaining.add(cast("asyncio.Task[_IngestResult]", task))
1752 return done, remaining
1755async def _collect_results(
1756 feed: _ResultFeed,
1757 added: dict[str, None],
1758 updated: dict[str, None],
1759 failed: dict[str, None],
1760 skipped: dict[str, None],
1761 *,
1762 window: int,
1763 on_progress: DetailedProgressCallback = noop_callback,
1764 progress: Progress | None = None,
1765 ptask: Any = None,
1766 flush_failed: set[str] | None = None,
1767 reasons: dict[str, str] | None = None,
1768) -> None:
1769 """Run *feed* through a bounded task window, batching writes and progress.
1771 At most *window* tasks exist at once: results are consumed as they complete
1772 and the window is refilled from the feed, so memory stays flat however many
1773 files a sync covers. Successful files are buffered and flushed to LanceDB in
1774 batches (one locked transaction per batch) rather than one write per file.
1775 The buffer is flushed on the way out too -- even on cancel -- so
1776 completed-but-unwritten work is persisted. On exception (typically
1777 asyncio.CancelledError from a user cancel), cancel every in-flight sibling
1778 and await them with ``return_exceptions=True`` so their pending
1779 CancelledErrors don't surface as "Task exception was never retrieved".
1780 """
1781 buffer: list[_IngestResult] = []
1782 buffered_chunks = 0
1783 completed_count = 0
1784 to_purge: list[str] = []
1785 in_flight: set[asyncio.Task[_IngestResult]] = set()
1786 try:
1787 await _refill_window(in_flight, feed, window)
1788 while in_flight:
1789 prefetch = feed.pull()
1790 done, in_flight = await _next_completions(in_flight, prefetch)
1791 saw_cancel = False
1792 for fut in done:
1793 if fut is prefetch:
1794 continue # a planned batch landing, not a file result
1795 try:
1796 result = fut.result()
1797 except asyncio.CancelledError:
1798 # A user cancel completes several futures together. Flag it but
1799 # keep draining `done` so a sibling that genuinely finished in
1800 # the same batch is still buffered and flushed (the
1801 # cancel-persists contract), then propagate after the loop. A
1802 # non-cancel exception still propagates immediately, as before,
1803 # so a genuine ingest bug surfaces and cancels the siblings.
1804 saw_cancel = True
1805 continue
1806 completed_count += 1
1807 status = _classify_result(result, added, updated, failed, skipped, reasons)
1808 if status is BatchStatus.INGESTED:
1809 buffered_chunks = await _buffer_and_maybe_flush(
1810 result,
1811 buffer,
1812 buffered_chunks,
1813 added,
1814 updated,
1815 failed,
1816 skipped,
1817 flush_failed,
1818 )
1819 elif status is BatchStatus.SKIPPED and result.needs_cleanup:
1820 # Zero-text result is never buffered; collect it for the
1821 # purge pass (see _purge_emptied_sources).
1822 to_purge.append(result.name)
1823 _report_file_progress(
1824 result,
1825 status,
1826 feed.resolved + completed_count,
1827 feed.corpus_total,
1828 on_progress,
1829 progress,
1830 ptask,
1831 )
1832 if saw_cancel:
1833 # Completed siblings in this batch are now buffered; propagate the
1834 # cancel so the finally flushes them and cancels still-running work.
1835 raise asyncio.CancelledError
1836 await _refill_window(in_flight, feed, window)
1837 finally:
1838 # The inner finally guarantees the sibling cancel even if the flush
1839 # itself raises (e.g. a cancellation landing on the to_thread await).
1840 try:
1841 await to_ingest_thread(
1842 _flush_writes, buffer, added, updated, failed, skipped, flush_failed
1843 )
1844 await to_ingest_thread(_purge_emptied_sources, to_purge)
1845 finally:
1846 try:
1847 await _cancel_in_flight(in_flight)
1848 finally:
1849 # Closing the feed stops the planner behind it, so a cancelled
1850 # sync does not keep hashing the rest of the corpus.
1851 await feed.aclose()
1854async def _collect_under_bar(
1855 feed: _ResultFeed,
1856 added: dict[str, None],
1857 updated: dict[str, None],
1858 failed: dict[str, None],
1859 skipped: dict[str, None],
1860 *,
1861 window: int,
1862 on_progress: DetailedProgressCallback = noop_callback,
1863 flush_failed: set[str] | None = None,
1864 reasons: dict[str, str] | None = None,
1865) -> None:
1866 """Run :func:`_collect_results` under a transient Rich progress bar."""
1867 with Progress(
1868 SpinnerColumn(),
1869 TextColumn("{task.description}"),
1870 BarColumn(),
1871 MofNCompleteColumn(),
1872 TimeElapsedColumn(),
1873 transient=True,
1874 ) as progress:
1875 # The corpus the discovery walk found, known before the first batch is
1876 # planned. None only when the caller supplied no plan to measure against.
1877 ptask = progress.add_task("Ingesting documents...", total=feed.corpus_total or None)
1878 # The bar advances once per file (in _collect_results), so a single
1879 # multi-page scanned PDF would freeze at "0/1" through its whole
1880 # OCR + embed phase. Drive the spinner's description off the same
1881 # EXTRACT (OCR page i/N) and EMBED (chunk i/N) events the TUI uses
1882 # so the row visibly moves while one file is being worked.
1883 await _collect_results(
1884 feed,
1885 added,
1886 updated,
1887 failed,
1888 skipped,
1889 window=window,
1890 on_progress=_phase_progress_callback(progress, ptask, on_progress),
1891 progress=progress,
1892 ptask=ptask,
1893 flush_failed=flush_failed,
1894 reasons=reasons,
1895 )
1898async def _cancel_in_flight(in_flight: set[asyncio.Task[_IngestResult]]) -> None:
1899 """Cancel still-running tasks and await them so their CancelledErrors are retrieved."""
1900 # Explicit loop: Nuitka miscompiled the comprehension form of this cleanup.
1901 still_pending = []
1902 for t in in_flight:
1903 if not t.done():
1904 still_pending.append(t)
1905 for task in still_pending:
1906 task.cancel()
1907 if still_pending:
1908 await asyncio.gather(*still_pending, return_exceptions=True)
1911async def _buffer_and_maybe_flush(
1912 result: _IngestResult,
1913 buffer: list[_IngestResult],
1914 buffered_chunks: int,
1915 added: dict[str, None],
1916 updated: dict[str, None],
1917 failed: dict[str, None],
1918 skipped: dict[str, None],
1919 flush_failed: set[str] | None,
1920) -> int:
1921 """Buffer one ingested file, flushing at the chunk threshold; returns the new count."""
1922 buffer.append(result)
1923 # Zero-chunk files count one unit so the buffer stays bounded.
1924 buffered_chunks += max(result.chunk_count, 1)
1925 if buffered_chunks >= _WRITE_FLUSH_CHUNKS:
1926 await to_ingest_thread(_flush_writes, buffer, added, updated, failed, skipped, flush_failed)
1927 buffered_chunks = 0
1928 return buffered_chunks
1931def _report_file_progress(
1932 result: _IngestResult,
1933 status: BatchStatus,
1934 completed_count: int,
1935 total: int,
1936 on_progress: DetailedProgressCallback,
1937 progress: Progress | None,
1938 ptask: Any,
1939) -> None:
1940 """Advance the Rich bar (when present) and emit one BATCH_PROGRESS event.
1942 *completed_count* counts every file the pass has disposed of and *total* is
1943 the corpus the discovery walk found, so the pair answers how much of the
1944 corpus is done rather than how much of the plan so far is.
1945 """
1946 if progress is not None and ptask is not None:
1947 desc = f"Ingested {result.name}" if result.error is None else f"Failed {result.name}"
1948 # Set, not advanced: files the plan resolved without ingest produce no
1949 # result of their own and would otherwise never reach the bar.
1950 progress.update(ptask, description=desc, completed=completed_count)
1951 with contextlib.suppress(TaskCancelledError):
1952 on_progress(
1953 EventType.BATCH_PROGRESS,
1954 BatchProgressEvent(
1955 file=result.name,
1956 status=status,
1957 current=completed_count,
1958 total=total,
1959 ),
1960 )
1963def _classify_result(
1964 result: _IngestResult,
1965 added: dict[str, None],
1966 updated: dict[str, None],
1967 failed: dict[str, None],
1968 skipped: dict[str, None],
1969 reasons: dict[str, str] | None = None,
1970) -> BatchStatus:
1971 """Record a completed file's outcome and return its batch status.
1973 Failures, refusals and zero-chunk files are tracked here; a successful file is
1974 reported as ``INGESTED`` and its chunks are persisted by the batched flush, so
1975 it stays in ``added`` / ``updated`` until then. When *reasons* is given, the
1976 human-readable cause is recorded there (filename → reason) for reporting.
1977 """
1978 if result.skip_reason is not None:
1979 added.pop(result.name, None)
1980 updated.pop(result.name, None)
1981 skipped[result.name] = None
1982 if reasons is not None:
1983 reasons[result.name] = result.skip_reason
1984 return BatchStatus.SKIPPED
1985 if result.error is not None:
1986 # A traceback here would bleed into the TUI chat pane; the full trace stays at DEBUG.
1987 log.warning("Failed to ingest %s: %s", result.name, result.error)
1988 log.debug("Traceback for failed ingest of %s", result.name, exc_info=result.error)
1989 added.pop(result.name, None)
1990 updated.pop(result.name, None)
1991 failed[result.name] = None
1992 if reasons is not None:
1993 reasons[result.name] = error_reason(result.error)
1994 return BatchStatus.FAILED
1995 if result.chunk_count == 0:
1996 # No searchable chunks: never report it as added/updated. With no page
1997 # texts either, nothing is persisted and the file retries next sync. With
1998 # page texts, it stays INGESTED so its pages persist (export/recon) and it
1999 # stops replanning, but it is reported as skipped since search can't see it.
2000 added.pop(result.name, None)
2001 updated.pop(result.name, None)
2002 skipped[result.name] = None
2003 if reasons is not None:
2004 reasons[result.name] = (
2005 "no text extracted (0 chunks)"
2006 if not result.page_texts
2007 else "stored page text only (0 searchable chunks)"
2008 )
2009 return BatchStatus.SKIPPED if not result.page_texts else BatchStatus.INGESTED
2010 return BatchStatus.INGESTED
2013# Back off briefly before the single flush retry: the usual contender is a
2014# search-triggered FTS optimize holding the store lock past its 30s timeout.
2015_FLUSH_RETRY_DELAY_SECONDS = 2.0
2018def _retry_after_lock_timeout(write: Callable[[], object]) -> None:
2019 """Run one store write, retrying once after a lock timeout."""
2020 try:
2021 write()
2022 except LockTimeoutError:
2023 log.warning(
2024 "Store write lock busy; retrying batch flush in %.0fs", _FLUSH_RETRY_DELAY_SECONDS
2025 )
2026 time.sleep(_FLUSH_RETRY_DELAY_SECONDS)
2027 write()
2030def _flush_batch(buffer: list[_IngestResult]) -> None:
2031 """Persist one flush unit in a single locked ``write_chunks_batch`` transaction.
2033 Page texts travel inside each :class:`ChunkWrite` so the store writes them
2034 after the cleanup delete (which clears the source's old page-text rows) and
2035 before the source row: a page-text failure leaves the row stale and the file
2036 replans next sync instead of losing its pages forever behind the stat
2037 short-circuit. The write retries once on a lock timeout.
2038 """
2039 store = get_services().store
2040 items: list[ChunkWrite] = []
2041 stale: list[str] = []
2042 for r in buffer:
2043 digest = r.file_hash or file_hash(r.path)
2044 items.append(
2045 ChunkWrite(
2046 source=r.name,
2047 file_hash=digest,
2048 records=cast(list[dict], r.records or []),
2049 needs_cleanup=r.needs_cleanup,
2050 stat=r.stat,
2051 page_texts=cast(list[dict], r.page_texts or []),
2052 meta=r.meta,
2053 )
2054 )
2055 if r.members is None:
2056 continue
2057 items.extend(_member_write(m, digest, r.stat) for m in r.members)
2058 current = {m.name for m in r.members}
2059 stale.extend(n for n in store.member_sources(r.name) if n not in current)
2060 if stale:
2061 store.remove_documents(stale)
2062 _retry_after_lock_timeout(lambda: store.write_chunks_batch(items))
2063 _flush_concept_records(buffer)
2064 _flush_entity_rows(buffer)
2067def _member_write(member: MemberRecords, digest: str, stat: SourceStat | None) -> ChunkWrite:
2068 """A member's write item: its own source row, keyed to the archive's hash and stat."""
2069 return ChunkWrite(
2070 source=member.name,
2071 file_hash=digest,
2072 records=cast(list[dict], member.records),
2073 needs_cleanup=True,
2074 stat=stat,
2075 page_texts=cast(list[dict], member.page_texts),
2076 meta=member.meta,
2077 )
2080def _flush_concept_records(buffer: list[_IngestResult]) -> None:
2081 """Write the flush unit's buffered concept rows in one batched pass.
2083 Runs after the chunk write so a failed flush (files moved to ``failed``
2084 and replanned) never lands concept rows for unwritten chunks. A concept
2085 write failure is logged and never fails the files, matching the
2086 per-file extraction failure semantics.
2087 """
2088 batches = [r.concept_records for r in buffer if r.concept_records is not None]
2089 if not batches:
2090 return
2091 try:
2092 get_services().concepts.write_concept_records(ConceptRecords.merged(batches))
2093 except Exception:
2094 log.warning("Concept indexing failed for %d-file batch", len(batches), exc_info=True)
2097def _flush_entity_rows(buffer: list[_IngestResult]) -> None:
2098 """Write the flush unit's buffered entity rows in one batched pass.
2100 Runs after the chunk write, which also performed the per-source deletes,
2101 so replacement never leaves a source's stale entity rows behind. A write
2102 failure is logged and never fails the files, matching concept semantics.
2103 """
2104 rows = [row for r in buffer if r.entity_rows for row in r.entity_rows]
2105 if not rows:
2106 return
2107 try:
2108 get_services().store.add_entities(rows)
2109 except Exception:
2110 log.warning("Entity indexing failed for a %d-row batch", len(rows), exc_info=True)
2113def _purge_emptied_sources(names: list[str]) -> None:
2114 """Remove the prior index entry for files that now extract to nothing.
2116 An already-indexed file edited to yield zero chunks and zero page texts is
2117 classified SKIPPED and never buffered, so the batched cleanup delete never
2118 runs and its old chunks and source row would linger in search results. Full
2119 removal here keeps the index consistent; ``remove_documents`` is a no-op for
2120 never-indexed (brand-new empty) files, so unindexed inputs cost nothing.
2121 """
2122 if not names:
2123 return
2124 get_services().store.remove_documents(names)
2127def _flush_writes(
2128 buffer: list[_IngestResult],
2129 added: dict[str, None],
2130 updated: dict[str, None],
2131 failed: dict[str, None],
2132 skipped: dict[str, None],
2133 flush_failed: set[str] | None = None,
2134) -> None:
2135 """Flush the buffered documents to the store; track a write failure.
2137 Each buffered file's page texts, chunks, cleanup delete, and source upsert
2138 are written by :func:`_flush_batch`. If that fails, every file in the batch
2139 is moved to ``failed`` since its source row did not land, and recorded in
2140 *flush_failed* so the caller replans them next sync instead of skip-marking
2141 them; the exception never escapes, so the caller's sibling-cancel and
2142 skip-marker path always runs. The buffer is cleared either way.
2143 """
2144 if not buffer:
2145 return
2146 try:
2147 _flush_batch(buffer)
2148 except Exception as exc:
2149 for r in buffer:
2150 log.warning("Failed to write %s: %s", r.name, exc)
2151 added.pop(r.name, None)
2152 updated.pop(r.name, None)
2153 # A page-text-only file was pre-marked skipped at classification; on a
2154 # flush failure it belongs in failed only, never both.
2155 skipped.pop(r.name, None)
2156 failed[r.name] = None
2157 if flush_failed is not None:
2158 flush_failed.add(r.name)
2159 finally:
2160 buffer.clear()