Coverage for src/lilbee/data/extract/document.py: 100%

1"""Document extraction: one xberg pass that natively extracts text and OCRs 

2scanned pages/images through the registered backend; chunk + embed the result.""" 

3 

4from __future__ import annotations 

5 

6import asyncio 

7import contextvars 

8import logging 

9import time 

10from collections.abc import AsyncGenerator, Callable, Generator, Sequence 

11from contextlib import asynccontextmanager, contextmanager 

12from dataclasses import dataclass, replace 

13from pathlib import Path 

14from typing import TYPE_CHECKING, Any, Protocol 

15 

16from lilbee.app.services import get_services 

17from lilbee.core.config import active_config 

18from lilbee.core.config.enums import OcrMode, OcrPageStrategy 

19from lilbee.core.web_names import is_web_source 

20from lilbee.data.offload import to_ingest_thread 

21from lilbee.data.store import ChunkType, PageTextRecord, SourceMeta 

22from lilbee.data.title import derive_title, source_meta_from_extraction 

23from lilbee.data.types import ( 

24 IMAGE_CONTENT_TYPE, 

25 MARKDOWN_OUTPUT, 

26 PDF_CONTENT_TYPE, 

27 ChunkRecord, 

28 DocumentRecords, 

29 ExtractMode, 

30 MemberRecords, 

31 OcrBackendName, 

32 OcrReport, 

33) 

34from lilbee.providers.base import aux_options 

35from lilbee.runtime.cancellation import CancelSignal 

36from lilbee.runtime.progress import ( 

37 DetailedProgressCallback, 

38 EventType, 

39 ExtractEvent, 

40 OcrBackendUsed, 

41 OcrStartEvent, 

42 noop_callback, 

43) 

44 

45from .backends.vision_ocr import backend_options_for, ocr_request 

46from .batch import active_extract_batcher 

47from .chunk import ChunkLimitError, build_chunking_config, chunk_text, enforce_chunk_limit 

48from .trace import ExtractionTrace, trace_extraction, trace_log 

49 

50if TYPE_CHECKING: 

51 from xberg import ( 

52 ExtractedDocument, 

53 ExtractionConfig, 

54 LayoutDetectionConfig, 

55 OcrConfig, 

56 OcrStrategy, 

57 PdfConfig, 

58 ) 

59 from xberg.progress import ProgressEvent as OcrPageEvent 

60 

61 from .batch import ExtractBatcher 

62 

63log = logging.getLogger(__name__) 

64 

65 

66class _ExtractedTable(Protocol): 

67 """The table fields lilbee indexes as dedicated chunks. 

68 

69 Structural, so it is satisfied by both xberg's public ``Table`` and the native 

70 type that ``ExtractedDocument.tables`` actually yields. 

71 """ 

72 

73 @property 

74 def markdown(self) -> str: ... 

75 

76 @property 

77 def page_number(self) -> int: ... 

78 

79 

80def content_type_to_mode(content_type: str) -> ExtractMode: 

81 """Map a content_type to the extraction mode (paginated for PDFs and images).""" 

82 if content_type in (PDF_CONTENT_TYPE, IMAGE_CONTENT_TYPE): 

83 return ExtractMode.PAGINATED 

84 return ExtractMode.MARKDOWN 

85 

86 

87def _page_text_record(source: str, page: int, text: str, content_type: str) -> PageTextRecord: 

88 """Build one per-page text row for the export dataset.""" 

89 return PageTextRecord(source=source, page=page, text=text, content_type=content_type) 

90 

91 

92_ocr_mode_override: contextvars.ContextVar[OcrMode | None] = contextvars.ContextVar( 

93 "lilbee_ocr_mode_override", default=None 

94) 

95_ocr_timeout_override: contextvars.ContextVar[float | None] = contextvars.ContextVar( 

96 "lilbee_ocr_timeout_override", default=None 

97) 

98 

99 

100# Per-file title for embedding input; a ContextVar (like the OCR overrides) so the 

101# call chains need no signature threading and concurrent ingests don't leak titles. 

102_embed_title: contextvars.ContextVar[str] = contextvars.ContextVar("lilbee_embed_title", default="") 

103 

104 

105@contextmanager 

106def _title_scope(title: str) -> Generator[None, None, None]: 

107 token = _embed_title.set(title or "") 

108 try: 

109 yield 

110 finally: 

111 _embed_title.reset(token) 

112 

113 

114def _embed_inputs(texts: list[str], title: str | None = None) -> list[str]: 

