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.runtime.cancellation import CancelSignal, TaskCancelledError 

16from lilbee.vision import resolve_ocr_prompt 

17 

18from .registry import BackendKind, XbergBinding, register_binding 

19 

20if TYPE_CHECKING: 

21 from collections.abc import Callable, Generator 

22 

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

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

25 from xberg import ExtractedDocument, OcrBackendType, OcrConfig 

26 

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

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

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

30 

31 

32@dataclass(frozen=True) 

33class OcrRequestContext: 

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

35 

36 timeout: float = 0.0 

37 cancel: CancelSignal | None = None 

38 

39 

40class _OcrRequestRegistry: 

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

42 

43 def __init__(self) -> None: 

44 self._lock = threading.Lock() 

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

46 

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

48 token = uuid.uuid4().hex 

49 with self._lock: 

50 self._by_token[token] = ctx 

51 return token 

52 

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

54 if token is None: 

55 return None 

56 with self._lock: 

57 return self._by_token.get(token) 

58 

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

60 with self._lock: 

61 self._by_token.pop(token, None) 

62 

63 

64ocr_requests = _OcrRequestRegistry() 

65 

66 

67@contextmanager 

68def ocr_request( 

69 *, timeout: float = 0.0, cancel: CancelSignal | None = None 

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

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

72 token = ocr_requests.register(OcrRequestContext(timeout=timeout, cancel=cancel)) 

73 try: 

74 yield token 

75 finally: 

76 ocr_requests.unregister(token) 

77 

78 

79def backend_options_for(token: str) -> str: 

80 """The OcrConfig.backend_options JSON that carries a request token to process_image.""" 

81 return json.dumps({_REQUEST_TOKEN_KEY: token}) 

82 

83 

84class _OcrConfigView: 

85 """Typed reader over the xberg OcrConfig passed to process_image.""" 

86 

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

88 self._config = config 

89 

90 @property 

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

92 return self._config.vlm_prompt 

93 

94 @property 

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

96 options = self._config.backend_options 

97 if options is None: 

98 return None 

99 try: 

100 decoded = json.loads(options) 

101 except json.JSONDecodeError: 

102 return None 

103 # backend_options is free-form JSON; only lilbee's own object carries a token. 

104 token = decoded.get(_REQUEST_TOKEN_KEY) if isinstance(decoded, dict) else None 

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

106 

107 

108class _OcrFn(Protocol): 

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

110 def __call__( 

111 self, 

112 image_bytes: bytes, 

113 model: str, 

114 prompt: str, 

115 /, 

116 *, 

117 timeout: float, 

118 cancel: CancelSignal | None, 

119 ) -> str: ... 

120 

121 

122def _lilbee_version() -> str: 

123 try: 

124 return version("lilbee") 

125 except PackageNotFoundError: 

126 return "0" 

127 

128 

129class VisionOcrBackend: 

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

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

132 

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

134 self._ocr_fn = ocr_fn 

135 self._model_ref_fn = model_ref_fn 

136 

137 def name(self) -> str: 

138 return OcrBackendName.LILBEE_VISION 

139 

140 def version(self) -> str: 

141 return _lilbee_version() 

142 

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

144 return [] 

145 

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

147 return True 

148 

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

150 

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

152 

153 def backend_type(self) -> OcrBackendType: 

154 from xberg import OcrBackendType 

155 

156 return OcrBackendType.CUSTOM 

157 

158 def supports_table_detection(self) -> bool: 

159 return False 

160 

161 def supports_document_processing(self) -> bool: 

162 return False 

163 

164 def emits_structured_markdown(self) -> bool: 

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

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

167 return False 

168 

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

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

171 from xberg import ExtractedDocument 

172 

173 view = _OcrConfigView(config) 

174 model = self._model_ref_fn() 

175 prompt = view.vlm_prompt or resolve_ocr_prompt(model) 

176 ctx = ocr_requests.get(view.request_token) 

177 cancel = ctx.cancel if ctx is not None else None 

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

179 raise TaskCancelledError 

180 timeout = ctx.timeout if ctx is not None else 0.0 

181 text = self._ocr_fn(image_bytes, model, prompt, timeout=timeout, cancel=cancel) 

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

183 

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

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

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

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

188 

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

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

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

192 raise NotImplementedError( 

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

194 ) 

195 

196 

197register_binding( 

198 XbergBinding( 

199 kind=BackendKind.OCR, 

200 name=OcrBackendName.LILBEE_VISION, 

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

202 make=lambda provider, cfg: VisionOcrBackend( 

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

204 ), 

205 ) 

206)