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