Coverage for src/dataknobs_llm/prompts/base/types.py: 88%

1"""Core type definitions for the prompt library system. 

2 

3This module defines: 

4- Validation levels and configuration 

5- Prompt template structures 

6- RAG configuration types 

7- Message index types 

8- Render result types 

9""" 

10 

11from enum import Enum 

12from typing import Any, Dict, List, Optional, TypedDict, Union 

13from dataclasses import dataclass, field 

14 

15 

16class ValidationLevel(Enum): 

17 """Validation strictness levels for template rendering. 

18 

19 Attributes: 

20 ERROR: Raise exception for missing required parameters 

21 WARN: Log warning for missing required parameters (default) 

22 IGNORE: Silently ignore missing required parameters 

23 """ 

24 ERROR = "error" 

25 WARN = "warn" 

26 IGNORE = "ignore" 

27 

28 

29class TemplateMode(Enum): 

30 """Template rendering mode. 

31 

32 Attributes: 

33 MIXED: Pre-process (( )) conditionals then render with Jinja2 (default) 

34 JINJA2: Pure Jinja2 rendering, skip (( )) preprocessing 

35 """ 

36 MIXED = "mixed" 

37 JINJA2 = "jinja2" 

38 

39 @classmethod 

40 def from_string(cls, value: str) -> "TemplateMode": 

41 """Parse mode from string. 

42 

43 Args: 

44 value: Mode string ("mixed" or "jinja2") 

45 

46 Returns: 

47 TemplateMode enum value 

48 

49 Raises: 

50 ValueError: If value is not a valid mode 

51 """ 

52 try: 

53 return cls(value.lower()) 

54 except ValueError: 

55 raise ValueError( 

56 f"Invalid template mode: {value}. " 

57 f"Valid modes: {', '.join(m.value for m in cls)}" 

58 ) 

59 

60 

61@dataclass 

62class ValidationConfig: 

63 """Configuration for template parameter validation. 

64 

65 Attributes: 

66 level: Validation strictness level (None to inherit from context) 

67 required_params: Set of parameter names that must be provided 

68 optional_params: Set of parameter names that are optional 

69 """ 

70 level: Optional[ValidationLevel] = None 

71 required_params: set[str] = field(default_factory=set) 

72 optional_params: set[str] = field(default_factory=set) 

73 

74 def __init__( 

75 self, 

76 level: Optional[ValidationLevel] = None, 

77 required_params: Optional[List[str]] = None, 

78 optional_params: Optional[List[str]] = None 

79 ): 

80 """Initialize validation configuration. 

81 

82 Args: 

83 level: Validation strictness level (None to inherit from context) 

84 required_params: List of required parameter names 

85 optional_params: List of optional parameter names 

86 """ 

87 self.level = level 

88 self.required_params = set(required_params or []) 

89 self.optional_params = set(optional_params or []) 

90 

91 

92class PromptTemplate(TypedDict, total=False): 

93 """Structure for a prompt template definition. 

94 

95 Attributes: 

96 template: The template string with {{variables}} and ((conditionals)) 

97 defaults: Default values for template parameters 

98 validation: Validation configuration for this template 

99 metadata: Additional metadata (author, version, etc.) 

100 sections: Section definitions for template composition 

101 extends: Name of base template to inherit from 

102 rag_config_refs: References to standalone RAG configurations 

103 rag_configs: Inline RAG configurations 

104 template_mode: Template rendering mode ("mixed" or "jinja2") 

105 """ 

106 template: str 

107 defaults: Dict[str, Any] 

108 validation: ValidationConfig 

109 metadata: Dict[str, Any] 

110 sections: Dict[str, str] 

111 extends: str 

112 rag_config_refs: List[str] 

113 rag_configs: List['RAGConfig'] 

114 template_mode: str 

115 

116 

117class RAGConfig(TypedDict, total=False): 

118 """Configuration for RAG (Retrieval-Augmented Generation) searches. 

119 

120 Attributes: 

121 adapter_name: Name of the adapter to use for RAG search 

122 query: RAG search query (may contain {{variables}}) 

123 k: Number of results to retrieve 

124 filters: Additional filters for the search 

125 placeholder: Placeholder name in template (e.g., "RAG_CONTENT") 

126 header: Header text to prepend to RAG results 

127 item_template: Template for formatting individual RAG items 

128 """ 

129 adapter_name: str 

130 query: str 

131 k: int 

132 filters: Dict[str, Any] 

133 placeholder: str 

134 header: str 

135 item_template: str 

136 

137 

138class MessageIndex(TypedDict, total=False): 

139 """Structure for a message index definition. 

140 

141 Message indexes define sequences of messages with roles and content, 

142 including support for RAG content injection. 

143 

144 Attributes: 

145 messages: List of message dictionaries with 'role' and 'content' 

146 rag_configs: RAG configurations for this message sequence 

147 metadata: Additional metadata for this message index 

148 """ 

149 messages: List[Dict[str, str]] 

150 rag_configs: List[RAGConfig] 

151 metadata: Dict[str, Any] 

152 

153 

154@dataclass 

155class RenderResult: 

156 """Result of rendering a prompt template. 

157 

158 Attributes: 

159 content: The final rendered content 

160 params_used: Parameters that were actually used in rendering 

161 params_missing: Required parameters that were missing 

162 validation_warnings: List of validation warnings (if level=WARN) 

163 metadata: Additional metadata about the rendering 

164 rag_metadata: Optional RAG metadata (if return_rag_metadata=True) 

165 Contains details about RAG searches executed during rendering 

166 """ 

167 content: str 

168 params_used: Dict[str, Any] = field(default_factory=dict) 

169 params_missing: List[str] = field(default_factory=list) 

170 validation_warnings: List[str] = field(default_factory=list) 

171 metadata: Dict[str, Any] = field(default_factory=dict) 

172 rag_metadata: Optional[Dict[str, Any]] = None 

173 

174 

175# Type aliases for convenience 

176TemplateDict = Dict[str, PromptTemplate] 

177MessageIndexDict = Dict[str, MessageIndex] 

178ParameterDict = Dict[str, Any] 

179AdapterDict = Dict[str, Any] # Will be refined in adapter modules