115 """Embedding inputs, title-prefixed when ``cfg.embed_titles`` is on. 

116 

117 Only the vector sees the title; the stored chunk text is unchanged. 

118 ``None`` falls back to the scoped per-file title (the OCR chains). 

119 """ 

120 effective = title if title is not None else _embed_title.get() 

121 if not effective or not active_config().embed_titles: 

122 return texts 

123 return [f"{effective}\n{text}" for text in texts] 

124 

125 

126# Contextual enrichment: characters of document head shown to the model, and 

127# the reply budget for the one situating sentence. 

128_ENRICH_HEAD_CHARS = 2000 

129_ENRICH_CHUNK_CHARS = 2000 

130_ENRICH_MAX_TOKENS = 60 

131_ENRICH_PROMPT = ( 

132 "Document beginning:\n{head}\n\nChunk from the same document:\n{chunk}\n\n" 

133 "Write one short sentence situating this chunk within the document, to " 

134 "improve search retrieval of the chunk. Answer with only the sentence." 

135) 

136 

137 

138def _enrich_texts(texts: list[str], doc_head: str, source_name: str) -> list[str]: 

139 """Embedding inputs with one LLM-written situating sentence per chunk. 

140 

141 Anthropic-style contextual retrieval, opt-in (``cfg.contextual_enrichment``): 

142 one generation per chunk, so ingest slows accordingly. Only the vector sees 

143 the sentence; stored chunk text and citations stay verbatim. Any failure 

144 keeps that chunk's bare text. 

145 """ 

146 if not active_config().contextual_enrichment or not texts: 

147 return texts 

148 from lilbee.retrieval.reasoning import strip_reasoning 

149 

150 provider = get_services().provider 

151 head = doc_head[:_ENRICH_HEAD_CHARS] 

152 enriched: list[str] = [] 

153 failed = 0 

154 for text in texts: 

155 prompt = _ENRICH_PROMPT.format(head=head, chunk=text[:_ENRICH_CHUNK_CHARS]) 

156 try: 

157 response = provider.chat( 

158 [{"role": "user", "content": prompt}], 

159 stream=False, 

160 options=aux_options(_ENRICH_MAX_TOKENS), 

161 ) 

162 lines = strip_reasoning(response.text).strip().splitlines() 

163 sentence = lines[0].strip() if lines else "" 

164 except Exception: 

165 failed += 1 

166 sentence = "" 

167 enriched.append(f"{sentence}\n{text}" if sentence else text) 

168 if failed: 

169 log.warning( 

170 "Contextual enrichment failed for %d of %d chunks in %s; those embed bare", 

171 failed, 

172 len(texts), 

173 source_name, 

174 ) 

175 return enriched 

176 

177 

178def _effective_ocr_mode() -> OcrMode: 

179 """``cfg.ocr`` unless a per-request OCR override is active. 

180 

181 The override is a ContextVar, not a global cfg mutation, so concurrent 

182 ingests on the shared HTTP daemon each see their own setting. 

183 """ 

184 override = _ocr_mode_override.get() 

185 return active_config().ocr if override is None else override 

186 

187 

188def _extraction_timeout_secs() -> int | None: 

189 """``cfg.extraction_timeout`` as xberg's per-file cap; None when uncapped. 

190 

191 xberg defaults this to 600s. Passing lilbee's own value on every call keeps 

192 the cap something a user can see and raise instead of an inherited default. 

193 """ 

194 timeout = active_config().extraction_timeout 

195 return timeout if timeout > 0 else None 

196 

197 

198def _effective_ocr_timeout() -> float: 

199 """``cfg.ocr_timeout`` unless a per-request OCR timeout override is active.""" 

200 override = _ocr_timeout_override.get() 

201 return active_config().ocr_timeout if override is None else override 

202 

203 

204@contextmanager 

205def ocr_override( 

206 ocr: OcrMode | None = None, ocr_timeout: float | None = None 

207) -> Generator[None, None, None]: 

208 """Scope per-request OCR settings without mutating the global cfg. 

209 

210 A ``None`` argument leaves that setting at its cfg default. Each override is 

211 isolated to the entering context, so overlapping ingests never clobber one 

212 another's OCR config. 

213 """ 

214 tokens: list[tuple[contextvars.ContextVar[Any], contextvars.Token[Any]]] = [] 

215 try: 

216 if ocr is not None: 

217 tokens.append((_ocr_mode_override, _ocr_mode_override.set(ocr))) 

