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, now)
37 .await
38 .unwrap();
39
40 let now = t0 + Duration::seconds(10);
41 db.record_usage(user_id, false, provider, model, 2000, 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 output_tokens_this_month: 0,
54 spending_this_month: 0,
55 lifetime_spending: 0,
56 }
57 );
58
59 let now = t0 + Duration::seconds(60);
60 let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
61 assert_eq!(
62 usage,
63 Usage {
64 requests_this_minute: 1,
65 tokens_this_minute: 2000,
66 tokens_this_day: 3000,
67 input_tokens_this_month: 3000,
68 output_tokens_this_month: 0,
69 spending_this_month: 0,
70 lifetime_spending: 0,
71 }
72 );
73
74 let now = t0 + Duration::seconds(60);
75 db.record_usage(user_id, false, provider, model, 3000, 0, now)
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: 5000,
85 tokens_this_day: 6000,
86 input_tokens_this_month: 6000,
87 output_tokens_this_month: 0,
88 spending_this_month: 0,
89 lifetime_spending: 0,
90 }
91 );
92
93 let t1 = t0 + Duration::hours(24);
94 let now = t1;
95 let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
96 assert_eq!(
97 usage,
98 Usage {
99 requests_this_minute: 0,
100 tokens_this_minute: 0,
101 tokens_this_day: 5000,
102 input_tokens_this_month: 6000,
103 output_tokens_this_month: 0,
104 spending_this_month: 0,
105 lifetime_spending: 0,
106 }
107 );
108
109 db.record_usage(user_id, false, provider, model, 4000, 0, now)
110 .await
111 .unwrap();
112
113 let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
114 assert_eq!(
115 usage,
116 Usage {
117 requests_this_minute: 1,
118 tokens_this_minute: 4000,
119 tokens_this_day: 9000,
120 input_tokens_this_month: 10000,
121 output_tokens_this_month: 0,
122 spending_this_month: 0,
123 lifetime_spending: 0,
124 }
125 );
126
127 let t2 = t0 + Duration::days(30);
128 let now = t2;
129 let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
130 assert_eq!(
131 usage,
132 Usage {
133 requests_this_minute: 0,
134 tokens_this_minute: 0,
135 tokens_this_day: 0,
136 input_tokens_this_month: 9000,
137 output_tokens_this_month: 0,
138 spending_this_month: 0,
139 lifetime_spending: 0,
140 }
141 );
142}