Coverage for src/lilbee/server/handlers/ingest.py: 100%
191 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""Sync and add-files handlers (SSE-streamed)."""
3from __future__ import annotations
5import asyncio
6import contextlib
7import functools
8import hashlib
9import logging
10import os
11import re
12from collections.abc import AsyncGenerator, Callable, Coroutine
13from pathlib import Path
14from typing import TYPE_CHECKING, Any, TypeVar
16from lilbee.app.ingest import register_sources
17from lilbee.app.services import get_services
18from lilbee.core.config import cfg, validate_ocr_timeout
19from lilbee.core.config.enums import OcrMode
20from lilbee.core.security import validate_path_within
21from lilbee.data.ingest.discovery import excluded_extension_reasons
22from lilbee.runtime.ingest_lock import IngestLockRegistry
23from lilbee.runtime.lock import refuse_where_a_sync_cannot_lock
24from lilbee.runtime.progress import SseEvent
25from lilbee.server.handlers.sse import SseStream, sse_event
26from lilbee.server.models import AddSummary, SyncSummary
28if TYPE_CHECKING:
29 from lilbee.app.dataset import ImportSummary
30 from lilbee.data.ingest import SyncResult
32log = logging.getLogger(__name__)
33# What a sync of this server is called in the notice of a sync that waits for it.
34_SYNC_STARTER = "lilbee serve"
36# Payload carried alongside each source key through _ingest_stream: a server
37# path (str) for /api/add, or an (filename, content) pair for /api/add/upload.
38_Payload = TypeVar("_Payload")
39_Result = TypeVar("_Result")
42async def _to_thread_unbroken(work: Callable[[], _Result]) -> _Result:
43 """Run blocking *work* in a thread to its end; a cancel that arrives meanwhile lands after it.
45 The caller is cancelled again when *work* has returned, so its next await raises.
46 """
47 running = asyncio.ensure_future(asyncio.to_thread(work))
48 cancelled = False
49 while not running.done():
50 try:
51 await asyncio.shield(running)
52 except asyncio.CancelledError:
53 cancelled = True
54 task = asyncio.current_task()
55 if cancelled and task is not None:
56 task.cancel()
57 return running.result()
60async def _run_sync_with_sentinel(
61 sse: SseStream,
62 ocr: OcrMode | None,
63 ocr_timeout: float | None = None,
64 force_rebuild: bool = False,
65 retry_skipped: bool = False,
66 prune_ignored: bool = False,
67) -> SyncResult:
68 """Run ingest.sync() and guarantee the drain sentinel is enqueued."""
69 from lilbee.app.ingest import temporary_ocr_config
70 from lilbee.data.ingest import sync
72 try:
73 with temporary_ocr_config(ocr, ocr_timeout):
74 return await sync(
75 quiet=True,
76 on_progress=sse.callback,
77 cancel=sse.cancel,
78 force_rebuild=force_rebuild,
79 retry_skipped=retry_skipped,
80 prune_ignored=prune_ignored,
81 started_by=_SYNC_STARTER,
82 )
83 finally:
84 sse.queue.put_nowait(None)
87async def sync_stream(
88 *,
89 ocr: OcrMode | None = None,
90 ocr_timeout: float | None = None,
91 force_rebuild: bool = False,
92 retry_skipped: bool = False,
93 prune_ignored: bool = False,
94 done_payload: Callable[[SyncResult], dict[str, Any]] = lambda result: result.model_dump(),
95) -> AsyncGenerator[str, None]:
96 """Trigger sync, yield SSE progress events, then done event.
98 When ``force_rebuild`` is true, the underlying sync re-ingests every file from
99 ``cfg.documents_dir`` over the index (the REST equivalent of ``lilbee rebuild``).
100 When ``retry_skipped`` is true, it clears the failed-file markers so files
101 that were skipped on a previous sync get another attempt, without dropping
102 the store. When ``prune_ignored`` is true, it also drops sources a
103 ``.lilbeeignore`` now excludes. ``done_payload`` builds the done event's data
104 from the sync's result.
105 """
106 sse = SseStream()
107 task = asyncio.create_task(
108 _run_sync_with_sentinel(sse, ocr, ocr_timeout, force_rebuild, retry_skipped, prune_ignored)
109 )
110 async for event in sse.drain(task, "Sync stream"):
111 yield event
112 frame = sse.terminal_frame(task, done_payload)
113 if frame is not None:
114 yield frame
117async def _run_add(
118 paths: list[str],
119 force: bool,
120 ocr: OcrMode | None,
121 ocr_timeout: float | None,
122 sse: SseStream,
123) -> AddSummary:
124 """Copy files and sync, returning the summary for the final done event."""
125 from lilbee.app.ingest import temporary_ocr_config
126 from lilbee.data.ingest import sync
128 try:
129 errors: list[str] = []
130 valid: list[Path] = []
131 for p_str in paths:
132 p = Path(p_str)
133 if not p.exists():
134 errors.append(p_str)
135 else:
136 valid.append(p)
138 # A stream that stops while the registration runs stops the sync, not the add.
139 stopped = sse.cancel.is_set()
140 reg_result = await _to_thread_unbroken(
141 functools.partial(register_sources, valid, force=force)
142 )
144 errors.extend(reg_result.refused)
145 if stopped:
146 return AddSummary(
147 copied=reg_result.registered,
148 name_taken=reg_result.name_taken,
149 overlapping=reg_result.overlapping,
150 absorbed=reg_result.absorbed,
151 tracked=reg_result.tracked,
152 errors=errors,
153 )
155 if not reg_result.reached_corpus:
156 return AddSummary(copied=[], name_taken=reg_result.name_taken, errors=errors)
158 with temporary_ocr_config(ocr, ocr_timeout):
159 sync_result = await sync(
160 quiet=True, on_progress=sse.callback, cancel=sse.cancel, started_by=_SYNC_STARTER
161 )
163 return AddSummary(
164 copied=reg_result.registered,
165 name_taken=reg_result.name_taken,
166 overlapping=reg_result.overlapping,
167 absorbed=reg_result.absorbed,
168 tracked=reg_result.tracked,
169 errors=errors,
170 sync=SyncSummary(**sync_result.model_dump()),
171 )
172 finally:
173 sse.queue.put_nowait(None)
176def validate_add_paths(
177 data: dict[str, Any],
178) -> tuple[list[str], bool, OcrMode | None, float | None]:
179 """Validate add-files input. Raises ValueError on bad input."""
180 paths = data.get("paths")
181 if not isinstance(paths, list) or not paths:
182 raise ValueError("'paths' must be a non-empty list of strings")
183 # No file-count cap: the resource guard is the app's size-based
184 # request_max_body_size; a count limit only breaks the point-lilbee-at-
185 # your-codebase use case (hundreds of small files).
187 for p_str in paths:
188 # The basename becomes the source root's label (its key prefix). It must
189 # be a clean single segment: Path(x).name cannot traverse, so this only
190 # rejects an empty name ("/", "a/") that would name no root.
191 name = Path(p_str).name
192 if not name:
193 raise ValueError(f"{p_str!r} does not name a file")
194 validate_path_within(cfg.documents_dir / name, cfg.documents_dir)
196 force = bool(data.get("force", False))
197 ocr, ocr_timeout = _parse_ocr_params(data)
198 return paths, force, ocr, ocr_timeout
201def parse_ocr_mode(value: Any) -> OcrMode | None:
202 """Decode a request's ``ocr`` value; raises ValueError naming the accepted modes."""
203 if value is None:
204 return None
205 try:
206 return OcrMode(value)
207 except ValueError as exc:
208 accepted = ", ".join(mode.value for mode in OcrMode)
209 raise ValueError(f"ocr must be one of {accepted}; got {value!r}") from exc
212def _parse_ocr_params(data: dict[str, Any]) -> tuple[OcrMode | None, float | None]:
213 """Extract, decode, and validate OCR parameters from a request dict.
215 Raises ValueError on an unknown ocr mode, or a non-numeric or
216 out-of-bound ocr_timeout.
217 """
218 ocr_timeout = data.get("ocr_timeout")
219 if ocr_timeout is not None:
220 ocr_timeout = float(ocr_timeout)
221 validate_ocr_timeout(ocr_timeout)
222 return parse_ocr_mode(data.get("ocr")), ocr_timeout
225async def _ingest_stream(
226 items: list[tuple[str, _Payload]],
227 run: Callable[[list[_Payload], SseStream], Coroutine[Any, Any, AddSummary]],
228 label: str,
229) -> AsyncGenerator[str, None]:
230 """Lock per source, run ``run`` over the acquired subset, and stream SSE.
232 Shared by /api/add (server paths) and /api/add/upload (uploaded content).
233 Each item is ``(lock_key, payload)``: the key is the source identifier used
234 for the per-source ingest lock; the payload is what ``run`` receives for the
235 subset whose lock was acquired. Contended sources emit ``already_ingesting``.
236 When every source is contended the stream closes with no ``done`` event,
237 signalling the client to wait rather than retry. When only some are, the
238 acquired subset still runs and the ``done`` summary names the contended ones
239 in ``already_ingesting``, so a partial batch never reads as a full success.
240 """
241 registry = get_services().ingest_lock_registry
242 acquired, busy = await registry.acquire([key for key, _payload in items])
243 try:
244 for name in busy:
245 log.info("Rejecting %s for %s: already ingesting", label, name)
246 yield sse_event(SseEvent.ALREADY_INGESTING, {"source": name})
248 if not acquired:
249 return
251 acquired_names = {name for name, _lock in acquired}
252 locked = [payload for key, payload in items if key in acquired_names]
253 sse = SseStream()
254 task = asyncio.create_task(run(locked, sse))
255 try:
256 async for event in sse.drain(task, label):
257 yield event
258 # already_ingesting names the sources this run never attempted, so
259 # a client reading only the terminal event still sees a partial batch.
260 frame = sse.terminal_frame(
261 task, lambda s: s.model_copy(update={"already_ingesting": list(busy)}).model_dump()
262 )
263 if frame is not None:
264 yield frame
265 finally:
266 if not task.done():
267 task.cancel()
268 with contextlib.suppress(asyncio.CancelledError, Exception):
269 await task
270 finally:
271 registry.release(acquired)
274async def add_files_stream(
275 paths: list[str],
276 *,
277 force: bool = False,
278 ocr: OcrMode | None = None,
279 ocr_timeout: float | None = None,
280) -> AsyncGenerator[str, None]:
281 """Copy server-side files, sync, and yield SSE progress events.
283 Takes the already-validated/parsed values from ``validate_add_paths`` so the
284 request dict is decoded once.
285 """
286 async for event in _ingest_stream(
287 [(IngestLockRegistry.canonical_source_name(p), p) for p in paths],
288 lambda locked, sse: _run_add(locked, force, ocr, ocr_timeout, sse),
289 "Add files stream",
290 ):
291 yield event
294def _clean_upload_name(name: str) -> str:
295 """Normalize one upload filename to a safe relative path inside the corpus.
297 Relative paths are preserved (a source tree keeps its layout instead of
298 colliding on basenames); absolute paths, drive letters, and ``..`` segments
299 are rejected. Raises ValueError on bad input.
300 """
301 normalized = name.replace("\\", "/")
302 if normalized.startswith("/") or re.match(r"^[A-Za-z]:", normalized):
303 raise ValueError(f"upload filename must be relative: {name!r}")
304 parts = [part for part in normalized.split("/") if part not in ("", ".")]
305 if not parts:
306 raise ValueError(f"invalid upload filename: {name!r}")
307 if ".." in parts:
308 raise ValueError(f"upload filename may not contain '..': {name!r}")
309 relative = "/".join(parts)
310 reason = excluded_extension_reasons().get(Path(relative).suffix.lower())
311 if reason is not None:
312 raise ValueError(f"{name!r}: {reason}")
313 validate_path_within(cfg.documents_dir / relative, cfg.documents_dir)
314 return relative
317def validate_upload_names(names: list[str | None]) -> list[str]:
318 """Validate uploaded filenames, returning the cleaned relative paths.
320 Names only, deliberately: the route validates before reading any part's
321 bytes, so a request that will be rejected never costs a full copy of the
322 payload in the server's own memory. Filenames keep their relative path,
323 validated to stay inside ``cfg.documents_dir``, so an uploaded source tree
324 preserves its layout instead of colliding on basenames. There is no
325 file-count cap; the resource guard is the app's request_max_body_size.
327 Raises ValueError on bad input.
328 """
329 if not names:
330 raise ValueError("no files uploaded")
331 # A multipart part is allowed to carry no filename at all; that is not a
332 # file upload, and the cleaner's message should name it as missing rather
333 # than crash on None.
334 return [_clean_upload_name(name if name is not None else "") for name in names]
337async def _run_upload(
338 files: list[tuple[str, bytes]],
339 sse: SseStream,
340 ocr: OcrMode | None = None,
341 ocr_timeout: float | None = None,
342) -> AddSummary:
343 """Write uploaded file bytes into ``cfg.documents_dir``, then sync.
345 The upload equivalent of :func:`_run_add`: instead of copying from a
346 server-readable path, the client's content is written straight into the
347 documents dir and the same ingest pipeline runs. This is what lets an
348 external-mode client (a remote lilbee / GPU box) ingest files that live only
349 on the client. Unchanged content is a no-op re-embed inside ``sync`` (it
350 hashes each source), so there is no separate force flag. A data root that
351 the sync cannot lock refuses the upload before any file is written.
352 """
353 from lilbee.app.ingest import temporary_ocr_config
354 from lilbee.data.ingest import sync
356 try:
357 await asyncio.to_thread(refuse_where_a_sync_cannot_lock, cfg.data_root)
358 cfg.documents_dir.mkdir(parents=True, exist_ok=True)
359 written: list[str] = []
360 for name, content in files:
361 dest = cfg.documents_dir / name
362 dest.parent.mkdir(parents=True, exist_ok=True)
363 if not _move_same_content(name, content, dest):
364 dest.write_bytes(content)
365 written.append(name)
366 with temporary_ocr_config(ocr, ocr_timeout):
367 sync_result = await sync(
368 quiet=True, on_progress=sse.callback, cancel=sse.cancel, started_by=_SYNC_STARTER
369 )
370 return AddSummary(
371 copied=written,
372 errors=[],
373 sync=SyncSummary(**sync_result.model_dump()),
374 )
375 finally:
376 sse.queue.put_nowait(None)
379def _move_same_content(name: str, content: bytes, dest: Path) -> bool:
380 """Move the one indexed file holding *content* to *dest*, so sync repoints its key."""
381 digest = hashlib.sha256(content).hexdigest()
382 matches = [
383 s["filename"]
384 for s in get_services().store.get_sources()
385 if s["file_hash"] == digest and s["filename"] != name
386 ]
387 if len(matches) != 1:
388 return False
389 old_path = cfg.documents_dir / matches[0]
390 if not old_path.is_file():
391 return False
392 os.replace(old_path, dest)
393 log.info("Moved %s to %s: the same content arrived under a new name", matches[0], name)
394 return True
397async def add_uploads_stream(
398 files: list[tuple[str, bytes]],
399 *,
400 ocr: OcrMode | None = None,
401 ocr_timeout: float | None = None,
402) -> AsyncGenerator[str, None]:
403 """Ingest uploaded file content, yielding the same SSE progress as add_files_stream.
405 Locks per source name (the validated relative path) so an upload never
406 races an in-flight add of the same source.
407 """
408 async for event in _ingest_stream(
409 [(name, (name, content)) for name, content in files],
410 lambda locked, sse: _run_upload(locked, sse, ocr, ocr_timeout),
411 "Add uploads stream",
412 ):
413 yield event
416async def _run_import_with_sentinel(sse: SseStream, data: bytes, fmt: str) -> ImportSummary:
417 """Run the dataset import and guarantee the drain sentinel is enqueued."""
418 from lilbee.app.dataset import import_from_bytes
420 try:
421 return await import_from_bytes(data, fmt, on_progress=sse.callback)
422 finally:
423 sse.queue.put_nowait(None)
426async def import_stream(data: bytes, fmt: str) -> AsyncGenerator[str, None]:
427 """Import a dataset, yield SSE embed-progress events, then a done event."""
428 sse = SseStream()
429 task = asyncio.create_task(_run_import_with_sentinel(sse, data, fmt))
430 async for event in sse.drain(task, "Import stream"):
431 yield event
432 frame = sse.terminal_frame(task, lambda result: result.model_dump())
433 if frame is not None:
434 yield frame