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              AND provider IS NOT NULL AND model IS NOT NULL
103            GROUP BY provider, model
104            ORDER BY COUNT(*) DESC
105            LIMIT $3
106            "#,
107            start,
108            end,
109            limit
110        )
111        .fetch_all(&*self.pool)
112        .await
113        .map_err(Into::into)
114    }
115
116    pub async fn get_requests_for_trends(
117        &self,
118        start: DateTime<Utc>,
119        end: DateTime<Utc>,
120    ) -> Result<Vec<RequestTrendRow>> {
121        sqlx::query_as!(
122            RequestTrendRow,
123            r#"
124            SELECT
125                created_at as "created_at!",
126                tokens_used,
127                cost_microdollars,
128                latency_ms
129            FROM ai_requests
130            WHERE created_at >= $1 AND created_at < $2
131            ORDER BY created_at
132            "#,
133            start,
134            end
135        )
136        .fetch_all(&*self.pool)
137        .await
138        .map_err(Into::into)
139    }
140
141    pub async fn list_requests(
142        &self,
143        start: DateTime<Utc>,
144        end: DateTime<Utc>,
145        limit: i64,
146        model_filter: Option<&str>,
147    ) -> Result<Vec<RequestListRow>> {
148        if let Some(model) = model_filter {
149            let pattern = format!("%{model}%");
150            sqlx::query_as!(
151                RequestListRow,
152                r#"
153                SELECT
154                    id as "id!",
155                    provider,
156                    model,
157                    input_tokens,
158                    output_tokens,
159                    cost_microdollars,
160                    latency_ms,
161                    cache_hit,
162                    created_at as "created_at!",
163                    status as "status!",
164                    error_message,
165                    user_id as "user_id!: UserId"
166                FROM ai_requests
167                WHERE created_at >= $1 AND created_at < $2
168                  AND model ILIKE $3
169                ORDER BY created_at DESC
170                LIMIT $4
171                "#,
172                start,
173                end,
174                pattern,
175                limit
176            )
177            .fetch_all(&*self.pool)
178            .await
179            .map_err(Into::into)
180        } else {
181            sqlx::query_as!(
182                RequestListRow,
183                r#"
184                SELECT
185                    id as "id!",
186                    provider,
187                    model,
188                    input_tokens,
189                    output_tokens,
190                    cost_microdollars,
191                    latency_ms,
192                    cache_hit,
193                    created_at as "created_at!",
194                    status as "status!",
195                    error_message,
196                    user_id as "user_id!: UserId"
197                FROM ai_requests
198                WHERE created_at >= $1 AND created_at < $2
199                ORDER BY created_at DESC
200                LIMIT $3
201                "#,
202                start,
203                end,
204                limit
205            )
206            .fetch_all(&*self.pool)
207            .await
208            .map_err(Into::into)
209        }
210    }
211}