systemprompt_analytics/repository/
costs.rs

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

systemprompt_analytics/repository/
costs.rs
