Coverage for src/dataknobs_llm/prompts/base/base_prompt_library.py: 22%

1"""Base implementation of prompt library with shared functionality.

2

3 This module provides BasePromptLibrary, a concrete implementation of

4 AbstractPromptLibrary that includes caching and common utilities.

5"""

6

7 from typing import Any, Dict, List, Optional, Union

8 import logging

9

10 from .abstract_prompt_library import AbstractPromptLibrary

11 from .types import PromptTemplate, MessageIndex, RAGConfig, ValidationConfig, ValidationLevel

12

13 logger = logging.getLogger(__name__)

14

15

16 class BasePromptLibrary(AbstractPromptLibrary):

17 """Base implementation with caching and common functionality.

18

19 This class provides:

20 - Optional caching of loaded prompts and message indexes

21 - Helper methods for cache management

22 - Shared metadata handling

23 - Default implementations of optional methods

24

25 Subclasses should implement the abstract methods to provide

26 the actual prompt loading logic.

27 """

28

29 def __init__(

30 self,

31 enable_cache: bool = True,

32 metadata: Optional[Dict[str, Any]] = None

33 ):

34 """Initialize the base prompt library.

35

36 Args:

37 enable_cache: Whether to cache loaded prompts (default: True)

38 metadata: Optional metadata dictionary

39 """

40 self._enable_cache = enable_cache

41 self._metadata = metadata or {}

42

43 # Caches for loaded prompts and indexes

44 self._system_prompt_cache: Dict[str, PromptTemplate] = {}

45 self._user_prompt_cache: Dict[str, PromptTemplate] = {}

46 self._message_index_cache: Dict[str, MessageIndex] = {}

47 self._rag_config_cache: Dict[str, RAGConfig] = {} # Standalone RAG configs

48 self._prompt_rag_cache: Dict[tuple, List[RAGConfig]] = {} # (name, type)

49

50 # ===== Cache Management =====

51

52 def clear_cache(self) -> None:

53 """Clear all cached prompts and indexes."""

54 self._system_prompt_cache.clear()

55 self._user_prompt_cache.clear()

56 self._message_index_cache.clear()

57 self._rag_config_cache.clear()

58 self._prompt_rag_cache.clear()

59 logger.debug(f"Cleared cache for {self.__class__.__name__}")

60

61 def reload(self) -> None:

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

63

64 Subclasses can override to perform additional reload logic.

65 """

66 if self._enable_cache:

67 self.clear_cache()

68 logger.info(f"Reloaded {self.__class__.__name__}")

69

70 # ===== Metadata =====

71

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

73 """Get metadata about this prompt library.

74

75 Returns:

76 Dictionary with library metadata

77 """

78 return {

79 "class": self.__class__.__name__,

80 "cache_enabled": self._enable_cache,

81 **self._metadata

82 }

83

84 # ===== Cache Helpers =====

85

86 def _get_cached_system_prompt(self, name: str) -> Optional[PromptTemplate]:

87 """Get system prompt from cache if caching is enabled.

88

89 Args:

90 name: System prompt identifier

91

92 Returns:

93 Cached PromptTemplate if found, None otherwise

94 """

95 if not self._enable_cache:

96 return None

97 return self._system_prompt_cache.get(name)

98

99 def _cache_system_prompt(self, name: str, template: PromptTemplate) -> None:

100 """Cache a system prompt if caching is enabled.

101

102 Args:

103 name: System prompt identifier

104 template: PromptTemplate to cache

105 """

106 if self._enable_cache:

107 self._system_prompt_cache[name] = template

108

109 def _get_cached_user_prompt(self, name: str) -> Optional[PromptTemplate]:

110 """Get user prompt from cache if caching is enabled.

111

112 Args:

113 name: User prompt identifier

114

115 Returns:

116 Cached PromptTemplate if found, None otherwise

117 """

118 if not self._enable_cache:

119 return None

120 return self._user_prompt_cache.get(name)

121

122 def _cache_user_prompt(self, name: str, template: PromptTemplate) -> None:

123 """Cache a user prompt if caching is enabled.

124

125 Args:

126 name: User prompt identifier

127 template: PromptTemplate to cache

128 """

129 if self._enable_cache:

130 self._user_prompt_cache[name] = template

131

132 def _get_cached_message_index(self, name: str) -> Optional[MessageIndex]:

133 """Get message index from cache if caching is enabled.

134

135 Args:

136 name: Message index identifier

137

138 Returns:

139 Cached MessageIndex if found, None otherwise

140 """

141 if not self._enable_cache:

142 return None

143 return self._message_index_cache.get(name)

144

145 def _cache_message_index(self, name: str, index: MessageIndex) -> None:

146 """Cache a message index if caching is enabled.

147

148 Args:

149 name: Message index identifier

150 index: MessageIndex to cache

