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