rate_limits/
reset_time.rs1use crate::convert;
2use crate::error::{Error, Result};
3use headers::HeaderValue;
4use time::format_description::well_known::{Iso8601, Rfc2822};
5use time::{Duration, OffsetDateTime, PrimitiveDateTime};
6
7#[derive(Copy, Clone, Debug, PartialEq)]
14#[non_exhaustive]
15pub enum ResetTimeKind {
16 Seconds,
18 Timestamp,
21 TimestampMillis,
24 ImfFixdate,
26 Iso8601,
28 OpenAIDuration,
30}
31
32#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd)]
37pub enum ResetTime {
38 Seconds(usize),
40 DateTime(OffsetDateTime),
42}
43
44impl ResetTime {
45 pub fn new(value: &HeaderValue, kind: ResetTimeKind) -> Result<Self> {
52 let value = value.to_str()?;
53 match kind {
54 ResetTimeKind::Seconds => Ok(ResetTime::Seconds(convert::to_usize(value)?)),
55 ResetTimeKind::Timestamp => Ok(Self::DateTime(
56 OffsetDateTime::from_unix_timestamp(convert::to_i64(value)?)
57 .map_err(Error::Time)?,
58 )),
59 ResetTimeKind::TimestampMillis => Ok(Self::DateTime(
60 OffsetDateTime::from_unix_timestamp_nanos(convert::to_i128(value)? * 1_000_000)
61 .map_err(Error::Time)?,
62 )),
63 ResetTimeKind::Iso8601 => {
64 let d = PrimitiveDateTime::parse(value, &Iso8601::PARSING).map_err(Error::Parse)?;
66 Ok(ResetTime::DateTime(d.assume_utc()))
67 }
68 ResetTimeKind::ImfFixdate => {
69 let d = PrimitiveDateTime::parse(value, &Rfc2822).map_err(Error::Parse)?;
70 Ok(ResetTime::DateTime(d.assume_utc()))
71 }
72 ResetTimeKind::OpenAIDuration => {
73 let seconds = parse_openai_duration_to_seconds(value)?;
74 Ok(ResetTime::Seconds(seconds))
75 }
76 }
77 }
78
79 #[must_use]
81 pub fn seconds(&self) -> usize {
82 match self {
83 ResetTime::Seconds(s) => *s,
84 #[allow(clippy::cast_possible_truncation)]
88 ResetTime::DateTime(d) => {
89 let diff = *d - OffsetDateTime::now_utc();
90 let seconds = diff.whole_seconds();
91 if seconds < 0 { 0 } else { seconds as usize }
92 }
93 }
94 }
95
96 #[must_use]
98 pub fn duration(&self) -> Duration {
99 match self {
100 ResetTime::Seconds(s) => Duration::seconds(*s as i64),
101 ResetTime::DateTime(d) => {
102 Duration::seconds((*d - OffsetDateTime::now_utc()).whole_seconds())
103 }
104 }
105 }
106}
107
108fn parse_openai_duration_to_seconds(value: &str) -> Result<usize> {
112 let value = value.trim();
113 if value.is_empty() {
114 return Err(Error::InvalidDuration(value.to_string()));
115 }
116
117 let mut total_seconds = 0;
118 let mut current_number = String::new();
119
120 let mut chars = value.chars().peekable();
121
122 while let Some(c) = chars.next() {
123 if c.is_ascii_digit() {
124 current_number.push(c);
125 } else {
126 if current_number.is_empty() {
127 return Err(Error::InvalidDuration(value.to_string()));
128 }
129 let n: usize = current_number
130 .parse()
131 .map_err(|_| Error::InvalidDuration(value.to_string()))?;
132 current_number.clear();
133
134 match c {
135 'd' => total_seconds += n * 86400,
136 'h' => total_seconds += n * 3600,
137 'm' => {
138 if let Some('s') = chars.peek() {
140 chars.next(); if n > 0 {
143 total_seconds += 1;
144 }
145 } else {
146 total_seconds += n * 60;
147 }
148 }
149 's' => total_seconds += n,
150 _ => return Err(Error::InvalidDuration(value.to_string())),
151 }
152 }
153 }
154
155 if !current_number.is_empty() {
156 return Err(Error::InvalidDuration(value.to_string()));
157 }
158
159 Ok(total_seconds)
160}
161
162#[cfg(test)]
163mod tests {
164 use super::*;
165 use headers::HeaderValue;
166
167 #[test]
168 fn test_parse_openai_duration() {
169 assert!(parse_openai_duration_to_seconds("").is_err());
171 assert!(parse_openai_duration_to_seconds("🤖").is_err());
172 assert!(parse_openai_duration_to_seconds("1").is_err());
173 assert!(parse_openai_duration_to_seconds("s").is_err());
174 assert!(parse_openai_duration_to_seconds("1x").is_err());
175 assert!(parse_openai_duration_to_seconds("1m30").is_err());
176 assert!(parse_openai_duration_to_seconds("around 30s").is_err());
177 assert!(parse_openai_duration_to_seconds("1m30s hello").is_err());
178
179 assert_eq!(parse_openai_duration_to_seconds("1s").unwrap(), 1);
180 assert_eq!(parse_openai_duration_to_seconds("1s ").unwrap(), 1);
181 assert_eq!(parse_openai_duration_to_seconds("1m").unwrap(), 60);
182 assert_eq!(parse_openai_duration_to_seconds("1h").unwrap(), 3600);
183 assert_eq!(parse_openai_duration_to_seconds("1d").unwrap(), 86400);
184
185 assert_eq!(parse_openai_duration_to_seconds("1m30s").unwrap(), 90);
187 assert_eq!(parse_openai_duration_to_seconds("1h1m1s").unwrap(), 3661);
188 assert_eq!(parse_openai_duration_to_seconds("6m0s").unwrap(), 360);
189
190 assert_eq!(parse_openai_duration_to_seconds("10ms").unwrap(), 1);
192 assert_eq!(parse_openai_duration_to_seconds("0ms").unwrap(), 0);
193 assert_eq!(parse_openai_duration_to_seconds("1000ms").unwrap(), 1);
194 }
195
196 #[test]
197 fn test_reset_time_new_openai_duration() {
198 let v = HeaderValue::from_str("1h30m").unwrap();
199 let rt = ResetTime::new(&v, ResetTimeKind::OpenAIDuration).unwrap();
200 assert_eq!(rt, ResetTime::Seconds(5400));
201 }
202}