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

1"""lilbee's embedder tokenizer exposed as a xberg plugin tokenizer backend. 

2 

3xberg's chunk sizer can budget in tokens, but only with tokenizers it loads itself; 

4lilbee's embedder is a GGUF model whose vocab isn't published that way. This lets 

5``chunk_size`` be a real token ceiling instead of a chars-per-token guess (off by 

62-4x on token-dense text). 

7""" 

8 

9from __future__ import annotations 

10 

11import logging 

12from typing import TYPE_CHECKING 

13 

14from lilbee.data.types import TokenizerBackendName 

15 

16from .registry import BackendKind, XbergBinding, register_binding 

17 

18if TYPE_CHECKING: 

19 from collections.abc import Callable 

20 

21 from lilbee.core.config.model import Config 

22 from lilbee.providers.base import LLMProvider 

23 

24log = logging.getLogger(__name__) 

25 

26# Over-counting fallback for when the exact count is unavailable: splits a touch 

27# early rather than emitting an over-length chunk. Mirrors providers.fleet.client. 

28_FALLBACK_CHARS_PER_TOKEN = 3 

29 

30# The text xberg's registration probe counts. Its bridge reports an exception as a zero count. 

31_XBERG_PROBE_TEXT = "a" 

32 

33 

34def _estimate_tokens(text: str) -> int: 

35 """Conservative token estimate from character length (ceiling division).""" 

36 return max(1, -(-len(text) // _FALLBACK_CHARS_PER_TOKEN)) 

37 

38 

39class LilbeeTokenizerBackend: 

40 """Counts chunk-sizing tokens with lilbee's embedder tokenizer. 

41 

42 ``count_fn`` is read live (an embedding-model swap needs no re-registration). 

43 Only a backend without a local tokenizer degrades to the character estimate; 

44 any other failure raises to the caller. xberg requires a non-zero count for 

45 non-empty text, so a zero count degrades to the estimate as well. 

46 After :meth:`probe`, xberg's registration probe is answered once without a 

47 provider call, so a provider that stops after the probe cannot fail the 

48 registration. Later counts of the same text go to the provider. 

49 """ 

50 

51 def __init__(self, *, count_fn: Callable[[str], int]) -> None: 

52 self._count_fn = count_fn 

53 self._probe_count: int | None = None 

54 

55 def name(self) -> str: 

56 return TokenizerBackendName.LILBEE 

57 

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

59 

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

61 

62 def probe(self) -> None: 

63 """Count xberg's probe text through the provider; a count failure raises here.""" 

64 self._probe_count = self._count(_XBERG_PROBE_TEXT) 

65 

66 def count_tokens(self, text: str) -> int: 

67 if not text: 

68 return 0 

69 if text == _XBERG_PROBE_TEXT: 

70 answer, self._probe_count = self._probe_count, None 

71 if answer is not None: 

72 return answer 

73 return self._count(text) 

74 

75 def _count(self, text: str) -> int: 

76 try: 

77 count = self._count_fn(text) 

78 except NotImplementedError: # SDK embedders expose no local tokenizer 

79 log.debug("exact token count failed; using character estimate", exc_info=True) 

80 return _estimate_tokens(text) 

81 return count if count > 0 else _estimate_tokens(text) 

82 

83 

84def _make_tokenizer_backend(provider: LLMProvider, cfg: Config) -> LilbeeTokenizerBackend: 

85 """Build the tokenizer backend, probing its count before xberg does.""" 

86 backend = LilbeeTokenizerBackend(count_fn=provider.count_tokens) 

87 backend.probe() 

88 return backend 

89 

90 

91register_binding( 

92 XbergBinding( 

93 kind=BackendKind.TOKENIZER, 

94 name=TokenizerBackendName.LILBEE, 

95 enabled=lambda cfg: cfg.token_sizing, 

96 make=_make_tokenizer_backend, 

97 ) 

98)