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