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

1"""lilbee's vision model exposed as a xberg custom OCR backend.""" 

2 

3from __future__ import annotations 

4 

5import json 

6import threading 

7import uuid 

8from contextlib import contextmanager 

9from dataclasses import dataclass 

10from importlib.metadata import PackageNotFoundError, version 

11from pathlib import Path 

12from typing import TYPE_CHECKING, Protocol 

13 

14from lilbee.data.types import MARKDOWN_MIME, OcrBackendName 

15from lilbee.vision import resolve_ocr_prompt 

16 

17from .registry import BackendKind, XbergBinding, register_binding 

18 

19if TYPE_CHECKING: 

20 from collections.abc import Callable, Generator 

21 

22 # OcrBackend types the config as the public xberg.OcrConfig; its 

23 # backend_options arrives as a JSON string at runtime (see _OcrConfigView). 

24 from xberg import ExtractedDocument, OcrBackendType, OcrConfig 

25 

26# xberg doesn't propagate contextvars into process_image, so per-request state 

27# travels as this token in OcrConfig.backend_options and resolves via the registry. 

28_REQUEST_TOKEN_KEY = "req" # noqa: S105 # JSON key name, not a secret 

29 

30 

31@dataclass(frozen=True) 

32class OcrRequestContext: 

33 """Per-extraction state the backend needs but xberg won't carry for it.""" 

34 

35 on_page: Callable[[], None] | None = None 

36 timeout: float = 0.0 

37 

38 

39class _OcrRequestRegistry: 

40 """Token-keyed request contexts; lock-guarded (process_image runs on xberg threads).""" 

41 

42 def __init__(self) -> None: 

43 self._lock = threading.Lock() 

44 self._by_token: dict[str, OcrRequestContext] = {} 

45 

46 def register(self, ctx: OcrRequestContext) -> str: 

47 token = uuid.uuid4().hex 

48 with self._lock: 

49 self._by_token[token] = ctx 

50 return token 

51 

52 def get(self, token: str | None) -> OcrRequestContext | None: 

53 if token is None: 

54 return None 

55 with self._lock: 

56 return self._by_token.get(token) 

57 

58 def unregister(self, token: str) -> None: 

59 with self._lock: 

60 self._by_token.pop(token, None) 

61 

62 

63ocr_requests = _OcrRequestRegistry() 

64 

65 

66@contextmanager 

67def ocr_request( 

68 *, on_page: Callable[[], None] | None = None, timeout: float = 0.0 

69) -> Generator[str, None, None]: 

70 """Register a per-extraction context and yield its token for OcrConfig.backend_options.""" 

71 token = ocr_requests.register(OcrRequestContext(on_page=on_page, timeout=timeout)) 

72 try: 

73 yield token 

74 finally: 

75 ocr_requests.unregister(token) 

76 

77 

78def backend_options_for(token: str) -> dict[str, str]: 

79 """Carry a request token in OcrConfig.backend_options for process_image to read.""" 

80 return {_REQUEST_TOKEN_KEY: token} 

81 

82 

83class _OcrConfigView: 

84 """Typed reader over the xberg OcrConfig passed to process_image. The native 

85 round-trip hands ``backend_options`` back as a JSON string, so ``request_token`` 

86 accepts both the dict and string shapes.""" 

87 

88 def __init__(self, config: OcrConfig) -> None: 

89 self._config = config 

90 

91 @property 

92 def vlm_prompt(self) -> str | None: 

93 return self._config.vlm_prompt 

94 

95 @property 

96 def request_token(self) -> str | None: 

97 options = self._config.backend_options 

98 if isinstance(options, str): 

99 try: 

100 options = json.loads(options) 

101 except json.JSONDecodeError: 

102 return None 

103 token = options.get(_REQUEST_TOKEN_KEY) if isinstance(options, dict) else None 

104 return token if isinstance(token, str) else None 

105 

106 

107class _OcrFn(Protocol): 

108 # Positional-only so the provider's vision_ocr (named png_bytes) matches structurally. 

109 def __call__( 

110 self, image_bytes: bytes, model: str, prompt: str, /, *, timeout: float 

111 ) -> str: ... 

112 

113 

114def _lilbee_version() -> str: 

115 try: 

116 return version("lilbee") 

117 except PackageNotFoundError: 

118 return "0" 

119 

120 

121class VisionOcrBackend: 

122 """Routes xberg OCR calls to lilbee's vision model through the injected 

123 ``ocr_fn`` (single-image OCR) and ``model_ref_fn`` (the active vision model).""" 

124 

125 def __init__(self, *, ocr_fn: _OcrFn, model_ref_fn: Callable[[], str]) -> None: 

126 self._ocr_fn = ocr_fn 

127 self._model_ref_fn = model_ref_fn 

128 

129 def name(self) -> str: 

130 return OcrBackendName.LILBEE_VISION 

131 

132 def version(self) -> str: 

133 return _lilbee_version() 

134 

135 def supported_languages(self) -> list[str]: 

136 return [] 

137 

138 def supports_language(self, lang: str) -> bool: 

139 return True 

140 

141 def initialize(self) -> None: ... 

142 

143 def shutdown(self) -> None: ... 

144 

145 def backend_type(self) -> OcrBackendType: 

146 from xberg import OcrBackendType 

147 

148 return OcrBackendType.CUSTOM 

149 

150 def supports_table_detection(self) -> bool: 

151 return False 

152 

153 def supports_document_processing(self) -> bool: 

154 return False 

155 

156 def emits_structured_markdown(self) -> bool: 

157 # False keeps xberg's layout reconstruction (the validated path). The model 

158 # does emit markdown, so True is a valid optimization but changes OCR output. 

159 return False 

160 

161 def process_image(self, image_bytes: bytes, config: OcrConfig) -> ExtractedDocument: 

162 # xberg passes a native OcrConfig and expects a native ExtractedDocument back. 

163 from xberg import ExtractedDocument 

164 

165 view = _OcrConfigView(config) 

166 model = self._model_ref_fn() 

167 prompt = view.vlm_prompt or resolve_ocr_prompt(model) 

168 ctx = ocr_requests.get(view.request_token) 

169 text = self._ocr_fn(image_bytes, model, prompt, timeout=ctx.timeout if ctx else 0.0) 

170 if ctx is not None and ctx.on_page is not None: 

171 ctx.on_page() 

172 return ExtractedDocument(content=text, mime_type=MARKDOWN_MIME) 

173 

174 def process_image_file(self, path: str, config: OcrConfig) -> ExtractedDocument: 

175 # OCR an image file by reading its bytes and delegating to process_image 

176 # (mirrors xberg's default OcrBackend::process_image_file). 

177 return self.process_image(Path(path).read_bytes(), config) 

178 

179 def process_document(self, _path: str, _config: OcrConfig) -> ExtractedDocument: 

180 # Image-only backend: xberg only calls this when supports_document_processing() 

181 # is True, which it never is here, so document-level OCR is unreachable. 

182 raise NotImplementedError( 

183 "Document-level OCR is not supported by the lilbee vision backend" 

184 ) 

185 

186 

187register_binding( 

188 XbergBinding( 

189 kind=BackendKind.OCR, 

190 name=OcrBackendName.LILBEE_VISION, 

191 enabled=lambda cfg: bool(cfg.vision_model), 

192 make=lambda provider, cfg: VisionOcrBackend( 

193 ocr_fn=provider.vision_ocr, model_ref_fn=lambda: cfg.vision_model 

194 ), 

195 ) 

196)