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

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 as e: 

85 logger.exception( 

86 f"Failed to initialize Jinja2 environment: {e}" 

87 ) 

88 return None 

89 

90 return cls._jinja_env 

91 

92 @classmethod 

93 def format( 

94 cls, 

95 event_type: EventType, 

96 context: Dict[str, Any], 

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

98 ) -> Dict[str, str]: 

99 """ 

100 Format a notification template with context data using Jinja2. 

101 

102 Args: 

103 event_type: Type of event 

104 context: Context data for template formatting 

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

106 

107 Returns: 

108 Dict with 'title' and 'body' keys 

109 """ 

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

111 if custom_template: 111 ↛ 112line 111 didn't jump to line 112 because the condition on line 111 was never true

112 try: 

113 title = custom_template["title"].format(**context) 

114 body = custom_template["body"].format(**context) 

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

116 except KeyError as e: 

117 return { 

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

119 "body": f"Template error: Missing variable {e}. Context: {context}", 

120 } 

121 

122 # Get template filename for this event type 

123 template_file = cls.TEMPLATE_FILES.get(event_type) 

124 if not template_file: 124 ↛ 125line 124 didn't jump to line 125 because the condition on line 124 was never true

125 return { 

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

127 "body": str(context), 

128 } 

129 

130 # Try to render with Jinja2 

131 jinja_env = cls._get_jinja_env() 

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

133 # Fallback to simple format if Jinja2 failed to initialize 

134 logger.warning( 

135 "Jinja2 not available, using fallback template format" 

136 ) 

137 return cls._get_fallback_template(event_type, context) 

138 

139 try: 

140 template = jinja_env.get_template(template_file) 

141 rendered_content = template.render(**context) 

142 

143 # Parse the rendered content into title and body 

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

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

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

147 

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

149 except Exception as e: 

150 logger.exception( 

151 f"Error rendering Jinja2 template {template_file}: {e}" 

152 ) 

153 # Fall back to simple format 

154 return cls._get_fallback_template(event_type, context) 

155 

156 @classmethod 

157 def _get_fallback_template( 

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

159 ) -> Dict[str, str]: 

160 """ 

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

162 

163 Args: 

164 event_type: Type of event 

165 context: Context data 

166 

167 Returns: 

168 Dict with 'title' and 'body' keys 

169 """ 

170 # Generic fallback that works for all event types 

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

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

173 

174 return { 

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

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

177 } 

178 

179 @classmethod 

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

181 """ 

182 Get required context variables for an event type. 

183 

184 Args: 

185 event_type: Type of event 

186 

187 Returns: 

188 List of required variable names 

189 """ 

190 template_file = cls.TEMPLATE_FILES.get(event_type) 

191 if not template_file: 

192 return [] 

193 

194 # Try to parse Jinja2 template to get variables 

195 jinja_env = cls._get_jinja_env() 

196 if jinja_env is None: 

197 # No Jinja2 environment available, return empty list 

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

199 return [] 

200 

201 try: 

202 template = jinja_env.get_template(template_file) 

203 # Get variables from the parsed Jinja2 template 

204 variables = set() 

205 if hasattr(template, "environment"): 

206 # Use Jinja2's meta module to find variables 

207 from jinja2 import meta 

208 

209 template_source = template.environment.loader.get_source( 

210 jinja_env, template_file 

211 )[0] 

212 ast = jinja_env.parse(template_source) 

213 variables = meta.find_undeclared_variables(ast) 

214 else: 

215 # Fallback: extract from template source 

216 template_source = template.environment.loader.get_source( 

217 jinja_env, template_file 

218 )[0] 

219 import re 

220 

221 variables.update( 

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

223 ) 

224 

225 return sorted(variables) 

226 except Exception as e: 

227 logger.exception( 

228 f"Error parsing Jinja2 template {template_file}: {e}" 

229 ) 

230 return []