1use crate::{
  2    llm::db::{queries::providers::ModelParams, queries::usages::Usage, LlmDatabase},
  3    test_llm_db,
  4};
  5use chrono::{Duration, Utc};
  6use pretty_assertions::assert_eq;
  7use rpc::LanguageModelProvider;
  8
  9test_llm_db!(test_tracking_usage, test_tracking_usage_postgres);
 10
 11async fn test_tracking_usage(db: &mut LlmDatabase) {
 12    let provider = LanguageModelProvider::Anthropic;
 13    let model = "claude-3-5-sonnet";
 14
 15    db.initialize().await.unwrap();
 16    db.insert_models(&[ModelParams {
 17        provider,
 18        name: model.to_string(),
 19        max_requests_per_minute: 5,
 20        max_tokens_per_minute: 10_000,
 21        max_tokens_per_day: 50_000,
 22        price_per_million_input_tokens: 50,
 23        price_per_million_output_tokens: 50,
 24    }])
 25    .await
 26    .unwrap();
 27
 28    let t0 = Utc::now();
 29    let user_id = 123;
 30
 31    let now = t0;
 32    db.record_usage(user_id, provider, model, 1000, 0, now)
 33        .await
 34        .unwrap();
 35
 36    let now = t0 + Duration::seconds(10);
 37    db.record_usage(user_id, provider, model, 2000, 0, now)
 38        .await
 39        .unwrap();
 40
 41    let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
 42    assert_eq!(
 43        usage,
 44        Usage {
 45            requests_this_minute: 2,
 46            tokens_this_minute: 3000,
 47            tokens_this_day: 3000,
 48            input_tokens_this_month: 3000,
 49            output_tokens_this_month: 0,
 50            spending_this_month: 0,
 51        }
 52    );
 53
 54    let now = t0 + Duration::seconds(60);
 55    let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
 56    assert_eq!(
 57        usage,
 58        Usage {
 59            requests_this_minute: 1,
 60            tokens_this_minute: 2000,
 61            tokens_this_day: 3000,
 62            input_tokens_this_month: 3000,
 63            output_tokens_this_month: 0,
 64            spending_this_month: 0,
 65        }
 66    );
 67
 68    let now = t0 + Duration::seconds(60);
 69    db.record_usage(user_id, provider, model, 3000, 0, now)
 70        .await
 71        .unwrap();
 72
 73    let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
 74    assert_eq!(
 75        usage,
 76        Usage {
 77            requests_this_minute: 2,
 78            tokens_this_minute: 5000,
 79            tokens_this_day: 6000,
 80            input_tokens_this_month: 6000,
 81            output_tokens_this_month: 0,
 82            spending_this_month: 0,
 83        }
 84    );
 85
 86    let t1 = t0 + Duration::hours(24);
 87    let now = t1;
 88    let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
 89    assert_eq!(
 90        usage,
 91        Usage {
 92            requests_this_minute: 0,
 93            tokens_this_minute: 0,
 94            tokens_this_day: 5000,
 95            input_tokens_this_month: 6000,
 96            output_tokens_this_month: 0,
 97            spending_this_month: 0,
 98        }
 99    );
100
101    db.record_usage(user_id, provider, model, 4000, 0, now)
102        .await
103        .unwrap();
104
105    let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
106    assert_eq!(
107        usage,
108        Usage {
109            requests_this_minute: 1,
110            tokens_this_minute: 4000,
111            tokens_this_day: 9000,
112            input_tokens_this_month: 10000,
113            output_tokens_this_month: 0,
114            spending_this_month: 0,
115        }
116    );
117
118    let t2 = t0 + Duration::days(30);
119    let now = t2;
120    let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
121    assert_eq!(
122        usage,
123        Usage {
124            requests_this_minute: 0,
125            tokens_this_minute: 0,
126            tokens_this_day: 0,
127            input_tokens_this_month: 9000,
128            output_tokens_this_month: 0,
129            spending_this_month: 0,
130        }
131    );
132}