Coverage for src/lilbee/server/handlers/ingest.py: 100%
166 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-04 17:08 +0000
1"""Sync and add-files handlers (SSE-streamed)."""
3from __future__ import annotations
5import asyncio
6import contextlib
7import hashlib
8import logging
9import os
10import re
11from collections.abc import AsyncGenerator, Callable, Coroutine
12from pathlib import Path
13from typing import TYPE_CHECKING, Any, TypeVar
15from lilbee.app.ingest import register_sources
16from lilbee.app.services import get_services
17from lilbee.core.config import cfg
18from lilbee.core.security import validate_path_within
19from lilbee.data.ingest.discovery import excluded_extension_reasons
20from lilbee.runtime.ingest_lock import IngestLockRegistry
21from lilbee.runtime.progress import SseEvent
22from lilbee.server.handlers.sse import SseStream, sse_event
23from lilbee.server.models import AddSummary, SyncSummary
25if TYPE_CHECKING:
26 from lilbee.app.dataset import ImportSummary
27 from lilbee.data.ingest import SyncResult
29log = logging.getLogger(__name__)
31# Payload carried alongside each source key through _ingest_stream: a server
32# path (str) for /api/add, or an (filename, content) pair for /api/add/upload.
33_Payload = TypeVar("_Payload")
36async def _run_sync_with_sentinel(
37 sse: SseStream,
38 enable_ocr: bool | None,
39 force_rebuild: bool = False,
40 retry_skipped: bool = False,
41 prune_ignored: bool = False,
42) -> SyncResult:
43 """Run ingest.sync() and guarantee the drain sentinel is enqueued."""
44 from lilbee.app.ingest import temporary_ocr_config
45 from lilbee.data.ingest import sync
47 try:
48 with temporary_ocr_config(enable_ocr):
49 return await sync(
50 quiet=True,
51 on_progress=sse.callback,
52 cancel=sse.cancel,
53 force_rebuild=force_rebuild,
54 retry_skipped=retry_skipped,
55 prune_ignored=prune_ignored,
56 )
57 finally:
58 sse.queue.put_nowait(None)
61async def sync_stream(
62 *,
63 enable_ocr: bool | None = None,
64 force_rebuild: bool = False,
65 retry_skipped: bool = False,
66 prune_ignored: bool = False,
67) -> AsyncGenerator[str, None]:
68 """Trigger sync, yield SSE progress events, then done event.
70 When ``force_rebuild`` is true, the underlying sync drops every table and
71 re-ingests from ``cfg.documents_dir`` (the REST equivalent of ``lilbee rebuild``).
72 When ``retry_skipped`` is true, it clears the failed-file markers so files
73 that were skipped on a previous sync get another attempt, without dropping
74 the store. When ``prune_ignored`` is true, it also drops sources a
75 ``.lilbeeignore`` now excludes.
76 """
77 sse = SseStream()
78 task = asyncio.create_task(
79 _run_sync_with_sentinel(sse, enable_ocr, force_rebuild, retry_skipped, prune_ignored)
80 )
81 async for event in sse.drain(task, "Sync stream"):
82 yield event
83 frame = sse.terminal_frame(task, lambda result: result.model_dump())
84 if frame is not None:
85 yield frame
88async def _run_add(
89 paths: list[str],
90 force: bool,
91 enable_ocr: bool | None,
92 ocr_timeout: float | None,
93 sse: SseStream,
94) -> AddSummary:
95 """Copy files and sync, returning the summary for the final done event."""
96 from lilbee.app.ingest import temporary_ocr_config
97 from lilbee.data.ingest import sync
99 try:
100 errors: list[str] = []
101 valid: list[Path] = []
102 for p_str in paths:
103 p = Path(p_str)
104 if not p.exists():
105 errors.append(p_str)
106 else:
107 valid.append(p)
109 reg_result = register_sources(valid, force=force)
111 errors.extend(reg_result.refused)
112 if sse.cancel.is_set():
113 return AddSummary(
114 copied=reg_result.registered,
115 skipped=reg_result.skipped,
116 tracked=reg_result.tracked,
117 errors=errors,
118 )
120 if not reg_result.registered and not reg_result.skipped and not reg_result.tracked:
121 # Nothing reached the corpus, and sync() is a whole-vault pass
122 # holding the ingest lock. A *tracked* or *skipped* file is not this
123 # case: it is already in the corpus but may never have been indexed,
124 # and a tracked one may have just had its skip marker cleared.
125 return AddSummary(copied=[], skipped=[], errors=errors)
127 with temporary_ocr_config(enable_ocr, ocr_timeout):
128 sync_result = await sync(quiet=True, on_progress=sse.callback, cancel=sse.cancel)
130 return AddSummary(
131 copied=reg_result.registered,
132 skipped=reg_result.skipped,
133 tracked=reg_result.tracked,
134 errors=errors,
135 sync=SyncSummary(**sync_result.model_dump()),
136 )
137 finally:
138 sse.queue.put_nowait(None)
141def validate_add_paths(
142 data: dict[str, Any],
143) -> tuple[list[str], bool, bool | None, float | None]:
144 """Validate add-files input. Raises ValueError on bad input."""
145 paths = data.get("paths")
146 if not isinstance(paths, list) or not paths:
147 raise ValueError("'paths' must be a non-empty list of strings")
148 # No file-count cap: the resource guard is the app's size-based
149 # request_max_body_size; a count limit only breaks the point-lilbee-at-
150 # your-codebase use case (hundreds of small files).
152 for p_str in paths:
153 # The basename becomes the source root's label (its key prefix). It must
154 # be a clean single segment: Path(x).name cannot traverse, so this only
155 # rejects an empty name ("/", "a/") that would name no root.
156 name = Path(p_str).name
157 if not name:
158 raise ValueError(f"{p_str!r} does not name a file")
159 validate_path_within(cfg.documents_dir / name, cfg.documents_dir)
161 force = bool(data.get("force", False))
162 enable_ocr, ocr_timeout = _parse_ocr_params(data)
163 return paths, force, enable_ocr, ocr_timeout
166def _parse_ocr_params(data: dict[str, Any]) -> tuple[bool | None, float | None]:
167 """Extract and coerce OCR parameters from a request dict."""
168 enable_ocr = data.get("enable_ocr")
169 ocr_timeout = data.get("ocr_timeout")
170 if enable_ocr is not None:
171 enable_ocr = bool(enable_ocr)
172 if ocr_timeout is not None:
173 ocr_timeout = float(ocr_timeout)
174 return enable_ocr, ocr_timeout
177async def _ingest_stream(
178 items: list[tuple[str, _Payload]],
179 run: Callable[[list[_Payload], SseStream], Coroutine[Any, Any, AddSummary]],
180 label: str,
181) -> AsyncGenerator[str, None]:
182 """Lock per source, run ``run`` over the acquired subset, and stream SSE.
184 Shared by /api/add (server paths) and /api/add/upload (uploaded content).
185 Each item is ``(lock_key, payload)``: the key is the source identifier used
186 for the per-source ingest lock; the payload is what ``run`` receives for the
187 subset whose lock was acquired. Contended sources emit ``already_ingesting``.
188 When every source is contended the stream closes with no ``done`` event,
189 signalling the client to wait rather than retry. When only some are, the
190 acquired subset still runs and the ``done`` summary names the contended ones
191 in ``already_ingesting``, so a partial batch never reads as a full success.
192 """
193 registry = get_services().ingest_lock_registry
194 acquired, busy = await registry.acquire([key for key, _payload in items])
195 try:
196 for name in busy:
197 log.info("Rejecting %s for %s: already ingesting", label, name)
198 yield sse_event(SseEvent.ALREADY_INGESTING, {"source": name})
200 if not acquired:
201 return
203 acquired_names = {name for name, _lock in acquired}
204 locked = [payload for key, payload in items if key in acquired_names]
205 sse = SseStream()
206 task = asyncio.create_task(run(locked, sse))
207 try:
208 async for event in sse.drain(task, label):
209 yield event
210 # already_ingesting names the sources this run never attempted, so
211 # a client reading only the terminal event still sees a partial batch.
212 frame = sse.terminal_frame(
213 task, lambda s: s.model_copy(update={"already_ingesting": list(busy)}).model_dump()
214 )
215 if frame is not None:
216 yield frame
217 finally:
218 if not task.done():
219 task.cancel()
220 with contextlib.suppress(asyncio.CancelledError, Exception):
221 await task
222 finally:
223 registry.release(acquired)
226async def add_files_stream(
227 paths: list[str],
228 *,
229 force: bool = False,
230 enable_ocr: bool | None = None,
231 ocr_timeout: float | None = None,
232) -> AsyncGenerator[str, None]:
233 """Copy server-side files, sync, and yield SSE progress events.
235 Takes the already-validated/parsed values from ``validate_add_paths`` so the
236 request dict is decoded once.
237 """
238 async for event in _ingest_stream(
239 [(IngestLockRegistry.canonical_source_name(p), p) for p in paths],
240 lambda locked, sse: _run_add(locked, force, enable_ocr, ocr_timeout, sse),
241 "Add files stream",
242 ):
243 yield event
246def _clean_upload_name(name: str) -> str:
247 """Normalize one upload filename to a safe relative path inside the corpus.
249 Relative paths are preserved (a source tree keeps its layout instead of
250 colliding on basenames); absolute paths, drive letters, and ``..`` segments
251 are rejected. Raises ValueError on bad input.
252 """
253 normalized = name.replace("\\", "/")
254 if normalized.startswith("/") or re.match(r"^[A-Za-z]:", normalized):
255 raise ValueError(f"upload filename must be relative: {name!r}")
256 parts = [part for part in normalized.split("/") if part not in ("", ".")]
257 if not parts:
258 raise ValueError(f"invalid upload filename: {name!r}")
259 if ".." in parts:
260 raise ValueError(f"upload filename may not contain '..': {name!r}")
261 relative = "/".join(parts)
262 reason = excluded_extension_reasons().get(Path(relative).suffix.lower())
263 if reason is not None:
264 raise ValueError(f"{name!r}: {reason}")
265 validate_path_within(cfg.documents_dir / relative, cfg.documents_dir)
266 return relative
269def validate_upload_names(names: list[str | None]) -> list[str]:
270 """Validate uploaded filenames, returning the cleaned relative paths.
272 Names only, deliberately: the route validates before reading any part's
273 bytes, so a request that will be rejected never costs a full copy of the
274 payload in the server's own memory. Filenames keep their relative path,
275 validated to stay inside ``cfg.documents_dir``, so an uploaded source tree
276 preserves its layout instead of colliding on basenames. There is no
277 file-count cap; the resource guard is the app's request_max_body_size.
279 Raises ValueError on bad input.
280 """
281 if not names:
282 raise ValueError("no files uploaded")
283 # A multipart part is allowed to carry no filename at all; that is not a
284 # file upload, and the cleaner's message should name it as missing rather
285 # than crash on None.
286 return [_clean_upload_name(name if name is not None else "") for name in names]
289async def _run_upload(files: list[tuple[str, bytes]], sse: SseStream) -> AddSummary:
290 """Write uploaded file bytes into ``cfg.documents_dir``, then sync.
292 The upload equivalent of :func:`_run_add`: instead of copying from a
293 server-readable path, the client's content is written straight into the
294 documents dir and the same ingest pipeline runs. This is what lets an
295 external-mode client (a remote lilbee / GPU box) ingest files that live only
296 on the client. Unchanged content is a no-op re-embed inside ``sync`` (it
297 hashes each source), so there is no separate force flag.
298 """
299 from lilbee.app.ingest import temporary_ocr_config
300 from lilbee.data.ingest import sync
302 try:
303 cfg.documents_dir.mkdir(parents=True, exist_ok=True)
304 written: list[str] = []
305 for name, content in files:
306 dest = cfg.documents_dir / name
307 dest.parent.mkdir(parents=True, exist_ok=True)
308 if not _move_same_content(name, content, dest):
309 dest.write_bytes(content)
310 written.append(name)
311 with temporary_ocr_config(None):
312 sync_result = await sync(quiet=True, on_progress=sse.callback, cancel=sse.cancel)
313 return AddSummary(
314 copied=written,
315 skipped=[],
316 errors=[],
317 sync=SyncSummary(**sync_result.model_dump()),
318 )
319 finally:
320 sse.queue.put_nowait(None)
323def _move_same_content(name: str, content: bytes, dest: Path) -> bool:
324 """Move the one indexed file holding *content* to *dest*, so sync repoints its key."""
325 digest = hashlib.sha256(content).hexdigest()
326 matches = [
327 s["filename"]
328 for s in get_services().store.get_sources()
329 if s["file_hash"] == digest and s["filename"] != name
330 ]
331 if len(matches) != 1:
332 return False
333 old_path = cfg.documents_dir / matches[0]
334 if not old_path.is_file():
335 return False
336 os.replace(old_path, dest)
337 log.info("Moved %s to %s: the same content arrived under a new name", matches[0], name)
338 return True
341async def add_uploads_stream(files: list[tuple[str, bytes]]) -> AsyncGenerator[str, None]:
342 """Ingest uploaded file content, yielding the same SSE progress as add_files_stream.
344 Locks per source name (the validated relative path) so an upload never
345 races an in-flight add of the same source.
346 """
347 async for event in _ingest_stream(
348 [(name, (name, content)) for name, content in files],
349 lambda locked, sse: _run_upload(locked, sse),
350 "Add uploads stream",
351 ):
352 yield event
355async def _run_import_with_sentinel(sse: SseStream, data: bytes, fmt: str) -> ImportSummary:
356 """Run the dataset import and guarantee the drain sentinel is enqueued."""
357 from lilbee.app.dataset import import_from_bytes
359 try:
360 return await import_from_bytes(data, fmt, on_progress=sse.callback)
361 finally:
362 sse.queue.put_nowait(None)
365async def import_stream(data: bytes, fmt: str) -> AsyncGenerator[str, None]:
366 """Import a dataset, yield SSE embed-progress events, then a done event."""
367 sse = SseStream()
368 task = asyncio.create_task(_run_import_with_sentinel(sse, data, fmt))
369 async for event in sse.drain(task, "Import stream"):
370 yield event
371 frame = sse.terminal_frame(task, lambda result: result.model_dump())
372 if frame is not None:
373 yield frame