Skip to main content

systemprompt_analytics/repository/costs/
platform.rs

1use super::CostAnalyticsRepository;
2use crate::Result;
3use chrono::{DateTime, Utc};
4
5use crate::models::cli::{CostBreakdownRow, CostSummaryRow, CostTrendRow, PreviousCostRow};
6
7impl CostAnalyticsRepository {
8    pub async fn get_summary(
9        &self,
10        start: DateTime<Utc>,
11        end: DateTime<Utc>,
12    ) -> Result<CostSummaryRow> {
13        sqlx::query_as!(
14            CostSummaryRow,
15            r#"
16            SELECT
17                COUNT(*)::bigint as "requests!",
18                SUM(cost_microdollars)::bigint as "cost",
19                SUM(tokens_used)::bigint as "tokens"
20            FROM ai_requests
21            WHERE created_at >= $1 AND created_at < $2
22            "#,
23            start,
24            end
25        )
26        .fetch_one(&*self.pool)
27        .await
28        .map_err(Into::into)
29    }
30
31    pub async fn get_previous_cost(
32        &self,
33        start: DateTime<Utc>,
34        end: DateTime<Utc>,
35    ) -> Result<PreviousCostRow> {
36        sqlx::query_as!(
37            PreviousCostRow,
38            r#"
39            SELECT SUM(cost_microdollars)::bigint as "cost"
40            FROM ai_requests
41            WHERE created_at >= $1 AND created_at < $2
42            "#,
43            start,
44            end
45        )
46        .fetch_one(&*self.pool)
47        .await
48        .map_err(Into::into)
49    }
50
51    pub async fn get_breakdown_by_model(
52        &self,
53        start: DateTime<Utc>,
54        end: DateTime<Utc>,
55        limit: i64,
56    ) -> Result<Vec<CostBreakdownRow>> {
57        sqlx::query_as!(
58            CostBreakdownRow,
59            r#"
60            SELECT
61                model as "name!",
62                COALESCE(SUM(cost_microdollars), 0)::bigint as "cost!",
63                COUNT(*)::bigint as "requests!",
64                COALESCE(SUM(tokens_used), 0)::bigint as "tokens!"
65            FROM ai_requests
66            WHERE created_at >= $1 AND created_at < $2
67            GROUP BY model
68            ORDER BY SUM(cost_microdollars) DESC NULLS LAST
69            LIMIT $3
70            "#,
71            start,
72            end,
73            limit
74        )
75        .fetch_all(&*self.pool)
76        .await
77        .map_err(Into::into)
78    }
79
80    pub async fn get_breakdown_by_provider(
81        &self,
82        start: DateTime<Utc>,
83        end: DateTime<Utc>,
84        limit: i64,
85    ) -> Result<Vec<CostBreakdownRow>> {
86        sqlx::query_as!(
87            CostBreakdownRow,
88            r#"
89            SELECT
90                provider as "name!",
91                COALESCE(SUM(cost_microdollars), 0)::bigint as "cost!",
92                COUNT(*)::bigint as "requests!",
93                COALESCE(SUM(tokens_used), 0)::bigint as "tokens!"
94            FROM ai_requests
95            WHERE created_at >= $1 AND created_at < $2
96            GROUP BY provider
97            ORDER BY SUM(cost_microdollars) DESC NULLS LAST
98            LIMIT $3
99            "#,
100            start,
101            end,
102            limit
103        )
104        .fetch_all(&*self.pool)
105        .await
106        .map_err(Into::into)
107    }
108
109    pub async fn get_breakdown_by_agent(
110        &self,
111        start: DateTime<Utc>,
112        end: DateTime<Utc>,
113        limit: i64,
114    ) -> Result<Vec<CostBreakdownRow>> {
115        sqlx::query_as!(
116            CostBreakdownRow,
117            r#"
118            (
119                SELECT
120                    at.agent_name as "name!",
121                    COALESCE(SUM(r.cost_microdollars), 0)::bigint as "cost!",
122                    COUNT(*)::bigint as "requests!",
123                    COALESCE(SUM(r.tokens_used), 0)::bigint as "tokens!"
124                FROM ai_requests r
125                INNER JOIN agent_tasks at ON at.task_id = r.task_id
126                WHERE r.created_at >= $1 AND r.created_at < $2
127                  AND at.agent_name IS NOT NULL
128                GROUP BY at.agent_name
129                ORDER BY SUM(r.cost_microdollars) DESC NULLS LAST
130                LIMIT $3
131            )
132            UNION ALL
133            (
134                SELECT
135                    'unattributed' as "name!",
136                    COALESCE(SUM(r.cost_microdollars), 0)::bigint as "cost!",
137                    COUNT(*)::bigint as "requests!",
138                    COALESCE(SUM(r.tokens_used), 0)::bigint as "tokens!"
139                FROM ai_requests r
140                LEFT JOIN agent_tasks at ON at.task_id = r.task_id
141                WHERE r.created_at >= $1 AND r.created_at < $2
142                  AND (r.task_id IS NULL OR at.agent_name IS NULL)
143                HAVING COUNT(*) > 0
144            )
145            "#,
146            start,
147            end,
148            limit
149        )
150        .fetch_all(&*self.pool)
151        .await
152        .map_err(Into::into)
153    }
154
155    pub async fn get_costs_for_trends(
156        &self,
157        start: DateTime<Utc>,
158        end: DateTime<Utc>,
159    ) -> Result<Vec<CostTrendRow>> {
160        sqlx::query_as!(
161            CostTrendRow,
162            r#"
163            SELECT
164                created_at as "created_at!",
165                cost_microdollars,
166                tokens_used
167            FROM ai_requests
168            WHERE created_at >= $1 AND created_at < $2
169            ORDER BY created_at
170            "#,
171            start,
172            end
173        )
174        .fetch_all(&*self.pool)
175        .await
176        .map_err(Into::into)
177    }
178}