Coverage for src/dataknobs_llm/prompts/implementations/config_library.py: 90%

1r"""Configuration-based prompt library implementation. 

2 

3This module provides a prompt library that loads prompts from Python dictionaries. 

4Useful for programmatic prompt definition and testing. 

5 

6Example: 

7 config = { 

8 "system": { 

9 "analyze_code": { 

10 "template": "Analyze this {{language}} code:\\n{{code}}", 

11 "defaults": {"language": "python"}, 

12 "validation": { 

13 "level": "error", 

14 "required_params": ["code"] 

15 } 

16 } 

17 }, 

18 "user": { 

19 "code_question": { 

20 "template": "Please analyze the following code...", 

21 "rag_configs": [...] 

22 }, 

23 "followup_question": { 

24 "template": "Additionally, check for..." 

25 } 

26 } 

27 } 

28 

29 library = ConfigPromptLibrary(config) 

30 template = library.get_system_prompt("analyze_code") 

31""" 

32 

33import logging 

34from typing import Any, Dict, List 

35 

36from ..base import ( 

37 BasePromptLibrary, 

38 PromptTemplateDict, 

39 RAGConfig, 

40 MessageIndex, 

41) 

42 

43logger = logging.getLogger(__name__) 

44 

45 

46class ConfigPromptLibrary(BasePromptLibrary): 

47 """Prompt library that loads prompts from configuration dictionaries. 

48 

49 Features: 

50 - Direct in-memory configuration 

51 - No filesystem dependencies 

52 - Useful for testing and programmatic definition 

53 - Supports all prompt types (system, user, messages, RAG) 

54 

55 Example: 

56 >>> config = { 

57 ... "system": {"greet": {"template": "Hello {{name}}!"}}, 

58 ... "user": {"ask": {"template": "Tell me about {{topic}}"}} 

59 ... } 

60 >>> library = ConfigPromptLibrary(config) 

61 >>> template = library.get_system_prompt("greet") 

62 """ 

63 

64 def __init__(self, config: Dict[str, Any] | None = None): 

65 """Initialize configuration-based prompt library. 

66 

67 Args: 

68 config: Configuration dictionary with structure: 

69 { 

70 "system": {name: PromptTemplateDict, ...}, 

71 "user": {name: PromptTemplateDict, ...}, 

72 "messages": {name: MessageIndex, ...}, 

73 "rag": {name: RAGConfig, ...} 

74 } 

75 """ 

76 super().__init__() 

77 self._config = config or {} 

78 

79 # Load prompts from config 

80 self._load_from_config() 

81 

82 def _load_from_config(self) -> None: 

83 """Load all prompts from the configuration dictionary.""" 

84 self._load_system_prompts() 

85 self._load_user_prompts() 

86 self._load_message_indexes() 

87 self._load_rag_configs() # Load standalone RAG configs 

88 

89 def _load_system_prompts(self) -> None: 

90 """Load system prompts from config.""" 

91 system_config = self._config.get("system", {}) 

92 

93 for name, data in system_config.items(): 

94 try: 

95 template = self._parse_prompt_template(data) 

96 self._cache_system_prompt(name, template) 

97 logger.debug(f"Loaded system prompt from config: {name}") 

98 except Exception as e: 

99 logger.error(f"Error loading system prompt {name}: {e}") 

100 

101 def _load_user_prompts(self) -> None: 

102 """Load user prompts from config.""" 

103 user_config = self._config.get("user", {}) 

104 

105 for name, data in user_config.items(): 

106 try: 

107 template = self._parse_prompt_template(data) 

108 self._cache_user_prompt(name, template) 

109 logger.debug(f"Loaded user prompt from config: {name}") 

110 except Exception as e: 

111 logger.error(f"Error loading user prompt {name}: {e}") 

112 

113 def _load_message_indexes(self) -> None: 

114 """Load message indexes from config.""" 

115 messages_config = self._config.get("messages", {}) 

116 

117 for name, data in messages_config.items(): 

118 try: 

119 message_index = self._parse_message_index(data) 

120 self._cache_message_index(name, message_index) 

121 logger.debug(f"Loaded message index from config: {name}") 

122 except Exception as e: 

123 logger.error(f"Error loading message index {name}: {e}") 

124 

125 def _load_rag_configs(self) -> None: 

126 """Load RAG configurations from config.""" 

127 rag_config = self._config.get("rag", {}) 

128 

129 for name, data in rag_config.items(): 

130 try: 

131 rag = self._parse_rag_config(data) 

132 self._cache_rag_config(name, rag) 

133 logger.debug(f"Loaded RAG config from config: {name}") 

134 except Exception as e: 

135 logger.error(f"Error loading RAG config {name}: {e}") 

136 

137 # Note: _parse_prompt_template(), _parse_validation_config(), and 

138 # _parse_rag_config() are now inherited from BasePromptLibrary 

139 

140 def _parse_message_index(self, data: Dict[str, Any]) -> MessageIndex: 

141 """Parse message index from config data. 

142 

143 Args: 

144 data: Message index data 

145 

146 Returns: 

147 MessageIndex dictionary 

148 """ 

149 message_index: MessageIndex = { 

150 "messages": data.get("messages", []), 

151 } 

152 

153 # Add optional fields 

154 if "rag_configs" in data: 