218 if ocr_timeout is not None: 

219 tokens.append((_ocr_timeout_override, _ocr_timeout_override.set(ocr_timeout))) 

220 yield 

221 finally: 

222 for var, token in reversed(tokens): 

223 var.reset(token) 

224 

225 

226def ocr_backend() -> OcrBackendUsed: 

227 """The OCR backend for this extraction: none when OCR is off, else vision or Tesseract.""" 

228 return OcrBackendUsed.chosen(_effective_ocr_mode(), active_config().vision_model) 

229 

230 

231def _ocr_config(ocr_token: str | None) -> OcrConfig | None: 

232 """xberg's OcrConfig for the backend ``ocr_backend`` picks, or None when OCR is off. 

233 

234 xberg OCRs each page without usable text; a page with usable text keeps it. 

235 """ 

236 from xberg import OcrConfig 

237 

238 config = active_config() 

239 backend = ocr_backend() 

240 if backend is OcrBackendUsed.NONE: 

241 return None 

242 if backend is OcrBackendUsed.VISION: 

243 options = backend_options_for(ocr_token) if ocr_token else None 

244 return OcrConfig( 

245 backend=OcrBackendName.LILBEE_VISION, 

246 backend_options=options, 

247 ) 

248 # xberg requires a non-empty language list (4.x defaulted to English; 

249 # xberg 1.0 errors on an empty one). cfg.ocr_language is validated non-empty. 

250 return OcrConfig( 

251 backend=OcrBackendName.TESSERACT, 

252 language=list(config.ocr_language), 

253 ) 

254 

255 

256def _ocr_strategy() -> OcrStrategy: 

257 """xberg's page-selection strategy for cfg.ocr_strategy; auto when OCR is off.""" 

258 from xberg import OcrStrategy 

259 

260 config = active_config() 

261 # xberg rejects scanned_pages when OCR is disabled. 

262 if ocr_backend() is OcrBackendUsed.NONE or config.ocr_strategy is OcrPageStrategy.AUTO: 

263 return OcrStrategy(OcrPageStrategy.AUTO.value) 

264 return OcrStrategy.scanned_pages(config.ocr_scan_confidence) 

265 

266 

267def _force_ocr_pages() -> list[int] | None: 

268 """cfg.force_ocr_pages for xberg; None when empty or OCR is off.""" 

269 if ocr_backend() is OcrBackendUsed.NONE: 

270 return None 

271 return list(active_config().force_ocr_pages) or None 

272 

273 

274# Header/footer band stripped when layout detection is on: outermost 5%. 

275_TOP_MARGIN_FRACTION = 0.05 

276_BOTTOM_MARGIN_FRACTION = 0.05 

277 

278 

279def _pdf_options() -> PdfConfig | None: 

280 """PdfConfig for the enabled opt-in features (tables, layout), or None when all off.""" 

281 config = active_config() 

282 if not (config.table_extraction or config.layout_detection): 

283 return None 

284 from xberg import PdfConfig 

285 

286 kwargs: dict[str, Any] = {} 

287 if config.table_extraction: 

288 kwargs["extract_tables"] = True 

289 if config.layout_detection: 

290 kwargs.update( 

291 reading_order=True, 

292 top_margin_fraction=_TOP_MARGIN_FRACTION, 

293 bottom_margin_fraction=_BOTTOM_MARGIN_FRACTION, 

294 ) 

295 return PdfConfig(**kwargs) 

296 

297 

298def warn_if_table_model_ignored() -> None: 

299 """Warn when table extraction runs with layout_detection off: xberg only 

300 applies table_model inside layout detection, so the model is silently ignored. 

301 """ 

302 config = active_config() 

303 if config.table_extraction and not config.layout_detection: 

304 log.warning( 

305 "table_model=%s is ignored while layout_detection is off: tables use " 

306 "the native extractor, not the structure model. Enable layout_detection " 

307 "to apply the table model.", 

308 config.table_model.value, 

309 ) 

310 

311 

312def _layout_config() -> LayoutDetectionConfig | None: 

313 """AUTO-strategy layout config when enabled, else None.""" 

314 config = active_config() 

315 if not config.layout_detection: 

316 return None 

317 from xberg import LayoutDetectionConfig, LayoutStrategy 

318 

319 return LayoutDetectionConfig(strategy=LayoutStrategy.AUTO, table_model=config.table_model) 

320 

321 

