Coverage for src / local_deep_research / notifications / templates.py: 79%

1""" 

2Notification templates for different event types. 

3""" 

4 

5from enum import Enum 

6from pathlib import Path 

7from typing import Dict, Any, Optional 

8 

9from jinja2 import Environment, FileSystemLoader, select_autoescape 

10from loguru import logger 

11 

12 

13class EventType(Enum): 

14 """Types of events that can trigger notifications.""" 

15 

16 # Research events 

17 RESEARCH_COMPLETED = "research_completed" 

18 RESEARCH_FAILED = "research_failed" 

19 RESEARCH_QUEUED = "research_queued" 

20 

21 # Subscription events 

22 SUBSCRIPTION_UPDATE = "subscription_update" 

23 SUBSCRIPTION_ERROR = "subscription_error" 

24 

25 # System events 

26 RATE_LIMIT_WARNING = "rate_limit_warning" 

27 API_QUOTA_WARNING = "api_quota_warning" 

28 AUTH_ISSUE = "auth_issue" 

29 

30 # Test event 

31 TEST = "test" 

32 

33 

34class NotificationTemplate: 

35 """ 

36 Manages notification message templates using Jinja2. 

37 

38 Uses Jinja2 templates for consistency with the rest of the project. 

39 Templates are stored in the notifications/templates/ directory. 

40 """ 

41 

42 # Map event types to template filenames 

43 TEMPLATE_FILES: Dict[EventType, str] = { 

44 EventType.RESEARCH_COMPLETED: "research_completed.jinja2", 

45 EventType.RESEARCH_FAILED: "research_failed.jinja2", 

46 EventType.RESEARCH_QUEUED: "research_queued.jinja2", 

47 EventType.SUBSCRIPTION_UPDATE: "subscription_update.jinja2", 

48 EventType.SUBSCRIPTION_ERROR: "subscription_error.jinja2", 

49 EventType.API_QUOTA_WARNING: "api_quota_warning.jinja2", 

50 EventType.AUTH_ISSUE: "auth_issue.jinja2", 

51 EventType.TEST: "test.jinja2", 

52 } 

53 

54 # Shared Jinja2 environment for all template rendering 

55 _jinja_env: Optional[Environment] = None 

56 

57 @classmethod 

58 def _get_jinja_env(cls) -> Environment: 

59 """ 

60 Get or create the Jinja2 environment. 

61 

62 Returns: 

63 Jinja2 Environment configured for notification templates 

64 """ 

65 if cls._jinja_env is None: 

66 # Get the templates directory relative to this file 

67 template_dir = Path(__file__).parent / "templates" 

68 

69 if not template_dir.exists(): 69 ↛ 70line 69 didn't jump to line 70 because the condition on line 69 was never true

70 logger.warning(f"Templates directory not found: {template_dir}") 

71 # Fall back to simple string formatting 

72 return None 

73 

74 try: 

75 cls._jinja_env = Environment( 

76 loader=FileSystemLoader(template_dir), 

77 autoescape=select_autoescape(["html", "xml"]), 

78 trim_blocks=True, 

79 lstrip_blocks=True, 

80 ) 

81 logger.debug( 

82 f"Jinja2 environment initialized with templates from: {template_dir}" 

83 ) 

84 except Exception: 

85 logger.exception("Failed to initialize Jinja2 environment") 

86 return None 

87 

88 return cls._jinja_env 

89 

90 @classmethod 

91 def format( 

92 cls, 

93 event_type: EventType, 

94 context: Dict[str, Any], 

95 custom_template: Optional[Dict[str, str]] = None, 

96 ) -> Dict[str, str]: 

97 """ 

98 Format a notification template with context data using Jinja2. 

99 

100 Args: 

101 event_type: Type of event 

102 context: Context data for template formatting 

103 custom_template: Optional custom template override (still uses string format) 

104 

105 Returns: 

106 Dict with 'title' and 'body' keys 

107 """ 

108 # If custom template is provided, use the old string format for backward compatibility 

109 # Security: Sanitize context values to prevent format string attacks 

110 if custom_template: 

111 try: 

112 # Only allow simple string substitution - convert all values to strings 

113 # This prevents format string attacks like {query.__class__} 