155 message_index["rag_configs"] = [ 

156 self._parse_rag_config(rag_data) 

157 for rag_data in data["rag_configs"] 

158 ] 

159 

160 if "metadata" in data: 

161 message_index["metadata"] = data["metadata"] 

162 

163 return message_index 

164 

165 def get_system_prompt(self, name: str, **kwargs: Any) -> PromptTemplateDict | None: 

166 """Get a system prompt by name. 

167 

168 Args: 

169 name: System prompt name 

170 **kwargs: Additional arguments (unused) 

171 

172 Returns: 

173 PromptTemplateDict if found, None otherwise 

174 """ 

175 return self._get_cached_system_prompt(name) 

176 

177 def get_user_prompt(self, name: str, **kwargs: Any) -> PromptTemplateDict | None: 

178 """Get a user prompt by name. 

179 

180 Args: 

181 name: User prompt name 

182 **kwargs: Additional arguments (unused) 

183 

184 Returns: 

185 PromptTemplateDict if found, None otherwise 

186 """ 

187 return self._get_cached_user_prompt(name) 

188 

189 def get_message_index(self, name: str, **kwargs: Any) -> MessageIndex | None: 

190 """Get a message index by name. 

191 

192 Args: 

193 name: Message index name 

194 **kwargs: Additional arguments (unused) 

195 

196 Returns: 

197 MessageIndex if found, None otherwise 

198 """ 

199 return self._get_cached_message_index(name) 

200 

201 def get_rag_config(self, name: str, **kwargs: Any) -> RAGConfig | None: 

202 """Get a standalone RAG configuration by name. 

203 

204 Args: 

205 name: RAG config name 

206 **kwargs: Additional arguments (unused) 

207 

208 Returns: 

209 RAGConfig if found, None otherwise 

210 """ 

211 return self._get_cached_rag_config(name) 

212 

213 def get_prompt_rag_configs( 

214 self, 

215 prompt_name: str, 

216 prompt_type: str = "user", 

217 **kwargs: Any 

218 ) -> List[RAGConfig]: 

219 """Get RAG configurations for a specific prompt. 

220 

221 Resolves both inline RAG configs and references to standalone configs. 

222 

223 Args: 

224 prompt_name: Prompt name 

225 prompt_type: Type of prompt ("user" or "system") 

226 **kwargs: Additional arguments (unused) 

227 

228 Returns: 

229 List of RAGConfig (empty if none defined) 

230 """ 

231 # Get the prompt template 

232 if prompt_type == "system": 

233 template = self.get_system_prompt(prompt_name) 

234 else: 

235 template = self.get_user_prompt(prompt_name) 

236 

237 if template is None: 

238 return [] 

239 

240 configs = [] 

241 

242 # Get inline RAG configs from template 

243 if "rag_configs" in template: 

244 configs.extend(template["rag_configs"]) 

245 

246 # Resolve RAG config references 

247 if "rag_config_refs" in template: 

248 for ref_name in template["rag_config_refs"]: 

249 ref_config = self.get_rag_config(ref_name) 

250 if ref_config: 

251 configs.append(ref_config) 

252 else: 

253 logger.warning( 

254 f"RAG config reference '{ref_name}' not found " 

255 f"for prompt '{prompt_name}'" 

256 ) 

257 

258 return configs 

259 

260 def add_system_prompt(self, name: str, template: PromptTemplateDict) -> None: 

261 """Add or update a system prompt. 

262 

263 Args: 

264 name: System prompt name 

265 template: Prompt template to add 

266 """ 

267 self._cache_system_prompt(name, template) 

268 logger.debug(f"Added/updated system prompt: {name}") 

269 

270 def add_user_prompt(self, name: str, template: PromptTemplateDict) -> None: 

271 """Add or update a user prompt. 

272 

273 Args: 

274 name: User prompt name 

275 template: Prompt template to add 

276 """ 

277 self._cache_user_prompt(name, template) 

278 logger.debug(f"Added/updated user prompt: {name}") 

279 

280 def add_message_index(self, name: str, message_index: MessageIndex) -> None: 

281 """Add or update a message index. 

282 

283 Args: 

284 name: Message index name 

285 message_index: Message index to add 

286 """ 

287 self._cache_message_index(name, message_index) 

288 logger.debug(f"Added/updated message index: {name}") 

289 

290 def add_rag_config(self, name: str, rag_config: RAGConfig) -> None: 

291 """Add or update a RAG configuration. 

292 

293 Args: 

294 name: RAG config name 

295 rag_config: RAG configuration to add 

296 """ 

297 self._cache_rag_config(name, rag_config) 

298 logger.debug(f"Added/updated RAG config: {name}") 

299 

300 def list_system_prompts(self) -> List[str]: 

301 """List all available system prompt names. 

302 

303 Returns: 

304 List of system prompt identifiers 

305 """ 

306 return list(self._system_prompt_cache.keys()) 

307 

308 def list_user_prompts(self) -> List[str]: 

309 """List available user prompts. 

310 

311 Returns: 

312 List of user prompt names 

313 """ 

314 return list(self._user_prompt_cache.keys()) 

315 

316 def list_message_indexes(self) -> List[str]: 

317 """List all available message index names. 

318 

319 Returns: 

320 List of message index identifiers 

321 """ 

322 return list(self._message_index_cache.keys())