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

1"""Binds lilbee's providers into xberg's process-global OCR/embedding/tokenizer backends. 

2 

3Each backend module declares an :class:`XbergBinding` and self-registers it here at 

4import. The bind is a locked unregister-then-register: a rebuild invalidates the 

5captured provider, and the lock avoids racing xberg's "already registered". 

6""" 

7 

8from __future__ import annotations 

9 

10import threading 

11from collections.abc import Callable 

12from dataclasses import dataclass 

13from enum import Enum, auto 

14from typing import TYPE_CHECKING, Any 

15 

16if TYPE_CHECKING: 

17 from lilbee.core.config.model import Config 

18 from lilbee.providers.base import LLMProvider 

19 

20 

21class BackendKind(Enum): 

22 """Which xberg registry a binding targets (selects the register trio).""" 

23 

24 OCR = auto() 

25 EMBEDDING = auto() 

26 TOKENIZER = auto() 

27 

28 

29@dataclass(frozen=True) 

30class XbergBinding: 

31 """How to bind one lilbee backend into xberg. Declared by each backend module. 

32 

33 ``enabled`` gates registration on live cfg (embedding is always on; OCR and 

34 the tokenizer are opt-in). ``make`` builds the backend from the current 

35 provider, capturing the provider callable the binding routes through. 

36 """ 

37 

38 kind: BackendKind 

39 name: str 

40 enabled: Callable[[Config], bool] 

41 make: Callable[[LLMProvider, Config], Any] 

42 

43 

44# Keyed by kind, not name: embedding and tokenizer both register as "lilbee". 

45_BINDINGS: dict[BackendKind, XbergBinding] = {} 

46_bind_lock = threading.Lock() 

47 

48 

49def register_binding(binding: XbergBinding) -> None: 

50 """Add a backend binding to the registry (called at backend-module import).""" 

51 _BINDINGS[binding.kind] = binding 

52 

53 

54def _registry_fns(kind: BackendKind) -> tuple[Any, Any, Any]: 

55 """xberg's (list, register, unregister) functions for *kind*, imported lazily 

56 so backend modules never import the heavy xberg package at module scope. 

57 """ 

58 import xberg 

59 

60 return { 

61 BackendKind.OCR: ( 

62 xberg.list_ocr_backends, 

63 xberg.register_ocr_backend, 

64 xberg.unregister_ocr_backend, 

65 ), 

66 BackendKind.EMBEDDING: ( 

67 xberg.list_embedding_backends, 

68 xberg.register_embedding_backend, 

69 xberg.unregister_embedding_backend, 

70 ), 

71 BackendKind.TOKENIZER: ( 

72 xberg.list_tokenizer_backends, 

73 xberg.register_tokenizer_backend, 

74 xberg.unregister_tokenizer_backend, 

75 ), 

76 }[kind] 

77 

78 

79def _sync(binding: XbergBinding, provider: LLMProvider, cfg: Config) -> None: 

80 list_fn, register_fn, unregister_fn = _registry_fns(binding.kind) 

81 with _bind_lock: 

82 present = binding.name in list_fn() 

83 if binding.enabled(cfg): 

84 # Re-register so the backend always binds to the current provider. 

85 if present: 

86 unregister_fn(binding.name) 

87 register_fn(binding.make(provider, cfg)) 

88 elif present: 

89 unregister_fn(binding.name) 

90 

91 

92def _load_bindings() -> None: 

93 """Import the backend modules so their ``register_binding`` calls have run.""" 

94 from . import embedding, tokenizer, vision_ocr # noqa: F401 

95 

96 

97def sync_xberg_backends(provider: LLMProvider) -> None: 

98 """(Re)bind every registered backend to *provider*.""" 

99 from lilbee.core.config import cfg 

100 

101 _load_bindings() 

102 for binding in _BINDINGS.values(): 

103 _sync(binding, provider, cfg) 

104 

105 

106def sync_xberg_backend(kind: BackendKind, provider: LLMProvider) -> None: 

107 """(Re)bind one backend by kind, after the setting that gates it changes.""" 

108 from lilbee.core.config import cfg 

109 

110 _load_bindings() 

111 _sync(_BINDINGS[kind], provider, cfg)