322def extraction_config(mode: ExtractMode, *, ocr_token: str | None = None) -> ExtractionConfig: 

323 """Build ExtractionConfig for the given extraction mode.""" 

324 from xberg import ExtractionConfig, PageConfig 

325 

326 # Files are extracted one per call; xberg parallelizes OCR across a document's 

327 # pages internally, and cross-file concurrency is the pipeline's semaphore. 

328 chunking = build_chunking_config() 

329 ocr = _ocr_config(ocr_token) 

330 # OCR off sends no OCR block and sets disable_ocr; xberg auto-OCRs a textless PDF otherwise. 

331 disable_ocr = ocr is None 

332 # ocr = all defeats xberg's text-layer short-circuit for either engine; xberg 

333 # rejects force_ocr together with disable_ocr. 

334 force_ocr = ocr is not None and _effective_ocr_mode() is OcrMode.ALL 

335 if mode is ExtractMode.PAGINATED: 

336 paginated = ExtractionConfig( 

337 chunking=chunking, 

338 pages=PageConfig(extract_pages=True, insert_page_markers=False), 

339 ocr=ocr, 

340 disable_ocr=disable_ocr, 

341 force_ocr=force_ocr, 

342 pdf_options=_pdf_options(), 

343 extraction_timeout_secs=_extraction_timeout_secs(), 

344 ocr_strategy=_ocr_strategy(), 

345 force_ocr_pages=_force_ocr_pages(), 

346 ) 

347 # The layout fields keep xberg's defaults when layout detection is off. 

348 layout = _layout_config() 

349 if layout is None: 

350 return paginated 

351 return replace(paginated, layout=layout) 

352 return ExtractionConfig( 

353 chunking=chunking, 

354 output_format=MARKDOWN_OUTPUT, 

355 ocr=ocr, 

356 disable_ocr=disable_ocr, 

357 force_ocr=force_ocr, 

358 extraction_timeout_secs=_extraction_timeout_secs(), 

359 ocr_strategy=_ocr_strategy(), 

360 force_ocr_pages=_force_ocr_pages(), 

361 ) 

362 

363 

364def analysis_config() -> ExtractionConfig: 

365 """Native extraction for corpus analysis: no OCR, no layout model, tables and languages on.""" 

366 from xberg import ExtractionConfig, LanguageDetectionConfig, PageConfig, PdfConfig 

367 

368 # No OCR block at all: xberg 1.2.7 OCRs a PDF's page images whenever one is present, 

369 # even disabled (fixed in 1.2.9). 

370 return ExtractionConfig( 

371 disable_ocr=True, 

372 pages=PageConfig(extract_pages=True), 

373 pdf_options=PdfConfig(extract_tables=True), 

374 language_detection=LanguageDetectionConfig(enabled=True, detect_multiple=True), 

375 extraction_timeout_secs=_extraction_timeout_secs(), 

376 ) 

377 

378 

379def make_extract_batcher() -> ExtractBatcher | None: 

380 """The extraction batcher for this ingest run, or None when batching is off.""" 

381 config = active_config() 

382 if not config.batch_extraction: 

383 return None 

384 from .batch import ExtractBatcher 

385 from .xberg import aextract_batch 

386 

387 return ExtractBatcher( 

388 size=config.batch_extraction_size, 

389 config_fn=extraction_config, 

390 ocr_fn=_ocr_config, 

391 batch_fn=aextract_batch, 

392 ) 

393 

394 

395@asynccontextmanager 

396async def extract_batching() -> AsyncGenerator[None]: 

397 """Activate extraction batching for the enclosed ingest, when the toggle is on. 

398 

399 The batcher is set before the block runs so the ingest tasks created inside it 

400 inherit it in their copied context; off (the default) is a no-op. 

401 """ 

402 batcher = make_extract_batcher() 

403 if batcher is None: 

404 yield 

405 return 

406 from .batch import reset_active_batcher, set_active_batcher 

407 

408 token = set_active_batcher(batcher) 

409 try: 

410 yield 

411 finally: 

412 await batcher.close() 

413 reset_active_batcher(token) 

414 

415 

416def _chunk_pages(page_texts: Sequence[tuple[int, str]]) -> list[tuple[int, str]]: 

417 """Chunk each page's text. Semantic chunking is off: a single page rarely spans 

418 multiple topics, so the semantic round-trip is not worth it.""" 

419 return [ 

420 (page_num, chunk) 

421 for page_num, text in page_texts 

422 for chunk in chunk_text(text, use_semantic=False) 

423 ] 

