1use crate::db::UserId;
2use crate::llm::Cents;
3use chrono::Datelike;
4use futures::StreamExt as _;
5use std::str::FromStr;
6use strum::IntoEnumIterator as _;
7
8use super::*;
9
10impl LlmDatabase {
11 pub async fn initialize_usage_measures(&mut self) -> Result<()> {
12 let all_measures = self
13 .transaction(|tx| async move {
14 let existing_measures = usage_measure::Entity::find().all(&*tx).await?;
15
16 let new_measures = UsageMeasure::iter()
17 .filter(|measure| {
18 !existing_measures
19 .iter()
20 .any(|m| m.name == measure.to_string())
21 })
22 .map(|measure| usage_measure::ActiveModel {
23 name: ActiveValue::set(measure.to_string()),
24 ..Default::default()
25 })
26 .collect::<Vec<_>>();
27
28 if !new_measures.is_empty() {
29 usage_measure::Entity::insert_many(new_measures)
30 .exec(&*tx)
31 .await?;
32 }
33
34 Ok(usage_measure::Entity::find().all(&*tx).await?)
35 })
36 .await?;
37
38 self.usage_measure_ids = all_measures
39 .into_iter()
40 .filter_map(|measure| {
41 UsageMeasure::from_str(&measure.name)
42 .ok()
43 .map(|um| (um, measure.id))
44 })
45 .collect();
46 Ok(())
47 }
48
49 pub async fn get_user_spending_for_month(
50 &self,
51 user_id: UserId,
52 now: DateTimeUtc,
53 ) -> Result<Cents> {
54 self.transaction(|tx| async move {
55 let month = now.date_naive().month() as i32;
56 let year = now.date_naive().year();
57
58 let mut monthly_usages = monthly_usage::Entity::find()
59 .filter(
60 monthly_usage::Column::UserId
61 .eq(user_id)
62 .and(monthly_usage::Column::Month.eq(month))
63 .and(monthly_usage::Column::Year.eq(year)),
64 )
65 .stream(&*tx)
66 .await?;
67 let mut monthly_spending = Cents::ZERO;
68
69 while let Some(usage) = monthly_usages.next().await {
70 let usage = usage?;
71 let Ok(model) = self.model_by_id(usage.model_id) else {
72 continue;
73 };
74
75 monthly_spending += calculate_spending(
76 model,
77 usage.input_tokens as usize,
78 usage.cache_creation_input_tokens as usize,
79 usage.cache_read_input_tokens as usize,
80 usage.output_tokens as usize,
81 );
82 }
83
84 Ok(monthly_spending)
85 })
86 .await
87 }
88}
89
90fn calculate_spending(
91 model: &model::Model,
92 input_tokens_this_month: usize,
93 cache_creation_input_tokens_this_month: usize,
94 cache_read_input_tokens_this_month: usize,
95 output_tokens_this_month: usize,
96) -> Cents {
97 let input_token_cost =
98 input_tokens_this_month * model.price_per_million_input_tokens as usize / 1_000_000;
99 let cache_creation_input_token_cost = cache_creation_input_tokens_this_month
100 * model.price_per_million_cache_creation_input_tokens as usize
101 / 1_000_000;
102 let cache_read_input_token_cost = cache_read_input_tokens_this_month
103 * model.price_per_million_cache_read_input_tokens as usize
104 / 1_000_000;
105 let output_token_cost =
106 output_tokens_this_month * model.price_per_million_output_tokens as usize / 1_000_000;
107 let spending = input_token_cost
108 + cache_creation_input_token_cost
109 + cache_read_input_token_cost
110 + output_token_cost;
111 Cents::new(spending as u32)
112}