1use crate::llm::FREE_TIER_MONTHLY_SPENDING_LIMIT;
2use crate::{
3 Cents,
4 db::UserId,
5 llm::db::{
6 LlmDatabase, TokenUsage,
7 queries::{providers::ModelParams, usages::Usage},
8 },
9 test_llm_db,
10};
11use chrono::{DateTime, Duration, Utc};
12use pretty_assertions::assert_eq;
13use rpc::LanguageModelProvider;
14
15test_llm_db!(test_tracking_usage, test_tracking_usage_postgres);
16
17async fn test_tracking_usage(db: &mut LlmDatabase) {
18 let provider = LanguageModelProvider::Anthropic;
19 let model = "claude-3-5-sonnet";
20
21 db.initialize().await.unwrap();
22 db.insert_models(&[ModelParams {
23 provider,
24 name: model.to_string(),
25 max_requests_per_minute: 5,
26 max_tokens_per_minute: 10_000,
27 max_tokens_per_day: 50_000,
28 price_per_million_input_tokens: 50,
29 price_per_million_output_tokens: 50,
30 }])
31 .await
32 .unwrap();
33
34 // We're using a fixed datetime to prevent flakiness based on the clock.
35 let t0 = DateTime::parse_from_rfc3339("2024-08-08T22:46:33Z")
36 .unwrap()
37 .with_timezone(&Utc);
38 let user_id = UserId::from_proto(123);
39
40 let now = t0;
41 db.record_usage(
42 user_id,
43 false,
44 provider,
45 model,
46 TokenUsage {
47 input: 1000,
48 input_cache_creation: 0,
49 input_cache_read: 0,
50 output: 0,
51 },
52 false,
53 Cents::ZERO,
54 FREE_TIER_MONTHLY_SPENDING_LIMIT,
55 now,
56 )
57 .await
58 .unwrap();
59
60 let now = t0 + Duration::seconds(10);
61 db.record_usage(
62 user_id,
63 false,
64 provider,
65 model,
66 TokenUsage {
67 input: 2000,
68 input_cache_creation: 0,
69 input_cache_read: 0,
70 output: 0,
71 },
72 false,
73 Cents::ZERO,
74 FREE_TIER_MONTHLY_SPENDING_LIMIT,
75 now,
76 )
77 .await
78 .unwrap();
79
80 let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
81 assert_eq!(
82 usage,
83 Usage {
84 requests_this_minute: 2,
85 tokens_this_minute: 3000,
86 tokens_this_day: 3000,
87 tokens_this_month: TokenUsage {
88 input: 3000,
89 input_cache_creation: 0,
90 input_cache_read: 0,
91 output: 0,
92 },
93 spending_this_month: Cents::ZERO,
94 lifetime_spending: Cents::ZERO,
95 }
96 );
97
98 let now = t0 + Duration::seconds(60);
99 let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
100 assert_eq!(
101 usage,
102 Usage {
103 requests_this_minute: 1,
104 tokens_this_minute: 2000,
105 tokens_this_day: 3000,
106 tokens_this_month: TokenUsage {
107 input: 3000,
108 input_cache_creation: 0,
109 input_cache_read: 0,
110 output: 0,
111 },
112 spending_this_month: Cents::ZERO,
113 lifetime_spending: Cents::ZERO,
114 }
115 );
116
117 let now = t0 + Duration::seconds(60);
118 db.record_usage(
119 user_id,
120 false,
121 provider,
122 model,
123 TokenUsage {
124 input: 3000,
125 input_cache_creation: 0,
126 input_cache_read: 0,
127 output: 0,
128 },
129 false,
130 Cents::ZERO,
131 FREE_TIER_MONTHLY_SPENDING_LIMIT,
132 now,
133 )
134 .await
135 .unwrap();
136
137 let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
138 assert_eq!(
139 usage,
140 Usage {
141 requests_this_minute: 2,
142 tokens_this_minute: 5000,
143 tokens_this_day: 6000,
144 tokens_this_month: TokenUsage {
145 input: 6000,
146 input_cache_creation: 0,
147 input_cache_read: 0,
148 output: 0,
149 },
150 spending_this_month: Cents::ZERO,
151 lifetime_spending: Cents::ZERO,
152 }
153 );
154
155 let t1 = t0 + Duration::hours(24);
156 let now = t1;
157 let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
158 assert_eq!(
159 usage,
160 Usage {
161 requests_this_minute: 0,
162 tokens_this_minute: 0,
163 tokens_this_day: 5000,
164 tokens_this_month: TokenUsage {
165 input: 6000,
166 input_cache_creation: 0,
167 input_cache_read: 0,
168 output: 0,
169 },
170 spending_this_month: Cents::ZERO,
171 lifetime_spending: Cents::ZERO,
172 }
173 );
174
175 db.record_usage(
176 user_id,
177 false,
178 provider,
179 model,
180 TokenUsage {
181 input: 4000,
182 input_cache_creation: 0,
183 input_cache_read: 0,
184 output: 0,
185 },
186 false,
187 Cents::ZERO,
188 FREE_TIER_MONTHLY_SPENDING_LIMIT,
189 now,
190 )
191 .await
192 .unwrap();
193
194 let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
195 assert_eq!(
196 usage,
197 Usage {
198 requests_this_minute: 1,
199 tokens_this_minute: 4000,
200 tokens_this_day: 9000,
201 tokens_this_month: TokenUsage {
202 input: 10000,
203 input_cache_creation: 0,
204 input_cache_read: 0,
205 output: 0,
206 },
207 spending_this_month: Cents::ZERO,
208 lifetime_spending: Cents::ZERO,
209 }
210 );
211
212 // We're using a fixed datetime to prevent flakiness based on the clock.
213 let now = DateTime::parse_from_rfc3339("2024-10-08T22:15:58Z")
214 .unwrap()
215 .with_timezone(&Utc);
216
217 // Test cache creation input tokens
218 db.record_usage(
219 user_id,
220 false,
221 provider,
222 model,
223 TokenUsage {
224 input: 1000,
225 input_cache_creation: 500,
226 input_cache_read: 0,
227 output: 0,
228 },
229 false,
230 Cents::ZERO,
231 FREE_TIER_MONTHLY_SPENDING_LIMIT,
232 now,
233 )
234 .await
235 .unwrap();
236
237 let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
238 assert_eq!(
239 usage,
240 Usage {
241 requests_this_minute: 1,
242 tokens_this_minute: 1500,
243 tokens_this_day: 1500,
244 tokens_this_month: TokenUsage {
245 input: 1000,
246 input_cache_creation: 500,
247 input_cache_read: 0,
248 output: 0,
249 },
250 spending_this_month: Cents::ZERO,
251 lifetime_spending: Cents::ZERO,
252 }
253 );
254
255 // Test cache read input tokens
256 db.record_usage(
257 user_id,
258 false,
259 provider,
260 model,
261 TokenUsage {
262 input: 1000,
263 input_cache_creation: 0,
264 input_cache_read: 300,
265 output: 0,
266 },
267 false,
268 Cents::ZERO,
269 FREE_TIER_MONTHLY_SPENDING_LIMIT,
270 now,
271 )
272 .await
273 .unwrap();
274
275 let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
276 assert_eq!(
277 usage,
278 Usage {
279 requests_this_minute: 2,
280 tokens_this_minute: 2800,
281 tokens_this_day: 2800,
282 tokens_this_month: TokenUsage {
283 input: 2000,
284 input_cache_creation: 500,
285 input_cache_read: 300,
286 output: 0,
287 },
288 spending_this_month: Cents::ZERO,
289 lifetime_spending: Cents::ZERO,
290 }
291 );
292}