424 

425 

426async def chunk_and_embed_pages( 

427 page_texts: Sequence[tuple[int, str]], 

428 source_name: str, 

429 content_type: str, 

430 on_progress: DetailedProgressCallback, 

431) -> list[ChunkRecord]: 

432 """Chunk per-page text and embed every chunk. Used by the dataset import path.""" 

433 if not page_texts: 

434 return [] 

435 

436 # chunk_text runs xberg's synchronous extractor; offload it so a long 

437 # document does not stall sibling files sharing this event loop. 

438 all_chunks = await to_ingest_thread(_chunk_pages, page_texts) 

439 if not all_chunks: 

440 return [] 

441 texts = [c for _, c in all_chunks] 

442 embed_texts = await to_ingest_thread(_enrich_texts, texts, page_texts[0][1], source_name) 

443 vectors = await to_ingest_thread( 

444 get_services().embedder.embed_batch, 

445 _embed_inputs(embed_texts), 

446 source=source_name, 

447 on_progress=on_progress, 

448 ) 

449 return [ 

450 ChunkRecord( 

451 source=source_name, 

452 content_type=content_type, 

453 chunk_type=ChunkType.RAW, 

454 page_start=page_num, 

455 page_end=page_num, 

456 line_start=0, 

457 line_end=0, 

458 chunk=text, 

459 chunk_index=i, 

460 vector=vec, 

461 ) 

462 for i, ((page_num, text), vec) in enumerate(zip(all_chunks, vectors, strict=True)) 

463 ] 

464 

465 

466def _capture_result_page_texts( 

467 doc: ExtractedDocument, 

468 source_name: str, 

469 content_type: str, 

470 page_texts_out: list[PageTextRecord] | None, 

471) -> None: 

472 """Append an extraction's page texts to the export accumulator. 

473 

474 Paginated documents yield one row per ``doc.pages`` entry; others have no 

475 page split, so the full ``doc.content`` is recorded as page 0. 

476 """ 

477 if page_texts_out is None: 

478 return 

479 if doc.pages: 

480 page_texts_out.extend( 

481 _page_text_record(source_name, page.page_number, page.content, content_type) 

482 for page in doc.pages 

483 ) 

484 elif doc.content.strip(): 

485 page_texts_out.append(_page_text_record(source_name, 0, doc.content, content_type)) 

486 

487 

488def _document_tables(doc: ExtractedDocument) -> list[_ExtractedTable]: 

489 """The result tables to index as dedicated chunks, when table extraction is on. 

490 

491 Each table becomes its own chunk carrying xberg's markdown serialization and 

492 page metadata. The table's flattened text also stays inside the content 

493 chunks: stripping it there would tear holes in reading-order prose and the 

494 page-text export, so both are indexed deliberately. The dedicated table 

495 chunk adds the structured serialization for targeted retrieval. 

496 """ 

497 if not active_config().table_extraction: 

498 return [] 

499 return [t for t in (doc.tables or []) if t.markdown and t.markdown.strip()] 

500 

501 

502def _warn_empty_ocr(source_name: str, media: str, backend: OcrBackendUsed) -> None: 

503 """Warn that extraction yielded no text, with advice that matches the backend that ran.""" 

504 if backend is OcrBackendUsed.NONE: 

505 log.warning( 

506 "Skipped %s: text extraction produced no usable text. " 

507 "OCR is off (ocr = off); set ocr to auto to read %s.", 

508 source_name, 

509 media, 

510 ) 

511 return 

512 if backend is OcrBackendUsed.VISION: 

513 log.warning( 

514 "Skipped %s: the vision model returned no usable text for %s.", 

515 source_name, 

516 media, 

517 ) 

518 return 

519 log.warning( 

520 "Skipped %s: text extraction produced no usable text. " 

521 "For better results on %s, configure a vision model via PUT /api/models/vision.", 

522 source_name, 

523 media, 

524 ) 

525 

526 

527def _warn_skipped_scanned_pages(doc: ExtractedDocument, source_name: str) -> None: 

528 """Warn that OCR being off left the scanned pages of an indexed PDF unread.""" 

529 skipped = _scanned_pages(doc) 

530 if skipped: 

531 log.warning( 

532 "Indexed %s without its scanned pages %s: OCR is off (ocr = off). " 

533 "Set ocr to auto to read them.", 

534 source_name, 

535 ", ".join(str(page) for page in skipped), 

536 ) 

537 

538 

