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 input_tokens_this_minute: 3000,
87 output_tokens_this_minute: 0,
88 tokens_this_day: 3000,
89 tokens_this_month: TokenUsage {
90 input: 3000,
91 input_cache_creation: 0,
92 input_cache_read: 0,
93 output: 0,
94 },
95 spending_this_month: Cents::ZERO,
96 lifetime_spending: Cents::ZERO,
97 }
98 );
99
100 let now = t0 + Duration::seconds(60);
101 let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
102 assert_eq!(
103 usage,
104 Usage {
105 requests_this_minute: 1,
106 tokens_this_minute: 2000,
107 input_tokens_this_minute: 2000,
108 output_tokens_this_minute: 0,
109 tokens_this_day: 3000,
110 tokens_this_month: TokenUsage {
111 input: 3000,
112 input_cache_creation: 0,
113 input_cache_read: 0,
114 output: 0,
115 },
116 spending_this_month: Cents::ZERO,
117 lifetime_spending: Cents::ZERO,
118 }
119 );
120
121 let now = t0 + Duration::seconds(60);
122 db.record_usage(
123 user_id,
124 false,
125 provider,
126 model,
127 TokenUsage {
128 input: 3000,
129 input_cache_creation: 0,
130 input_cache_read: 0,
131 output: 0,
132 },
133 false,
134 Cents::ZERO,
135 FREE_TIER_MONTHLY_SPENDING_LIMIT,
136 now,
137 )
138 .await
139 .unwrap();
140
141 let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
142 assert_eq!(
143 usage,
144 Usage {
145 requests_this_minute: 2,
146 tokens_this_minute: 5000,
147 input_tokens_this_minute: 5000,
148 output_tokens_this_minute: 0,
149 tokens_this_day: 6000,
150 tokens_this_month: TokenUsage {
151 input: 6000,
152 input_cache_creation: 0,
153 input_cache_read: 0,
154 output: 0,
155 },
156 spending_this_month: Cents::ZERO,
157 lifetime_spending: Cents::ZERO,
158 }
159 );
160
161 let t1 = t0 + Duration::hours(24);
162 let now = t1;
163 let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
164 assert_eq!(
165 usage,
166 Usage {
167 requests_this_minute: 0,
168 tokens_this_minute: 0,
169 input_tokens_this_minute: 0,
170 output_tokens_this_minute: 0,
171 tokens_this_day: 5000,
172 tokens_this_month: TokenUsage {
173 input: 6000,
174 input_cache_creation: 0,
175 input_cache_read: 0,
176 output: 0,
177 },
178 spending_this_month: Cents::ZERO,
179 lifetime_spending: Cents::ZERO,
180 }
181 );
182
183 db.record_usage(
184 user_id,
185 false,
186 provider,
187 model,
188 TokenUsage {
189 input: 4000,
190 input_cache_creation: 0,
191 input_cache_read: 0,
192 output: 0,
193 },
194 false,
195 Cents::ZERO,
196 FREE_TIER_MONTHLY_SPENDING_LIMIT,
197 now,
198 )
199 .await
200 .unwrap();
201
202 let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
203 assert_eq!(
204 usage,
205 Usage {
206 requests_this_minute: 1,
207 tokens_this_minute: 4000,
208 input_tokens_this_minute: 4000,
209 output_tokens_this_minute: 0,
210 tokens_this_day: 9000,
211 tokens_this_month: TokenUsage {
212 input: 10000,
213 input_cache_creation: 0,
214 input_cache_read: 0,
215 output: 0,
216 },
217 spending_this_month: Cents::ZERO,
218 lifetime_spending: Cents::ZERO,
219 }
220 );
221
222 // We're using a fixed datetime to prevent flakiness based on the clock.
223 let now = DateTime::parse_from_rfc3339("2024-10-08T22:15:58Z")
224 .unwrap()
225 .with_timezone(&Utc);
226
227 // Test cache creation input tokens
228 db.record_usage(
229 user_id,
230 false,
231 provider,
232 model,
233 TokenUsage {
234 input: 1000,
235 input_cache_creation: 500,
236 input_cache_read: 0,
237 output: 0,
238 },
239 false,
240 Cents::ZERO,
241 FREE_TIER_MONTHLY_SPENDING_LIMIT,
242 now,
243 )
244 .await
245 .unwrap();
246
247 let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
248 assert_eq!(
249 usage,
250 Usage {
251 requests_this_minute: 1,
252 tokens_this_minute: 1500,
253 input_tokens_this_minute: 1500,
254 output_tokens_this_minute: 0,
255 tokens_this_day: 1500,
256 tokens_this_month: TokenUsage {
257 input: 1000,
258 input_cache_creation: 500,
259 input_cache_read: 0,
260 output: 0,
261 },
262 spending_this_month: Cents::ZERO,
263 lifetime_spending: Cents::ZERO,
264 }
265 );
266
267 // Test cache read input tokens
268 db.record_usage(
269 user_id,
270 false,
271 provider,
272 model,
273 TokenUsage {
274 input: 1000,
275 input_cache_creation: 0,
276 input_cache_read: 300,
277 output: 0,
278 },
279 false,
280 Cents::ZERO,
281 FREE_TIER_MONTHLY_SPENDING_LIMIT,
282 now,
283 )
284 .await
285 .unwrap();
286
287 let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
288 assert_eq!(
289 usage,
290 Usage {
291 requests_this_minute: 2,
292 tokens_this_minute: 2800,
293 input_tokens_this_minute: 2500,
294 output_tokens_this_minute: 0,
295 tokens_this_day: 2800,
296 tokens_this_month: TokenUsage {
297 input: 2000,
298 input_cache_creation: 500,
299 input_cache_read: 300,
300 output: 0,
301 },
302 spending_this_month: Cents::ZERO,
303 lifetime_spending: Cents::ZERO,
304 }
305 );
306}