Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-rag/src/lexigram/ai/rag/pipeline/stages/synthesis_registry.py: 47%

1"""Synthesis strategy registry for RAG pipeline.""" 

2 

3from __future__ import annotations 

4 

5from typing import TYPE_CHECKING, Any, Protocol 

6 

7if TYPE_CHECKING: 

8 from lexigram.ai.rag.synthesis.types import SynthesisStrategy 

9 from lexigram.contracts.ai import LLMClientProtocol 

10 

11 

12class SynthesisStrategyHandler(Protocol): 

13 """Protocol for synthesis strategy handlers.""" 

14 

15 def can_handle(self, strategy: SynthesisStrategy) -> bool: 

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

17 ... 

18 

19 def create_synthesizer( 

20 self, config: Any, llm_client: LLMClientProtocol | None 

21 ) -> Any: 

22 """Create a synthesizer instance.""" 

23 ... 

24 

25 

26class DirectSynthesisStrategyHandler: 

27 """Handler for DIRECT synthesis strategy.""" 

28 

29 def can_handle(self, strategy: SynthesisStrategy) -> bool: 

30 from lexigram.ai.rag.pipeline.stages.synthesis import SynthesisStrategy 

31 

32 return strategy == SynthesisStrategy.DIRECT 

33 

34 def create_synthesizer( 

35 self, config: Any, llm_client: LLMClientProtocol | None 

36 ) -> Any: 

37 from lexigram.ai.rag.pipeline.stages.synthesis import DirectSynthesizer 

38 

39 return DirectSynthesizer( 

40 separator="\n\n", 

41 max_chunks=None, 

42 include_sources=config.include_citations, 

43 ) 

44 

45 

46class ExtractiveSynthesisStrategyHandler: 

47 """Handler for EXTRACTIVE synthesis strategy.""" 

48 

49 def can_handle(self, strategy: SynthesisStrategy) -> bool: 

50 from lexigram.ai.rag.pipeline.stages.synthesis import SynthesisStrategy 

51 

52 return strategy == SynthesisStrategy.EXTRACTIVE 

53 

54 def create_synthesizer( 

55 self, config: Any, llm_client: LLMClientProtocol | None 

56 ) -> Any: 

57 from lexigram.ai.rag.pipeline.stages.synthesis import ExtractiveSynthesizer 

58 

59 return ExtractiveSynthesizer( 

60 max_sentences=10, 

61 min_sentence_length=20, 

62 ) 

63 

64 

65class AbstractiveSynthesisStrategyHandler: 

66 """Handler for ABSTRACTIVE synthesis strategy.""" 

67 

68 def can_handle(self, strategy: SynthesisStrategy) -> bool: 

69 from lexigram.ai.rag.pipeline.stages.synthesis import SynthesisStrategy 

70 

71 return strategy == SynthesisStrategy.ABSTRACTIVE 

72 

73 def create_synthesizer( 

74 self, config: Any, llm_client: LLMClientProtocol | None 

75 ) -> Any: 

76 from lexigram.ai.rag.pipeline.stages.synthesis import ( 

77 AbstractiveSynthesizer, 

78 ExtractiveSynthesizer, 

79 ) 

80 from lexigram.logging import get_logger 

81 

82 logger = get_logger(__name__) 

83 if llm_client is None: 

84 logger.warning( 

85 "LLM client not provided for abstractive synthesis, " 

86 "falling back to extractive", 

87 ) 

88 return ExtractiveSynthesizer( 

89 max_sentences=10, 

90 min_sentence_length=20, 

91 ) 

92 return AbstractiveSynthesizer( 

93 llm_client=llm_client, 

94 max_context_chunks=5, 

95 include_citations=config.include_citations, 

96 ) 

97 

98 

99class HybridSynthesisStrategyHandler: 

100 """Handler for HYBRID synthesis strategy.""" 

101 

102 def can_handle(self, strategy: SynthesisStrategy) -> bool: 

103 from lexigram.ai.rag.pipeline.stages.synthesis import SynthesisStrategy 

104 

105 return strategy == SynthesisStrategy.HYBRID 

106 

107 def create_synthesizer( 

108 self, config: Any, llm_client: LLMClientProtocol | None 

109 ) -> Any: 

110 from lexigram.ai.rag.pipeline.stages.synthesis import ( 

111 ExtractiveSynthesizer, 

112 HybridSynthesizer, 

113 ) 

114 from lexigram.logging import get_logger 

115 

116 logger = get_logger(__name__) 

117 if llm_client is None: 

118 logger.warning( 

119 "LLM client not provided for hybrid synthesis, " 

120 "falling back to extractive", 

121 ) 

122 return ExtractiveSynthesizer( 

123 max_sentences=10, 

124 min_sentence_length=20, 

125 ) 

126 return HybridSynthesizer( 

127 llm_client=llm_client, 

128 max_extractive_sentences=5, 

129 ) 

130 

131 

132class SynthesisStrategyRegistry: 

133 """Central registry for synthesis strategy handlers.""" 

134 

135 def __init__(self) -> None: 

136 self._handlers: list[SynthesisStrategyHandler] = [] 

137 

138 @classmethod 

139 def with_defaults(cls) -> SynthesisStrategyRegistry: 

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

141 registry = cls() 

142 registry._handlers = [ 

143 DirectSynthesisStrategyHandler(), 

144 ExtractiveSynthesisStrategyHandler(), 

145 AbstractiveSynthesisStrategyHandler(), 

146 HybridSynthesisStrategyHandler(), 

147 ] 

148 return registry 

149 

150 def register(self, handler: SynthesisStrategyHandler) -> None: 

151 """Register a new strategy handler.""" 

152 self._handlers.insert(0, handler) 

153 

154 def create_synthesizer( 

155 self, 

156 strategy: SynthesisStrategy, 

157 config: Any, 

158 llm_client: LLMClientProtocol | None, 

159 ) -> Any: 

160 """Create a synthesizer for the given strategy.""" 

161 for handler in self._handlers: 

162 if handler.can_handle(strategy): 

163 return handler.create_synthesizer(config, llm_client) 

164 raise ValueError(f"Unknown synthesis strategy: {strategy}")