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 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}