Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-rag/src/lexigram/ai/rag/chunking/strategy_registry.py: 53%

1"""Chunking strategy registry for RAG document splitting. 

2 

3Maps :class:`~lexigram.ai.rag.chunking.types.ChunkingStrategy` enum values 

4to chunker *classes* and instantiates on demand. 

5""" 

6 

7from __future__ import annotations 

8 

9from typing import TYPE_CHECKING, Any 

10 

11from lexigram.ai.rag.chunking.types import ChunkingConfig, ChunkingStrategy 

12from lexigram.logging import ( 

13 get_logger, 

14) 

15from lexigram.primitives.registry import StrategyRegistry 

16 

17if TYPE_CHECKING: 

18 from lexigram.ai.rag.chunking.base import AbstractChunker 

19 

20logger = get_logger(__name__) 

21 

22 

23class ChunkingStrategyRegistry(StrategyRegistry): 

24 """Registry mapping chunking strategy names to chunker classes. 

25 

26 Usage:: 

27 

28 registry = ChunkingStrategyRegistry.with_defaults() 

29 chunker = registry.instantiate( 

30 ChunkingStrategy.FIXED_SIZE, chunk_size=512, overlap=100, 

31 ) 

32 """ 

33 

34 def __init__(self) -> None: 

35 super().__init__(name="chunking.strategies", allow_overwrite=True) 

36 

37 @classmethod 

38 def with_defaults(cls) -> ChunkingStrategyRegistry: 

39 """Create a registry pre-populated with the built-in chunking strategies. 

40 

41 Returns: 

42 A new registry with all built-in chunking strategies registered. 

43 """ 

44 instance = cls() 

45 instance._register_defaults() 

46 return instance 

47 

48 def _register_defaults(self) -> None: 

49 """Populate with built-in chunker classes.""" 

50 from lexigram.ai.rag.chunking.strategies.fixed_size import FixedSizeChunker 

51 from lexigram.ai.rag.chunking.strategies.recursive import RecursiveChunker 

52 from lexigram.ai.rag.chunking.strategies.semantic import SemanticChunker 

53 from lexigram.ai.rag.chunking.strategies.sliding_window import ( 

54 SlidingWindowChunker, 

55 ) 

56 from lexigram.ai.rag.chunking.strategies.token import TokenChunker 

57 

58 self.register(ChunkingStrategy.FIXED_SIZE, FixedSizeChunker) 

59 self.register(ChunkingStrategy.RECURSIVE, RecursiveChunker) 

60 self.register(ChunkingStrategy.SEMANTIC, SemanticChunker) 

61 self.register(ChunkingStrategy.SLIDING_WINDOW, SlidingWindowChunker) 

62 self.register(ChunkingStrategy.TOKEN, TokenChunker) 

63 

64 def create_chunker( 

65 self, 

66 strategy: ChunkingStrategy = ChunkingStrategy.FIXED_SIZE, 

67 config: ChunkingConfig | None = None, 

68 **kwargs: Any, 

69 ) -> AbstractChunker: 

70 """Create a chunker instance for the given strategy. 

71 

72 Merges *config* defaults with explicit *kwargs* (kwargs win). 

73 

74 Args: 

75 strategy: Which chunking strategy to use. 

76 config: Optional chunking config; fields used as defaults. 

77 **kwargs: Constructor overrides forwarded to the chunker class. 

78 

79 Returns: 

80 A configured :class:`Chunker` instance. 

81 """ 

82 config = config or ChunkingConfig() 

83 merged = self._merge_config(strategy, config, kwargs) 

84 return self.instantiate(strategy, **merged) 

85 

86 @staticmethod 

87 def _merge_config( 

88 strategy: ChunkingStrategy, 

89 config: ChunkingConfig, 

90 overrides: dict[str, Any], 

91 ) -> dict[str, Any]: 

92 """Build kwargs from config with overrides applied.""" 

93 base: dict[str, Any] = {} 

94 if strategy == ChunkingStrategy.FIXED_SIZE: 

95 base = { 

96 "chunk_size": overrides.get("chunk_size", config.chunk_size), 

97 "overlap": overrides.get("overlap", config.overlap), 

98 } 

99 elif strategy == ChunkingStrategy.RECURSIVE: 

100 base = { 

101 "chunk_size": overrides.get("chunk_size", config.chunk_size), 

102 "overlap": overrides.get("overlap", config.overlap), 

103 "separators": overrides.get("separators", config.separators), 

104 } 

105 elif strategy == ChunkingStrategy.SEMANTIC: 

106 base = { 

107 "max_chunk_size": overrides.get("chunk_size", config.chunk_size), 

108 "min_chunk_size": overrides.get( 

109 "min_chunk_size", config.min_chunk_size 

110 ), 

111 } 

112 elif strategy == ChunkingStrategy.SLIDING_WINDOW: 

113 stride = overrides.get("stride") 

114 if stride is None: 

115 overlap = overrides.get("overlap", config.overlap) 

116 chunk_size = overrides.get("chunk_size", config.chunk_size) 

117 stride = max(1, chunk_size - overlap) 

118 base = { 

119 "window_size": overrides.get("chunk_size", config.chunk_size), 

120 "stride": stride, 

121 } 

122 elif strategy == ChunkingStrategy.TOKEN: 

123 base = { 

124 "chunk_size": overrides.get("chunk_size", config.chunk_size), 

125 "overlap": overrides.get("overlap", config.overlap), 

126 "encoding_name": overrides.get("encoding_name", config.encoding_name), 

127 } 

128 return base 

129 

130 def default_strategies(self) -> dict[str, type]: 

131 """Return built-in strategy key → class mapping.""" 

132 from lexigram.ai.rag.chunking.strategies.fixed_size import FixedSizeChunker 

133 from lexigram.ai.rag.chunking.strategies.recursive import RecursiveChunker 

134 from lexigram.ai.rag.chunking.strategies.semantic import SemanticChunker 

135 from lexigram.ai.rag.chunking.strategies.sliding_window import ( 

136 SlidingWindowChunker, 

137 ) 

138 from lexigram.ai.rag.chunking.strategies.token import TokenChunker 

139 

140 return { 

141 ChunkingStrategy.FIXED_SIZE: FixedSizeChunker, 

142 ChunkingStrategy.RECURSIVE: RecursiveChunker, 

143 ChunkingStrategy.SEMANTIC: SemanticChunker, 

144 ChunkingStrategy.SLIDING_WINDOW: SlidingWindowChunker, 

145 ChunkingStrategy.TOKEN: TokenChunker, 

146 }