Skip to main content

systemprompt_analytics/services/
throttle.rs

1use chrono::{DateTime, Duration, Utc};
2use serde::{Deserialize, Serialize};
3
4const BEHAVIORAL_BOT_SCORE_THRESHOLD: i32 = 50;
5const HIGH_REQUESTS_PER_MINUTE_THRESHOLD: f64 = 30.0;
6const HIGH_ERROR_RATE_THRESHOLD: f64 = 0.5;
7const MIN_REQUESTS_FOR_ERROR_ESCALATION: i64 = 20;
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
10#[repr(i32)]
11pub enum ThrottleLevel {
12    Normal = 0,
13    Warning = 1,
14    Severe = 2,
15    Blocked = 3,
16}
17
18impl From<i32> for ThrottleLevel {
19    fn from(value: i32) -> Self {
20        match value {
21            1 => Self::Warning,
22            2 => Self::Severe,
23            3 => Self::Blocked,
24            _ => Self::Normal,
25        }
26    }
27}
28
29impl From<ThrottleLevel> for i32 {
30    fn from(level: ThrottleLevel) -> Self {
31        level as Self
32    }
33}
34
35impl ThrottleLevel {
36    pub const fn rate_multiplier(self) -> f64 {
37        match self {
38            Self::Normal => 1.0,
39            Self::Warning => 0.5,
40            Self::Severe => 0.25,
41            Self::Blocked => 0.0,
42        }
43    }
44
45    pub const fn allows_requests(self) -> bool {
46        !matches!(self, Self::Blocked)
47    }
48
49    pub const fn escalate(self) -> Self {
50        match self {
51            Self::Normal => Self::Warning,
52            Self::Warning => Self::Severe,
53            Self::Severe | Self::Blocked => Self::Blocked,
54        }
55    }
56
57    pub const fn deescalate(self) -> Self {
58        match self {
59            Self::Normal | Self::Warning => Self::Normal,
60            Self::Severe => Self::Warning,
61            Self::Blocked => Self::Severe,
62        }
63    }
64}
65
66#[derive(Debug, Clone)]
67pub struct EscalationCriteria {
68    pub behavioral_bot_score: i32,
69    pub request_count: i64,
70    pub error_rate: f64,
71    pub requests_per_minute: f64,
72}
73
74impl Copy for EscalationCriteria {}
75
76#[derive(Debug, Clone, Copy, Default)]
77pub struct ThrottleService;
78
79impl ThrottleService {
80    pub const fn new() -> Self {
81        Self
82    }
83
84    pub fn should_escalate(criteria: &EscalationCriteria, current_level: ThrottleLevel) -> bool {
85        if current_level == ThrottleLevel::Blocked {
86            return false;
87        }
88
89        if criteria.behavioral_bot_score >= BEHAVIORAL_BOT_SCORE_THRESHOLD {
90            return true;
91        }
92
93        if criteria.requests_per_minute > HIGH_REQUESTS_PER_MINUTE_THRESHOLD {
94            return true;
95        }
96
97        if criteria.error_rate > HIGH_ERROR_RATE_THRESHOLD
98            && criteria.request_count > MIN_REQUESTS_FOR_ERROR_ESCALATION
99        {
100            return true;
101        }
102
103        false
104    }
105
106    pub fn adjusted_rate_limit(base_rate: u64, level: ThrottleLevel) -> u64 {
107        let multiplier = level.rate_multiplier();
108        ((base_rate as f64) * multiplier).max(1.0) as u64
109    }
110
111    pub fn can_deescalate(
112        current_level: ThrottleLevel,
113        last_escalation: Option<DateTime<Utc>>,
114        cooldown_minutes: i64,
115    ) -> bool {
116        if current_level == ThrottleLevel::Normal {
117            return false;
118        }
119
120        last_escalation.is_none_or(|escalated_at| {
121            Utc::now() > escalated_at + Duration::minutes(cooldown_minutes)
122        })
123    }
124}