usages.rs

  1use crate::db::UserId;
  2use crate::llm::Cents;
  3use chrono::{Datelike, Duration};
  4use futures::StreamExt as _;
  5use rpc::LanguageModelProvider;
  6use sea_orm::QuerySelect;
  7use std::{iter, str::FromStr};
  8use strum::IntoEnumIterator as _;
  9
 10use super::*;
 11
 12#[derive(Debug, PartialEq, Clone, Copy, Default, serde::Serialize)]
 13pub struct TokenUsage {
 14    pub input: usize,
 15    pub input_cache_creation: usize,
 16    pub input_cache_read: usize,
 17    pub output: usize,
 18}
 19
 20impl TokenUsage {
 21    pub fn total(&self) -> usize {
 22        self.input + self.input_cache_creation + self.input_cache_read + self.output
 23    }
 24}
 25
 26#[derive(Debug, PartialEq, Clone, Copy, serde::Serialize)]
 27pub struct Usage {
 28    pub requests_this_minute: usize,
 29    pub tokens_this_minute: usize,
 30    pub input_tokens_this_minute: usize,
 31    pub output_tokens_this_minute: usize,
 32    pub tokens_this_day: usize,
 33    pub tokens_this_month: TokenUsage,
 34    pub spending_this_month: Cents,
 35    pub lifetime_spending: Cents,
 36}
 37
 38#[derive(Debug, PartialEq, Clone)]
 39pub struct ApplicationWideUsage {
 40    pub provider: LanguageModelProvider,
 41    pub model: String,
 42    pub requests_this_minute: usize,
 43    pub tokens_this_minute: usize,
 44    pub input_tokens_this_minute: usize,
 45    pub output_tokens_this_minute: usize,
 46}
 47
 48#[derive(Clone, Copy, Debug, Default)]
 49pub struct ActiveUserCount {
 50    pub users_in_recent_minutes: usize,
 51    pub users_in_recent_days: usize,
 52}
 53
 54impl LlmDatabase {
 55    pub async fn initialize_usage_measures(&mut self) -> Result<()> {
 56        let all_measures = self
 57            .transaction(|tx| async move {
 58                let existing_measures = usage_measure::Entity::find().all(&*tx).await?;
 59
 60                let new_measures = UsageMeasure::iter()
 61                    .filter(|measure| {
 62                        !existing_measures
 63                            .iter()
 64                            .any(|m| m.name == measure.to_string())
 65                    })
 66                    .map(|measure| usage_measure::ActiveModel {
 67                        name: ActiveValue::set(measure.to_string()),
 68                        ..Default::default()
 69                    })
 70                    .collect::<Vec<_>>();
 71
 72                if !new_measures.is_empty() {
 73                    usage_measure::Entity::insert_many(new_measures)
 74                        .exec(&*tx)
 75                        .await?;
 76                }
 77
 78                Ok(usage_measure::Entity::find().all(&*tx).await?)
 79            })
 80            .await?;
 81
 82        self.usage_measure_ids = all_measures
 83            .into_iter()
 84            .filter_map(|measure| {
 85                UsageMeasure::from_str(&measure.name)
 86                    .ok()
 87                    .map(|um| (um, measure.id))
 88            })
 89            .collect();
 90        Ok(())
 91    }
 92
 93    pub async fn get_application_wide_usages_by_model(
 94        &self,
 95        now: DateTimeUtc,
 96    ) -> Result<Vec<ApplicationWideUsage>> {
 97        self.transaction(|tx| async move {
 98            let past_minute = now - Duration::minutes(1);
 99            let requests_per_minute = self.usage_measure_ids[&UsageMeasure::RequestsPerMinute];
100            let tokens_per_minute = self.usage_measure_ids[&UsageMeasure::TokensPerMinute];
101            let input_tokens_per_minute =
102                self.usage_measure_ids[&UsageMeasure::InputTokensPerMinute];
103            let output_tokens_per_minute =
104                self.usage_measure_ids[&UsageMeasure::OutputTokensPerMinute];
105
106            let mut results = Vec::new();
107            for ((provider, model_name), model) in self.models.iter() {
108                let mut usages = usage::Entity::find()
109                    .filter(
110                        usage::Column::Timestamp
111                            .gte(past_minute.naive_utc())
112                            .and(usage::Column::IsStaff.eq(false))
113                            .and(usage::Column::ModelId.eq(model.id))
114                            .and(
115                                usage::Column::MeasureId
116                                    .eq(requests_per_minute)
117                                    .or(usage::Column::MeasureId.eq(tokens_per_minute)),
118                            ),
119                    )
120                    .stream(&*tx)
121                    .await?;
122
123                let mut requests_this_minute = 0;
124                let mut tokens_this_minute = 0;
125                let mut input_tokens_this_minute = 0;
126                let mut output_tokens_this_minute = 0;
127                while let Some(usage) = usages.next().await {
128                    let usage = usage?;
129                    if usage.measure_id == requests_per_minute {
130                        requests_this_minute += Self::get_live_buckets(
131                            &usage,
132                            now.naive_utc(),
133                            UsageMeasure::RequestsPerMinute,
134                        )
135                        .0
136                        .iter()
137                        .copied()
138                        .sum::<i64>() as usize;
139                    } else if usage.measure_id == tokens_per_minute {
140                        tokens_this_minute += Self::get_live_buckets(
141                            &usage,
142                            now.naive_utc(),
143                            UsageMeasure::TokensPerMinute,
144                        )
145                        .0
146                        .iter()
147                        .copied()
148                        .sum::<i64>() as usize;
149                    } else if usage.measure_id == input_tokens_per_minute {
150                        input_tokens_this_minute += Self::get_live_buckets(
151                            &usage,
152                            now.naive_utc(),
153                            UsageMeasure::InputTokensPerMinute,
154                        )
155                        .0
156                        .iter()
157                        .copied()
158                        .sum::<i64>() as usize;
159                    } else if usage.measure_id == output_tokens_per_minute {
160                        output_tokens_this_minute += Self::get_live_buckets(
161                            &usage,
162                            now.naive_utc(),
163                            UsageMeasure::OutputTokensPerMinute,
164                        )
165                        .0
166                        .iter()
167                        .copied()
168                        .sum::<i64>() as usize;
169                    }
170                }
171
172                results.push(ApplicationWideUsage {
173                    provider: *provider,
174                    model: model_name.clone(),
175                    requests_this_minute,
176                    tokens_this_minute,
177                    input_tokens_this_minute,
178                    output_tokens_this_minute,
179                })
180            }
181
182            Ok(results)
183        })
184        .await
185    }
186
187    pub async fn get_user_spending_for_month(
188        &self,
189        user_id: UserId,
190        now: DateTimeUtc,
191    ) -> Result<Cents> {
192        self.transaction(|tx| async move {
193            let month = now.date_naive().month() as i32;
194            let year = now.date_naive().year();
195
196            let mut monthly_usages = monthly_usage::Entity::find()
197                .filter(
198                    monthly_usage::Column::UserId
199                        .eq(user_id)
200                        .and(monthly_usage::Column::Month.eq(month))
201                        .and(monthly_usage::Column::Year.eq(year)),
202                )
203                .stream(&*tx)
204                .await?;
205            let mut monthly_spending = Cents::ZERO;
206
207            while let Some(usage) = monthly_usages.next().await {
208                let usage = usage?;
209                let Ok(model) = self.model_by_id(usage.model_id) else {
210                    continue;
211                };
212
213                monthly_spending += calculate_spending(
214                    model,
215                    usage.input_tokens as usize,
216                    usage.cache_creation_input_tokens as usize,
217                    usage.cache_read_input_tokens as usize,
218                    usage.output_tokens as usize,
219                );
220            }
221
222            Ok(monthly_spending)
223        })
224        .await
225    }
226
227    pub async fn get_usage(
228        &self,
229        user_id: UserId,
230        provider: LanguageModelProvider,
231        model_name: &str,
232        now: DateTimeUtc,
233    ) -> Result<Usage> {
234        self.transaction(|tx| async move {
235            let model = self
236                .models
237                .get(&(provider, model_name.to_string()))
238                .ok_or_else(|| anyhow!("unknown model {provider}:{model_name}"))?;
239
240            let usages = usage::Entity::find()
241                .filter(
242                    usage::Column::UserId
243                        .eq(user_id)
244                        .and(usage::Column::ModelId.eq(model.id)),
245                )
246                .all(&*tx)
247                .await?;
248
249            let month = now.date_naive().month() as i32;
250            let year = now.date_naive().year();
251            let monthly_usage = monthly_usage::Entity::find()
252                .filter(
253                    monthly_usage::Column::UserId
254                        .eq(user_id)
255                        .and(monthly_usage::Column::ModelId.eq(model.id))
256                        .and(monthly_usage::Column::Month.eq(month))
257                        .and(monthly_usage::Column::Year.eq(year)),
258                )
259                .one(&*tx)
260                .await?;
261            let lifetime_usage = lifetime_usage::Entity::find()
262                .filter(
263                    lifetime_usage::Column::UserId
264                        .eq(user_id)
265                        .and(lifetime_usage::Column::ModelId.eq(model.id)),
266                )
267                .one(&*tx)
268                .await?;
269
270            let requests_this_minute =
271                self.get_usage_for_measure(&usages, now, UsageMeasure::RequestsPerMinute)?;
272            let tokens_this_minute =
273                self.get_usage_for_measure(&usages, now, UsageMeasure::TokensPerMinute)?;
274            let input_tokens_this_minute =
275                self.get_usage_for_measure(&usages, now, UsageMeasure::InputTokensPerMinute)?;
276            let output_tokens_this_minute =
277                self.get_usage_for_measure(&usages, now, UsageMeasure::OutputTokensPerMinute)?;
278            let tokens_this_day =
279                self.get_usage_for_measure(&usages, now, UsageMeasure::TokensPerDay)?;
280            let spending_this_month = if let Some(monthly_usage) = &monthly_usage {
281                calculate_spending(
282                    model,
283                    monthly_usage.input_tokens as usize,
284                    monthly_usage.cache_creation_input_tokens as usize,
285                    monthly_usage.cache_read_input_tokens as usize,
286                    monthly_usage.output_tokens as usize,
287                )
288            } else {
289                Cents::ZERO
290            };
291            let lifetime_spending = if let Some(lifetime_usage) = &lifetime_usage {
292                calculate_spending(
293                    model,
294                    lifetime_usage.input_tokens as usize,
295                    lifetime_usage.cache_creation_input_tokens as usize,
296                    lifetime_usage.cache_read_input_tokens as usize,
297                    lifetime_usage.output_tokens as usize,
298                )
299            } else {
300                Cents::ZERO
301            };
302
303            Ok(Usage {
304                requests_this_minute,
305                tokens_this_minute,
306                input_tokens_this_minute,
307                output_tokens_this_minute,
308                tokens_this_day,
309                tokens_this_month: TokenUsage {
310                    input: monthly_usage
311                        .as_ref()
312                        .map_or(0, |usage| usage.input_tokens as usize),
313                    input_cache_creation: monthly_usage
314                        .as_ref()
315                        .map_or(0, |usage| usage.cache_creation_input_tokens as usize),
316                    input_cache_read: monthly_usage
317                        .as_ref()
318                        .map_or(0, |usage| usage.cache_read_input_tokens as usize),
319                    output: monthly_usage
320                        .as_ref()
321                        .map_or(0, |usage| usage.output_tokens as usize),
322                },
323                spending_this_month,
324                lifetime_spending,
325            })
326        })
327        .await
328    }
329
330    pub async fn record_usage(
331        &self,
332        user_id: UserId,
333        is_staff: bool,
334        provider: LanguageModelProvider,
335        model_name: &str,
336        tokens: TokenUsage,
337        has_llm_subscription: bool,
338        max_monthly_spend: Cents,
339        free_tier_monthly_spending_limit: Cents,
340        now: DateTimeUtc,
341    ) -> Result<Usage> {
342        self.transaction(|tx| async move {
343            let model = self.model(provider, model_name)?;
344
345            let usages = usage::Entity::find()
346                .filter(
347                    usage::Column::UserId
348                        .eq(user_id)
349                        .and(usage::Column::ModelId.eq(model.id)),
350                )
351                .all(&*tx)
352                .await?;
353
354            let requests_this_minute = self
355                .update_usage_for_measure(
356                    user_id,
357                    is_staff,
358                    model.id,
359                    &usages,
360                    UsageMeasure::RequestsPerMinute,
361                    now,
362                    1,
363                    &tx,
364                )
365                .await?;
366            let tokens_this_minute = self
367                .update_usage_for_measure(
368                    user_id,
369                    is_staff,
370                    model.id,
371                    &usages,
372                    UsageMeasure::TokensPerMinute,
373                    now,
374                    tokens.total(),
375                    &tx,
376                )
377                .await?;
378            let input_tokens_this_minute = self
379                .update_usage_for_measure(
380                    user_id,
381                    is_staff,
382                    model.id,
383                    &usages,
384                    UsageMeasure::InputTokensPerMinute,
385                    now,
386                    // Cache read input tokens are not counted for the purposes of rate limits (but they are still billed).
387                    tokens.input + tokens.input_cache_creation,
388                    &tx,
389                )
390                .await?;
391            let output_tokens_this_minute = self
392                .update_usage_for_measure(
393                    user_id,
394                    is_staff,
395                    model.id,
396                    &usages,
397                    UsageMeasure::OutputTokensPerMinute,
398                    now,
399                    tokens.output,
400                    &tx,
401                )
402                .await?;
403            let tokens_this_day = self
404                .update_usage_for_measure(
405                    user_id,
406                    is_staff,
407                    model.id,
408                    &usages,
409                    UsageMeasure::TokensPerDay,
410                    now,
411                    tokens.total(),
412                    &tx,
413                )
414                .await?;
415
416            let month = now.date_naive().month() as i32;
417            let year = now.date_naive().year();
418
419            // Update monthly usage
420            let monthly_usage = monthly_usage::Entity::find()
421                .filter(
422                    monthly_usage::Column::UserId
423                        .eq(user_id)
424                        .and(monthly_usage::Column::ModelId.eq(model.id))
425                        .and(monthly_usage::Column::Month.eq(month))
426                        .and(monthly_usage::Column::Year.eq(year)),
427                )
428                .one(&*tx)
429                .await?;
430
431            let monthly_usage = match monthly_usage {
432                Some(usage) => {
433                    monthly_usage::Entity::update(monthly_usage::ActiveModel {
434                        id: ActiveValue::unchanged(usage.id),
435                        input_tokens: ActiveValue::set(usage.input_tokens + tokens.input as i64),
436                        cache_creation_input_tokens: ActiveValue::set(
437                            usage.cache_creation_input_tokens + tokens.input_cache_creation as i64,
438                        ),
439                        cache_read_input_tokens: ActiveValue::set(
440                            usage.cache_read_input_tokens + tokens.input_cache_read as i64,
441                        ),
442                        output_tokens: ActiveValue::set(usage.output_tokens + tokens.output as i64),
443                        ..Default::default()
444                    })
445                    .exec(&*tx)
446                    .await?
447                }
448                None => {
449                    monthly_usage::ActiveModel {
450                        user_id: ActiveValue::set(user_id),
451                        model_id: ActiveValue::set(model.id),
452                        month: ActiveValue::set(month),
453                        year: ActiveValue::set(year),
454                        input_tokens: ActiveValue::set(tokens.input as i64),
455                        cache_creation_input_tokens: ActiveValue::set(
456                            tokens.input_cache_creation as i64,
457                        ),
458                        cache_read_input_tokens: ActiveValue::set(tokens.input_cache_read as i64),
459                        output_tokens: ActiveValue::set(tokens.output as i64),
460                        ..Default::default()
461                    }
462                    .insert(&*tx)
463                    .await?
464                }
465            };
466
467            let spending_this_month = calculate_spending(
468                model,
469                monthly_usage.input_tokens as usize,
470                monthly_usage.cache_creation_input_tokens as usize,
471                monthly_usage.cache_read_input_tokens as usize,
472                monthly_usage.output_tokens as usize,
473            );
474
475            if !is_staff
476                && spending_this_month > free_tier_monthly_spending_limit
477                && has_llm_subscription
478                && (spending_this_month - free_tier_monthly_spending_limit) <= max_monthly_spend
479            {
480                billing_event::ActiveModel {
481                    id: ActiveValue::not_set(),
482                    idempotency_key: ActiveValue::not_set(),
483                    user_id: ActiveValue::set(user_id),
484                    model_id: ActiveValue::set(model.id),
485                    input_tokens: ActiveValue::set(tokens.input as i64),
486                    input_cache_creation_tokens: ActiveValue::set(
487                        tokens.input_cache_creation as i64,
488                    ),
489                    input_cache_read_tokens: ActiveValue::set(tokens.input_cache_read as i64),
490                    output_tokens: ActiveValue::set(tokens.output as i64),
491                }
492                .insert(&*tx)
493                .await?;
494            }
495
496            // Update lifetime usage
497            let lifetime_usage = lifetime_usage::Entity::find()
498                .filter(
499                    lifetime_usage::Column::UserId
500                        .eq(user_id)
501                        .and(lifetime_usage::Column::ModelId.eq(model.id)),
502                )
503                .one(&*tx)
504                .await?;
505
506            let lifetime_usage = match lifetime_usage {
507                Some(usage) => {
508                    lifetime_usage::Entity::update(lifetime_usage::ActiveModel {
509                        id: ActiveValue::unchanged(usage.id),
510                        input_tokens: ActiveValue::set(usage.input_tokens + tokens.input as i64),
511                        cache_creation_input_tokens: ActiveValue::set(
512                            usage.cache_creation_input_tokens + tokens.input_cache_creation as i64,
513                        ),
514                        cache_read_input_tokens: ActiveValue::set(
515                            usage.cache_read_input_tokens + tokens.input_cache_read as i64,
516                        ),
517                        output_tokens: ActiveValue::set(usage.output_tokens + tokens.output as i64),
518                        ..Default::default()
519                    })
520                    .exec(&*tx)
521                    .await?
522                }
523                None => {
524                    lifetime_usage::ActiveModel {
525                        user_id: ActiveValue::set(user_id),
526                        model_id: ActiveValue::set(model.id),
527                        input_tokens: ActiveValue::set(tokens.input as i64),
528                        cache_creation_input_tokens: ActiveValue::set(
529                            tokens.input_cache_creation as i64,
530                        ),
531                        cache_read_input_tokens: ActiveValue::set(tokens.input_cache_read as i64),
532                        output_tokens: ActiveValue::set(tokens.output as i64),
533                        ..Default::default()
534                    }
535                    .insert(&*tx)
536                    .await?
537                }
538            };
539
540            let lifetime_spending = calculate_spending(
541                model,
542                lifetime_usage.input_tokens as usize,
543                lifetime_usage.cache_creation_input_tokens as usize,
544                lifetime_usage.cache_read_input_tokens as usize,
545                lifetime_usage.output_tokens as usize,
546            );
547
548            Ok(Usage {
549                requests_this_minute,
550                tokens_this_minute,
551                input_tokens_this_minute,
552                output_tokens_this_minute,
553                tokens_this_day,
554                tokens_this_month: TokenUsage {
555                    input: monthly_usage.input_tokens as usize,
556                    input_cache_creation: monthly_usage.cache_creation_input_tokens as usize,
557                    input_cache_read: monthly_usage.cache_read_input_tokens as usize,
558                    output: monthly_usage.output_tokens as usize,
559                },
560                spending_this_month,
561                lifetime_spending,
562            })
563        })
564        .await
565    }
566
567    /// Returns the active user count for the specified model.
568    pub async fn get_active_user_count(
569        &self,
570        provider: LanguageModelProvider,
571        model_name: &str,
572        now: DateTimeUtc,
573    ) -> Result<ActiveUserCount> {
574        self.transaction(|tx| async move {
575            let minute_since = now - Duration::minutes(5);
576            let day_since = now - Duration::days(5);
577
578            let model = self
579                .models
580                .get(&(provider, model_name.to_string()))
581                .ok_or_else(|| anyhow!("unknown model {provider}:{model_name}"))?;
582
583            let tokens_per_minute = self.usage_measure_ids[&UsageMeasure::TokensPerMinute];
584
585            let users_in_recent_minutes = usage::Entity::find()
586                .filter(
587                    usage::Column::ModelId
588                        .eq(model.id)
589                        .and(usage::Column::MeasureId.eq(tokens_per_minute))
590                        .and(usage::Column::Timestamp.gte(minute_since.naive_utc()))
591                        .and(usage::Column::IsStaff.eq(false)),
592                )
593                .select_only()
594                .column(usage::Column::UserId)
595                .group_by(usage::Column::UserId)
596                .count(&*tx)
597                .await? as usize;
598
599            let users_in_recent_days = usage::Entity::find()
600                .filter(
601                    usage::Column::ModelId
602                        .eq(model.id)
603                        .and(usage::Column::MeasureId.eq(tokens_per_minute))
604                        .and(usage::Column::Timestamp.gte(day_since.naive_utc()))
605                        .and(usage::Column::IsStaff.eq(false)),
606                )
607                .select_only()
608                .column(usage::Column::UserId)
609                .group_by(usage::Column::UserId)
610                .count(&*tx)
611                .await? as usize;
612
613            Ok(ActiveUserCount {
614                users_in_recent_minutes,
615                users_in_recent_days,
616            })
617        })
618        .await
619    }
620
621    async fn update_usage_for_measure(
622        &self,
623        user_id: UserId,
624        is_staff: bool,
625        model_id: ModelId,
626        usages: &[usage::Model],
627        usage_measure: UsageMeasure,
628        now: DateTimeUtc,
629        usage_to_add: usize,
630        tx: &DatabaseTransaction,
631    ) -> Result<usize> {
632        let now = now.naive_utc();
633        let measure_id = *self
634            .usage_measure_ids
635            .get(&usage_measure)
636            .ok_or_else(|| anyhow!("usage measure {usage_measure} not found"))?;
637
638        let mut id = None;
639        let mut timestamp = now;
640        let mut buckets = vec![0_i64];
641
642        if let Some(old_usage) = usages.iter().find(|usage| usage.measure_id == measure_id) {
643            id = Some(old_usage.id);
644            let (live_buckets, buckets_since) =
645                Self::get_live_buckets(old_usage, now, usage_measure);
646            if !live_buckets.is_empty() {
647                buckets.clear();
648                buckets.extend_from_slice(live_buckets);
649                buckets.extend(iter::repeat(0).take(buckets_since));
650                timestamp =
651                    old_usage.timestamp + (usage_measure.bucket_duration() * buckets_since as i32);
652            }
653        }
654
655        *buckets.last_mut().unwrap() += usage_to_add as i64;
656        let total_usage = buckets.iter().sum::<i64>() as usize;
657
658        let mut model = usage::ActiveModel {
659            user_id: ActiveValue::set(user_id),
660            is_staff: ActiveValue::set(is_staff),
661            model_id: ActiveValue::set(model_id),
662            measure_id: ActiveValue::set(measure_id),
663            timestamp: ActiveValue::set(timestamp),
664            buckets: ActiveValue::set(buckets),
665            ..Default::default()
666        };
667
668        if let Some(id) = id {
669            model.id = ActiveValue::unchanged(id);
670            model.update(tx).await?;
671        } else {
672            usage::Entity::insert(model)
673                .exec_without_returning(tx)
674                .await?;
675        }
676
677        Ok(total_usage)
678    }
679
680    fn get_usage_for_measure(
681        &self,
682        usages: &[usage::Model],
683        now: DateTimeUtc,
684        usage_measure: UsageMeasure,
685    ) -> Result<usize> {
686        let now = now.naive_utc();
687        let measure_id = *self
688            .usage_measure_ids
689            .get(&usage_measure)
690            .ok_or_else(|| anyhow!("usage measure {usage_measure} not found"))?;
691        let Some(usage) = usages.iter().find(|usage| usage.measure_id == measure_id) else {
692            return Ok(0);
693        };
694
695        let (live_buckets, _) = Self::get_live_buckets(usage, now, usage_measure);
696        Ok(live_buckets.iter().sum::<i64>() as _)
697    }
698
699    fn get_live_buckets(
700        usage: &usage::Model,
701        now: chrono::NaiveDateTime,
702        measure: UsageMeasure,
703    ) -> (&[i64], usize) {
704        let seconds_since_usage = (now - usage.timestamp).num_seconds().max(0);
705        let buckets_since_usage =
706            seconds_since_usage as f32 / measure.bucket_duration().num_seconds() as f32;
707        let buckets_since_usage = buckets_since_usage.ceil() as usize;
708        let mut live_buckets = &[] as &[i64];
709        if buckets_since_usage < measure.bucket_count() {
710            let expired_bucket_count =
711                (usage.buckets.len() + buckets_since_usage).saturating_sub(measure.bucket_count());
712            live_buckets = &usage.buckets[expired_bucket_count..];
713            while live_buckets.first() == Some(&0) {
714                live_buckets = &live_buckets[1..];
715            }
716        }
717        (live_buckets, buckets_since_usage)
718    }
719}
720
721fn calculate_spending(
722    model: &model::Model,
723    input_tokens_this_month: usize,
724    cache_creation_input_tokens_this_month: usize,
725    cache_read_input_tokens_this_month: usize,
726    output_tokens_this_month: usize,
727) -> Cents {
728    let input_token_cost =
729        input_tokens_this_month * model.price_per_million_input_tokens as usize / 1_000_000;
730    let cache_creation_input_token_cost = cache_creation_input_tokens_this_month
731        * model.price_per_million_cache_creation_input_tokens as usize
732        / 1_000_000;
733    let cache_read_input_token_cost = cache_read_input_tokens_this_month
734        * model.price_per_million_cache_read_input_tokens as usize
735        / 1_000_000;
736    let output_token_cost =
737        output_tokens_this_month * model.price_per_million_output_tokens as usize / 1_000_000;
738    let spending = input_token_cost
739        + cache_creation_input_token_cost
740        + cache_read_input_token_cost
741        + output_token_cost;
742    Cents::new(spending as u32)
743}
744
745const MINUTE_BUCKET_COUNT: usize = 12;
746const DAY_BUCKET_COUNT: usize = 48;
747
748impl UsageMeasure {
749    fn bucket_count(&self) -> usize {
750        match self {
751            UsageMeasure::RequestsPerMinute => MINUTE_BUCKET_COUNT,
752            UsageMeasure::TokensPerMinute
753            | UsageMeasure::InputTokensPerMinute
754            | UsageMeasure::OutputTokensPerMinute => MINUTE_BUCKET_COUNT,
755            UsageMeasure::TokensPerDay => DAY_BUCKET_COUNT,
756        }
757    }
758
759    fn total_duration(&self) -> Duration {
760        match self {
761            UsageMeasure::RequestsPerMinute => Duration::minutes(1),
762            UsageMeasure::TokensPerMinute
763            | UsageMeasure::InputTokensPerMinute
764            | UsageMeasure::OutputTokensPerMinute => Duration::minutes(1),
765            UsageMeasure::TokensPerDay => Duration::hours(24),
766        }
767    }
768
769    fn bucket_duration(&self) -> Duration {
770        self.total_duration() / self.bucket_count() as i32
771    }
772}