systemprompt_analytics/repository/
queries.rs

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

systemprompt_analytics/repository/
queries.rs
