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

1"""Compression strategy registry for RAG context compression.""" 

2 

3from __future__ import annotations 

4 

5from typing import Any, Protocol 

6 

7from lexigram.ai.rag.context_compression.abstractive import AbstractiveCompressor 

8from lexigram.ai.rag.context_compression.extractive import ExtractiveSummaryCompressor 

9from lexigram.ai.rag.context_compression.hybrid import HybridCompressor 

10from lexigram.ai.rag.context_compression.semantic_dedup import ( 

11 SemanticDeduplicationCompressor, 

12) 

13from lexigram.ai.rag.context_compression.token_limit import TokenLimitCompressor 

14from lexigram.ai.rag.context_compression.types import CompressionStrategy 

15 

16 

17class CompressionStrategyHandler(Protocol): 

18 """Protocol for compression strategy handlers.""" 

19 

20 def can_handle(self, strategy: Any) -> bool: 

21 """Check if this handler can handle the strategy.""" 

22 ... 

23 

24 async def create_and_compress( 

25 self, 

26 strategy: Any, 

27 context: Any, 

28 query: Any, 

29 kwargs: dict[str, Any], 

30 ) -> Any: 

31 """Create compressor and compress context.""" 

32 ... 

33 

34 

35class ExtractiveStrategyHandler: 

36 """Handler for extractive compression strategy.""" 

37 

38 def can_handle(self, strategy: Any) -> bool: 

39 return strategy == CompressionStrategy.EXTRACTIVE 

40 

41 async def create_and_compress( 

42 self, 

43 strategy: Any, 

44 context: Any, 

45 query: Any, 

46 kwargs: dict[str, Any], 

47 ) -> Any: 

48 compressor = ExtractiveSummaryCompressor(**kwargs) 

49 return await compressor.compress(context, query=query) 

50 

51 

52class AbstractiveStrategyHandler: 

53 """Handler for abstractive compression strategy.""" 

54 

55 def can_handle(self, strategy: Any) -> bool: 

56 return strategy == CompressionStrategy.ABSTRACTIVE 

57 

58 async def create_and_compress( 

59 self, 

60 strategy: Any, 

61 context: Any, 

62 query: Any, 

63 kwargs: dict[str, Any], 

64 ) -> Any: 

65 if "llm_client" not in kwargs: 

66 msg = "AbstractiveCompressor requires 'llm_client' parameter" 

67 raise ValueError(msg) 

68 compressor = AbstractiveCompressor(**kwargs) 

69 return await compressor.compress(context, query=query) 

70 

71 

72class TokenLimitStrategyHandler: 

73 """Handler for token limit compression strategy.""" 

74 

75 def can_handle(self, strategy: Any) -> bool: 

76 return strategy == CompressionStrategy.TOKEN_LIMIT 

77 

78 async def create_and_compress( 

79 self, 

80 strategy: Any, 

81 context: Any, 

82 query: Any, 

83 kwargs: dict[str, Any], 

84 ) -> Any: 

85 compressor = TokenLimitCompressor(**kwargs) 

86 return await compressor.compress(context, query=query) 

87 

88 

89class SemanticDedupStrategyHandler: 

90 """Handler for semantic deduplication compression strategy.""" 

91 

92 def can_handle(self, strategy: Any) -> bool: 

93 return strategy == CompressionStrategy.SEMANTIC_DEDUP 

94 

95 async def create_and_compress( 

96 self, 

97 strategy: Any, 

98 context: Any, 

99 query: Any, 

100 kwargs: dict[str, Any], 

101 ) -> Any: 

102 compressor = SemanticDeduplicationCompressor(**kwargs) 

103 return await compressor.compress(context, query=query) 

104 

105 

106class HybridStrategyHandler: 

107 """Handler for hybrid compression strategy.""" 

108 

109 def can_handle(self, strategy: Any) -> bool: 

110 return strategy == CompressionStrategy.HYBRID 

111 

112 async def create_and_compress( 

113 self, 

114 strategy: Any, 

115 context: Any, 

116 query: Any, 

117 kwargs: dict[str, Any], 

118 ) -> Any: 

119 if "compressors" not in kwargs: 

120 msg = "HybridCompressor requires 'compressors' parameter" 

121 raise ValueError(msg) 

122 compressor = HybridCompressor(**kwargs) 

123 return await compressor.compress(context, query=query) 

124 

125 

126class CompressionStrategyRegistry: 

127 """Central registry for compression strategy handlers.""" 

128 

129 def __init__(self) -> None: 

130 self._handlers: list[CompressionStrategyHandler] = [] 

131 

132 @classmethod 

133 def with_defaults(cls) -> CompressionStrategyRegistry: 

134 """Create a registry pre-populated with all built-in strategy handlers.""" 

135 registry = cls() 

136 registry._handlers = [ 

137 ExtractiveStrategyHandler(), 

138 AbstractiveStrategyHandler(), 

139 TokenLimitStrategyHandler(), 

140 SemanticDedupStrategyHandler(), 

141 HybridStrategyHandler(), 

142 ] 

143 return registry 

144 

145 def register(self, handler: CompressionStrategyHandler) -> None: 

146 """Register a new strategy handler.""" 

147 self._handlers.insert(0, handler) 

148 

149 async def compress( 

150 self, 

151 strategy: Any, 

152 context: Any, 

153 query: Any, 

154 kwargs: dict[str, Any], 

155 ) -> Any: 

156 """Compress context using the appropriate strategy.""" 

157 for handler in self._handlers: 

158 if handler.can_handle(strategy): 

159 return await handler.create_and_compress( 

160 strategy, 

161 context, 

162 query, 

163 kwargs, 

164 ) 

165 msg = f"Unknown compression strategy: {strategy}" 

166 raise ValueError(msg)