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(11);
49
50 // Record usage that brings us close to the limit but doesn't exceed it
51 // Let's say we use $10.50 worth of tokens
52 let tokens_to_use = 210_000_000; // This will cost $10.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
60 // Verify that before we record any usage, there are 0 billing events
61 let billing_events = db.get_billing_events().await.unwrap();
62 assert_eq!(billing_events.len(), 0);
63
64 db.record_usage(
65 user_id,
66 false,
67 provider,
68 model,
69 usage,
70 true,
71 max_monthly_spend,
72 now,
73 )
74 .await
75 .unwrap();
76
77 // Verify the recorded usage and spending
78 let recorded_usage = db.get_usage(user_id, provider, model, now).await.unwrap();
79
80 // Verify that we exceeded the free tier usage
81 assert!(
82 recorded_usage.spending_this_month > FREE_TIER_MONTHLY_SPENDING_LIMIT,
83 "Expected spending to exceed free tier limit"
84 );
85
86 assert_eq!(
87 recorded_usage,
88 Usage {
89 requests_this_minute: 1,
90 tokens_this_minute: tokens_to_use,
91 tokens_this_day: tokens_to_use,
92 tokens_this_month: TokenUsage {
93 input: tokens_to_use,
94 input_cache_creation: 0,
95 input_cache_read: 0,
96 output: 0,
97 },
98 spending_this_month: Cents::new(1050),
99 lifetime_spending: Cents::new(1050),
100 }
101 );
102
103 // Verify that there is one `billing_event` record
104 let billing_events = db.get_billing_events().await.unwrap();
105 assert_eq!(billing_events.len(), 1);
106
107 let (billing_event, _model) = &billing_events[0];
108 assert_eq!(billing_event.user_id, user_id);
109 assert_eq!(billing_event.input_tokens, tokens_to_use as i64);
110 assert_eq!(billing_event.input_cache_creation_tokens, 0);
111 assert_eq!(billing_event.input_cache_read_tokens, 0);
112 assert_eq!(billing_event.output_tokens, 0);
113
114 let tokens_to_exceed = 20_000_000; // This will cost $1.00 more, pushing us from $10.50 to $11.50, which is over the $11 monthly maximum limit
115 let usage_exceeding = TokenUsage {
116 input: tokens_to_exceed,
117 input_cache_creation: 0,
118 input_cache_read: 0,
119 output: 0,
120 };
121
122 // This should still create a billing event as it's the first request that exceeds the limit
123 db.record_usage(
124 user_id,
125 false,
126 provider,
127 model,
128 usage_exceeding,
129 true,
130 max_monthly_spend,
131 now,
132 )
133 .await
134 .unwrap();
135
136 // Verify that there is still one billing record
137 let billing_events = db.get_billing_events().await.unwrap();
138 assert_eq!(billing_events.len(), 1);
139
140 // Verify the updated usage and spending
141 let updated_usage = db.get_usage(user_id, provider, model, now).await.unwrap();
142 assert_eq!(
143 updated_usage,
144 Usage {
145 requests_this_minute: 2,
146 tokens_this_minute: tokens_to_use + tokens_to_exceed,
147 tokens_this_day: tokens_to_use + tokens_to_exceed,
148 tokens_this_month: TokenUsage {
149 input: tokens_to_use + tokens_to_exceed,
150 input_cache_creation: 0,
151 input_cache_read: 0,
152 output: 0,
153 },
154 spending_this_month: Cents::new(1150),
155 lifetime_spending: Cents::new(1150),
156 }
157 );
158}