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. A 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, field 

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 at services init on live cfg (embedding is 

34 always on; OCR is opt-in; the tokenizer only unregisters when off, and the 

35 chunker binds it on demand). ``make`` builds the backend from 

36 the current provider, capturing the provider callable the binding routes through. 

37 """ 

38 

39 kind: BackendKind 

40 name: str 

41 enabled: Callable[[Config], bool] 

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

43 

44 

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

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

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

48 """ 

49 import xberg 

50 

51 return { 

52 BackendKind.OCR: ( 

53 xberg.list_ocr_backends, 

54 xberg.register_ocr_backend, 

55 xberg.unregister_ocr_backend, 

56 ), 

57 BackendKind.EMBEDDING: ( 

58 xberg.list_embedding_backends, 

59 xberg.register_embedding_backend, 

60 xberg.unregister_embedding_backend, 

61 ), 

62 BackendKind.TOKENIZER: ( 

63 xberg.list_tokenizer_backends, 

64 xberg.register_tokenizer_backend, 

65 xberg.unregister_tokenizer_backend, 

66 ), 

67 }[kind] 

68 

69 

70@dataclass 

71class _BackendRegistry: 

72 """The binding table, the bind lock, and the provider each kind was last bound to.""" 

73 

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

75 bindings: dict[BackendKind, XbergBinding] = field(default_factory=dict) 

76 bound: dict[BackendKind, LLMProvider] = field(default_factory=dict) 

77 lock: threading.Lock = field(default_factory=threading.Lock) 

78 

79 def bind(self, kind: BackendKind, provider: LLMProvider, cfg: Config) -> None: 

80 """Register *kind* for *provider*; a no-op when xberg already holds that binding.""" 

81 binding = self.bindings[kind] 

82 list_fn, register_fn, unregister_fn = _registry_fns(kind) 

83 with self.lock: 

84 present = binding.name in list_fn() 

85 if present and self.bound.get(kind) is provider: 

86 return 

87 if present: 

88 unregister_fn(binding.name) 

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

90 self.bound[kind] = provider 

91 

92 def unbind(self, kind: BackendKind) -> None: 

93 """Unregister *kind* from xberg when present.""" 

94 binding = self.bindings[kind] 

95 list_fn, _register_fn, unregister_fn = _registry_fns(kind) 

96 with self.lock: 

97 if binding.name in list_fn(): 

98 unregister_fn(binding.name) 

99 

100 def sync(self, kind: BackendKind, provider: LLMProvider, cfg: Config) -> None: 

101 """Bind or unbind *kind* as its ``enabled`` gate says under *cfg*. 

102 

103 An enabled tokenizer is left to the chunker's on-demand bind: binding it 

104 counts through the provider, which must not run while services are built. 

105 """ 

106 if not self.bindings[kind].enabled(cfg): 

107 self.unbind(kind) 

108 elif kind is not BackendKind.TOKENIZER: 

109 self.bind(kind, provider, cfg) 

110 

111 

112_registry = _BackendRegistry() 

113 

114 

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

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

117 _registry.bindings[binding.kind] = binding 

118 

119 

120def _load_bindings() -> None: 

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

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

123 

124 

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

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

127 from lilbee.core.config import cfg 

128 

129 _load_bindings() 

130 for kind in _registry.bindings: 

131 _registry.sync(kind, provider, cfg) 

132 

133 

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

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

136 from lilbee.core.config import cfg 

137 

138 _load_bindings() 

139 _registry.sync(kind, provider, cfg) 

140 

141 

142def bind_backend(kind: BackendKind, provider: LLMProvider) -> None: 

143 """Bind one backend to *provider* on demand, whatever its ``enabled`` gate says.""" 

144 from lilbee.core.config import cfg 

145 

146 _load_bindings() 

147 _registry.bind(kind, provider, cfg)