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