1use crate::{
  2    llm::db::{queries::providers::ModelRateLimits, 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(&[(
 17        provider,
 18        model.to_string(),
 19        ModelRateLimits {
 20            max_requests_per_minute: 5,
 21            max_tokens_per_minute: 10_000,
 22            max_tokens_per_day: 50_000,
 23        },
 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, now)
 33        .await
 34        .unwrap();
 35
 36    let now = t0 + Duration::seconds(10);
 37    db.record_usage(user_id, provider, model, 2000, 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            tokens_this_month: 3000,
 49        }
 50    );
 51
 52    let now = t0 + Duration::seconds(60);
 53    let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
 54    assert_eq!(
 55        usage,
 56        Usage {
 57            requests_this_minute: 1,
 58            tokens_this_minute: 2000,
 59            tokens_this_day: 3000,
 60            tokens_this_month: 3000,
 61        }
 62    );
 63
 64    let now = t0 + Duration::seconds(60);
 65    db.record_usage(user_id, provider, model, 3000, now)
 66        .await
 67        .unwrap();
 68
 69    let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
 70    assert_eq!(
 71        usage,
 72        Usage {
 73            requests_this_minute: 2,
 74            tokens_this_minute: 5000,
 75            tokens_this_day: 6000,
 76            tokens_this_month: 6000,
 77        }
 78    );
 79
 80    let t1 = t0 + Duration::hours(24);
 81    let now = t1;
 82    let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
 83    assert_eq!(
 84        usage,
 85        Usage {
 86            requests_this_minute: 0,
 87            tokens_this_minute: 0,
 88            tokens_this_day: 5000,
 89            tokens_this_month: 6000,
 90        }
 91    );
 92
 93    db.record_usage(user_id, provider, model, 4000, now)
 94        .await
 95        .unwrap();
 96
 97    let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
 98    assert_eq!(
 99        usage,
100        Usage {
101            requests_this_minute: 1,
102            tokens_this_minute: 4000,
103            tokens_this_day: 9000,
104            tokens_this_month: 10000,
105        }
106    );
107
108    let t2 = t0 + Duration::days(30);
109    let now = t2;
110    let usage = db.get_usage(user_id, provider, model, now).await.unwrap();
111    assert_eq!(
112        usage,
113        Usage {
114            requests_this_minute: 0,
115            tokens_this_minute: 0,
116            tokens_this_day: 0,
117            tokens_this_month: 9000,
118        }
119    );
120}