151 """

152 if self._enable_cache:

153 self._message_index_cache[name] = index

154

155 def _get_cached_rag_config(self, name: str) -> Optional[RAGConfig]:

156 """Get standalone RAG config from cache if caching is enabled.

157

158 Args:

159 name: RAG config identifier

160

161 Returns:

162 Cached RAGConfig if found, None otherwise

163 """

164 if not self._enable_cache:

165 return None

166 return self._rag_config_cache.get(name)

167

168 def _cache_rag_config(self, name: str, config: RAGConfig) -> None:

169 """Cache a standalone RAG config if caching is enabled.

170

171 Args:

172 name: RAG config identifier

173 config: RAGConfig to cache

174 """

175 if self._enable_cache:

176 self._rag_config_cache[name] = config

177

178 def _get_cached_prompt_rag_configs(

179 self,

180 prompt_name: str,

181 prompt_type: str

182 ) -> Optional[List[RAGConfig]]:

183 """Get prompt RAG configs from cache if caching is enabled.

184

185 Args:

186 prompt_name: Prompt identifier

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

188

189 Returns:

190 Cached list of RAGConfig if found, None otherwise

191 """

192 if not self._enable_cache:

193 return None

194 return self._prompt_rag_cache.get((prompt_name, prompt_type))

195

196 def _cache_prompt_rag_configs(

197 self,

198 prompt_name: str,

199 prompt_type: str,

200 configs: List[RAGConfig]

201 ) -> None:

202 """Cache prompt RAG configurations if caching is enabled.

203

204 Args:

205 prompt_name: Prompt identifier

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

207 configs: List of RAGConfig to cache

208 """

209 if self._enable_cache:

210 self._prompt_rag_cache[(prompt_name, prompt_type)] = configs

211

212 # ===== Common Parsing Methods =====

213

214 def _parse_validation_config(self, data: Union[Dict, ValidationConfig]) -> ValidationConfig:

215 """Parse validation configuration from dict or ValidationConfig.

216

217 This method is shared by all library implementations for consistent

218 validation config parsing.

219

220 Args:

221 data: Validation data (dict or ValidationConfig instance)

222

223 Returns:

224 ValidationConfig instance

225

226 Raises:

227 ValueError: If data type is invalid

228 """

229 if isinstance(data, ValidationConfig):

230 return data

231

232 if not isinstance(data, dict):

233 raise ValueError(

234 f"Invalid validation config: expected dict or ValidationConfig, "

235 f"got {type(data)}"

236 )

237

238 # Parse level

239 level = None

240 if "level" in data:

241 level_data = data["level"]

242 if isinstance(level_data, str):

243 level = ValidationLevel(level_data.lower())

244 elif isinstance(level_data, ValidationLevel):

245 level = level_data

246

247 # Parse params

248 required_params = data.get("required_params", [])

249 optional_params = data.get("optional_params", [])

250

251 return ValidationConfig(

252 level=level,

253 required_params=required_params,

254 optional_params=optional_params

255 )

256

257 def _parse_rag_config(self, data: Dict[str, Any]) -> RAGConfig:

258 """Parse RAG configuration from dict.

259

260 This method is shared by all library implementations for consistent

261 RAG config parsing.

262

263 Args:

264 data: RAG config data dictionary

265

266 Returns:

267 RAGConfig dictionary

268 """

269 rag_config: RAGConfig = {

270 "adapter_name": data.get("adapter_name", ""),

271 "query": data.get("query", ""),

272 }

273

274 # Add optional fields

275 if "k" in data:

276 rag_config["k"] = data["k"]

277

278 if "filters" in data:

279 rag_config["filters"] = data["filters"]

280

281 if "placeholder" in data:

282 rag_config["placeholder"] = data["placeholder"]

283

284 if "header" in data:

285 rag_config["header"] = data["header"]

286

287 if "item_template" in data:

288 rag_config["item_template"] = data["item_template"]

289

290 return rag_config

291

292 def _parse_prompt_template(self, data: Any) -> PromptTemplate:

293 """Parse prompt template from various formats.

294

295 This method is shared by all library implementations for consistent

296 template parsing. Supports:

297 - String templates (converted to {"template": string})

298 - Dict with "template" key

299 - Dict with "extends" key but no "template" (template inherited)

300 - Empty dict (treated as {"template": ""})

301

302 Args:

303 data: Prompt template data (string or dict)

304

305 Returns:

306 PromptTemplate dictionary

307

308 Raises:

309 ValueError: If data format is invalid

