Coverage for src/lilbee/data/extract/document.py: 100%
231 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""Document extraction: one xberg pass that natively extracts text and OCRs
2scanned pages/images through the registered backend; chunk + embed the result."""
4from __future__ import annotations
6import contextvars
7import logging
8import time
9from collections.abc import AsyncGenerator, Generator, Sequence
10from contextlib import asynccontextmanager, contextmanager
11from pathlib import Path
12from typing import TYPE_CHECKING, Any, Protocol
14from lilbee.app.services import get_services
15from lilbee.core.config import active_config
16from lilbee.data.offload import to_ingest_thread
17from lilbee.data.store import ChunkType, PageTextRecord, SourceMeta
18from lilbee.data.title import derive_title, source_meta_from_extraction
19from lilbee.data.types import (
20 IMAGE_CONTENT_TYPE,
21 MARKDOWN_OUTPUT,
22 PDF_CONTENT_TYPE,
23 ChunkRecord,
24 ExtractMode,
25 OcrBackendName,
26)
27from lilbee.runtime.progress import (
28 DetailedProgressCallback,
29 EventType,
30 ExtractEvent,
31 noop_callback,
32)
34from .backends.vision_ocr import backend_options_for, ocr_request
35from .batch import active_extract_batcher
36from .chunk import build_chunking_config, chunk_text
37from .trace import ExtractionTrace, trace_extraction, trace_log
39if TYPE_CHECKING:
40 from xberg import (
41 ExtractedDocument,
42 ExtractionConfig,
43 LayoutDetectionConfig,
44 OcrConfig,
45 PdfConfig,
46 )
48 from .batch import ExtractBatcher
50log = logging.getLogger(__name__)
53class _ExtractedTable(Protocol):
54 """The table fields lilbee indexes as dedicated chunks.
56 Structural, so it is satisfied by both xberg's public ``Table`` and the native
57 type that ``ExtractedDocument.tables`` actually yields.
58 """
60 @property
61 def markdown(self) -> str: ...
63 @property
64 def page_number(self) -> int: ...
67def content_type_to_mode(content_type: str) -> ExtractMode:
68 """Map a content_type to the extraction mode (paginated for PDFs and images)."""
69 if content_type in (PDF_CONTENT_TYPE, IMAGE_CONTENT_TYPE):
70 return ExtractMode.PAGINATED
71 return ExtractMode.MARKDOWN
74def _page_text_record(source: str, page: int, text: str, content_type: str) -> PageTextRecord:
75 """Build one per-page text row for the export dataset."""
76 return PageTextRecord(source=source, page=page, text=text, content_type=content_type)
79_ocr_enable_override: contextvars.ContextVar[bool | None] = contextvars.ContextVar(
80 "lilbee_ocr_enable_override", default=None
81)
82_ocr_timeout_override: contextvars.ContextVar[float | None] = contextvars.ContextVar(
83 "lilbee_ocr_timeout_override", default=None
84)
87# Per-file title for embedding input; a ContextVar (like the OCR overrides) so the
88# call chains need no signature threading and concurrent ingests don't leak titles.
89_embed_title: contextvars.ContextVar[str] = contextvars.ContextVar("lilbee_embed_title", default="")
92@contextmanager
93def _title_scope(title: str) -> Generator[None, None, None]:
94 token = _embed_title.set(title or "")
95 try:
96 yield
97 finally:
98 _embed_title.reset(token)
101def _embed_inputs(texts: list[str], title: str | None = None) -> list[str]:
102 """Embedding inputs, title-prefixed when ``cfg.embed_titles`` is on.
104 Only the vector sees the title; the stored chunk text is unchanged.
105 ``None`` falls back to the scoped per-file title (the OCR chains).
106 """
107 effective = title if title is not None else _embed_title.get()
108 if not effective or not active_config().embed_titles:
109 return texts
110 return [f"{effective}\n{text}" for text in texts]
113# Contextual enrichment: characters of document head shown to the model, and
114# the reply budget for the one situating sentence.
115_ENRICH_HEAD_CHARS = 2000
116_ENRICH_CHUNK_CHARS = 2000
117_ENRICH_MAX_TOKENS = 60
118_ENRICH_PROMPT = (
119 "Document beginning:\n{head}\n\nChunk from the same document:\n{chunk}\n\n"
120 "Write one short sentence situating this chunk within the document, to "
121 "improve search retrieval of the chunk. Answer with only the sentence."
122)
125def _enrich_texts(texts: list[str], doc_head: str, source_name: str) -> list[str]:
126 """Embedding inputs with one LLM-written situating sentence per chunk.
128 Anthropic-style contextual retrieval, opt-in (``cfg.contextual_enrichment``):
129 one generation per chunk, so ingest slows accordingly. Only the vector sees
130 the sentence; stored chunk text and citations stay verbatim. Any failure
131 keeps that chunk's bare text.
132 """
133 if not active_config().contextual_enrichment or not texts:
134 return texts
135 from lilbee.retrieval.reasoning import strip_reasoning
137 provider = get_services().provider
138 head = doc_head[:_ENRICH_HEAD_CHARS]
139 enriched: list[str] = []
140 failed = 0
141 for text in texts:
142 prompt = _ENRICH_PROMPT.format(head=head, chunk=text[:_ENRICH_CHUNK_CHARS])
143 try:
144 response = provider.chat(
145 [{"role": "user", "content": prompt}],
146 stream=False,
147 options={"num_predict": _ENRICH_MAX_TOKENS},
148 )
149 lines = strip_reasoning(response.text).strip().splitlines()
150 sentence = lines[0].strip() if lines else ""
151 except Exception:
152 failed += 1
153 sentence = ""
154 enriched.append(f"{sentence}\n{text}" if sentence else text)
155 if failed:
156 log.warning(
157 "Contextual enrichment failed for %d of %d chunks in %s; those embed bare",
158 failed,
159 len(texts),
160 source_name,
161 )
162 return enriched
165def _effective_enable_ocr() -> bool | None:
166 """``cfg.enable_ocr`` unless a per-request OCR override is active.
168 The override is a ContextVar, not a global cfg mutation, so concurrent
169 ingests on the shared HTTP daemon each see their own setting.
170 """
171 override = _ocr_enable_override.get()
172 return active_config().enable_ocr if override is None else override
175def _effective_ocr_timeout() -> float:
176 """``cfg.ocr_timeout`` unless a per-request OCR timeout override is active."""
177 override = _ocr_timeout_override.get()
178 return active_config().ocr_timeout if override is None else override
181@contextmanager
182def ocr_override(
183 enable_ocr: bool | None = None, ocr_timeout: float | None = None
184) -> Generator[None, None, None]:
185 """Scope per-request OCR settings without mutating the global cfg.
187 A ``None`` argument leaves that setting at its cfg default. Each override is
188 isolated to the entering context, so overlapping ingests never clobber one
189 another's OCR config.
190 """
191 tokens: list[tuple[contextvars.ContextVar[Any], contextvars.Token[Any]]] = []
192 try:
193 if enable_ocr is not None:
194 tokens.append((_ocr_enable_override, _ocr_enable_override.set(enable_ocr)))
195 if ocr_timeout is not None:
196 tokens.append((_ocr_timeout_override, _ocr_timeout_override.set(ocr_timeout)))
197 yield
198 finally:
199 for var, token in reversed(tokens):
200 var.reset(token)
203def _ocr_config(ocr_token: str | None) -> OcrConfig:
204 """Pick the OCR backend for this extraction.
206 Mirrors the prior fallback policy: OCR off when ``enable_ocr`` is False; lilbee's
207 vision backend when a vision model is configured; otherwise xberg's tesseract.
208 xberg auto-OCRs only the pages that lack a text layer.
209 """
210 from xberg import OcrConfig
212 config = active_config()
213 if _effective_enable_ocr() is False:
214 return OcrConfig(enabled=False)
215 if config.vision_model:
216 options = backend_options_for(ocr_token) if ocr_token else None
217 return OcrConfig(
218 backend=OcrBackendName.LILBEE_VISION,
219 backend_options=options,
220 )
221 # xberg requires a non-empty language list (4.x defaulted to English;
222 # xberg 1.0 errors on an empty one). cfg.ocr_language is validated non-empty.
223 return OcrConfig(backend=OcrBackendName.TESSERACT, language=list(config.ocr_language))
226def _ocr_force_requested() -> bool:
227 """Whether LILBEE_OCR_FORCE forces vision OCR on every page (targeted re-ingest lever)."""
228 import os
230 return os.environ.get("LILBEE_OCR_FORCE", "").strip().lower() in {"1", "true", "yes"}
233# Header/footer band stripped when layout detection is on: outermost 5%.
234_TOP_MARGIN_FRACTION = 0.05
235_BOTTOM_MARGIN_FRACTION = 0.05
238def _pdf_options() -> PdfConfig | None:
239 """PdfConfig for the enabled opt-in features (tables, layout), or None when all off."""
240 config = active_config()
241 if not (config.table_extraction or config.layout_detection):
242 return None
243 from xberg import PdfConfig
245 kwargs: dict[str, Any] = {}
246 if config.table_extraction:
247 kwargs["extract_tables"] = True
248 if config.layout_detection:
249 kwargs.update(
250 reading_order=True,
251 top_margin_fraction=_TOP_MARGIN_FRACTION,
252 bottom_margin_fraction=_BOTTOM_MARGIN_FRACTION,
253 )
254 return PdfConfig(**kwargs)
257def warn_if_table_model_ignored() -> None:
258 """Warn when table extraction runs with layout_detection off: xberg only
259 applies table_model inside layout detection, so the model is silently ignored.
260 """
261 config = active_config()
262 if config.table_extraction and not config.layout_detection:
263 log.warning(
264 "table_model=%s is ignored while layout_detection is off: tables use "
265 "the native extractor, not the structure model. Enable layout_detection "
266 "to apply the table model.",
267 config.table_model.value,
268 )
271def _layout_config() -> LayoutDetectionConfig | None:
272 """AUTO-strategy layout config when enabled, else None."""
273 config = active_config()
274 if not config.layout_detection:
275 return None
276 from xberg import LayoutDetectionConfig, LayoutStrategy
278 return LayoutDetectionConfig(strategy=LayoutStrategy.AUTO, table_model=config.table_model)
281def extraction_config(mode: ExtractMode, *, ocr_token: str | None = None) -> ExtractionConfig:
282 """Build ExtractionConfig for the given extraction mode."""
283 from xberg import ExtractionConfig, PageConfig
285 # Files are extracted one per call; xberg parallelizes OCR across a document's
286 # pages internally, and cross-file concurrency is the pipeline's semaphore.
287 chunking = build_chunking_config()
288 ocr = _ocr_config(ocr_token)
289 # Defeats xberg's text-layer short-circuit; vision path only (GPU re-OCR lever).
290 force_ocr = _ocr_force_requested() and ocr.backend == OcrBackendName.LILBEE_VISION
291 if mode is ExtractMode.PAGINATED:
292 paginated = ExtractionConfig(
293 chunking=chunking,
294 pages=PageConfig(extract_pages=True, insert_page_markers=False),
295 ocr=ocr,
296 force_ocr=force_ocr,
297 pdf_options=_pdf_options(),
298 )
299 # Set only when on: the keys are absent rather than None, leaving xberg's
300 # defaults in place when layout detection is off.
301 layout = _layout_config()
302 if layout is not None:
303 paginated["layout"] = layout
304 paginated["use_layout_for_markdown"] = True
305 return paginated
306 return ExtractionConfig(
307 chunking=chunking,
308 output_format=MARKDOWN_OUTPUT,
309 ocr=ocr,
310 force_ocr=force_ocr,
311 )
314def make_extract_batcher() -> ExtractBatcher | None:
315 """The extraction batcher for this ingest run, or None when batching is off."""
316 config = active_config()
317 if not config.batch_extraction:
318 return None
319 from .batch import ExtractBatcher
320 from .xberg import aextract_batch
322 return ExtractBatcher(
323 size=config.batch_extraction_size,
324 config_fn=extraction_config,
325 ocr_fn=_ocr_config,
326 batch_fn=aextract_batch,
327 )
330@asynccontextmanager
331async def extract_batching() -> AsyncGenerator[None]:
332 """Activate extraction batching for the enclosed ingest, when the toggle is on.
334 The batcher is set before the block runs so the ingest tasks created inside it
335 inherit it in their copied context; off (the default) is a no-op.
336 """
337 batcher = make_extract_batcher()
338 if batcher is None:
339 yield
340 return
341 from .batch import reset_active_batcher, set_active_batcher
343 token = set_active_batcher(batcher)
344 try:
345 yield
346 finally:
347 await batcher.close()
348 reset_active_batcher(token)
351def _chunk_pages(page_texts: Sequence[tuple[int, str]]) -> list[tuple[int, str]]:
352 """Chunk each page's text. Semantic chunking is off: a single page rarely spans
353 multiple topics, so the semantic round-trip is not worth it."""
354 return [
355 (page_num, chunk)
356 for page_num, text in page_texts
357 for chunk in chunk_text(text, use_semantic=False)
358 ]
361async def chunk_and_embed_pages(
362 page_texts: Sequence[tuple[int, str]],
363 source_name: str,
364 content_type: str,
365 on_progress: DetailedProgressCallback,
366) -> list[ChunkRecord]:
367 """Chunk per-page text and embed every chunk. Used by the dataset import path."""
368 if not page_texts:
369 return []
371 # chunk_text runs xberg's synchronous extractor; offload it so a long
372 # document does not stall sibling files sharing this event loop.
373 all_chunks = await to_ingest_thread(_chunk_pages, page_texts)
374 if not all_chunks:
375 return []
376 texts = [c for _, c in all_chunks]
377 embed_texts = await to_ingest_thread(_enrich_texts, texts, page_texts[0][1], source_name)
378 vectors = await to_ingest_thread(
379 get_services().embedder.embed_batch,
380 _embed_inputs(embed_texts),
381 source=source_name,
382 on_progress=on_progress,
383 )
384 return [
385 ChunkRecord(
386 source=source_name,
387 content_type=content_type,
388 chunk_type=ChunkType.RAW,
389 page_start=page_num,
390 page_end=page_num,
391 line_start=0,
392 line_end=0,
393 chunk=text,
394 chunk_index=i,
395 vector=vec,
396 )
397 for i, ((page_num, text), vec) in enumerate(zip(all_chunks, vectors, strict=True))
398 ]
401def _capture_result_page_texts(
402 doc: ExtractedDocument,
403 source_name: str,
404 content_type: str,
405 page_texts_out: list[PageTextRecord] | None,
406) -> None:
407 """Append an extraction's page texts to the export accumulator.
409 Paginated documents yield one row per ``doc.pages`` entry; others have no
410 page split, so the full ``doc.content`` is recorded as page 0.
411 """
412 if page_texts_out is None:
413 return
414 if doc.pages:
415 page_texts_out.extend(
416 _page_text_record(source_name, page.page_number, page.content, content_type)
417 for page in doc.pages
418 )
419 elif doc.content.strip():
420 page_texts_out.append(_page_text_record(source_name, 0, doc.content, content_type))
423def _document_tables(doc: ExtractedDocument) -> list[_ExtractedTable]:
424 """The result tables to index as dedicated chunks, when table extraction is on.
426 Each table becomes its own chunk carrying xberg's markdown serialization and
427 page metadata. The table's flattened text also stays inside the content
428 chunks: stripping it there would tear holes in reading-order prose and the
429 page-text export, so both are indexed deliberately. The dedicated table
430 chunk adds the structured serialization for targeted retrieval.
431 """
432 if not active_config().table_extraction:
433 return []
434 return [t for t in (doc.tables or []) if t.markdown and t.markdown.strip()]
437def _warn_empty_ocr(source_name: str, media: str) -> None:
438 """Warn that extraction yielded no text and point to the vision-model remedy."""
439 log.warning(
440 "Skipped %s: text extraction produced no usable text. "
441 "For better results on %s, configure a vision model "
442 "via PUT /api/models/vision or set LILBEE_ENABLE_OCR=true.",
443 source_name,
444 media,
445 )
448async def ingest_document(
449 path: Path,
450 source_name: str,
451 content_type: str,
452 *,
453 quiet: bool = False,
454 on_progress: DetailedProgressCallback = noop_callback,
455 page_texts_out: list[PageTextRecord] | None = None,
456) -> tuple[list[ChunkRecord], SourceMeta]:
457 """Extract, chunk, and embed a document in a single xberg pass, with its metadata.
459 xberg extracts native text and, where a page has none, OCRs it through the
460 registered backend (lilbee's vision model, or tesseract). Per-page OCR progress
461 is streamed as a running count via ``ocr_request``. ``quiet`` is accepted for
462 pipeline call compatibility. The returned metadata carries the document's
463 extraction title/authors/date and is derived even when extraction yields nothing.
464 """
465 del quiet
466 from .xberg import aextract_document
468 page_seen = 0
470 def _tick() -> None:
471 nonlocal page_seen
472 page_seen += 1
473 on_progress(
474 EventType.EXTRACT,
475 ExtractEvent(file=source_name, page=page_seen, total_pages=0),
476 )
478 trace_log.debug("extract-start source=%r type=%s", source_name, content_type)
479 started = time.perf_counter()
480 with ocr_request(on_page=_tick, timeout=_effective_ocr_timeout()) as token:
481 mode = content_type_to_mode(content_type)
482 batcher = active_extract_batcher()
483 if batcher is not None:
484 doc = await batcher.submit(mode, path.read_bytes(), path.name, token)
485 else:
486 config = extraction_config(mode, ocr_token=token)
487 # xberg's extract is async; awaiting it keeps the OCR page loop off this thread.
488 doc = await aextract_document(path.read_bytes(), filename=path.name, config=config)
489 elapsed = time.perf_counter() - started
491 # One trace line per extraction (filename, timing, counts, OCR pages), plus a
492 # vision line for scanned files. Emitted for empty results too (a slow file
493 # that yields nothing is worth surfacing).
494 trace_extraction(
495 ExtractionTrace(
496 source=source_name,
497 content_type=content_type,
498 elapsed_s=elapsed,
499 page_count=len(doc.pages or []) or len(doc.chunks or []),
500 chunk_count=len(doc.chunks or []),
501 ocr_pages=page_seen,
502 vision_configured=bool(active_config().vision_model),
503 )
504 )
506 # Derived before the empty-result return so a scan's title/authors survive zero chunks.
507 meta = source_meta_from_extraction(doc.metadata, source_name)
509 tables = _document_tables(doc)
510 if not doc.chunks and not tables:
511 if content_type in (PDF_CONTENT_TYPE, IMAGE_CONTENT_TYPE):
512 _warn_empty_ocr(source_name, "scanned documents")
513 return [], meta
515 _capture_result_page_texts(doc, source_name, content_type, page_texts_out)
517 # One EXTRACT event per file so progress subscribers show "extracted N pages"
518 # before embedding; result.pages, or the chunk count for non-paginated docs.
519 page_count = len(doc.pages or []) or len(doc.chunks or [])
520 on_progress(
521 EventType.EXTRACT,
522 ExtractEvent(file=source_name, page=page_count, total_pages=page_count),
523 )
525 # Content chunks and table serializations share one embed batch; the vector
526 # list is split back apart below by position.
527 texts = [chunk.content for chunk in doc.chunks or []]
528 table_texts = [table.markdown for table in tables]
529 embed_texts = await to_ingest_thread(
530 _enrich_texts, texts + table_texts, texts[0] if texts else "", source_name
531 )
532 vectors = await to_ingest_thread(
533 get_services().embedder.embed_batch,
534 _embed_inputs(embed_texts, meta.title),
535 source=source_name,
536 on_progress=on_progress,
537 )
538 records = [
539 ChunkRecord(
540 source=source_name,
541 content_type=content_type,
542 chunk_type=ChunkType.RAW,
543 page_start=chunk.metadata.first_page or 0,
544 page_end=chunk.metadata.last_page or 0,
545 line_start=0,
546 line_end=0,
547 chunk=text,
548 chunk_index=chunk.metadata.chunk_index,
549 vector=vec,
550 )
551 for chunk, text, vec in zip(doc.chunks or [], texts, vectors[: len(texts)], strict=True)
552 ]
553 # Table chunk indices continue after the content chunks so a source's
554 # (source, chunk_index) pairs stay unique.
555 records.extend(
556 ChunkRecord(
557 source=source_name,
558 content_type=content_type,
559 chunk_type=ChunkType.TABLE,
560 page_start=table.page_number,
561 page_end=table.page_number,
562 line_start=0,
563 line_end=0,
564 chunk=text,
565 chunk_index=len(texts) + i,
566 vector=vec,
567 )
568 for i, (table, text, vec) in enumerate(
569 zip(tables, table_texts, vectors[len(texts) :], strict=True)
570 )
571 )
572 return records, meta
575def _markdown_h1(text: str) -> str | None:
576 """The document's leading ``# Heading``, the best title a note carries.
578 Only a top-level ATX heading counts; ``##`` and deeper are sections, not the
579 document title. None when the note opens without one.
580 """
581 for line in text.splitlines():
582 stripped = line.strip()
583 if not stripped:
584 continue
585 if stripped.startswith("# "):
586 return stripped[2:].strip() or None
587 return None
588 return None
591async def ingest_markdown(
592 path: Path,
593 source_name: str,
594 on_progress: DetailedProgressCallback = noop_callback,
595 page_texts_out: list[PageTextRecord] | None = None,
596) -> tuple[list[ChunkRecord], SourceMeta]:
597 """Chunk a markdown file with heading context prepended to each chunk.
599 Each chunk gets the heading hierarchy path (e.g. "# Setup > ## Install")
600 prepended for better retrieval context. When ``page_texts_out`` is given,
601 the full text is appended as page 0 for export. The returned metadata's
602 title is the note's leading ``# Heading`` when it has one, else the stem.
603 """
604 raw_text = await to_ingest_thread(path.read_text, encoding="utf-8", errors="replace")
605 meta = SourceMeta(title=derive_title(source_name, _markdown_h1(raw_text)))
606 if not raw_text.strip():
607 return [], meta
609 # chunk_text runs xberg's synchronous extractor; offload it so a large
610 # markdown doc does not stall sibling files sharing this event loop.
611 texts = await to_ingest_thread(
612 chunk_text, raw_text, mime_type="text/markdown", heading_context=True
613 )
614 if not texts:
615 return [], meta
617 if page_texts_out is not None:
618 page_texts_out.append(_page_text_record(source_name, 0, raw_text, "text"))
620 embed_texts = await to_ingest_thread(_enrich_texts, texts, raw_text, source_name)
621 vectors = await to_ingest_thread(
622 get_services().embedder.embed_batch,
623 _embed_inputs(embed_texts, meta.title),
624 source=source_name,
625 on_progress=on_progress,
626 )
627 records = [
628 ChunkRecord(
629 source=source_name,
630 content_type="text",
631 chunk_type=ChunkType.RAW,
632 page_start=0,
633 page_end=0,
634 line_start=0,
635 line_end=0,
636 chunk=t,
637 chunk_index=idx,
638 vector=vec,
639 )
640 for idx, (t, vec) in enumerate(zip(texts, vectors, strict=True))
641 ]
642 return records, meta