539async def ingest_document( 

540 path: Path, 

541 source_name: str, 

542 content_type: str, 

543 *, 

544 quiet: bool = False, 

545 on_progress: DetailedProgressCallback = noop_callback, 

546 page_texts_out: list[PageTextRecord] | None = None, 

547 cancel: CancelSignal | None = None, 

548) -> DocumentRecords: 

549 """Extract, chunk, and embed a document in a single xberg pass, with its metadata. 

550 

551 xberg extracts native text and, where a page has none, OCRs it through the 

552 registered backend (lilbee's vision model, or tesseract). xberg reports each 

553 OCR'd page of a PDF or image as it completes, and a scanned file under Tesseract 

554 also gets one OCR_START event before its extraction. ``quiet`` is accepted for 

555 pipeline call compatibility. The returned metadata carries the document's 

556 extraction title/authors/date and is derived even when extraction yields nothing; 

557 the OCR report says which backend the extraction ran and how many pages it OCR'd. 

558 """ 

559 del quiet 

560 doc, ocr = await _extract_document( 

561 path, source_name, content_type, content_type_to_mode(content_type), on_progress, cancel 

562 ) 

563 records, meta = await _records_from_document( 

564 doc, 

565 source_name, 

566 content_type, 

567 on_progress=on_progress, 

568 page_texts_out=page_texts_out, 

569 ocr_backend=ocr.backend, 

570 ) 

571 return DocumentRecords(records, meta, ocr) 

572 

573 

574async def ingest_archive( 

575 path: Path, 

576 source_name: str, 

577 content_type: str, 

578 *, 

579 on_progress: DetailedProgressCallback = noop_callback, 

580 cancel: CancelSignal | None = None, 

581) -> list[MemberRecords]: 

582 """Extract an archive once and build records for every member, nested archives included. 

583 

584 The archive itself contributes no chunks. Each member is its own source named 

585 ``<archive>/<member path>``. xberg unpacks to ``max_archive_depth`` under its 

586 zip-bomb limits, so depth and size are enforced before this runs. 

587 """ 

588 doc, ocr = await _extract_document( 

589 path, source_name, content_type, ExtractMode.PAGINATED, on_progress, cancel 

590 ) 

591 members: list[MemberRecords] = [] 

592 await _collect_members(doc, source_name, members, on_progress, ocr.backend) 

593 return members 

594 

595 

596async def _collect_members( 

597 doc: ExtractedDocument, 

598 prefix: str, 

599 members: list[MemberRecords], 

600 on_progress: DetailedProgressCallback, 

601 ocr_backend: OcrBackendUsed, 

602) -> None: 

603 # circular: document -> ingest.discovery via the ingest package's pipeline import 

604 from lilbee.data.ingest.discovery import archive_content_types, classify_file 

605 

606 unsupported: list[str] = [] 

607 for entry in doc.children or []: 

608 name = f"{prefix}/{entry.path}" 

609 # The same gate discovery applies on disk: a member of an unsupported format 

610 # is never chunked, so an archive of machine output cannot flood the index. 

611 content_type = classify_file(Path(entry.path)) 

612 if content_type is None: 

613 unsupported.append(entry.path) 

614 continue 

615 if content_type in archive_content_types(): 

616 await _collect_members(entry.result, name, members, on_progress, ocr_backend) 

617 continue 

618 page_texts: list[PageTextRecord] = [] 

619 try: 

620 records, meta = await _records_from_document( 

621 entry.result, 

622 name, 

623 content_type, 

624 on_progress=on_progress, 

625 page_texts_out=page_texts, 

626 ocr_backend=ocr_backend, 

627 ) 

628 except ChunkLimitError as exc: 

629 raise ChunkLimitError(exc.count, exc.limit, member=name) from None 

630 members.append(MemberRecords(name, content_type, records, page_texts, meta)) 

631 if unsupported: 

632 log.info( 

633 "Skipped %d member(s) of %s, unsupported format: %s", 

634 len(unsupported), 

635 prefix, 

636 ", ".join(sorted(unsupported)), 

637 ) 

638 

639 

640def _page_count_config() -> ExtractionConfig: 

641 """A metadata-only ExtractionConfig: OCR and page bodies off, structure kept.""" 

642 from xberg import ExtractionConfig, PageConfig 

643 

644 # The pass is metadata-only: no OCR block, disable_ocr set, no page bodies. 

