1use crate::{
2 db::UserId,
3 llm::db::{
4 queries::{providers::ModelParams, usages::Usage},
5 LlmDatabase,
6 },
7 test_llm_db,
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(user_id, false, provider, model, 1000, 0, 0, 0, now)
40 .await
41 .unwrap();
42
43 let now = t0 + Duration::seconds(10);
44 db.record_usage(user_id, false, provider, model, 2000, 0, 0, 0, now)
45 .await
46 .unwrap();
47
48 let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
49 assert_eq!(
50 usage,
51 Usage {
52 requests_this_minute: 2,
53 tokens_this_minute: 3000,
54 tokens_this_day: 3000,
55 input_tokens_this_month: 3000,
56 cache_creation_input_tokens_this_month: 0,
57 cache_read_input_tokens_this_month: 0,
58 output_tokens_this_month: 0,
59 spending_this_month: 0,
60 lifetime_spending: 0,
61 }
62 );
63
64 let now = t0 + Duration::seconds(60);
65 let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
66 assert_eq!(
67 usage,
68 Usage {
69 requests_this_minute: 1,
70 tokens_this_minute: 2000,
71 tokens_this_day: 3000,
72 input_tokens_this_month: 3000,
73 cache_creation_input_tokens_this_month: 0,
74 cache_read_input_tokens_this_month: 0,
75 output_tokens_this_month: 0,
76 spending_this_month: 0,
77 lifetime_spending: 0,
78 }
79 );
80
81 let now = t0 + Duration::seconds(60);
82 db.record_usage(user_id, false, provider, model, 3000, 0, 0, 0, now)
83 .await
84 .unwrap();
85
86 let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
87 assert_eq!(
88 usage,
89 Usage {
90 requests_this_minute: 2,
91 tokens_this_minute: 5000,
92 tokens_this_day: 6000,
93 input_tokens_this_month: 6000,
94 cache_creation_input_tokens_this_month: 0,
95 cache_read_input_tokens_this_month: 0,
96 output_tokens_this_month: 0,
97 spending_this_month: 0,
98 lifetime_spending: 0,
99 }
100 );
101
102 let t1 = t0 + Duration::hours(24);
103 let now = t1;
104 let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
105 assert_eq!(
106 usage,
107 Usage {
108 requests_this_minute: 0,
109 tokens_this_minute: 0,
110 tokens_this_day: 5000,
111 input_tokens_this_month: 6000,
112 cache_creation_input_tokens_this_month: 0,
113 cache_read_input_tokens_this_month: 0,
114 output_tokens_this_month: 0,
115 spending_this_month: 0,
116 lifetime_spending: 0,
117 }
118 );
119
120 db.record_usage(user_id, false, provider, model, 4000, 0, 0, 0, now)
121 .await
122 .unwrap();
123
124 let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
125 assert_eq!(
126 usage,
127 Usage {
128 requests_this_minute: 1,
129 tokens_this_minute: 4000,
130 tokens_this_day: 9000,
131 input_tokens_this_month: 10000,
132 cache_creation_input_tokens_this_month: 0,
133 cache_read_input_tokens_this_month: 0,
134 output_tokens_this_month: 0,
135 spending_this_month: 0,
136 lifetime_spending: 0,
137 }
138 );
139
140 // We're using a fixed datetime to prevent flakiness based on the clock.
141 let now = DateTime::parse_from_rfc3339("2024-10-08T22:15:58Z")
142 .unwrap()
143 .with_timezone(&Utc);
144
145 // Test cache creation input tokens
146 db.record_usage(user_id, false, provider, model, 1000, 500, 0, 0, now)
147 .await
148 .unwrap();
149
150 let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
151 assert_eq!(
152 usage,
153 Usage {
154 requests_this_minute: 1,
155 tokens_this_minute: 1500,
156 tokens_this_day: 1500,
157 input_tokens_this_month: 1000,
158 cache_creation_input_tokens_this_month: 500,
159 cache_read_input_tokens_this_month: 0,
160 output_tokens_this_month: 0,
161 spending_this_month: 0,
162 lifetime_spending: 0,
163 }
164 );
165
166 // Test cache read input tokens
167 db.record_usage(user_id, false, provider, model, 1000, 0, 300, 0, now)
168 .await
169 .unwrap();
170
171 let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
172 assert_eq!(
173 usage,
174 Usage {
175 requests_this_minute: 2,
176 tokens_this_minute: 2800,
177 tokens_this_day: 2800,
178 input_tokens_this_month: 2000,
179 cache_creation_input_tokens_this_month: 500,
180 cache_read_input_tokens_this_month: 300,
181 output_tokens_this_month: 0,
182 spending_this_month: 0,
183 lifetime_spending: 0,
184 }
185 );
186}