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 Chunk, 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# xberg's default sizer: chunk_size counts characters. 

20_CHARACTER_SIZING = "characters" 

21# Markdown heading path rendered into a chunk: "# Setup > ## Install". 

22_HEADING_MARK = "#" 

23_BREADCRUMB_SEPARATOR = " > " 

24 

25 

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

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

28 config = active_config() 

29 max_chars = config.chunk_size * CHARS_PER_TOKEN 

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

31 return max_chars, max_overlap 

32 

33 

34def _size_params() -> tuple[int, int, ChunkSizing | str]: 

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

36 

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

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

39 token ceiling. Otherwise the character heuristic with xberg's default 

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

41 characters and ignores ChunkSizing.""" 

42 config = active_config() 

43 if config.token_sizing: 

44 from xberg import ChunkSizing 

45 

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

47 return ( 

48 config.chunk_size, 

49 overlap, 

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

51 ) 

52 max_chars, max_overlap = _char_budget() 

53 return max_chars, max_overlap, _CHARACTER_SIZING 

54 

55 

56def _semantic_embedding_config() -> EmbeddingConfig: 

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

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

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

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

61 from xberg import EmbeddingConfig, EmbeddingModelType 

62 

63 model = EmbeddingModelType.plugin(EmbeddingBackendName.LILBEE) 

64 return EmbeddingConfig(model=model) 

65 

66 

67def _table_chunking() -> TableChunkingMode | None: 

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

69 xberg's default. 

70 

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

72 no chunk holds headerless rows. 

73 """ 

74 config = active_config() 

75 if not config.table_extraction: 

76 return None 

77 from xberg import TableChunkingMode 

78 

79 return TableChunkingMode.REPEAT_HEADER 

80 

81 

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

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

84 from xberg import ChunkingConfig 

85 

86 config = active_config() 

87 if use_semantic and config.semantic_chunking: 

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

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

90 max_chars, max_overlap = _char_budget() 

91 chunking = ChunkingConfig( 

92 chunker_type=_SEMANTIC_CHUNKER, 

93 embedding=_semantic_embedding_config(), 

94 topic_threshold=config.topic_threshold, 

95 max_characters=max_chars, 

96 overlap=max_overlap, 

97 ) 

98 else: 

99 max_size, max_overlap, sizing = _size_params() 

100 chunking = ChunkingConfig( 

101 max_characters=max_size, 

102 overlap=max_overlap, 

103 sizing=sizing, 

104 ) 

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

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

107 mode = _table_chunking() 

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

109 

110 

111def chunk_text( 

112 text: str, 

113 *, 

114 mime_type: str = "text/plain", 

115 heading_context: bool = False, 

116 use_semantic: bool = True, 

117) -> list[str]: 

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

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

120 return [] 

121 

122 from xberg import ChunkingConfig, ExtractionConfig 

123 

124 from .xberg import extract_document 

125 

126 if heading_context: 

127 max_size, max_overlap, sizing = _size_params() 

128 chunking = ChunkingConfig( 

129 max_characters=max_size, 

130 overlap=max_overlap, 

131 sizing=sizing, 

132 chunker_type=_MARKDOWN_CHUNKER, 

133 ) 

134 else: 

135 chunking = build_chunking_config(use_semantic=use_semantic) 

136 

137 config = ExtractionConfig(chunking=chunking) 

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

139 if not doc.chunks: 

140 return [] 

141 if heading_context: 

142 return [_with_heading_breadcrumb(c) for c in doc.chunks] 

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

144 

145 

146def _with_heading_breadcrumb(chunk: Chunk) -> str: 

147 """Prefix a markdown chunk with its heading path, e.g. ``# Setup > ## Install``. 

148 

149 xberg's ``render_heading_breadcrumb`` is Rust-only; the Python binding exposes 

150 the headings as ``metadata.heading_context``. 

151 """ 

152 context = chunk.metadata.heading_context 

153 if context is None or not context.headings: 

154 return chunk.content 

155 breadcrumb = _BREADCRUMB_SEPARATOR.join( 

156 f"{_HEADING_MARK * h.level} {h.text}" for h in context.headings 

157 ) 

158 return f"{breadcrumb}\n\n{chunk.content}" 

159 

160 

161class ChunkLimitError(Exception): 

162 """One file produced more chunks than ``cfg.max_chunks_per_file`` allows.""" 

163 

164 def __init__(self, count: int, limit: int, member: str | None = None) -> None: 

165 prefix = f"{member}: " if member else "" 

166 super().__init__( 

167 f"{prefix}{count} chunks exceed the per-file limit of {limit}; " 

168 f"raise max_chunks_per_file (0 = no limit), then retry skipped files" 

169 ) 

170 self.count = count 

171 self.member = member 

172 self.limit = limit 

173 

174 

175def enforce_chunk_limit(count: int) -> None: 

176 """Refuse a file whose *count* chunks exceed the per-file limit; a limit of 0 accepts any.""" 

177 limit = active_config().max_chunks_per_file 

178 if limit and count > limit: 

179 raise ChunkLimitError(count, limit)