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, on_progress)`` coroutine; lilbee extracts 

4a single in-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 from xberg.progress import ProgressCallback, ProgressEvent 

30 

31 

32@dataclass(frozen=True) 

33class BatchItem: 

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

35 

36 data: bytes 

37 mime: str | None 

38 filename: str | None 

39 ocr: OcrConfig | None 

40 on_progress: ProgressCallback | None = None 

41 

42 

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

44 from xberg import ExtractInput, ExtractInputKind 

45 

46 return ExtractInput( 

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

48 ) 

49 

50 

51def _concurrency_config() -> ConcurrencyConfig: 

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

53 from xberg import ConcurrencyConfig 

54 

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

56 

57 

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

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

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

61 

62 

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

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

65 

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

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

68 the timeout or unsupported-format it is reporting. 

69 """ 

70 if result.results: 

71 return result.results[0] 

72 if result.errors: 

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

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

75 

76 

77async def aextract_document( 

78 data: bytes, 

79 mime_type: str | None = None, 

80 *, 

81 filename: str | None = None, 

82 config: ExtractionConfig, 

83 on_progress: ProgressCallback | None = None, 

84) -> ExtractedDocument: 

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

86 

87 xberg calls *on_progress* once per OCR'd page, from a worker thread, and 

88 swallows whatever it raises. 

89 """ 

90 from xberg.progress import extract 

91 

92 return _first( 

93 await extract(_input(data, mime_type, filename), _with_concurrency(config), on_progress) 

94 ) 

95 

96 

97def _route_by_input(items: list[BatchItem]) -> ProgressCallback | None: 

98 """One callback handing each batch progress event to its own input's callback.""" 

99 callbacks: dict[int | None, ProgressCallback] = { 

100 index: item.on_progress for index, item in enumerate(items) if item.on_progress is not None 

101 } 

102 if not callbacks: 

103 return None 

104 

105 def _route(event: ProgressEvent) -> None: 

106 callback = callbacks.get(event.input_index) 

107 if callback is not None: 

108 callback(event) 

109 

110 return _route 

111 

112 

113async def aextract_batch( 

114 items: list[BatchItem], config: ExtractionConfig 

115) -> list[ExtractedDocument | Exception]: 

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

117 

118 Each item's OCR config overrides the batch default for that file, and each item's 

119 ``on_progress`` receives only that file's OCR page events. xberg compacts 

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

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

122 """ 

123 from xberg import ExtractInput, ExtractInputKind, FileExtractionConfig 

124 from xberg.progress import extract_batch 

125 

126 inputs = [ 

127 ExtractInput( 

128 kind=ExtractInputKind.BYTES, 

129 bytes=item.data, 

130 mime_type=item.mime, 

131 filename=item.filename, 

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

133 ) 

134 for item in items 

135 ] 

136 result = await extract_batch(inputs, _with_concurrency(config), _route_by_input(items)) 

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

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

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

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

141 ) 

142 by_index.update(failed) 

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

144 

145 

146def extract_document( 

147 data: bytes, 

148 mime_type: str | None = None, 

149 *, 

150 filename: str | None = None, 

151 config: ExtractionConfig, 

152) -> ExtractedDocument: 

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

154 

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

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

157 """ 

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

159 

160 

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

162 try: 

163 asyncio.get_running_loop() 

164 except RuntimeError: 

165 return asyncio.run(coro) 

166 with ThreadPoolExecutor(max_workers=1) as pool: 

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