Skip to main content

systemprompt_analytics/repository/
queries.rs

1use anyhow::{anyhow, Result};
2use serde::Serialize;
3use std::sync::Arc;
4use systemprompt_database::{DatabaseProvider, DbPool, JsonRow, ToDbValue};
5
6#[derive(Debug, Clone)]
7pub struct AnalyticsQueryRepository {
8    db_pool: DbPool,
9}
10
11impl AnalyticsQueryRepository {
12    #[allow(clippy::unnecessary_wraps)]
13    pub fn new(db_pool: &DbPool) -> Result<Self> {
14        Ok(Self {
15            db_pool: Arc::clone(db_pool),
16        })
17    }
18
19    pub async fn get_ai_provider_usage(
20        &self,
21        days: i32,
22        user_id: Option<&str>,
23    ) -> Result<Vec<ProviderUsage>> {
24        let base_query = r"
25            SELECT
26                provider,
27                model,
28                COUNT(*) as request_count,
29                SUM(tokens_used) as total_tokens,
30                SUM(cost_microdollars) as total_cost_microdollars,
31                AVG(latency_ms) as avg_latency_ms,
32                COUNT(DISTINCT user_id) as unique_users,
33                COUNT(DISTINCT session_id) as unique_sessions
34            FROM ai_requests
35            WHERE created_at >= NOW() - INTERVAL '1 day' * $1
36            ";
37
38        let mut query = base_query.to_string();
39        let mut params: Vec<Box<dyn ToDbValue>> = vec![Box::new(days)];
40        let mut param_index = 2;
41
42        let placeholder = |idx: &mut i32| {
43            let placeholder = format!("${idx}");
44            *idx += 1;
45            placeholder
46        };
47
48        if let Some(uid) = user_id {
49            query.push_str(&format!(" AND user_id = {}", placeholder(&mut param_index)));
50            params.push(Box::new(uid.to_string()));
51        }
52
53        query.push_str(" GROUP BY provider, model ORDER BY request_count DESC");
54
55        let param_refs: Vec<&dyn ToDbValue> = params.iter().map(|p| &**p).collect();
56
57        let rows = self.db_pool.as_ref().fetch_all(&query, &param_refs).await?;
58
59        rows.iter()
60            .map(ProviderUsage::from_json_row)
61            .collect::<Result<Vec<_>>>()
62    }
63}
64
65#[derive(Debug, Serialize)]
66pub struct ProviderUsage {
67    pub provider: String,
68    pub model: String,
69    pub request_count: i32,
70    pub total_tokens: Option<i32>,
71    pub total_cost_microdollars: Option<i32>,
72    pub avg_latency_ms: Option<f64>,
73    pub unique_users: i32,
74    pub unique_sessions: i32,
75}
76
77impl ProviderUsage {
78    pub fn from_json_row(row: &JsonRow) -> Result<Self> {
79        let provider = row
80            .get("provider")
81            .and_then(|v| v.as_str())
82            .ok_or_else(|| anyhow!("Missing provider"))?
83            .to_string();
84
85        let model = row
86            .get("model")
87            .and_then(|v| v.as_str())
88            .ok_or_else(|| anyhow!("Missing model"))?
89            .to_string();
90
91        let request_count = row
92            .get("request_count")
93            .and_then(serde_json::Value::as_i64)
94            .ok_or_else(|| anyhow!("Missing request_count"))? as i32;
95
96        let total_tokens = row
97            .get("total_tokens")
98            .and_then(serde_json::Value::as_i64)
99            .map(|i| i as i32);
100
101        let total_cost_microdollars = row
102            .get("total_cost_microdollars")
103            .and_then(serde_json::Value::as_i64)
104            .map(|i| i as i32);
105
106        let avg_latency_ms = row
107            .get("avg_latency_ms")
108            .and_then(serde_json::Value::as_f64);
109
110        let unique_users = row
111            .get("unique_users")
112            .and_then(serde_json::Value::as_i64)
113            .ok_or_else(|| anyhow!("Missing unique_users"))? as i32;
114
115        let unique_sessions =
116            row.get("unique_sessions")
117                .and_then(serde_json::Value::as_i64)
118                .ok_or_else(|| anyhow!("Missing unique_sessions"))? as i32;
119
120        Ok(Self {
121            provider,
122            model,
123            request_count,
124            total_tokens,
125            total_cost_microdollars,
126            avg_latency_ms,
127            unique_users,
128            unique_sessions,
129        })
130    }
131}