Coverage for src/dataknobs_llm/prompts/implementations/composite_library.py: 78%

1"""Composite prompt library implementation. 

2 

3This module provides a library that combines multiple prompt libraries with 

4fallback behavior. Searches libraries in order and returns the first match. 

5 

6Example: 

7 # Create layered configuration with fallbacks 

8 base_library = FileSystemPromptLibrary(Path("prompts/base/")) 

9 override_library = ConfigPromptLibrary({ 

10 "system": { 

11 "analyze_code": {"template": "Custom analysis prompt..."} 

12 } 

13 }) 

14 

15 # Composite library searches override first, then base 

16 library = CompositePromptLibrary( 

17 libraries=[override_library, base_library], 

18 names=["override", "base"] 

19 ) 

20 

21 # Gets from override if available, otherwise from base 

22 template = library.get_system_prompt("analyze_code") 

23""" 

24 

25import logging 

26from typing import Any, Dict, List 

27 

28from ..base import ( 

29 AbstractPromptLibrary, 

30 PromptTemplateDict, 

31 RAGConfig, 

32 MessageIndex, 

33) 

34 

35logger = logging.getLogger(__name__) 

36 

37 

38class CompositePromptLibrary(AbstractPromptLibrary): 

39 """Composite library that combines multiple prompt libraries with fallback. 

40 

41 Features: 

42 - Searches libraries in order (first match wins) 

43 - Useful for layered configurations (overrides + defaults) 

44 - Supports all prompt types 

45 - Optional library naming for debugging 

46 

47 Example: 

48 >>> base = FileSystemPromptLibrary(Path("base/")) 

49 >>> custom = ConfigPromptLibrary(custom_config) 

50 >>> library = CompositePromptLibrary([custom, base]) 

51 >>> # Searches custom first, then base 

52 >>> template = library.get_system_prompt("analyze") 

53 """ 

54 

55 def __init__( 

56 self, 

57 libraries: List[AbstractPromptLibrary] | None = None, 

58 names: List[str] | None = None 

59 ): 

60 """Initialize composite prompt library. 

61 

62 Args: 

63 libraries: List of libraries to search (in order of priority) 

64 names: Optional names for libraries (for debugging/logging) 

65 """ 

66 self._libraries = libraries or [] 

67 self._names = names or [f"library_{i}" for i in range(len(self._libraries))] 

68 

69 if len(self._names) != len(self._libraries): 

70 raise ValueError( 

71 f"Number of names ({len(self._names)}) must match " 

72 f"number of libraries ({len(self._libraries)})" 

73 ) 

74 

75 def add_library( 

76 self, 

77 library: AbstractPromptLibrary, 

78 name: str | None = None, 

79 priority: int = -1 

80 ) -> None: 

81 """Add a library to the composite. 

82 

83 Args: 

84 library: Library to add 

85 name: Optional name for the library 

86 priority: Position to insert library (default: -1 = end) 

87 0 = highest priority (searched first) 

88 """ 

89 if name is None: 

90 name = f"library_{len(self._libraries)}" 

91 

92 if priority < 0 or priority >= len(self._libraries): 

93 # Append to end 

94 self._libraries.append(library) 

95 self._names.append(name) 

96 else: 

97 # Insert at specific position 

98 self._libraries.insert(priority, library) 

99 self._names.insert(priority, name) 

100 

101 logger.debug( 

102 f"Added library '{name}' at priority {priority if priority >= 0 else len(self._libraries) - 1}" 

103 ) 

104 

105 def remove_library(self, name: str) -> bool: 

106 """Remove a library by name. 

107 

108 Args: 

109 name: Name of the library to remove 

110 

111 Returns: 

112 True if library was removed, False if not found 

113 """ 

114 try: 

115 index = self._names.index(name) 

116 self._libraries.pop(index) 

117 self._names.pop(index) 

118 logger.debug(f"Removed library '{name}'") 

119 return True 

120 except ValueError: 

121 logger.warning(f"Library '{name}' not found") 

122 return False 

123 

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

125 """Get a system prompt by name, searching libraries in order. 

126 

127 Args: 

128 name: System prompt name 

129 **kwargs: Additional arguments passed to libraries 

130 

131 Returns: 

132 PromptTemplateDict from first library that has it, or None 

133 """ 

134 for lib, lib_name in zip(self._libraries, self._names, strict=True): 

135 template = lib.get_system_prompt(name, **kwargs) 

136 if template is not None: 

137 logger.debug(f"Found system prompt '{name}' in library '{lib_name}'") 

138 return template 

139 

140 logger.debug(f"System prompt '{name}' not found in any library") 

141 return None 

142 

143 def get_user_prompt( 

144 self, 

145 name: str, 

146 **kwargs 

147 ) -> PromptTemplateDict | None: 

148 """Get a user prompt by name, searching libraries in order. 

149 

150 Args: 

151 name: User prompt name 

152 **kwargs: Additional arguments passed to libraries 

153 

154 Returns: 

155 PromptTemplateDict from first library that has it, or None 

156 """ 

