Skip to main content

systemprompt_analytics/repository/
requests.rs

1//! AI-request analytics over the `ai_requests` table.
2//!
3//! [`RequestAnalyticsRepository`] reports token, cost, latency, and
4//! cache-hit stats, per-model usage breakdowns, trend series, and a request
5//! list, each optionally filtered by a model substring.
6
7use crate::Result;
8use chrono::{DateTime, Utc};
9use sqlx::PgPool;
10use std::sync::Arc;
11use systemprompt_database::DbPool;
12use systemprompt_identifiers::UserId;
13
14use crate::models::cli::{ModelUsageRow, RequestListRow, RequestStatsRow, RequestTrendRow};
15
16#[derive(Debug)]
17pub struct RequestAnalyticsRepository {
18    pool: Arc<PgPool>,
19}
20
21impl RequestAnalyticsRepository {
22    pub fn new(db: &DbPool) -> Result<Self> {
23        let pool = db.pool_arc()?;
24        Ok(Self { pool })
25    }
26
27    pub async fn get_stats(
28        &self,
29        start: DateTime<Utc>,
30        end: DateTime<Utc>,
31        model_filter: Option<&str>,
32    ) -> Result<RequestStatsRow> {
33        if let Some(model) = model_filter {
34            let pattern = format!("%{}%", model);
35            sqlx::query_as!(
36                RequestStatsRow,
37                r#"
38                SELECT
39                    COUNT(*)::bigint as "total!",
40                    SUM(tokens_used)::bigint as "total_tokens",
41                    SUM(input_tokens)::bigint as "input_tokens",
42                    SUM(output_tokens)::bigint as "output_tokens",
43                    SUM(cost_microdollars)::bigint as "cost",
44                    AVG(latency_ms)::float8 as "avg_latency",
45                    COUNT(*) FILTER (WHERE cache_hit = true)::bigint as "cache_hits!"
46                FROM ai_requests
47                WHERE created_at >= $1 AND created_at < $2
48                  AND model ILIKE $3
49                "#,
50                start,
51                end,
52                pattern
53            )
54            .fetch_one(&*self.pool)
55            .await
56            .map_err(Into::into)
57        } else {
58            sqlx::query_as!(
59                RequestStatsRow,
60                r#"
61                SELECT
62                    COUNT(*)::bigint as "total!",
63                    SUM(tokens_used)::bigint as "total_tokens",
64                    SUM(input_tokens)::bigint as "input_tokens",
65                    SUM(output_tokens)::bigint as "output_tokens",
66                    SUM(cost_microdollars)::bigint as "cost",
67                    AVG(latency_ms)::float8 as "avg_latency",
68                    COUNT(*) FILTER (WHERE cache_hit = true)::bigint as "cache_hits!"
69                FROM ai_requests
70                WHERE created_at >= $1 AND created_at < $2
71                "#,
72                start,
73                end
74            )
75            .fetch_one(&*self.pool)
76            .await
77            .map_err(Into::into)
78        }
79    }
80
81    pub async fn list_models(
82        &self,
83        start: DateTime<Utc>,
84        end: DateTime<Utc>,
85        limit: i64,
86    ) -> Result<Vec<ModelUsageRow>> {
87        sqlx::query_as!(
88            ModelUsageRow,
89            r#"
90            SELECT
91                provider as "provider!",
92                model as "model!",
93                COUNT(*)::bigint as "request_count!",
94                SUM(tokens_used)::bigint as "total_tokens",
95                SUM(cost_microdollars)::bigint as "total_cost",
96                AVG(latency_ms)::float8 as "avg_latency"
97            FROM ai_requests
98            WHERE created_at >= $1 AND created_at < $2
99            GROUP BY provider, model
100            ORDER BY COUNT(*) DESC
101            LIMIT $3
102            "#,
103            start,
104            end,
105            limit
106        )
107        .fetch_all(&*self.pool)
108        .await
109        .map_err(Into::into)
110    }
111
112    pub async fn get_requests_for_trends(
113        &self,
114        start: DateTime<Utc>,
115        end: DateTime<Utc>,
116    ) -> Result<Vec<RequestTrendRow>> {
117        sqlx::query_as!(
118            RequestTrendRow,
119            r#"
120            SELECT
121                created_at as "created_at!",
122                tokens_used,
123                cost_microdollars,
124                latency_ms
125            FROM ai_requests
126            WHERE created_at >= $1 AND created_at < $2
127            ORDER BY created_at
128            "#,
129            start,
130            end
131        )
132        .fetch_all(&*self.pool)
133        .await
134        .map_err(Into::into)
135    }
136
137    pub async fn list_requests(
138        &self,
139        start: DateTime<Utc>,
140        end: DateTime<Utc>,
141        limit: i64,
142        model_filter: Option<&str>,
143    ) -> Result<Vec<RequestListRow>> {
144        if let Some(model) = model_filter {
145            let pattern = format!("%{}%", model);
146            sqlx::query_as!(
147                RequestListRow,
148                r#"
149                SELECT
150                    id as "id!",
151                    provider as "provider!",
152                    model as "model!",
153                    input_tokens,
154                    output_tokens,
155                    cost_microdollars,
156                    latency_ms,
157                    cache_hit,
158                    created_at as "created_at!",
159                    status as "status!",
160                    error_message,
161                    user_id as "user_id!: UserId"
162                FROM ai_requests
163                WHERE created_at >= $1 AND created_at < $2
164                  AND model ILIKE $3
165                ORDER BY created_at DESC
166                LIMIT $4
167                "#,
168                start,
169                end,
170                pattern,
171                limit
172            )
173            .fetch_all(&*self.pool)
174            .await
175            .map_err(Into::into)
176        } else {
177            sqlx::query_as!(
178                RequestListRow,
179                r#"
180                SELECT
181                    id as "id!",
182                    provider as "provider!",
183                    model as "model!",
184                    input_tokens,
185                    output_tokens,
186                    cost_microdollars,
187                    latency_ms,
188                    cache_hit,
189                    created_at as "created_at!",
190                    status as "status!",
191                    error_message,
192                    user_id as "user_id!: UserId"
193                FROM ai_requests
194                WHERE created_at >= $1 AND created_at < $2
195                ORDER BY created_at DESC
196                LIMIT $3
197                "#,
198                start,
199                end,
200                limit
201            )
202            .fetch_all(&*self.pool)
203            .await
204            .map_err(Into::into)
205        }
206    }
207}