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

1"""Synthesis stage for response generation.""" 

2 

3from __future__ import annotations 

4 

5from typing import TYPE_CHECKING, cast 

6 

7from lexigram.ai.rag.synthesis.synthesizers.base import ResponseSynthesizerProtocol 

8 

9if TYPE_CHECKING: 

10 from lexigram.contracts.ai import LLMClientProtocol 

11 

12from lexigram.ai.rag.config import SynthesisConfig 

13from lexigram.ai.rag.pipeline.types import PipelineContext 

14from lexigram.ai.rag.synthesis import ( 

15 AbstractiveSynthesizer, 

16 DirectSynthesizer, 

17 ExtractiveSynthesizer, 

18 HybridSynthesizer, 

19 ResponseSynthesizerProtocol, 

20 SynthesisStrategy, 

21) 

22from lexigram.logging import ( 

23 get_logger, 

24) 

25 

26logger = get_logger(__name__) 

27 

28 

29class SynthesisStage: 

30 """Pipeline stage for response synthesis. 

31 

32 This stage generates a response from the retrieved context chunks 

33 using the configured synthesis strategy. 

34 """ 

35 

36 def __init__( 

37 self, 

38 config: SynthesisConfig, 

39 llm_client: LLMClientProtocol | None = None, 

40 ): 

41 """Initialize the synthesis stage. 

42 

43 Args: 

44 config: Synthesis configuration 

45 llm_client: Optional LLM client for abstractive/hybrid synthesis 

46 """ 

47 self.config = config 

48 self.llm_client = llm_client 

49 self._synthesizer: ResponseSynthesizerProtocol | None = None 

50 

51 @property 

52 def name(self) -> str: 

53 """Get stage name.""" 

54 return "synthesis" 

55 

56 async def process(self, context: PipelineContext) -> PipelineContext: 

57 """Process the synthesis stage. 

58 

59 Args: 

60 context: Pipeline context with retrieved chunks 

61 

62 Returns: 

63 Updated context with synthesis result 

64 """ 

65 if not self.config.enabled: 

66 logger.info("Synthesis stage disabled, skipping") 

67 return context 

68 

69 # Check if we have chunks to synthesize from 

70 chunks_to_use = context.optimized_chunks or context.retrieved_chunks 

71 if not chunks_to_use: 

72 logger.warning("No chunks available for synthesis") 

73 context.add_warning("No chunks available for synthesis") 

74 return context 

75 

76 # Get or create synthesizer 

77 synthesizer = self._get_synthesizer() 

78 

79 try: 

80 # Synthesize response 

81 logger.info( 

82 "Synthesizing response", 

83 extra={ 

84 "request_id": context.request_id, 

85 "strategy": self.config.strategy.value, 

86 "num_chunks": len(chunks_to_use), 

87 }, 

88 ) 

89 

90 result = await synthesizer._synthesize_internal( 

91 query=context.query, 

92 context_chunks=chunks_to_use, 

93 ) 

94 

95 # Store result in context 

96 context.synthesis_result = result 

97 

98 logger.info( 

99 "Synthesis completed", 

100 extra={ 

101 "request_id": context.request_id, 

102 "response_length": len(result.response), 

103 "chunks_used": result.num_chunks_used, 

104 "sources": len(result.sources), 

105 }, 

106 ) 

107 

108 except Exception as e: 

109 logger.exception( 

110 "Synthesis failed", 

111 extra={ 

112 "request_id": context.request_id, 

113 "strategy": self.config.strategy.value, 

114 "error": str(e), 

115 }, 

116 ) 

117 

118 # Try fallback if configured 

119 if ( 

120 self.config.error_strategy == "fallback" 

121 and self.config.strategy != self.config.fallback_strategy 

122 ): 

123 logger.info( 

124 "Attempting fallback synthesis strategy", 

125 extra={ 

126 "request_id": context.request_id, 

127 "fallback_strategy": self.config.fallback_strategy, 

128 }, 

129 ) 

130 

131 try: 

132 fallback_synthesizer = self._get_fallback_synthesizer() 

133 result = await fallback_synthesizer._synthesize_internal( 

134 query=context.query, 

135 context_chunks=chunks_to_use, 

136 ) 

137 context.synthesis_result = result 

138 context.add_warning( 

139 f"Used fallback synthesis strategy: {self.config.fallback_strategy}", 

140 ) 

141 except Exception as fallback_error: 

142 logger.exception( 

143 "Fallback synthesis also failed", 

144 extra={ 

145 "request_id": context.request_id, 

146 "error": str(fallback_error), 

147 }, 

148 ) 

149 raise e from fallback_error # Raise original error with fallback cause 

150 else: 

151 raise 

152 

153 return context 

154 

155 def _get_synthesizer(self) -> ResponseSynthesizerProtocol: 

156 """Get or create the configured synthesizer. 

157 

158 Returns: 

159 Response synthesizer instance 

160 """ 

161 if self._synthesizer is not None: 

162 return self._synthesizer 

163 

164 from lexigram.ai.rag.pipeline.stages.synthesis_registry import ( 

165 SynthesisStrategyRegistry, 

166 ) 

167 

168 strategy = self.config.strategy 

169 registry = SynthesisStrategyRegistry.with_defaults() 

170 self._synthesizer = registry.create_synthesizer( 

171 strategy, 

172 self.config, 

173 self.llm_client, 

174 ) 

175 

176 return cast("ResponseSynthesizerProtocol", self._synthesizer) 

177 

178 def _get_fallback_synthesizer(self) -> ResponseSynthesizerProtocol: 

179 """Get the fallback synthesizer. 

180 

181 Returns: 

182 Fallback synthesizer instance 

183 """ 

184 strategy = self.config.fallback_strategy 

185 

186 if strategy == SynthesisStrategy.DIRECT: 

187 return DirectSynthesizer( 

188 separator="\n\n", 

189 max_chunks=None, 

190 include_sources=self.config.include_citations, 

191 ) 

192 

193 if strategy == SynthesisStrategy.EXTRACTIVE: 

194 return ExtractiveSynthesizer( 

195 max_sentences=10, 

196 min_sentence_length=20, 

197 ) 

198 

199 if strategy == SynthesisStrategy.ABSTRACTIVE: 

200 if self.llm_client is None: 

201 # Fall back to extractive if no LLM 

202 return ExtractiveSynthesizer( 

203 max_sentences=10, 

204 min_sentence_length=20, 

205 ) 

206 return AbstractiveSynthesizer( 

207 llm_client=self.llm_client, 

208 max_context_chunks=5, 

209 include_citations=self.config.include_citations, 

210 ) 

211 

212 if strategy == SynthesisStrategy.HYBRID: 

213 if self.llm_client is None: 

214 # Fall back to extractive if no LLM 

215 return ExtractiveSynthesizer( 

216 max_sentences=10, 

217 min_sentence_length=20, 

218 ) 

219 return HybridSynthesizer( 

220 llm_client=self.llm_client, 

221 extraction_ratio=0.5, 

222 max_extractive_sentences=8, 

223 ) 

224 

225 msg = f"Unknown fallback synthesis strategy: {strategy}" 

226 raise ValueError(msg)