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

1"""Bridge to xberg's async-only ``extract`` for lilbee's call sites. 

2 

3xberg exposes one ``extract(input, config)`` coroutine; lilbee extracts a single 

4in-memory document at a time, from both async and sync callers. 

5""" 

6 

7from __future__ import annotations 

8 

9import asyncio 

10import contextvars 

11from concurrent.futures import ThreadPoolExecutor 

12from dataclasses import dataclass, replace 

13from typing import TYPE_CHECKING 

14 

15from lilbee.core.config import active_config 

16from lilbee.runtime.cpu import cpu_quota 

17 

18if TYPE_CHECKING: 

19 from collections.abc import Coroutine 

20 

21 from xberg import ( 

22 ConcurrencyConfig, 

23 ExtractedDocument, 

24 ExtractInput, 

25 ExtractionConfig, 

26 ExtractionResult, 

27 OcrConfig, 

28 ) 

29 

30 

31@dataclass(frozen=True) 

32class BatchItem: 

33 """One input for :func:`aextract_batch`, with its per-file OCR override.""" 

34 

35 data: bytes 

36 mime: str | None 

37 filename: str | None 

38 ocr: OcrConfig | None 

39 

40 

41def _input(data: bytes, mime_type: str | None, filename: str | None) -> ExtractInput: 

42 from xberg import ExtractInput, ExtractInputKind 

43 

44 return ExtractInput( 

45 kind=ExtractInputKind.BYTES, bytes=data, mime_type=mime_type, filename=filename 

46 ) 

47 

48 

49def _concurrency_config() -> ConcurrencyConfig: 

50 """xberg's thread budget from the active config.""" 

51 from xberg import ConcurrencyConfig 

52 

53 return ConcurrencyConfig(max_threads=active_config().extraction_threads or cpu_quota()) 

54 

55 

56def _with_concurrency(config: ExtractionConfig) -> ExtractionConfig: 

57 """*config* carrying the active concurrency; xberg latches the first call's pools.""" 

58 return replace(config, concurrency=_concurrency_config()) 

59 

60 

61def _first(result: ExtractionResult) -> ExtractedDocument: 

62 """Return the single extracted document, or raise on an extraction error. 

63 

64 The error item carries the reason in ``message``; it has no ``__str__``, so 

65 formatting the item itself would hand the caller an object repr instead of 

66 the timeout or unsupported-format it is reporting. 

67 """ 

68 if result.results: 

69 return result.results[0] 

70 if result.errors: 

71 raise RuntimeError(result.errors[0].message) 

72 raise RuntimeError("xberg extraction returned no document") 

73 

74 

75async def aextract_document( 

76 data: bytes, 

77 mime_type: str | None = None, 

78 *, 

79 filename: str | None = None, 

80 config: ExtractionConfig, 

81) -> ExtractedDocument: 

82 """Extract one in-memory document. For callers already on the event loop.""" 

83 from xberg import extract 

84 

85 return _first(await extract(_input(data, mime_type, filename), _with_concurrency(config))) 

86 

87 

88async def aextract_batch( 

89 items: list[BatchItem], config: ExtractionConfig 

90) -> list[ExtractedDocument | Exception]: 

91 """Extract many inputs in one call, returning one document-or-error per input. 

92 

93 Each item's OCR config overrides the batch default for that file. xberg compacts 

94 ``results`` to successes in input order and reports failures in ``errors`` by 

95 input index; this remaps them back to one slot per input. 

96 """ 

97 from xberg import ExtractInput, ExtractInputKind, FileExtractionConfig, extract_batch 

98 

99 inputs = [ 

100 ExtractInput( 

101 kind=ExtractInputKind.BYTES, 

102 bytes=item.data, 

103 mime_type=item.mime, 

104 filename=item.filename, 

105 config=FileExtractionConfig(ocr=item.ocr) if item.ocr is not None else None, 

106 ) 

107 for item in items 

108 ] 

109 result = await extract_batch(inputs, _with_concurrency(config)) 

110 failed: dict[int, Exception] = {e.index: RuntimeError(e.message) for e in result.errors} 

111 success_indices = [i for i in range(len(items)) if i not in failed] 

112 by_index: dict[int, ExtractedDocument | Exception] = dict( 

113 zip(success_indices, result.results, strict=True) 

114 ) 

115 by_index.update(failed) 

116 return [by_index[i] for i in range(len(items))] 

117 

118 

119def extract_document( 

120 data: bytes, 

121 mime_type: str | None = None, 

122 *, 

123 filename: str | None = None, 

124 config: ExtractionConfig, 

125) -> ExtractedDocument: 

126 """Extract one in-memory document from synchronous code. 

127 

128 Uses ``asyncio.run``; if a loop is already running on this thread, drives the 

129 coroutine on a fresh worker thread so it never re-enters that loop. 

130 """ 

131 return _run(aextract_document(data, mime_type, filename=filename, config=config)) 

132 

133 

134def _run(coro: Coroutine[None, None, ExtractedDocument]) -> ExtractedDocument: 

135 try: 

136 asyncio.get_running_loop() 

137 except RuntimeError: 

138 return asyncio.run(coro) 

139 with ThreadPoolExecutor(max_workers=1) as pool: 

140 return pool.submit(contextvars.copy_context().run, asyncio.run, coro).result()