1use crate::error::TypeError;
2use crate::{
3 dispatch::AlertDispatchType, AlertDispatchConfig, AlertThreshold, CommonCrons,
4 DispatchAlertDescription, OpsGenieDispatchConfig, PyHelperFuncs, SlackDispatchConfig,
5 ValidateAlertConfig,
6};
7use crate::{AlertCondition, AlertMap};
8use core::fmt::Debug;
9use pyo3::prelude::*;
10use pyo3::types::PyString;
11
12use serde::{Deserialize, Serialize};
13use std::collections::HashMap;
14
15#[pyclass]
16#[derive(Debug, Serialize, Deserialize, Clone)]
17pub struct CustomMetric {
18 #[pyo3(get, set)]
19 pub name: String,
20
21 #[pyo3(get, set)]
22 pub baseline_value: f64,
23
24 #[pyo3(get, set)]
25 pub alert_condition: AlertCondition,
26}
27
28#[pymethods]
29impl CustomMetric {
30 #[new]
31 #[pyo3(signature = (name, baseline_value, alert_threshold, delta=None))]
32 pub fn new(
33 name: &str,
34 baseline_value: f64,
35 alert_threshold: AlertThreshold,
36 delta: Option<f64>,
37 ) -> Result<Self, TypeError> {
38 let custom_condition = AlertCondition::new(baseline_value, alert_threshold, delta);
39
40 Ok(Self {
41 name: name.to_lowercase(),
42 baseline_value,
43 alert_condition: custom_condition,
44 })
45 }
46
47 pub fn __str__(&self) -> String {
48 PyHelperFuncs::__str__(self)
50 }
51
52 #[getter]
53 pub fn alert_threshold(&self) -> AlertThreshold {
54 self.alert_condition.alert_threshold.clone()
55 }
56
57 #[getter]
58 pub fn delta(&self) -> Option<f64> {
59 self.alert_condition.delta
60 }
61}
62
63#[pyclass]
64#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
65pub struct CustomMetricAlertConfig {
66 pub dispatch_config: AlertDispatchConfig,
67
68 #[pyo3(get, set)]
69 pub schedule: String,
70
71 #[pyo3(get, set)]
72 pub alert_conditions: Option<HashMap<String, AlertCondition>>,
73}
74
75impl CustomMetricAlertConfig {
76 pub fn set_alert_conditions(&mut self, metrics: &[CustomMetric]) {
77 self.alert_conditions = Some(
78 metrics
79 .iter()
80 .map(|m| (m.name.clone(), m.alert_condition.clone()))
81 .collect(),
82 );
83 }
84}
85
86impl ValidateAlertConfig for CustomMetricAlertConfig {}
87
88#[pymethods]
89impl CustomMetricAlertConfig {
90 #[new]
91 #[pyo3(signature = (schedule=None, dispatch_config=None))]
92 pub fn new(
93 schedule: Option<&Bound<'_, PyAny>>,
94 dispatch_config: Option<&Bound<'_, PyAny>>,
95 ) -> Result<Self, TypeError> {
96 let alert_dispatch_config = match dispatch_config {
97 None => AlertDispatchConfig::default(),
98 Some(config) => {
99 if config.is_instance_of::<SlackDispatchConfig>() {
100 AlertDispatchConfig::Slack(config.extract::<SlackDispatchConfig>()?)
101 } else if config.is_instance_of::<OpsGenieDispatchConfig>() {
102 AlertDispatchConfig::OpsGenie(config.extract::<OpsGenieDispatchConfig>()?)
103 } else {
104 AlertDispatchConfig::default()
105 }
106 }
107 };
108
109 let schedule = match schedule {
110 Some(schedule) => {
111 if schedule.is_instance_of::<PyString>() {
112 schedule.to_string()
113 } else if schedule.is_instance_of::<CommonCrons>() {
114 schedule.extract::<CommonCrons>().unwrap().cron()
115 } else {
116 return Err(TypeError::InvalidScheduleError)?;
117 }
118 }
119 None => CommonCrons::EveryDay.cron(),
120 };
121
122 let schedule = Self::resolve_schedule(&schedule);
123
124 Ok(Self {
125 schedule,
126 dispatch_config: alert_dispatch_config,
127 alert_conditions: None,
128 })
129 }
130
131 #[getter]
132 pub fn dispatch_type(&self) -> AlertDispatchType {
133 self.dispatch_config.dispatch_type()
134 }
135
136 #[getter]
137 pub fn dispatch_config<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
138 self.dispatch_config.config(py)
139 }
140}
141
142impl Default for CustomMetricAlertConfig {
143 fn default() -> CustomMetricAlertConfig {
144 Self {
145 dispatch_config: AlertDispatchConfig::default(),
146 schedule: CommonCrons::EveryDay.cron(),
147 alert_conditions: None,
148 }
149 }
150}
151
152#[derive(Serialize, Deserialize, Debug, Default, Clone)]
153pub struct ComparisonMetricAlert {
154 pub metric_name: String,
155 pub baseline_value: f64,
156 pub observed_value: f64,
157 pub delta: Option<f64>,
158 pub alert_threshold: AlertThreshold,
159}
160
161impl From<ComparisonMetricAlert> for AlertMap {
162 fn from(val: ComparisonMetricAlert) -> Self {
163 AlertMap::Custom(val)
164 }
165}
166
167impl ComparisonMetricAlert {
168 fn alert_description_header(&self) -> String {
169 let below_threshold = |delta: Option<f64>| match delta {
170 Some(b) => format!(
171 "The observed {} metric value has dropped below the threshold (initial value - {})",
172 self.metric_name, b
173 ),
174 None => format!(
175 "The {} metric value has dropped below the initial value",
176 self.metric_name
177 ),
178 };
179
180 let above_threshold = |delta: Option<f64>| match delta {
181 Some(b) => format!(
182 "The {} metric value has increased beyond the threshold (initial value + {})",
183 self.metric_name, b
184 ),
185 None => format!(
186 "The {} metric value has increased beyond the initial value",
187 self.metric_name
188 ),
189 };
190
191 let outside_threshold = |delta: Option<f64>| match delta {
192 Some(b) => format!(
193 "The {} metric value has fallen outside the threshold (initial value ± {})",
194 self.metric_name, b,
195 ),
196 None => format!(
197 "The metric value has fallen outside the initial value for {}",
198 self.metric_name
199 ),
200 };
201
202 match self.alert_threshold {
203 AlertThreshold::Below => below_threshold(self.delta),
204 AlertThreshold::Above => above_threshold(self.delta),
205 AlertThreshold::Outside => outside_threshold(self.delta),
206 }
207 }
208}
209
210impl DispatchAlertDescription for ComparisonMetricAlert {
211 fn create_alert_description(&self, _dispatch_type: AlertDispatchType) -> String {
213 let mut alert_description = String::new();
214 let header = format!("{}\n", self.alert_description_header());
215 alert_description.push_str(&header);
216
217 let current_metric = format!("Current Metric Value: {}\n", self.observed_value);
218 let historical_metric = format!("Initial Metric Value: {}\n", self.baseline_value);
219
220 alert_description.push_str(&historical_metric);
221 alert_description.push_str(¤t_metric);
222
223 alert_description
224 }
225}
226
227#[cfg(test)]
228mod tests {
229 use super::*;
230
231 #[test]
232 fn test_alert_config() {
233 let dispatch_config = AlertDispatchConfig::OpsGenie(OpsGenieDispatchConfig {
235 team: "test-team".to_string(),
236 priority: "P5".to_string(),
237 });
238 let schedule = "0 0 * * * *".to_string();
239 let mut alert_config = CustomMetricAlertConfig {
240 dispatch_config,
241 schedule,
242 ..Default::default()
243 };
244 assert_eq!(alert_config.dispatch_type(), AlertDispatchType::OpsGenie);
245
246 let custom_metrics = vec![
247 CustomMetric::new("mae", 12.4, AlertThreshold::Above, Some(2.3)).unwrap(),
248 CustomMetric::new("accuracy", 0.85, AlertThreshold::Below, None).unwrap(),
249 ];
250
251 alert_config.set_alert_conditions(&custom_metrics);
252
253 if let Some(alert_conditions) = alert_config.alert_conditions.as_ref() {
254 assert_eq!(
255 alert_conditions["mae"].alert_threshold,
256 AlertThreshold::Above
257 );
258 assert_eq!(alert_conditions["mae"].delta, Some(2.3));
259 assert_eq!(
260 alert_conditions["accuracy"].alert_threshold,
261 AlertThreshold::Below
262 );
263 assert_eq!(alert_conditions["accuracy"].delta, None);
264 } else {
265 panic!("alert_conditions should not be None");
266 }
267 }
268
269 #[test]
270 fn test_create_alert_description() {
271 let alert_above_threshold = ComparisonMetricAlert {
272 metric_name: "mse".to_string(),
273 baseline_value: 12.5,
274 observed_value: 14.0,
275 delta: Some(1.0),
276 alert_threshold: AlertThreshold::Above,
277 };
278
279 let description =
280 alert_above_threshold.create_alert_description(AlertDispatchType::Console);
281 assert!(description.contains(
282 "The mse metric value has increased beyond the threshold (initial value + 1)"
283 ));
284 assert!(description.contains("Initial Metric Value: 12.5"));
285 assert!(description.contains("Current Metric Value: 14"));
286
287 let alert_below_threshold = ComparisonMetricAlert {
288 metric_name: "accuracy".to_string(),
289 baseline_value: 0.9,
290 observed_value: 0.7,
291 delta: None,
292 alert_threshold: AlertThreshold::Below,
293 };
294
295 let description =
296 alert_below_threshold.create_alert_description(AlertDispatchType::Console);
297 assert!(
298 description.contains("The accuracy metric value has dropped below the initial value")
299 );
300 assert!(description.contains("Initial Metric Value: 0.9"));
301 assert!(description.contains("Current Metric Value: 0.7"));
302
303 let alert_outside_threshold = ComparisonMetricAlert {
304 metric_name: "mae".to_string(),
305 baseline_value: 12.5,
306 observed_value: 22.0,
307 delta: Some(2.0),
308 alert_threshold: AlertThreshold::Outside,
309 };
310
311 let description =
312 alert_outside_threshold.create_alert_description(AlertDispatchType::Console);
313 assert!(description
314 .contains("The mae metric value has fallen outside the threshold (initial value ± 2)"));
315 assert!(description.contains("Initial Metric Value: 12.5"));
316 assert!(description.contains("Current Metric Value: 22"));
317 }
318}