1use crate::{
2 db::UserId,
3 llm::{
4 db::{
5 queries::{providers::ModelParams, usages::Usage},
6 LlmDatabase, TokenUsage,
7 },
8 FREE_TIER_MONTHLY_SPENDING_LIMIT,
9 },
10 test_llm_db, Cents,
11};
12use chrono::{DateTime, Utc};
13use pretty_assertions::assert_eq;
14use rpc::LanguageModelProvider;
15
16test_llm_db!(
17 test_billing_limit_exceeded,
18 test_billing_limit_exceeded_postgres
19);
20
21async fn test_billing_limit_exceeded(db: &mut LlmDatabase) {
22 let provider = LanguageModelProvider::Anthropic;
23 let model = "fake-claude-limerick";
24 const PRICE_PER_MILLION_INPUT_TOKENS: i32 = 5;
25 const PRICE_PER_MILLION_OUTPUT_TOKENS: i32 = 5;
26
27 // Initialize the database and insert the model
28 db.initialize().await.unwrap();
29 db.insert_models(&[ModelParams {
30 provider,
31 name: model.to_string(),
32 max_requests_per_minute: 5,
33 max_tokens_per_minute: 10_000,
34 max_tokens_per_day: 50_000,
35 price_per_million_input_tokens: PRICE_PER_MILLION_INPUT_TOKENS,
36 price_per_million_output_tokens: PRICE_PER_MILLION_OUTPUT_TOKENS,
37 }])
38 .await
39 .unwrap();
40
41 // Set a fixed datetime for consistent testing
42 let now = DateTime::parse_from_rfc3339("2024-08-08T22:46:33Z")
43 .unwrap()
44 .with_timezone(&Utc);
45
46 let user_id = UserId::from_proto(123);
47
48 let max_monthly_spend = Cents::from_dollars(10);
49
50 // Record usage that brings us close to the limit but doesn't exceed it
51 // Let's say we use $9.50 worth of tokens
52 let tokens_to_use = 190_000_000; // This will cost $9.50 at $0.05 per 1 million tokens
53 let usage = TokenUsage {
54 input: tokens_to_use,
55 input_cache_creation: 0,
56 input_cache_read: 0,
57 output: 0,
58 };
59 let cost = Cents::new(tokens_to_use as u32 / 1_000_000 * PRICE_PER_MILLION_INPUT_TOKENS as u32);
60
61 assert_eq!(
62 cost,
63 Cents::new(950),
64 "expected the cost to be $9.50, based on the inputs, but it wasn't"
65 );
66
67 // Verify that before we record any usage, there are 0 billing events
68 let billing_events = db.get_billing_events().await.unwrap();
69 assert_eq!(billing_events.len(), 0);
70
71 db.record_usage(
72 user_id,
73 false,
74 provider,
75 model,
76 usage,
77 true,
78 max_monthly_spend,
79 now,
80 )
81 .await
82 .unwrap();
83
84 // Verify the recorded usage and spending
85 let recorded_usage = db.get_usage(user_id, provider, model, now).await.unwrap();
86
87 // Verify that we exceeded the free tier usage
88 assert!(
89 recorded_usage.spending_this_month > FREE_TIER_MONTHLY_SPENDING_LIMIT,
90 "Expected spending to exceed free tier limit"
91 );
92
93 assert_eq!(
94 recorded_usage,
95 Usage {
96 requests_this_minute: 1,
97 tokens_this_minute: tokens_to_use,
98 tokens_this_day: tokens_to_use,
99 tokens_this_month: TokenUsage {
100 input: tokens_to_use,
101 input_cache_creation: 0,
102 input_cache_read: 0,
103 output: 0,
104 },
105 spending_this_month: Cents::new(950),
106 lifetime_spending: Cents::new(950),
107 }
108 );
109
110 // Verify that there is one `billing_event` record
111 let billing_events = db.get_billing_events().await.unwrap();
112 assert_eq!(billing_events.len(), 1);
113
114 let (billing_event, _model) = &billing_events[0];
115 assert_eq!(billing_event.user_id, user_id);
116 assert_eq!(billing_event.input_tokens, tokens_to_use as i64);
117 assert_eq!(billing_event.input_cache_creation_tokens, 0);
118 assert_eq!(billing_event.input_cache_read_tokens, 0);
119 assert_eq!(billing_event.output_tokens, 0);
120
121 let tokens_to_exceed = 20_000_000; // This will cost $1.00 more, pushing us from $9.50 to $10.50, which is over the $10 monthly maximum limit
122 let usage_exceeding = TokenUsage {
123 input: tokens_to_exceed,
124 input_cache_creation: 0,
125 input_cache_read: 0,
126 output: 0,
127 };
128
129 // This should still create a billing event as it's the first request that exceeds the limit
130 db.record_usage(
131 user_id,
132 false,
133 provider,
134 model,
135 usage_exceeding,
136 true,
137 max_monthly_spend,
138 now,
139 )
140 .await
141 .unwrap();
142
143 // Verify that there is still one billing record
144 let billing_events = db.get_billing_events().await.unwrap();
145 assert_eq!(billing_events.len(), 1);
146
147 // Verify the updated usage and spending
148 let updated_usage = db.get_usage(user_id, provider, model, now).await.unwrap();
149 assert_eq!(
150 updated_usage,
151 Usage {
152 requests_this_minute: 2,
153 tokens_this_minute: tokens_to_use + tokens_to_exceed,
154 tokens_this_day: tokens_to_use + tokens_to_exceed,
155 tokens_this_month: TokenUsage {
156 input: tokens_to_use + tokens_to_exceed,
157 input_cache_creation: 0,
158 input_cache_read: 0,
159 output: 0,
160 },
161 spending_this_month: Cents::new(1050),
162 lifetime_spending: Cents::new(1050),
163 }
164 );
165}