645 return ExtractionConfig( 

646 pages=PageConfig(extract_pages=False, insert_page_markers=False), 

647 disable_ocr=True, 

648 enable_quality_processing=False, 

649 ) 

650 

651 

652@dataclass(frozen=True) 

653class _PageProbe: 

654 """What the metadata-only pass read: the page count and whether any page is a scan.""" 

655 

656 pages: int = 0 

657 has_scanned_pages: bool = False 

658 

659 

660def _scanned_pages(doc: ExtractedDocument) -> list[int]: 

661 """The PDF pages xberg found with no usable text layer.""" 

662 fmt = doc.metadata.format 

663 pdf = fmt.pdf if fmt is not None else None 

664 return list(pdf.scanned_pages or []) if pdf is not None else [] 

665 

666 

667async def _probe_pages(data: bytes, filename: str) -> _PageProbe: 

668 """Read *data*'s page count and scanned pages in a metadata-only xberg pass. 

669 

670 Costs a second structural parse of *data* with OCR and page bodies both off. 

671 Any failure here returns an empty probe rather than raising, so the caller's 

672 real extraction still runs and reports its own error. 

673 """ 

674 from .xberg import aextract_document 

675 

676 try: 

677 doc = await aextract_document(data, filename=filename, config=_page_count_config()) 

678 except Exception: 

679 log.debug("Page-count probe failed for %s; its OCR start is not announced", filename) 

680 return _PageProbe() 

681 return _PageProbe(pages=doc.counts.pages, has_scanned_pages=bool(_scanned_pages(doc))) 

682 

683 

684async def _announce_tesseract_ocr( 

685 data: bytes, 

686 filename: str, 

687 source_name: str, 

688 content_type: str, 

689 on_progress: DetailedProgressCallback, 

690) -> None: 

691 """Emit OCR_START when Tesseract will OCR a scanned page of this PDF.""" 

692 if ocr_backend() is not OcrBackendUsed.TESSERACT or not _has_own_pages(content_type): 

693 return 

694 probe = await _probe_pages(data, filename) 

695 if probe.has_scanned_pages: 

696 on_progress(EventType.OCR_START, OcrStartEvent(file=source_name, total_pages=probe.pages)) 

697 

698 

699def _has_own_pages(content_type: str) -> bool: 

700 """Whether the file is itself paginated (a PDF or an image), which an archive is not.""" 

701 # The content type, not the extraction mode: ingest_archive always requests PAGINATED. 

702 return content_type_to_mode(content_type) is ExtractMode.PAGINATED 

703 

704 

705def _ocr_page_reporter( 

706 source_name: str, backend: OcrBackendUsed, on_progress: DetailedProgressCallback 

707) -> Callable[[OcrPageEvent], None]: 

708 """An xberg progress callback that reports each OCR'd page as an EXTRACT event. 

709 

710 ``completed`` counts distinct pages, so a page xberg sends to the backend 

711 more than once still counts once. xberg calls this from a worker thread. 

712 """ 

713 

714 def _report(event: OcrPageEvent) -> None: 

715 on_progress( 

716 EventType.EXTRACT, 

717 ExtractEvent( 

718 file=source_name, 

719 page=event.completed, 

720 total_pages=event.total, 

721 ocr_backend=backend, 

722 ), 

723 ) 

724 

725 return _report 

726 

727 

728async def _extract_document( 

729 path: Path, 

730 source_name: str, 

731 content_type: str, 

732 mode: ExtractMode, 

733 on_progress: DetailedProgressCallback, 

734 cancel: CancelSignal | None, 

735) -> tuple[ExtractedDocument, OcrReport]: 

736 """Run one xberg pass over *path*, with per-page OCR progress and the extraction trace. 

737 

738 A set *cancel* stops vision OCR before its next page, and a pass that 

739 ends under a set *cancel* raises ``CancelledError`` instead of returning a 

740 partial document. 

741 """ 

742 from .xberg import aextract_document 

743 

744 data = path.read_bytes() 

745 await _announce_tesseract_ocr(data, path.name, source_name, content_type, on_progress) 

746 backend = ocr_backend() 

747 # An archive gets no reporter: xberg counts its pages across members, each total per member. 

748 on_page = ( 

749 _ocr_page_reporter(source_name, backend, on_progress) 

750 if _has_own_pages(content_type) 

751 else None 

752 ) 

753 

754 trace_log.debug("extract-start source=%r type=%s", source_name, content_type) 

755 started = time.perf_counter() 

756 with ocr_request(timeout=_effective_ocr_timeout(), cancel=cancel) as token: 

