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

1"""Text chunking with optional heading-aware and topic-aware splitting.""" 

2 

3from __future__ import annotations 

4 

5from dataclasses import replace 

6from typing import TYPE_CHECKING 

7 

8from lilbee.core.config import active_config 

9from lilbee.data.types import EmbeddingBackendName, TokenizerBackendName 

10 

11if TYPE_CHECKING: 

12 from xberg import ChunkingConfig, ChunkSizing, EmbeddingConfig, TableChunkingMode 

13 

14# Char->token ratio for English. 

15CHARS_PER_TOKEN = 4 

16 

17_SEMANTIC_CHUNKER = "semantic" 

18_MARKDOWN_CHUNKER = "markdown" 

19 

20 

21def _char_budget() -> tuple[int, int]: 

22 """Return (max_chars, max_overlap) in characters from the token-based cfg.""" 

23 config = active_config() 

24 max_chars = config.chunk_size * CHARS_PER_TOKEN 

25 max_overlap = min(config.chunk_overlap * CHARS_PER_TOKEN, max_chars // 2) 

26 return max_chars, max_overlap 

27 

28 

29def _size_params() -> tuple[int, int, ChunkSizing | None]: 

30 """Return (max, overlap, sizing) for the plain and heading chunkers. 

31 

32 With ``cfg.token_sizing`` on, the budget is a raw token count and ``sizing`` 

33 routes to lilbee's registered tokenizer backend, so ``chunk_size`` is a real 

34 token ceiling. Otherwise the character heuristic with no sizing (xberg's default 

35 character sizer). The semantic chunker does not use this -- it sizes by 

36 characters and ignores ChunkSizing.""" 

37 config = active_config() 

38 if config.token_sizing: 

39 from xberg import ChunkSizing 

40 

41 overlap = min(config.chunk_overlap, config.chunk_size // 2) 

42 return ( 

43 config.chunk_size, 

44 overlap, 

45 ChunkSizing(type="tokenizer", model=TokenizerBackendName.LILBEE), 

46 ) 

47 max_chars, max_overlap = _char_budget() 

48 return max_chars, max_overlap, None 

49 

50 

51def _semantic_embedding_config() -> EmbeddingConfig: 

52 """EmbeddingConfig for semantic chunking. Boundary-detection embeddings route to 

53 lilbee's embedder, registered as xberg's plugin backend in 

54 ``lilbee.data.extract.backends.registry`` (embedding binding), so the model that 

55 vectorizes chunks for retrieval is the one that decides where they split.""" 

56 from xberg import EmbeddingConfig, EmbeddingModelType 

57 

58 model = EmbeddingModelType.plugin(EmbeddingBackendName.LILBEE) 

59 return EmbeddingConfig(model=model) 

60 

61 

62def _table_chunking() -> TableChunkingMode | None: 

63 """Header-repeating table splits when table extraction is on, else None for 

64 xberg's default. 

65 

66 REPEAT_HEADER carries the header row into every piece of a long table, so 

67 no chunk holds headerless rows. 

68 """ 

69 config = active_config() 

70 if not config.table_extraction: 

71 return None 

72 from xberg import TableChunkingMode 

73 

74 return TableChunkingMode.REPEAT_HEADER 

75 

76 

77def build_chunking_config(*, use_semantic: bool = True) -> ChunkingConfig: 

78 """Build an xberg ChunkingConfig from the current cfg.""" 

79 from xberg import ChunkingConfig 

80 

81 config = active_config() 

82 if use_semantic and config.semantic_chunking: 

83 # The semantic chunker sizes by characters and ignores ChunkSizing, so it 

84 # stays on the character budget regardless of cfg.token_sizing. 

85 max_chars, max_overlap = _char_budget() 

86 chunking = ChunkingConfig( 

87 chunker_type=_SEMANTIC_CHUNKER, 

88 embedding=_semantic_embedding_config(), 

89 topic_threshold=config.topic_threshold, 

90 max_characters=max_chars, 

91 overlap=max_overlap, 

92 ) 

93 else: 

94 max_size, max_overlap, sizing = _size_params() 

95 chunking = ChunkingConfig( 

96 max_characters=max_size, 

97 overlap=max_overlap, 

98 sizing=sizing, 

99 ) 

100 # table_chunking has no "unset" value on xberg's frozen ChunkingConfig, so the 

101 # field is left at its default rather than overwritten with None. 

102 mode = _table_chunking() 

103 return chunking if mode is None else replace(chunking, table_chunking=mode) 

104 

105 

106def chunk_text( 

107 text: str, 

108 *, 

109 mime_type: str = "text/plain", 

110 heading_context: bool = False, 

111 use_semantic: bool = True, 

112) -> list[str]: 

113 """Split text into chunks; heading_context wins over use_semantic wins over char-budget.""" 

114 if not text or not text.strip(): 

115 return [] 

116 

117 from xberg import ChunkingConfig, ExtractionConfig 

118 

119 from .xberg import extract_document 

120 

121 if heading_context: 

122 max_size, max_overlap, sizing = _size_params() 

123 chunking = ChunkingConfig( 

124 max_characters=max_size, 

125 overlap=max_overlap, 

126 sizing=sizing, 

127 chunker_type=_MARKDOWN_CHUNKER, 

128 prepend_heading_context=True, 

129 ) 

130 else: 

131 chunking = build_chunking_config(use_semantic=use_semantic) 

132 

133 config = ExtractionConfig(chunking=chunking) 

134 doc = extract_document(text.encode("utf-8"), mime_type, config=config) 

135 if doc.chunks: 

136 return [c.content for c in doc.chunks] 

137 return []