Skip to main content

systemprompt_analytics/repository/
queries.rs

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