1use crate::{
2 db::UserId,
3 llm::{
4 db::{queries::providers::ModelParams, LlmDatabase, TokenUsage},
5 FREE_TIER_MONTHLY_SPENDING_LIMIT,
6 },
7 test_llm_db, Cents,
8};
9use chrono::{DateTime, Utc};
10use pretty_assertions::assert_eq;
11use rpc::LanguageModelProvider;
12
13test_llm_db!(
14 test_billing_limit_exceeded,
15 test_billing_limit_exceeded_postgres
16);
17
18async fn test_billing_limit_exceeded(db: &mut LlmDatabase) {
19 let provider = LanguageModelProvider::Anthropic;
20 let model = "fake-claude-limerick";
21 const PRICE_PER_MILLION_INPUT_TOKENS: i32 = 5;
22 const PRICE_PER_MILLION_OUTPUT_TOKENS: i32 = 5;
23
24 // Initialize the database and insert the model
25 db.initialize().await.unwrap();
26 db.insert_models(&[ModelParams {
27 provider,
28 name: model.to_string(),
29 max_requests_per_minute: 5,
30 max_tokens_per_minute: 10_000,
31 max_tokens_per_day: 50_000,
32 price_per_million_input_tokens: PRICE_PER_MILLION_INPUT_TOKENS,
33 price_per_million_output_tokens: PRICE_PER_MILLION_OUTPUT_TOKENS,
34 }])
35 .await
36 .unwrap();
37
38 // Set a fixed datetime for consistent testing
39 let now = DateTime::parse_from_rfc3339("2024-08-08T22:46:33Z")
40 .unwrap()
41 .with_timezone(&Utc);
42
43 let user_id = UserId::from_proto(123);
44
45 let max_monthly_spend = Cents::from_dollars(11);
46
47 // Record usage that brings us close to the limit but doesn't exceed it
48 // Let's say we use $10.50 worth of tokens
49 let tokens_to_use = 210_000_000; // This will cost $10.50 at $0.05 per 1 million tokens
50 let usage = TokenUsage {
51 input: tokens_to_use,
52 input_cache_creation: 0,
53 input_cache_read: 0,
54 output: 0,
55 };
56
57 // Verify that before we record any usage, there are 0 billing events
58 let billing_events = db.get_billing_events().await.unwrap();
59 assert_eq!(billing_events.len(), 0);
60
61 db.record_usage(
62 user_id,
63 false,
64 provider,
65 model,
66 usage,
67 true,
68 max_monthly_spend,
69 FREE_TIER_MONTHLY_SPENDING_LIMIT,
70 now,
71 )
72 .await
73 .unwrap();
74
75 // Verify the recorded usage and spending
76 let recorded_usage = db.get_usage(user_id, provider, model, now).await.unwrap();
77 // Verify that we exceeded the free tier usage
78 assert_eq!(recorded_usage.spending_this_month, Cents::new(1050));
79 assert!(recorded_usage.spending_this_month > FREE_TIER_MONTHLY_SPENDING_LIMIT);
80
81 // Verify that there is one `billing_event` record
82 let billing_events = db.get_billing_events().await.unwrap();
83 assert_eq!(billing_events.len(), 1);
84
85 let (billing_event, _model) = &billing_events[0];
86 assert_eq!(billing_event.user_id, user_id);
87 assert_eq!(billing_event.input_tokens, tokens_to_use as i64);
88 assert_eq!(billing_event.input_cache_creation_tokens, 0);
89 assert_eq!(billing_event.input_cache_read_tokens, 0);
90 assert_eq!(billing_event.output_tokens, 0);
91
92 // Record usage that puts us at $20.50
93 let usage_2 = TokenUsage {
94 input: 200_000_000, // This will cost $10 more, pushing us from $10.50 to $20.50,
95 input_cache_creation: 0,
96 input_cache_read: 0,
97 output: 0,
98 };
99 db.record_usage(
100 user_id,
101 false,
102 provider,
103 model,
104 usage_2,
105 true,
106 max_monthly_spend,
107 FREE_TIER_MONTHLY_SPENDING_LIMIT,
108 now,
109 )
110 .await
111 .unwrap();
112
113 // Verify the updated usage and spending
114 let updated_usage = db.get_usage(user_id, provider, model, now).await.unwrap();
115 assert_eq!(updated_usage.spending_this_month, Cents::new(2050));
116
117 // Verify that there are now two billing events
118 let billing_events = db.get_billing_events().await.unwrap();
119 assert_eq!(billing_events.len(), 2);
120
121 let tokens_to_exceed = 20_000_000; // This will cost $1.00 more, pushing us from $20.50 to $21.50, which is over the $11 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 FREE_TIER_MONTHLY_SPENDING_LIMIT,
138 max_monthly_spend,
139 now,
140 )
141 .await
142 .unwrap();
143 // Verify the updated usage and spending
144 let updated_usage = db.get_usage(user_id, provider, model, now).await.unwrap();
145 assert_eq!(updated_usage.spending_this_month, Cents::new(2150));
146
147 // Verify that we never exceed the user max spending for the user
148 // and avoid charging them.
149 let billing_events = db.get_billing_events().await.unwrap();
150 assert_eq!(billing_events.len(), 2);
151}