Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-rag/src/lexigram/ai/rag/hyde/base.py: 89%

1"""Base HyDE generator with common functionality.""" 

2 

3from __future__ import annotations 

4 

5from abc import ABC, abstractmethod 

6from typing import Any 

7 

8from lexigram.ai.rag.hyde.protocols import EmbeddingClientProtocol 

9from lexigram.ai.rag.hyde.types import HyDEResult, HypotheticalDocument 

10from lexigram.contracts import ( 

11 ChatMessage, 

12 LLMClientProtocol, 

13) 

14 

15 

16class AbstractHyDEGenerator(ABC): 

17 """Base class for HyDE generators.""" 

18 

19 def __init__( 

20 self, 

21 llm_client: LLMClientProtocol, 

22 embedding_client: EmbeddingClientProtocol | None = None, 

23 ): 

24 """Initialize generator. 

25 

26 Args: 

27 llm_client: Client for generating hypothetical documents 

28 embedding_client: Optional client for generating embeddings 

29 """ 

30 self.llm_client = llm_client 

31 self.embedding_client = embedding_client 

32 

33 @abstractmethod 

34 async def generate( 

35 self, 

36 query: str, 

37 num_documents: int = 1, 

38 **kwargs: Any, 

39 ) -> HyDEResult: 

40 """Generate hypothetical documents for query. 

41 

42 Args: 

43 query: User query 

44 num_documents: Number of hypothetical documents to generate 

45 **kwargs: Additional parameters 

46 

47 Returns: 

48 HyDE result with hypothetical documents 

49 """ 

50 

51 def _build_prompt( 

52 self, 

53 query: str, 

54 context: str = "", 

55 domain: str = "", 

56 ) -> str: 

57 """Build prompt for generating hypothetical document. 

58 

59 Args: 

60 query: User query 

61 context: Optional context 

62 domain: Optional domain specification 

63 

64 Returns: 

65 Formatted prompt 

66 """ 

67 prompt = ( 

68 "Generate a detailed, informative passage that would answer this query:\n\n" 

69 ) 

70 prompt += f"Query: {query}\n\n" 

71 

72 if domain: 

73 prompt += f"Domain: {domain}\n\n" 

74 

75 if context: 

76 prompt += f"Context: {context}\n\n" 

77 

78 prompt += ( 

79 "Write a passage that directly answers this query with specific details, " 

80 "facts, and relevant information. Do not include the query itself in your response." 

81 ) 

82 

83 return prompt 

84 

85 async def _generate_single_document( 

86 self, 

87 query: str, 

88 temperature: float = 0.7, 

89 max_tokens: int = 200, 

90 **kwargs: Any, 

91 ) -> str: 

92 """Generate a single hypothetical document. 

93 

94 Args: 

95 query: User query 

96 temperature: Sampling temperature 

97 max_tokens: Maximum tokens to generate 

98 **kwargs: Additional parameters 

99 

100 Returns: 

101 Generated hypothetical document content 

102 """ 

103 prompt = self._build_prompt( 

104 query, 

105 context=kwargs.get("context", ""), 

106 domain=kwargs.get("domain", ""), 

107 ) 

108 

109 messages = [ChatMessage(role="user", content=prompt)] 

110 

111 result = await self.llm_client.complete( 

112 messages, 

113 temperature=temperature, 

114 max_tokens=max_tokens, 

115 ) 

116 if result.is_err(): 

117 raise result.unwrap_err() 

118 response = result.unwrap() 

119 

120 # Extract content from response 

121 content = self._extract_content(response) 

122 return content.strip() 

123 

124 def _extract_content(self, response: Any) -> str: 

125 """Extract content from LLM response. 

126 

127 Args: 

128 response: LLM response object 

129 

130 Returns: 

131 Extracted content string 

132 """ 

133 # Handle different response formats 

134 if hasattr(response, "content"): 

135 return response.content 

136 if hasattr(response, "text"): 

137 return response.text 

138 if hasattr(response, "choices"): 

139 return response.choices[0].message.content 

140 if isinstance(response, str): 

141 return response 

142 return str(response) 

143 

144 async def _embed_documents( 

145 self, 

146 documents: list[HypotheticalDocument], 

147 ) -> list[list[float]]: 

148 """Generate embeddings for hypothetical documents. 

149 

150 Args: 

151 documents: List of hypothetical documents 

152 

153 Returns: 

154 List of embeddings 

155 """ 

156 if not self.embedding_client: 

157 msg = "Embedding client required for embedding documents" 

158 raise ValueError(msg) 

159 

160 texts = [doc.content for doc in documents] 

161 return await self.embedding_client.embed(texts) 

162 

163 def _aggregate_embeddings( 

164 self, 

165 embeddings: list[list[float]], 

166 weights: list[float] | None = None, 

167 ) -> list[float]: 

168 """Aggregate multiple embeddings into one. 

169 

170 Args: 

171 embeddings: List of embedding vectors 

172 weights: Optional weights for each embedding 

173 

174 Returns: 

175 Aggregated embedding vector 

176 """ 

177 if not embeddings: 

178 return [] 

179 

180 if weights is None: 

181 weights = [1.0] * len(embeddings) 

182 

183 # Normalize weights 

184 total_weight = sum(weights) 

185 normalized_weights = [w / total_weight for w in weights] 

186 

187 # Weighted average 

188 embedding_dim = len(embeddings[0]) 

189 aggregated = [0.0] * embedding_dim 

190 

191 for embedding, weight in zip(embeddings, normalized_weights, strict=False): 

192 for i, val in enumerate(embedding): 

193 aggregated[i] += val * weight 

194 

195 # Normalize vector 

196 magnitude = sum(x * x for x in aggregated) ** 0.5 

197 if magnitude > 0: 

198 aggregated = [x / magnitude for x in aggregated] 

199 

200 return aggregated