Coverage for src/qdrant_loader/core/chunking/docling/chunker.py: 96%

1"""The chunker seam: the abstraction the rest of the codebase depends on. 

2 

3A :class:`DocumentChunker` turns a converted document into engine-neutral 

4:class:`~.outcome.Chunk` objects. The codebase depends on this Protocol, not on 

5docling's ``HybridChunker`` — the same swappability discipline the conversion layer 

6uses in its ``engine.py``. :func:`build_chunker` is the composition root that 

7selects and constructs the implementation, importing docling lazily. 

8 

9This consumes the :class:`~..conversion.ConvertedDocument` the conversion layer 

10produces — the structured ``DoclingDocument``, never a re-parsed markdown string — 

11which is what makes structure-aware chunking possible. 

12""" 

13 

14from __future__ import annotations 

15 

16import threading 

17from enum import StrEnum 

18from typing import TYPE_CHECKING, Protocol, runtime_checkable 

19 

20if TYPE_CHECKING: 

21 from ...conversion import ConvertedDocument 

22 from .config import ChunkingConfig 

23 from .outcome import Chunk 

24 

25 

26@runtime_checkable 

27class DocumentChunker(Protocol): 

28 """ConvertedDocument -> list[Chunk]. The codebase depends on this, not docling.""" 

29 

30 def chunk(self, document: ConvertedDocument) -> list[Chunk]: ... 

31 

32 

33class ChunkerKind(StrEnum): 

34 DOCLING = "docling" 

35 

36 

37# One chunker (and thus one HybridChunker + tokenizer) per distinct config for the 

38# whole process, mirroring entity_extraction's _EXTRACTOR_REGISTRY: ChunkingService 

39# builds a fresh DoclingChunkingStrategy per document, and without this the docling 

40# HybridChunker and its tokenizer (HuggingFaceTokenizer.from_pretrained / tiktoken) 

41# would be rebuilt for every document. ChunkingConfig is a frozen, hashable 

42# dataclass, so equal configs key the same cached instance. 

43_CHUNKER_REGISTRY: dict[tuple[ChunkerKind, ChunkingConfig], DocumentChunker] = {} 

44_REGISTRY_LOCK = threading.Lock() 

45 

46 

47def build_chunker(kind: ChunkerKind, config: ChunkingConfig) -> DocumentChunker: 

48 """Composition root: select and construct the chunker. 

49 

50 Memoized per ``(kind, config)`` so the same config shares one chunker instance 

51 process-wide — the per-document strategy construction in ChunkingService thus 

52 reuses one already-built (and pre-warmed) HybridChunker + tokenizer instead of 

53 rebuilding it every document. docling is imported lazily so importing this 

54 package never drags in the chunker stack, and so the import also breaks the 

55 chunker <-> docling_chunker cycle. 

56 """ 

57 key = (kind, config) 

58 with _REGISTRY_LOCK: 

59 chunker = _CHUNKER_REGISTRY.get(key) 

60 if chunker is None: 

61 chunker = _build_chunker(kind, config) 

62 _CHUNKER_REGISTRY[key] = chunker 

63 return chunker 

64 

65 

66def _build_chunker(kind: ChunkerKind, config: ChunkingConfig) -> DocumentChunker: 

67 match kind: 

68 case ChunkerKind.DOCLING: 

69 from .docling_chunker import DoclingChunker 

70 

71 return DoclingChunker(config) 

72 raise ValueError(f"unknown chunker kind: {kind!r}")