Coverage for src/qdrant_loader/core/text_processing/chunk_enricher.py: 100%

1"""Shared chunk NLP enrichment. 

2 

3The single definition of what semantic metadata a chunk gets, so the markdown and 

4docling chunking strategies produce identical fields. Owns one ``SemanticAnalyzer`` 

5and turns ``(content, doc_id)`` into the enrichment metadata dict. 

6""" 

7 

8from __future__ import annotations 

9 

10from typing import TYPE_CHECKING, Any 

11 

12from qdrant_loader.core.text_processing.semantic_analyzer import SemanticAnalyzer 

13from qdrant_loader.utils.logging import LoggingConfig 

14 

15if TYPE_CHECKING: 

16 from qdrant_loader.config import Settings 

17 

18logger = LoggingConfig.get_logger(__name__) 

19 

20 

21def _empty() -> dict[str, Any]: 

22 """The always-present enrichment keys with no data — fresh lists each call.""" 

23 return {"entities": [], "topics": [], "key_phrases": []} 

24 

25 

26class ChunkEnricher: 

27 """Owns a ``SemanticAnalyzer`` and maps ``(content, doc_id)`` to chunk metadata.""" 

28 

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

30 chunking = settings.global_config.chunking 

31 self._enhanced = bool(chunking.enable_enhanced_semantic_analysis) 

32 self._analyzer: SemanticAnalyzer | None = None 

33 if chunking.enable_semantic_analysis: 

34 sa = settings.global_config.semantic_analysis 

35 self._analyzer = SemanticAnalyzer( 

36 spacy_model=sa.spacy_model, 

37 num_topics=sa.num_topics, 

38 passes=sa.lda_passes, 

39 ) 

40 

41 @property 

42 def enabled(self) -> bool: 

43 return self._analyzer is not None 

44 

45 def fit_topics(self, contents: list[str]) -> None: 

46 """Front-load the document-level topic model on a document's chunk contents. 

47 

48 Training one model over all chunks lets each subsequent :meth:`enrich` call 

49 infer the chunk's topics against it, instead of training a degenerate 

50 single-chunk model each time. No-op when semantic analysis is disabled; a 

51 fit failure is swallowed so chunks simply fall back to per-chunk topics. 

52 """ 

53 if self._analyzer is None: 

54 return 

55 try: 

56 self._analyzer.fit_topic_model(contents) 

57 except Exception: 

58 logger.warning( 

59 "Topic model fit failed; chunks will fall back to per-chunk topics", 

60 exc_info=True, 

61 ) 

62 

63 def enrich(self, content: str, doc_id: str) -> dict[str, Any]: 

64 if self._analyzer is None: 

65 return _empty() 

66 try: 

67 result = self._analyzer.analyze_text( 

68 content, doc_id=doc_id, include_enhanced=self._enhanced 

69 ) 

70 except Exception: 

71 logger.warning( 

72 "Chunk semantic enrichment failed; emitting empty enrichment", 

73 doc_id=doc_id, 

74 exc_info=True, 

75 ) 

76 return _empty() 

77 enriched: dict[str, Any] = { 

78 "entities": result.entities, 

79 "topics": result.topics, 

80 "key_phrases": result.key_phrases, 

81 } 

82 if self._enhanced: 

83 enriched["pos_tags"] = result.pos_tags 

84 enriched["dependencies"] = result.dependencies 

85 enriched["document_similarity"] = result.document_similarity 

86 return enriched 

87 

88 def shutdown(self) -> None: 

89 if self._analyzer is not None: 

90 self._analyzer.shutdown() 

91 self._analyzer = None