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 

21log = logging.getLogger(__name__) 

22 

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

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

25_FALLBACK_CHARS_PER_TOKEN = 3 

26 

27 

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

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

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

31 

32 

33class LilbeeTokenizerBackend: 

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

35 

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

37 xberg requires a non-zero count for non-empty text, so any failure degrades to 

38 the character estimate rather than raising: the embedder may be unloaded (the 

39 registration probe, or chunking outside an ingest), and a raise or zero would 

40 abort extraction or make every span look within budget. 

41 """ 

42 

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

44 self._count_fn = count_fn 

45 

46 def name(self) -> str: 

47 return TokenizerBackendName.LILBEE 

48 

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

50 

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

52 

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

54 if not text: 

55 return 0 

56 try: 

57 count = self._count_fn(text) 

58 except Exception: # embedder unreachable/erroring: degrade, never crash chunking 

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

60 return _estimate_tokens(text) 

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

62 

63 

64register_binding( 

65 XbergBinding( 

66 kind=BackendKind.TOKENIZER, 

67 name=TokenizerBackendName.LILBEE, 

68 enabled=lambda cfg: cfg.token_sizing, 

69 make=lambda provider, cfg: LilbeeTokenizerBackend(count_fn=provider.count_tokens), 

70 ) 

71)