1use anyhow::{Context, Result};
 2use serde::Serialize;
 3
 4use crate::clickhouse::write_to_table;
 5
 6#[derive(Serialize, Debug, clickhouse::Row)]
 7pub struct LlmUsageEventRow {
 8    pub time: i64,
 9    pub user_id: i32,
10    pub is_staff: bool,
11    pub plan: String,
12    pub model: String,
13    pub provider: String,
14    pub input_token_count: u64,
15    pub cache_creation_input_token_count: u64,
16    pub cache_read_input_token_count: u64,
17    pub output_token_count: u64,
18    pub requests_this_minute: u64,
19    pub tokens_this_minute: u64,
20    pub tokens_this_day: u64,
21    pub input_tokens_this_month: u64,
22    pub cache_creation_input_tokens_this_month: u64,
23    pub cache_read_input_tokens_this_month: u64,
24    pub output_tokens_this_month: u64,
25    pub spending_this_month: u64,
26    pub lifetime_spending: u64,
27}
28
29#[derive(Serialize, Debug, clickhouse::Row)]
30pub struct LlmRateLimitEventRow {
31    pub time: i64,
32    pub user_id: i32,
33    pub is_staff: bool,
34    pub plan: String,
35    pub model: String,
36    pub provider: String,
37    pub usage_measure: String,
38    pub requests_this_minute: u64,
39    pub tokens_this_minute: u64,
40    pub tokens_this_day: u64,
41    pub users_in_recent_minutes: u64,
42    pub users_in_recent_days: u64,
43    pub max_requests_per_minute: u64,
44    pub max_tokens_per_minute: u64,
45    pub max_tokens_per_day: u64,
46}
47
48pub async fn report_llm_usage(client: &clickhouse::Client, row: LlmUsageEventRow) -> Result<()> {
49    const LLM_USAGE_EVENTS_TABLE: &str = "llm_usage_events";
50    write_to_table(LLM_USAGE_EVENTS_TABLE, &[row], client)
51        .await
52        .with_context(|| format!("failed to upload to table '{LLM_USAGE_EVENTS_TABLE}'"))?;
53    Ok(())
54}
55
56pub async fn report_llm_rate_limit(
57    client: &clickhouse::Client,
58    row: LlmRateLimitEventRow,
59) -> Result<()> {
60    const LLM_RATE_LIMIT_EVENTS_TABLE: &str = "llm_rate_limit_events";
61    write_to_table(LLM_RATE_LIMIT_EVENTS_TABLE, &[row], client)
62        .await
63        .with_context(|| format!("failed to upload to table '{LLM_RATE_LIMIT_EVENTS_TABLE}'"))?;
64    Ok(())
65}