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_cents)::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_cents)::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_cents), 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_cents) 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_cents), 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_cents) 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            SELECT
131                COALESCE(at.agent_name, 'unknown') as "name!",
132                COALESCE(SUM(r.cost_cents), 0)::bigint as "cost!",
133                COUNT(*)::bigint as "requests!",
134                COALESCE(SUM(r.tokens_used), 0)::bigint as "tokens!"
135            FROM ai_requests r
136            LEFT JOIN agent_tasks at ON at.task_id = r.task_id
137            WHERE r.created_at >= $1 AND r.created_at < $2
138            GROUP BY at.agent_name
139            ORDER BY SUM(r.cost_cents) DESC NULLS LAST
140            LIMIT $3
141            "#,
142            start,
143            end,
144            limit
145        )
146        .fetch_all(&*self.pool)
147        .await
148        .map_err(Into::into)
149    }
150
151    pub async fn get_costs_for_trends(
152        &self,
153        start: DateTime<Utc>,
154        end: DateTime<Utc>,
155    ) -> Result<Vec<CostTrendRow>> {
156        sqlx::query_as!(
157            CostTrendRow,
158            r#"
159            SELECT
160                created_at as "created_at!",
161                cost_cents,
162                tokens_used
163            FROM ai_requests
164            WHERE created_at >= $1 AND created_at < $2
165            ORDER BY created_at
166            "#,
167            start,
168            end
169        )
170        .fetch_all(&*self.pool)
171        .await
172        .map_err(Into::into)
173    }
174}
systemprompt_analytics/repository/costs.rs

systemprompt_analytics/repository/
costs.rs
