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

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

2 

3xberg's semantic chunker needs embeddings to detect topic boundaries. This routes 

4them to lilbee's own embedder, so the same model vectorizes and splits chunks; 

5without it xberg falls back to its bundled ONNX preset. 

6""" 

7 

8from __future__ import annotations 

9 

10from typing import TYPE_CHECKING 

11 

12from lilbee.data.types import EmbeddingBackendName 

13 

14from .registry import BackendKind, XbergBinding, register_binding 

15 

16if TYPE_CHECKING: 

17 from collections.abc import Callable 

18 

19 from lilbee.core.vectors import Vector 

20 

21 

22class LilbeeEmbeddingBackend: 

23 """Routes xberg's boundary-detection embeddings to lilbee's embedder. 

24 

25 ``embed_fn``/``dim_fn`` are read live, so an embedding-model swap needs no 

26 re-registration. xberg calls these sync methods on its own worker threads. 

27 """ 

28 

29 def __init__( 

30 self, 

31 *, 

32 embed_fn: Callable[[list[str]], list[Vector]], 

33 dim_fn: Callable[[], int], 

34 ) -> None: 

35 self._embed_fn = embed_fn 

36 self._dim_fn = dim_fn 

37 

38 def name(self) -> str: 

39 return EmbeddingBackendName.LILBEE 

40 

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

42 

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

44 

45 def dimensions(self) -> int: 

46 return self._dim_fn() 

47 

48 def embed(self, texts: list[str]) -> list[Vector]: 

49 return self._embed_fn(texts) 

50 

51 

52register_binding( 

53 XbergBinding( 

54 kind=BackendKind.EMBEDDING, 

55 name=EmbeddingBackendName.LILBEE, 

56 enabled=lambda cfg: True, 

57 make=lambda provider, cfg: LilbeeEmbeddingBackend( 

58 embed_fn=provider.embed, dim_fn=lambda: cfg.embedding_dim 

59 ), 

60 ) 

61)