Coverage for src/lilbee/data/extract/chunk.py: 100%
91 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""Text chunking with optional heading-aware and topic-aware splitting."""
3from __future__ import annotations
5from dataclasses import replace
6from typing import TYPE_CHECKING
8from lilbee.core.config import active_config
9from lilbee.data.types import EmbeddingBackendName, TokenizerBackendName
11if TYPE_CHECKING:
12 from xberg import (
13 ChunkingConfig,
14 ChunkSizing,
15 EmbeddingConfig,
16 HeadingContext,
17 TableChunkingMode,
18 )
20# Char->token ratio for English.
21CHARS_PER_TOKEN = 4
23_SEMANTIC_CHUNKER = "semantic"
24_MARKDOWN_CHUNKER = "markdown"
25# xberg's default sizer: chunk_size counts characters.
26_CHARACTER_SIZING = "characters"
27# Markdown heading path rendered into a chunk: "# Setup > ## Install".
28_HEADING_MARK = "#"
29_BREADCRUMB_SEPARATOR = " > "
32def _embed_token_cap() -> int | None:
33 """Tokens the embedder truncates one input to; None when it has no fixed cap."""
34 # circular: chunk -> app.services -> retrieval.embedder -> chunk via CHARS_PER_TOKEN
35 from lilbee.app.services import get_services
37 return get_services().provider.embed_token_cap()
40def token_sizing_in_effect(cap: int | None) -> bool:
41 """Whether chunks are sized in the embedder's tokens: opted in, or its window binds.
43 The window binds when *cap* is below the character budget ``chunk_size``
44 allows, so a character-sized chunk could exceed it at embedding time.
45 """
46 config = active_config()
47 if config.token_sizing:
48 return True
49 return cap is not None and cap < config.chunk_size * CHARS_PER_TOKEN
52def _char_budget() -> tuple[int, int]:
53 """Return (max_chars, max_overlap) in characters from the token-based cfg."""
54 config = active_config()
55 max_chars = config.chunk_size * CHARS_PER_TOKEN
56 max_overlap = min(config.chunk_overlap * CHARS_PER_TOKEN, max_chars // 2)
57 return max_chars, max_overlap
60def _tokenizer_sizing() -> ChunkSizing:
61 """xberg sizing through lilbee's tokenizer backend, bound to the provider on demand.
63 xberg fails extraction when a sizing names an unregistered tokenizer, so the
64 bind precedes every sizing that routes to it. The bind is a no-op once the
65 current provider holds the binding.
66 """
67 from xberg import ChunkSizing
69 from lilbee.app.services import get_services
71 from .backends import BackendKind, bind_backend
73 bind_backend(BackendKind.TOKENIZER, get_services().provider)
74 return ChunkSizing(type="tokenizer", model=TokenizerBackendName.LILBEE)
77def _size_params() -> tuple[int, int, ChunkSizing | str]:
78 """Return (max, overlap, sizing) for the plain and heading chunkers.
80 Under token sizing (``cfg.token_sizing``, or an embed window below the
81 character budget) the budget is a token count no larger than the embedder's
82 cap and ``sizing`` routes to lilbee's tokenizer backend, so no chunk loses
83 text at embedding time. Otherwise the character heuristic with xberg's
84 default character sizer. The semantic chunker does not use this -- it sizes
85 by characters and ignores ChunkSizing."""
86 config = active_config()
87 cap = _embed_token_cap()
88 if not token_sizing_in_effect(cap):
89 max_chars, max_overlap = _char_budget()
90 return max_chars, max_overlap, _CHARACTER_SIZING
91 max_tokens = config.chunk_size if cap is None else min(config.chunk_size, cap)
92 overlap = min(config.chunk_overlap, max_tokens // 2)
93 return max_tokens, overlap, _tokenizer_sizing()
96def _semantic_embedding_config() -> EmbeddingConfig:
97 """EmbeddingConfig for semantic chunking. Boundary-detection embeddings route to
98 lilbee's embedder, registered as xberg's plugin backend in
99 ``lilbee.data.extract.backends.registry`` (embedding binding), so the model that
100 vectorizes chunks for retrieval is the one that decides where they split."""
101 from xberg import EmbeddingConfig, EmbeddingModelType
103 model = EmbeddingModelType.plugin(EmbeddingBackendName.LILBEE)
104 return EmbeddingConfig(model=model)
107def _table_chunking() -> TableChunkingMode | None:
108 """Header-repeating table splits when table extraction is on, else None for
109 xberg's default.
111 REPEAT_HEADER carries the header row into every piece of a long table, so
112 no chunk holds headerless rows.
113 """
114 config = active_config()
115 if not config.table_extraction:
116 return None
117 from xberg import TableChunkingMode
119 return TableChunkingMode.REPEAT_HEADER
122def build_chunking_config(*, use_semantic: bool = True) -> ChunkingConfig:
123 """Build an xberg ChunkingConfig from the current cfg."""
124 from xberg import ChunkingConfig
126 config = active_config()
127 if use_semantic and config.semantic_chunking:
128 # The semantic chunker sizes by characters and ignores ChunkSizing, so it
129 # stays on the character budget regardless of cfg.token_sizing.
130 max_chars, max_overlap = _char_budget()
131 chunking = ChunkingConfig(
132 chunker_type=_SEMANTIC_CHUNKER,
133 embedding=_semantic_embedding_config(),
134 topic_threshold=config.topic_threshold,
135 max_characters=max_chars,
136 overlap=max_overlap,
137 )
138 else:
139 max_size, max_overlap, sizing = _size_params()
140 chunking = ChunkingConfig(
141 max_characters=max_size,
142 overlap=max_overlap,
143 sizing=sizing,
144 )
145 # table_chunking has no "unset" value on xberg's frozen ChunkingConfig, so the
146 # field is left at its default rather than overwritten with None.
147 mode = _table_chunking()
148 return chunking if mode is None else replace(chunking, table_chunking=mode)
151def chunk_text(
152 text: str,
153 *,
154 mime_type: str = "text/plain",
155 heading_context: bool = False,
156 use_semantic: bool = True,
157) -> list[str]:
158 """Split text into chunks; heading_context wins over use_semantic wins over char-budget."""
159 if not text or not text.strip():
160 return []
162 from xberg import ChunkingConfig, ExtractionConfig
164 from .xberg import extract_document
166 if heading_context:
167 max_size, max_overlap, sizing = _size_params()
168 chunking = ChunkingConfig(
169 max_characters=max_size,
170 overlap=max_overlap,
171 sizing=sizing,
172 chunker_type=_MARKDOWN_CHUNKER,
173 )
174 else:
175 chunking = build_chunking_config(use_semantic=use_semantic)
177 config = ExtractionConfig(chunking=chunking)
178 doc = extract_document(text.encode("utf-8"), mime_type, config=config)
179 if not doc.chunks:
180 return []
181 if heading_context:
182 return [_with_heading_breadcrumb(c.content, c.metadata.heading_context) for c in doc.chunks]
183 return [c.content for c in doc.chunks]
186def _with_heading_breadcrumb(content: str, context: HeadingContext | None) -> str:
187 """Prefix a markdown chunk with its heading path, e.g. ``# Setup > ## Install``.
189 xberg's ``render_heading_breadcrumb`` is Rust-only; the Python binding exposes
190 the headings as the chunk's ``metadata.heading_context``.
191 """
192 if context is None or not context.headings:
193 return content
194 breadcrumb = _BREADCRUMB_SEPARATOR.join(
195 f"{_HEADING_MARK * h.level} {h.text}" for h in context.headings
196 )
197 return f"{breadcrumb}\n\n{content}"
200class ChunkLimitError(Exception):
201 """One file produced more chunks than ``cfg.max_chunks_per_file`` allows."""
203 def __init__(self, count: int, limit: int, member: str | None = None) -> None:
204 prefix = f"{member}: " if member else ""
205 super().__init__(
206 f"{prefix}{count} chunks exceed the per-file limit of {limit}; "
207 f"raise max_chunks_per_file (0 = no limit), then retry skipped files"
208 )
209 self.count = count
210 self.member = member
211 self.limit = limit
214def enforce_chunk_limit(count: int) -> None:
215 """Refuse a file whose *count* chunks exceed the per-file limit; a limit of 0 accepts any."""
216 limit = active_config().max_chunks_per_file
217 if limit and count > limit:
218 raise ChunkLimitError(count, limit)