157 for lib, lib_name in zip(self._libraries, self._names, strict=True): 

158 template = lib.get_user_prompt(name, **kwargs) 

159 if template is not None: 

160 logger.debug( 

161 f"Found user prompt '{name}' in library '{lib_name}'" 

162 ) 

163 return template 

164 

165 logger.debug(f"User prompt '{name}' not found in any library") 

166 return None 

167 

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

169 """Get a message index by name, searching libraries in order. 

170 

171 Args: 

172 name: Message index name 

173 **kwargs: Additional arguments passed to libraries 

174 

175 Returns: 

176 MessageIndex from first library that has it, or None 

177 """ 

178 for lib, lib_name in zip(self._libraries, self._names, strict=True): 

179 message_index = lib.get_message_index(name, **kwargs) 

180 if message_index is not None: 

181 logger.debug(f"Found message index '{name}' in library '{lib_name}'") 

182 return message_index 

183 

184 logger.debug(f"Message index '{name}' not found in any library") 

185 return None 

186 

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

188 """Get a standalone RAG configuration by name, searching libraries in order. 

189 

190 Args: 

191 name: RAG config name 

192 **kwargs: Additional arguments passed to libraries 

193 

194 Returns: 

195 RAGConfig from first library that has it, or None 

196 """ 

197 for lib, lib_name in zip(self._libraries, self._names, strict=True): 

198 rag_config = lib.get_rag_config(name, **kwargs) 

199 if rag_config is not None: 

200 logger.debug(f"Found RAG config '{name}' in library '{lib_name}'") 

201 return rag_config 

202 

203 logger.debug(f"RAG config '{name}' not found in any library") 

204 return None 

205 

206 def get_prompt_rag_configs( 

207 self, 

208 prompt_name: str, 

209 prompt_type: str = "user", 

210 **kwargs 

211 ) -> List[RAGConfig]: 

212 """Get RAG configurations for a prompt, searching libraries in order. 

213 

214 Args: 

215 prompt_name: Prompt name 

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

217 **kwargs: Additional arguments passed to libraries 

218 

219 Returns: 

220 List of RAGConfig from first library that has the prompt 

221 """ 

222 for lib, lib_name in zip(self._libraries, self._names, strict=True): 

223 configs = lib.get_prompt_rag_configs(prompt_name, prompt_type, **kwargs) 

224 if configs: 

225 logger.debug( 

226 f"Found {len(configs)} RAG config(s) for prompt '{prompt_name}' " 

227 f"in library '{lib_name}'" 

228 ) 

229 return configs 

230 

231 logger.debug(f"No RAG configs found for prompt '{prompt_name}' in any library") 

232 return [] 

233 

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

235 """List all available system prompt names from all libraries. 

236 

237 Returns: 

238 Combined list of system prompt identifiers 

239 """ 

240 prompts = set() 

241 for lib in self._libraries: 

242 prompts.update(lib.list_system_prompts()) 

243 return sorted(prompts) 

244 

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

246 """List available user prompts from all libraries. 

247 

248 Returns: 

249 Combined list of user prompt names or indices 

250 """ 

251 prompts = set() 

252 for lib in self._libraries: 

253 prompts.update(lib.list_user_prompts()) 

254 return sorted(prompts) 

255 

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

257 """List all available message index names from all libraries. 

258 

259 Returns: 

260 Combined list of message index identifiers 

261 """ 

262 indexes = set() 

263 for lib in self._libraries: 

264 indexes.update(lib.list_message_indexes()) 

265 return sorted(indexes) 

266 

267 def get_metadata(self) -> Dict[str, Any]: 

268 """Get metadata about this composite library. 

269 

270 Returns: 

271 Dictionary with library metadata 

272 """ 

273 return { 

274 "class": self.__class__.__name__, 

275 "num_libraries": len(self._libraries), 

276 "library_names": self._names, 

277 } 

278 

279 def reload(self) -> None: 

280 """Reload the library by clearing the cache. 

281 

282 Subclasses can override to perform additional reload logic. 

283 """ 

284 for lib in self._libraries: 

285 lib.reload() 

286 

287 @property 

288 def libraries(self) -> List[AbstractPromptLibrary]: 

289 """Get list of libraries in priority order.""" 

290 return self._libraries.copy() 

291 

292 @property 

293 def library_names(self) -> List[str]: 

294 """Get list of library names in priority order.""" 

295 return self._names.copy() 

296 

297 def get_library_by_name(self, name: str) -> AbstractPromptLibrary | None: 

298 """Get a specific library by name. 

299 

300 Args: 

301 name: Library name 

302 

303 Returns: 

304 Library if found, None otherwise 

305 """ 

306 try: 

307 index = self._names.index(name) 

308 return self._libraries[index] 

309 except ValueError: 

310 return None