310 """

311 # If just a string, treat as template

312 if isinstance(data, str):

313 return {"template": data}

314

315 # If empty dict, treat as empty template

316 if isinstance(data, dict) and len(data) == 0:

317 return {"template": ""}

318

319 # If dict with template field

320 if isinstance(data, dict) and "template" in data:

321 template: PromptTemplate = {

322 "template": data["template"],

323 }

324

325 # Add optional fields

326 if "defaults" in data:

327 template["defaults"] = data["defaults"]

328

329 if "validation" in data:

330 template["validation"] = self._parse_validation_config(data["validation"])

331

332 if "metadata" in data:

333 template["metadata"] = data["metadata"]

334

335 # Add template mode field

336 if "template_mode" in data:

337 template["template_mode"] = data["template_mode"]

338

339 # Add composition fields

340 if "sections" in data:

341 template["sections"] = data["sections"]

342

343 if "extends" in data:

344 template["extends"] = data["extends"]

345

346 # Add RAG configuration fields

347 if "rag_config_refs" in data:

348 template["rag_config_refs"] = data["rag_config_refs"]

349

350 if "rag_configs" in data:

351 template["rag_configs"] = [

352 self._parse_rag_config(rag_data)

353 for rag_data in data["rag_configs"]

354 ]

355

356 return template

357

358 # If dict with extends field but no template (template will be inherited)

359 elif isinstance(data, dict) and "extends" in data:

360 template: PromptTemplate = {}

361

362 # Template will be inherited from base

363 template["extends"] = data["extends"]

364

365 # Add optional override fields

366 if "defaults" in data:

367 template["defaults"] = data["defaults"]

368

369 if "validation" in data:

370 template["validation"] = self._parse_validation_config(data["validation"])

371

372 if "metadata" in data:

373 template["metadata"] = data["metadata"]

374

375 # Add template mode field

376 if "template_mode" in data:

377 template["template_mode"] = data["template_mode"]

378

379 if "sections" in data:

380 template["sections"] = data["sections"]

381

382 # Add RAG configuration fields

383 if "rag_config_refs" in data:

384 template["rag_config_refs"] = data["rag_config_refs"]

385

386 if "rag_configs" in data:

387 template["rag_configs"] = [

388 self._parse_rag_config(rag_data)

389 for rag_data in data["rag_configs"]

390 ]

391

392 return template

393

394 else:

395 raise ValueError(

396 f"Invalid prompt template data: expected dict with 'template' or 'extends' key, "

397 f"or string, got {type(data)}"

398 )

399

400 # ===== Abstract Methods (must be implemented by subclasses) =====

401

402 def get_system_prompt(

403 self,

404 name: str,

405 **kwargs: Any

406 ) -> Optional[PromptTemplate]:

407 """Retrieve a system prompt template by name.

408

409 Subclasses must implement this method.

410 """

411 raise NotImplementedError(

412 f"{self.__class__.__name__} must implement get_system_prompt()"

413 )

414

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

416 """List all available system prompt names.

417

418 Subclasses must implement this method.

419 """

420 raise NotImplementedError(

421 f"{self.__class__.__name__} must implement list_system_prompts()"

422 )

423

424 def get_user_prompt(

425 self,

426 name: str,

427 index: int = 0,

428 **kwargs: Any

429 ) -> Optional[PromptTemplate]:

430 """Retrieve a user prompt template by name and index.

431

432 Subclasses must implement this method.

433 """

434 raise NotImplementedError(

435 f"{self.__class__.__name__} must implement get_user_prompt()"

436 )

437

438 def list_user_prompts(self, name: Optional[str] = None) -> List[str]:

439 """List available user prompts.

440

441 Subclasses must implement this method.

442 """

443 raise NotImplementedError(

444 f"{self.__class__.__name__} must implement list_user_prompts()"

445 )

446

447 def get_message_index(

448 self,

449 name: str,

450 **kwargs: Any

451 ) -> Optional[MessageIndex]:

452 """Retrieve a message index by name.

453

454 Subclasses must implement this method.

455 """

456 raise NotImplementedError(

457 f"{self.__class__.__name__} must implement get_message_index()"

458 )

459

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

461 """List all available message index names.

462

463 Subclasses must implement this method.

464 """

465 raise NotImplementedError(

466 f"{self.__class__.__name__} must implement list_message_indexes()"

467 )

468

469 def get_rag_config(

470 self,

471 name: str,

472 **kwargs: Any

473 ) -> Optional[RAGConfig]:

474 """Retrieve a standalone RAG configuration by name.

475

476 Subclasses must implement this method.

477 """

478 raise NotImplementedError(

479 f"{self.__class__.__name__} must implement get_rag_config()"

480 )

481

482 def get_prompt_rag_configs(

483 self,

484 prompt_name: str,

485 prompt_type: str = "user",

486 index: int = 0,

487 **kwargs: Any

488 ) -> List[RAGConfig]:

489 """Retrieve RAG configurations for a specific prompt.

490

491 Subclasses must implement this method.

492 """

493 raise NotImplementedError(

494 f"{self.__class__.__name__} must implement get_prompt_rag_configs()"

495 )