Coverage for src/qdrant_loader/core/chunking/strategy/docling_strategy.py: 97%

1"""Structure-aware chunking strategy backed by docling's ``HybridChunker``. 

2 

3This is the chunking half of full Option B: instead of re-parsing a markdown string, 

4it consumes the structured ``DoclingDocument`` the docling conversion engine produced 

5(carried on :attr:`Document.converted_document`) and chunks it natively — so heading 

6paths, page spans, bounding boxes and element labels survive into the payload via the 

7engine-neutral ``metadata.structure`` block. 

8 

9It composes the seam in :mod:`qdrant_loader.core.chunking.docling`: 

10``build_chunker`` (the docling ``HybridChunker`` facade) plus the docling-free 

11:class:`ChunkDocumentMapper` (chunks -> chunk ``Document``s). The chunk-size budget 

12and contextual-embedding toggle are read from ``chunking.strategies.docling`` (falling 

13back to the *embedding* settings), and the token counter is aligned to the embedding 

14model's tokenizer — see :meth:`DoclingChunkingStrategy._build_config`. 

15 

16NLP enrichment parity with the markdown path is provided by the shared 

17:class:`~qdrant_loader.core.text_processing.chunk_enricher.ChunkEnricher`, which writes 

18the entities/topics/key_phrases (+ enhanced) keys onto each chunk ``Document``. One 

19intentional divergence: this path front-loads a per-*document* LDA topic model 

20(``ChunkEnricher.fit_topics``) before enriching, whereas markdown still trains a 

21per-*chunk* model — see :meth:`chunk_document`. Parity is the floor; docling is allowed 

22to exceed it on topic quality. 

23""" 

24 

25from __future__ import annotations 

26 

27from typing import TYPE_CHECKING 

28 

29from qdrant_loader.core.chunking.docling import ( 

30 ChunkDocumentMapper, 

31 ChunkerKind, 

32 ChunkingConfig, 

33 ChunkingError, 

34 TableSerialization, 

35 build_chunker, 

36) 

37from qdrant_loader.core.chunking.strategy.base_strategy import BaseChunkingStrategy 

38from qdrant_loader.core.text_processing.chunk_enricher import ChunkEnricher 

39from qdrant_loader.utils.logging import LoggingConfig 

40 

41logger = LoggingConfig.get_logger(__name__) 

42 

43if TYPE_CHECKING: 

44 from qdrant_loader.config import Settings 

45 from qdrant_loader.config.chunking import DoclingStrategyConfig 

46 from qdrant_loader.config.embedding import EmbeddingConfig 

47 from qdrant_loader.core.chunking.docling import DocumentChunker 

48 from qdrant_loader.core.document import Document 

49 

50 

51class DoclingChunkingStrategy(BaseChunkingStrategy): 

52 """Chunks a docling-converted document via its structured ``DoclingDocument``.""" 

53 

54 # This strategy enriches via ChunkEnricher (SemanticAnalyzer), not the base 

55 # TextProcessor — so don't let BaseChunkingStrategy load a second spaCy model. 

56 _uses_base_text_processor = False 

57 

58 def __init__(self, settings: Settings) -> None: 

59 super().__init__(settings) 

60 self._config: ChunkingConfig = self._build_config( 

61 settings.global_config.chunking.strategies.docling, 

62 settings.global_config.embedding, 

63 ) 

64 self._chunker: DocumentChunker = build_chunker( 

65 ChunkerKind.DOCLING, self._config 

66 ) 

67 self._mapper = ChunkDocumentMapper(self._config) 

68 self._enricher = ChunkEnricher(settings) 

69 

70 @staticmethod 

71 def _build_config( 

72 docling: DoclingStrategyConfig, embedding: EmbeddingConfig 

73 ) -> ChunkingConfig: 

74 """Bridge the YAML knobs into the frozen engine config. 

75 

76 The chunk-size budget is ``chunking.strategies.docling.max_tokens`` when set, 

77 otherwise the embedding model's ``max_tokens_per_chunk``. The tokenizer 

78 identity (how tokens are counted) always comes from ``embedding.tokenizer`` — 

79 the override changes the budget, never the counter. 

80 """ 

81 max_tokens = ( 

82 docling.max_tokens 

83 if docling.max_tokens is not None 

84 else embedding.max_tokens_per_chunk 

85 ) 

86 return ChunkingConfig.from_embedding( 

87 tokenizer=embedding.tokenizer, 

88 max_tokens=max_tokens, 

89 include_context_in_embed=docling.include_context_in_embed, 

90 table_serialization=TableSerialization(docling.table_serialization), 

91 ) 

92 

93 def chunk_document(self, document: Document) -> list[Document]: 

94 """Chunk the structured artifact, then map chunks into chunk ``Document``s. 

95 

96 Requires the structured ``DoclingDocument`` produced by the docling engine. 

97 Its absence is a wiring bug, not a soft-degrade case, so it fails loud rather 

98 than silently emitting one giant chunk. 

99 """ 

100 converted = document.converted_document 

101 if converted is None: 

102 raise ChunkingError( 

103 "docling chunking requires a converted DoclingDocument, but " 

104 f"document {document.id!r} carries none — was it converted with the " 

105 "docling engine? (conversion_method=" 

106 f"{document.metadata.get('conversion_method')!r})" 

107 ) 

108 

109 chunks = self._chunker.chunk(converted) 

110 

111 # Configuration-driven safety limit, in parity with the other strategies: a 

112 # huge document must not emit unbounded chunks and overwhelm embedding/upsert. 

113 max_chunks = self.settings.global_config.chunking.max_chunks_per_document 

114 if len(chunks) > max_chunks: 

115 logger.warning( 

116 f"Docling chunking produced {len(chunks)} chunks, limiting to " 

117 f"{max_chunks} per max_chunks_per_document. Consider raising " 

118 f"max_chunks_per_document or the token budget. Document: {document.title}" 

119 ) 

120 chunks = chunks[:max_chunks] 

121 

122 documents = self._mapper.to_documents(chunks, document) 

123 

124 # Front-load the topic model on this document's chunks, so each chunk's 

125 # topics are inferred against one model instead of a degenerate per-chunk 

126 # single-document LDA. Must run before the enrich loop below. 

127 # 

128 # NOTE: this intentionally diverges from the markdown path, which still 

129 # trains a per-chunk LDA. 

130 self._enricher.fit_topics([chunk_doc.content for chunk_doc in documents]) 

131 

132 # NLP enrichment parity with the markdown path: the shared ChunkEnricher 

133 # writes entities/topics/key_phrases (+ enhanced when enabled), or the 

134 # empty-shape keys when disabled. Runs on the stored chunk content. 

135 for chunk_doc in documents: 

136 chunk_doc.metadata.update( 

137 self._enricher.enrich(chunk_doc.content, doc_id=chunk_doc.id) 

138 ) 

139 

140 return documents 

141 

142 def shutdown(self) -> None: 

143 """Release the enricher's analyzer resources (spaCy/LDA).""" 

144 if getattr(self, "_enricher", None) is not None: 

145 self._enricher.shutdown() 

146 

147 def __del__(self) -> None: 

148 """Release per-document NLP resources on GC. 

149 

150 The service builds a fresh strategy per document (see 

151 ``ChunkingService._get_strategy``), so without this the spaCy/LDA model 

152 each instance loads would linger until GC ran its finalizer-free path. 

153 Mirrors the markdown strategy's ``__del__`` so both paths free the 

154 enricher the same way. 

155 """ 

156 self.shutdown()