Coverage for src/lilbee/server/handlers/ingest.py: 100%
166 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-17 10:02 +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 name_taken=reg_result.name_taken,
116 overlapping=reg_result.overlapping,
117 tracked=reg_result.tracked,
118 errors=errors,
119 )
121 if not reg_result.reached_corpus:
122 return AddSummary(copied=[], name_taken=reg_result.name_taken, errors=errors)
124 with temporary_ocr_config(enable_ocr, ocr_timeout):
125 sync_result = await sync(quiet=True, on_progress=sse.callback, cancel=sse.cancel)
127 return AddSummary(
128 copied=reg_result.registered,
129 name_taken=reg_result.name_taken,
130 overlapping=reg_result.overlapping,
131 tracked=reg_result.tracked,
132 errors=errors,
133 sync=SyncSummary(**sync_result.model_dump()),
134 )
135 finally:
136 sse.queue.put_nowait(None)
139def validate_add_paths(
140 data: dict[str, Any],
141) -> tuple[list[str], bool, bool | None, float | None]:
142 """Validate add-files input. Raises ValueError on bad input."""
143 paths = data.get("paths")
144 if not isinstance(paths, list) or not paths:
145 raise ValueError("'paths' must be a non-empty list of strings")
146 # No file-count cap: the resource guard is the app's size-based
147 # request_max_body_size; a count limit only breaks the point-lilbee-at-
148 # your-codebase use case (hundreds of small files).
150 for p_str in paths:
151 # The basename becomes the source root's label (its key prefix). It must
152 # be a clean single segment: Path(x).name cannot traverse, so this only
153 # rejects an empty name ("/", "a/") that would name no root.
154 name = Path(p_str).name
155 if not name:
156 raise ValueError(f"{p_str!r} does not name a file")
157 validate_path_within(cfg.documents_dir / name, cfg.documents_dir)
159 force = bool(data.get("force", False))
160 enable_ocr, ocr_timeout = _parse_ocr_params(data)
161 return paths, force, enable_ocr, ocr_timeout
164def _parse_ocr_params(data: dict[str, Any]) -> tuple[bool | None, float | None]:
165 """Extract and coerce OCR parameters from a request dict."""
166 enable_ocr = data.get("enable_ocr")
167 ocr_timeout = data.get("ocr_timeout")
168 if enable_ocr is not None:
169 enable_ocr = bool(enable_ocr)
170 if ocr_timeout is not None:
171 ocr_timeout = float(ocr_timeout)
172 return enable_ocr, ocr_timeout
175async def _ingest_stream(
176 items: list[tuple[str, _Payload]],
177 run: Callable[[list[_Payload], SseStream], Coroutine[Any, Any, AddSummary]],
178 label: str,
179) -> AsyncGenerator[str, None]:
180 """Lock per source, run ``run`` over the acquired subset, and stream SSE.
182 Shared by /api/add (server paths) and /api/add/upload (uploaded content).
183 Each item is ``(lock_key, payload)``: the key is the source identifier used
184 for the per-source ingest lock; the payload is what ``run`` receives for the
185 subset whose lock was acquired. Contended sources emit ``already_ingesting``.
186 When every source is contended the stream closes with no ``done`` event,
187 signalling the client to wait rather than retry. When only some are, the
188 acquired subset still runs and the ``done`` summary names the contended ones
189 in ``already_ingesting``, so a partial batch never reads as a full success.
190 """
191 registry = get_services().ingest_lock_registry
192 acquired, busy = await registry.acquire([key for key, _payload in items])
193 try:
194 for name in busy:
195 log.info("Rejecting %s for %s: already ingesting", label, name)
196 yield sse_event(SseEvent.ALREADY_INGESTING, {"source": name})
198 if not acquired:
199 return
201 acquired_names = {name for name, _lock in acquired}
202 locked = [payload for key, payload in items if key in acquired_names]
203 sse = SseStream()
204 task = asyncio.create_task(run(locked, sse))
205 try:
206 async for event in sse.drain(task, label):
207 yield event
208 # already_ingesting names the sources this run never attempted, so
209 # a client reading only the terminal event still sees a partial batch.
210 frame = sse.terminal_frame(
211 task, lambda s: s.model_copy(update={"already_ingesting": list(busy)}).model_dump()
212 )
213 if frame is not None:
214 yield frame
215 finally:
216 if not task.done():
217 task.cancel()
218 with contextlib.suppress(asyncio.CancelledError, Exception):
219 await task
220 finally:
221 registry.release(acquired)
224async def add_files_stream(
225 paths: list[str],
226 *,
227 force: bool = False,
228 enable_ocr: bool | None = None,
229 ocr_timeout: float | None = None,
230) -> AsyncGenerator[str, None]:
231 """Copy server-side files, sync, and yield SSE progress events.
233 Takes the already-validated/parsed values from ``validate_add_paths`` so the
234 request dict is decoded once.
235 """
236 async for event in _ingest_stream(
237 [(IngestLockRegistry.canonical_source_name(p), p) for p in paths],
238 lambda locked, sse: _run_add(locked, force, enable_ocr, ocr_timeout, sse),
239 "Add files stream",
240 ):
241 yield event
244def _clean_upload_name(name: str) -> str:
245 """Normalize one upload filename to a safe relative path inside the corpus.
247 Relative paths are preserved (a source tree keeps its layout instead of
248 colliding on basenames); absolute paths, drive letters, and ``..`` segments
249 are rejected. Raises ValueError on bad input.
250 """
251 normalized = name.replace("\\", "/")
252 if normalized.startswith("/") or re.match(r"^[A-Za-z]:", normalized):
253 raise ValueError(f"upload filename must be relative: {name!r}")
254 parts = [part for part in normalized.split("/") if part not in ("", ".")]
255 if not parts:
256 raise ValueError(f"invalid upload filename: {name!r}")
257 if ".." in parts:
258 raise ValueError(f"upload filename may not contain '..': {name!r}")
259 relative = "/".join(parts)
260 reason = excluded_extension_reasons().get(Path(relative).suffix.lower())
261 if reason is not None:
262 raise ValueError(f"{name!r}: {reason}")
263 validate_path_within(cfg.documents_dir / relative, cfg.documents_dir)
264 return relative
267def validate_upload_names(names: list[str | None]) -> list[str]:
268 """Validate uploaded filenames, returning the cleaned relative paths.
270 Names only, deliberately: the route validates before reading any part's
271 bytes, so a request that will be rejected never costs a full copy of the
272 payload in the server's own memory. Filenames keep their relative path,
273 validated to stay inside ``cfg.documents_dir``, so an uploaded source tree
274 preserves its layout instead of colliding on basenames. There is no
275 file-count cap; the resource guard is the app's request_max_body_size.
277 Raises ValueError on bad input.
278 """
279 if not names:
280 raise ValueError("no files uploaded")
281 # A multipart part is allowed to carry no filename at all; that is not a
282 # file upload, and the cleaner's message should name it as missing rather
283 # than crash on None.
284 return [_clean_upload_name(name if name is not None else "") for name in names]
287async def _run_upload(files: list[tuple[str, bytes]], sse: SseStream) -> AddSummary:
288 """Write uploaded file bytes into ``cfg.documents_dir``, then sync.
290 The upload equivalent of :func:`_run_add`: instead of copying from a
291 server-readable path, the client's content is written straight into the
292 documents dir and the same ingest pipeline runs. This is what lets an
293 external-mode client (a remote lilbee / GPU box) ingest files that live only
294 on the client. Unchanged content is a no-op re-embed inside ``sync`` (it
295 hashes each source), so there is no separate force flag.
296 """
297 from lilbee.app.ingest import temporary_ocr_config
298 from lilbee.data.ingest import sync
300 try:
301 cfg.documents_dir.mkdir(parents=True, exist_ok=True)
302 written: list[str] = []
303 for name, content in files:
304 dest = cfg.documents_dir / name
305 dest.parent.mkdir(parents=True, exist_ok=True)
306 if not _move_same_content(name, content, dest):
307 dest.write_bytes(content)
308 written.append(name)
309 with temporary_ocr_config(None):
310 sync_result = await sync(quiet=True, on_progress=sse.callback, cancel=sse.cancel)
311 return AddSummary(
312 copied=written,
313 errors=[],
314 sync=SyncSummary(**sync_result.model_dump()),
315 )
316 finally:
317 sse.queue.put_nowait(None)
320def _move_same_content(name: str, content: bytes, dest: Path) -> bool:
321 """Move the one indexed file holding *content* to *dest*, so sync repoints its key."""
322 digest = hashlib.sha256(content).hexdigest()
323 matches = [
324 s["filename"]
325 for s in get_services().store.get_sources()
326 if s["file_hash"] == digest and s["filename"] != name
327 ]
328 if len(matches) != 1:
329 return False
330 old_path = cfg.documents_dir / matches[0]
331 if not old_path.is_file():
332 return False
333 os.replace(old_path, dest)
334 log.info("Moved %s to %s: the same content arrived under a new name", matches[0], name)
335 return True
338async def add_uploads_stream(files: list[tuple[str, bytes]]) -> AsyncGenerator[str, None]:
339 """Ingest uploaded file content, yielding the same SSE progress as add_files_stream.
341 Locks per source name (the validated relative path) so an upload never
342 races an in-flight add of the same source.
343 """
344 async for event in _ingest_stream(
345 [(name, (name, content)) for name, content in files],
346 lambda locked, sse: _run_upload(locked, sse),
347 "Add uploads stream",
348 ):
349 yield event
352async def _run_import_with_sentinel(sse: SseStream, data: bytes, fmt: str) -> ImportSummary:
353 """Run the dataset import and guarantee the drain sentinel is enqueued."""
354 from lilbee.app.dataset import import_from_bytes
356 try:
357 return await import_from_bytes(data, fmt, on_progress=sse.callback)
358 finally:
359 sse.queue.put_nowait(None)
362async def import_stream(data: bytes, fmt: str) -> AsyncGenerator[str, None]:
363 """Import a dataset, yield SSE embed-progress events, then a done event."""
364 sse = SseStream()
365 task = asyncio.create_task(_run_import_with_sentinel(sse, data, fmt))
366 async for event in sse.drain(task, "Import stream"):
367 yield event
368 frame = sse.terminal_frame(task, lambda result: result.model_dump())
369 if frame is not None:
370 yield frame