Coverage for src/qdrant_loader/core/chunking/docling/config.py: 97%

1"""Engine-agnostic chunking configuration. 

2 

3Frozen dataclasses describing *our* chunking knobs — never docling types. The 

4tokenizer factory in :mod:`.tokenizer` is the single place that turns this into a 

5docling tokenizer, which keeps the chunker behind the seam in :mod:`.chunker`. 

6 

7The defaults encode the chunking design: size the token budget against the embedding 

8model's own tokenizer, let HybridChunker merge undersized peers, optionally embed the 

9heading-path-prefixed text, and stamp a ``chunk_schema_version`` so a future re-index 

10is detectable. 

11""" 

12 

13from __future__ import annotations 

14 

15import dataclasses 

16from dataclasses import dataclass, field 

17from enum import StrEnum 

18from typing import Any 

19 

20 

21class TokenizerKind(StrEnum): 

22 """Which docling tokenizer wrapper backs the token budget. 

23 

24 The string values double as the type-safe parse target for a YAML config 

25 field, mirroring :class:`~..conversion.config.ConversionProfile`. 

26 """ 

27 

28 OPENAI = "openai" # tiktoken — pairs with an OpenAI-style embedding model 

29 HUGGINGFACE = "huggingface" # transformers AutoTokenizer for a HF model id 

30 

31 

32class TableSerialization(StrEnum): 

33 """How table items are rendered into chunk text. 

34 

35 The string values double as the type-safe parse target for the YAML field 

36 ``chunking.strategies.docling.table_serialization``. 

37 """ 

38 

39 # docling's chunking default: "row, column = value" triplets — more 

40 # semantically meaningful embeddings than grid layouts (per docling's docs) 

41 TRIPLETS = "triplets" 

42 # GitHub-flavored markdown rows — better LLM/display readability; very wide 

43 # tables may overflow the token budget (docling#3428) 

44 MARKDOWN = "markdown" 

45 

46 

47@dataclass(frozen=True, slots=True) 

48class TokenizerConfig: 

49 """How to count tokens, aligned to the *embedding* model. 

50 

51 ``HybridChunker`` requires a real tokenizer to size chunks; counting against a 

52 tokenizer unrelated to the embedder is the character-budget mistake this layer 

53 avoids. ``model`` is a tiktoken encoding name for ``OPENAI`` or a HF model id for 

54 ``HUGGINGFACE``. 

55 """ 

56 

57 kind: TokenizerKind = TokenizerKind.OPENAI 

58 model: str = "cl100k_base" 

59 max_tokens: int = 512 # the embedding model's per-chunk budget (a hard ceiling) 

60 # True when this is a *counting proxy*, not the embedding model's own tokenizer 

61 # (the model exposes none — e.g. Ollama). Counts are approximate; the factory 

62 # warns once so the choice is observable, never silent. 

63 approximate: bool = False 

64 

65 

66@dataclass(frozen=True, slots=True) 

67class ChunkingConfig: 

68 """The complete, immutable chunking configuration for one chunker instance.""" 

69 

70 tokenizer: TokenizerConfig = field(default_factory=TokenizerConfig) 

71 

72 # ── HybridChunker behaviour ── 

73 merge_peers: bool = True # recombine undersized adjacent chunks sharing headings 

74 delim: str = "\n" # serialization delimiter between merged doc items 

75 # how table items become chunk text; YAML-exposed via 

76 # ``chunking.strategies.docling.table_serialization`` 

77 table_serialization: TableSerialization = TableSerialization.TRIPLETS 

78 

79 # ── contract shaping ── 

80 # Contextual embedding: when on, the chunk's contextualized text (heading_path 

81 # + body) becomes the embedded text AND the stored content. Opt-in (default off) via 

82 # ``chunking.strategies.docling.include_context_in_embed``; the seam is wired through 

83 # the chunker (contextualize) and the mapper (content = embed_text). 

84 include_context_in_embed: bool = False 

85 # Caveat: contextualize() can push a chunk past max_tokens (HybridChunker only 

86 # guarantees the budget for the bare body). When True, the chunker logs a warning on 

87 # overflow rather than failing — the embedder may truncate oversized input. 

88 enforce_token_budget: bool = True 

89 chunk_schema_version: str = "1" # lets a future re-index be detected 

90 

91 @classmethod 

92 def from_embedding( 

93 cls, *, tokenizer: str, max_tokens: int, **overrides: Any 

94 ) -> ChunkingConfig: 

95 """Derive a config from the loader's embedding settings. 

96 

97 ``tokenizer`` is the ``embedding.tokenizer`` string the loader already 

98 carries; ``max_tokens`` is its ``max_tokens_per_chunk``. Top-level overrides 

99 replace whole fields (pass a built :class:`TokenizerConfig`, not its parts). 

100 """ 

101 baseline = cls(tokenizer=cls._resolve_tokenizer(tokenizer, max_tokens)) 

102 return dataclasses.replace(baseline, **overrides) if overrides else baseline 

103 

104 @staticmethod 

105 def _resolve_tokenizer(tokenizer: str, max_tokens: int) -> TokenizerConfig: 

106 """Map an ``embedding.tokenizer`` string onto a :class:`TokenizerConfig`. 

107 

108 ``"none"`` (the Ollama-local default, which exposes no tokenizer) falls back 

109 to tiktoken ``cl100k_base`` as an *approximate counting proxy* — flagged so 

110 the factory warns once, never a silent guess. A slashed value (``org/model``) 

111 is a HF model id; anything else is a tiktoken encoding name, and is treated as 

112 the embedding model's real tokenizer (not approximate). 

113 """ 

114 if not tokenizer or tokenizer == "none": 

115 return TokenizerConfig( 

116 kind=TokenizerKind.OPENAI, 

117 model="cl100k_base", 

118 max_tokens=max_tokens, 

119 approximate=True, 

120 ) 

121 if "/" in tokenizer: 

122 return TokenizerConfig( 

123 kind=TokenizerKind.HUGGINGFACE, model=tokenizer, max_tokens=max_tokens 

124 ) 

125 return TokenizerConfig( 

126 kind=TokenizerKind.OPENAI, model=tokenizer, max_tokens=max_tokens 

127 )