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