Skip to main content

systemprompt_analytics/repository/costs/
per_user.rs

1use super::CostAnalyticsRepository;
2use crate::Result;
3use chrono::{DateTime, Utc};
4use systemprompt_identifiers::{ContextId, UserId};
5
6use crate::models::cli::{
7    ContextGroupRow, ContextSummaryRow, CostBreakdownRow, CostSummaryRow, PreviousCostRow,
8    RecentContextRow,
9};
10
11impl CostAnalyticsRepository {
12    pub async fn get_summary_for_user(
13        &self,
14        user_id: &UserId,
15        start: DateTime<Utc>,
16        end: DateTime<Utc>,
17    ) -> Result<CostSummaryRow> {
18        sqlx::query_as!(
19            CostSummaryRow,
20            r#"
21            SELECT
22                COUNT(*)::bigint as "requests!",
23                SUM(cost_microdollars)::bigint as "cost",
24                SUM(tokens_used)::bigint as "tokens"
25            FROM ai_requests
26            WHERE created_at >= $1 AND created_at < $2 AND user_id = $3
27            "#,
28            start,
29            end,
30            user_id.as_str()
31        )
32        .fetch_one(&*self.pool)
33        .await
34        .map_err(Into::into)
35    }
36
37    pub async fn get_previous_cost_for_user(
38        &self,
39        user_id: &UserId,
40        start: DateTime<Utc>,
41        end: DateTime<Utc>,
42    ) -> Result<PreviousCostRow> {
43        sqlx::query_as!(
44            PreviousCostRow,
45            r#"
46            SELECT SUM(cost_microdollars)::bigint as "cost"
47            FROM ai_requests
48            WHERE created_at >= $1 AND created_at < $2 AND user_id = $3
49            "#,
50            start,
51            end,
52            user_id.as_str()
53        )
54        .fetch_one(&*self.pool)
55        .await
56        .map_err(Into::into)
57    }
58
59    pub async fn get_breakdown_by_model_for_user(
60        &self,
61        user_id: &UserId,
62        start: DateTime<Utc>,
63        end: DateTime<Utc>,
64        limit: i64,
65    ) -> Result<Vec<CostBreakdownRow>> {
66        sqlx::query_as!(
67            CostBreakdownRow,
68            r#"
69            SELECT
70                model as "name!",
71                COALESCE(SUM(cost_microdollars), 0)::bigint as "cost!",
72                COUNT(*)::bigint as "requests!",
73                COALESCE(SUM(tokens_used), 0)::bigint as "tokens!"
74            FROM ai_requests
75            WHERE created_at >= $1 AND created_at < $2 AND user_id = $4
76            GROUP BY model
77            ORDER BY SUM(cost_microdollars) DESC NULLS LAST
78            LIMIT $3
79            "#,
80            start,
81            end,
82            limit,
83            user_id.as_str()
84        )
85        .fetch_all(&*self.pool)
86        .await
87        .map_err(Into::into)
88    }
89
90    pub async fn get_context_summary_for_user(
91        &self,
92        user_id: &UserId,
93        start: DateTime<Utc>,
94        end: DateTime<Utc>,
95    ) -> Result<ContextSummaryRow> {
96        sqlx::query_as!(
97            ContextSummaryRow,
98            r#"
99            SELECT
100                COUNT(DISTINCT context_id)::bigint as "conversations!",
101                COUNT(*)::bigint as "ai_requests!"
102            FROM ai_requests
103            WHERE created_at >= $1 AND created_at < $2
104              AND user_id = $3
105              AND context_id IS NOT NULL
106            "#,
107            start,
108            end,
109            user_id.as_str()
110        )
111        .fetch_one(&*self.pool)
112        .await
113        .map_err(Into::into)
114    }
115
116    pub async fn get_contexts_by_model_for_user(
117        &self,
118        user_id: &UserId,
119        start: DateTime<Utc>,
120        end: DateTime<Utc>,
121        limit: i64,
122    ) -> Result<Vec<ContextGroupRow>> {
123        sqlx::query_as!(
124            ContextGroupRow,
125            r#"
126            SELECT
127                model as "name!",
128                COUNT(DISTINCT context_id)::bigint as "conversations!",
129                COUNT(*)::bigint as "ai_requests!"
130            FROM ai_requests
131            WHERE created_at >= $1 AND created_at < $2
132              AND user_id = $3
133              AND context_id IS NOT NULL
134            GROUP BY model
135            ORDER BY COUNT(DISTINCT context_id) DESC
136            LIMIT $4
137            "#,
138            start,
139            end,
140            user_id.as_str(),
141            limit
142        )
143        .fetch_all(&*self.pool)
144        .await
145        .map_err(Into::into)
146    }
147
148    pub async fn get_contexts_by_agent_for_user(
149        &self,
150        user_id: &UserId,
151        start: DateTime<Utc>,
152        end: DateTime<Utc>,
153        limit: i64,
154    ) -> Result<Vec<ContextGroupRow>> {
155        sqlx::query_as!(
156            ContextGroupRow,
157            r#"
158            SELECT
159                COALESCE(at.agent_name, 'unattributed') as "name!",
160                COUNT(DISTINCT r.context_id)::bigint as "conversations!",
161                COUNT(*)::bigint as "ai_requests!"
162            FROM ai_requests r
163            LEFT JOIN agent_tasks at ON at.task_id = r.task_id
164            WHERE r.created_at >= $1 AND r.created_at < $2
165              AND r.user_id = $3
166              AND r.context_id IS NOT NULL
167            GROUP BY COALESCE(at.agent_name, 'unattributed')
168            ORDER BY COUNT(DISTINCT r.context_id) DESC
169            LIMIT $4
170            "#,
171            start,
172            end,
173            user_id.as_str(),
174            limit
175        )
176        .fetch_all(&*self.pool)
177        .await
178        .map_err(Into::into)
179    }
180
181    pub async fn get_recent_contexts_for_user(
182        &self,
183        user_id: &UserId,
184        end: DateTime<Utc>,
185        limit: i64,
186    ) -> Result<Vec<RecentContextRow>> {
187        sqlx::query_as!(
188            RecentContextRow,
189            r#"
190            SELECT
191                ctx.context_id as "context_id!: ContextId",
192                ctx.last_activity as "last_activity!",
193                ctx.ai_requests as "ai_requests!",
194                last_req.model,
195                last_task.agent_name
196            FROM (
197                SELECT
198                    r.context_id,
199                    MAX(r.created_at) AS last_activity,
200                    COUNT(*) AS ai_requests
201                FROM ai_requests r
202                WHERE r.user_id = $1
203                  AND r.created_at < $2
204                  AND r.context_id IS NOT NULL
205                GROUP BY r.context_id
206                ORDER BY MAX(r.created_at) DESC
207                LIMIT $3
208            ) ctx
209            LEFT JOIN LATERAL (
210                SELECT model, task_id FROM ai_requests
211                WHERE context_id = ctx.context_id
212                ORDER BY created_at DESC
213                LIMIT 1
214            ) last_req ON TRUE
215            LEFT JOIN agent_tasks last_task ON last_task.task_id = last_req.task_id
216            ORDER BY ctx.last_activity DESC
217            "#,
218            user_id.as_str(),
219            end,
220            limit
221        )
222        .fetch_all(&*self.pool)
223        .await
224        .map_err(Into::into)
225    }
226}