114 safe_context = { 

115 k: str(v) if v is not None else "" 

116 for k, v in context.items() 

117 } 

118 title = custom_template["title"].format(**safe_context) 

119 body = custom_template["body"].format(**safe_context) 

120 return {"title": title, "body": body} 

121 except KeyError as e: 

122 return { 

123 "title": f"Notification: {event_type.value}", 

124 "body": f"Template error: Missing variable {e}", 

125 } 

126 

127 # Get template filename for this event type 

128 template_file = cls.TEMPLATE_FILES.get(event_type) 

129 if not template_file: 

130 return { 

131 "title": f"Notification: {event_type.value}", 

132 "body": str(context), 

133 } 

134 

135 # Try to render with Jinja2 

136 jinja_env = cls._get_jinja_env() 

137 if jinja_env is None: 137 ↛ 139line 137 didn't jump to line 139 because the condition on line 137 was never true

138 # Fallback to simple format if Jinja2 failed to initialize 

139 logger.warning( 

140 "Jinja2 not available, using fallback template format" 

141 ) 

142 return cls._get_fallback_template(event_type, context) 

143 

144 try: 

145 template = jinja_env.get_template(template_file) 

146 rendered_content = template.render(**context) 

147 

148 # Parse the rendered content into title and body 

149 lines = rendered_content.strip().split("\n") 

150 title = lines[0].strip() if lines else "Notification" 

151 body = "\n".join(lines[1:]).strip() if len(lines) > 1 else "" 

152 

153 return {"title": title, "body": body} 

154 except Exception: 

155 logger.exception(f"Error rendering Jinja2 template {template_file}") 

156 # Fall back to simple format 

157 return cls._get_fallback_template(event_type, context) 

158 

159 @classmethod 

160 def _get_fallback_template( 

161 cls, event_type: EventType, context: Dict[str, Any] 

162 ) -> Dict[str, str]: 

163 """ 

164 Get a simple fallback template when Jinja2 is not available. 

165 

166 Args: 

167 event_type: Type of event 

168 context: Context data 

169 

170 Returns: 

171 Dict with 'title' and 'body' keys 

172 """ 

173 # Generic fallback that works for all event types 

174 # Reduces maintenance burden - no need to update for new events 

175 event_display_name = event_type.value.replace("_", " ").title() 

176 

177 return { 

178 "title": f"Notification: {event_display_name}", 

179 "body": f"Event: {event_display_name}\n\nDetails: {str(context)}\n\nPlease check the application for complete information.", 

180 } 

181 

182 @classmethod 

183 def get_required_context(cls, event_type: EventType) -> list[str]: 

184 """ 

185 Get required context variables for an event type. 

186 

187 Args: 

188 event_type: Type of event 

189 

190 Returns: 

191 List of required variable names 

192 """ 

193 template_file = cls.TEMPLATE_FILES.get(event_type) 

194 if not template_file: 

195 return [] 

196 

197 # Try to parse Jinja2 template to get variables 

198 jinja_env = cls._get_jinja_env() 

199 if jinja_env is None: 199 ↛ 202line 199 didn't jump to line 202 because the condition on line 199 was never true

200 # No Jinja2 environment available, return empty list 

201 # With simplified fallback approach, we don't need to track required variables 

202 return [] 

203 

204 try: 

205 template = jinja_env.get_template(template_file) 

206 # Get variables from the parsed Jinja2 template 

207 variables = set() 

208 if hasattr(template, "environment"): 208 ↛ 219line 208 didn't jump to line 219 because the condition on line 208 was always true

209 # Use Jinja2's meta module to find variables 

210 from jinja2 import meta 

211 

212 template_source = template.environment.loader.get_source( 

213 jinja_env, template_file 

214 )[0] 

215 ast = jinja_env.parse(template_source) 

216 variables = meta.find_undeclared_variables(ast) 

217 else: 

218 # Fallback: extract from template source 

219 template_source = template.environment.loader.get_source( 

220 jinja_env, template_file 

221 )[0] 

222 import re 

223 

224 variables.update( 

225 re.findall(r"\{\{\s*(\w+)\s*\}\}", template_source) 

226 ) 

227 

228 return sorted(variables) 

229 except Exception: 

230 logger.exception(f"Error parsing Jinja2 template {template_file}") 

231 return []