systemprompt_analytics/repository/
requests.rs1use 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}