757 batcher = active_extract_batcher() 

758 if batcher is not None: 

759 doc = await batcher.submit(mode, data, path.name, token, on_page) 

760 else: 

761 config = extraction_config(mode, ocr_token=token) 

762 # xberg's extract is async; awaiting it keeps the OCR page loop off this thread. 

763 doc = await aextract_document( 

764 data, filename=path.name, config=config, on_progress=on_page 

765 ) 

766 if cancel is not None and cancel.is_set(): 

767 raise asyncio.CancelledError 

768 elapsed = time.perf_counter() - started 

769 ocr = OcrReport(backend=backend, pages=_ocr_page_count(doc)) 

770 

771 # One trace line per extraction (filename, timing, counts, OCR pages), plus a 

772 # vision line for scanned files. Emitted for empty results too (a slow file 

773 # that yields nothing is worth surfacing). 

774 trace_extraction( 

775 ExtractionTrace( 

776 source=source_name, 

777 content_type=content_type, 

778 elapsed_s=elapsed, 

779 page_count=len(doc.pages or []) or len(doc.chunks or []), 

780 chunk_count=len(doc.chunks or []), 

781 ocr=ocr, 

782 ) 

783 ) 

784 

785 return doc, ocr 

786 

787 

788def _ocr_page_count(doc: ExtractedDocument) -> int: 

789 """Pages xberg OCR'd, which are the pages it attaches an OCR confidence to.""" 

790 return sum(1 for page in doc.pages or [] if page.ocr_confidence is not None) 

791 

792 

793async def _records_from_document( 

794 doc: ExtractedDocument, 

795 source_name: str, 

796 content_type: str, 

797 *, 

798 on_progress: DetailedProgressCallback, 

799 page_texts_out: list[PageTextRecord] | None, 

800 ocr_backend: OcrBackendUsed, 

801) -> tuple[list[ChunkRecord], SourceMeta]: 

802 """Chunk-cap, page-capture, and embed one extracted document into its records.""" 

803 # Derived before the empty-result return so a scan's title/authors survive zero chunks. 

804 meta = source_meta_from_extraction(doc.metadata, source_name) 

805 

806 tables = _document_tables(doc) 

807 if not doc.chunks and not tables: 

808 if content_type in (PDF_CONTENT_TYPE, IMAGE_CONTENT_TYPE): 

809 _warn_empty_ocr(source_name, "scanned documents", ocr_backend) 

810 return [], meta 

811 

812 if ocr_backend is OcrBackendUsed.NONE: 

813 _warn_skipped_scanned_pages(doc, source_name) 

814 enforce_chunk_limit(len(doc.chunks or []) + len(tables)) 

815 _capture_result_page_texts(doc, source_name, content_type, page_texts_out) 

816 

817 # One EXTRACT event per file so progress subscribers show "extracted N pages" 

818 # before embedding; result.pages, or the chunk count for non-paginated docs. 

819 page_count = len(doc.pages or []) or len(doc.chunks or []) 

820 ran = ocr_backend if _ocr_page_count(doc) else OcrBackendUsed.NONE 

821 on_progress( 

822 EventType.EXTRACT, 

823 ExtractEvent(file=source_name, page=page_count, total_pages=page_count, ocr_backend=ran), 

824 ) 

825 

826 # Content chunks and table serializations share one embed batch; the vector 

827 # list is split back apart below by position. 

828 texts = [chunk.content for chunk in doc.chunks or []] 

829 table_texts = [table.markdown for table in tables] 

830 embed_texts = await to_ingest_thread( 

831 _enrich_texts, texts + table_texts, texts[0] if texts else "", source_name 

832 ) 

833 vectors = await to_ingest_thread( 

834 get_services().embedder.embed_batch, 

835 _embed_inputs(embed_texts, meta.title), 

836 source=source_name, 

837 on_progress=on_progress, 

838 ) 

839 records = [ 

840 ChunkRecord( 

841 source=source_name, 

842 content_type=content_type, 

843 chunk_type=ChunkType.RAW, 

844 page_start=chunk.metadata.first_page or 0, 

845 page_end=chunk.metadata.last_page or 0, 

846 line_start=0, 

847 line_end=0, 

848 chunk=text, 

849 chunk_index=chunk.metadata.chunk_index, 

850 vector=vec, 

851 ) 

852 for chunk, text, vec in zip(doc.chunks or [], texts, vectors[: len(texts)], strict=True) 

853 ]