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