systemprompt_runtime/trace/repository/
request_stats.rs

1//! AI-request aggregate queries: totals, per-provider and per-model rollups.
2//!
3//! Copyright (c) systemprompt.io — Business Source License 1.1.
4//! See <https://systemprompt.io> for licensing details.
5
6use chrono::{DateTime, Utc};
7
8use super::{Result, TraceRepository};
9use crate::trace::models::{AiRequestStats, ModelStatsRow, ProviderStatsRow};
10
11struct TotalRow {
12    request_count: Option<i64>,
13    total_input_tokens: Option<i64>,
14    total_output_tokens: Option<i64>,
15    total_cost_microdollars: Option<i64>,
16    avg_latency_ms: Option<i64>,
17}
18
19struct ProviderRow {
20    provider: String,
21    request_count: Option<i64>,
22    total_tokens: Option<i64>,
23    total_cost_microdollars: Option<i64>,
24    avg_latency_ms: Option<i64>,
25}
26
27struct ModelRow {
28    model: String,
29    provider: String,
30    request_count: Option<i64>,
31    total_tokens: Option<i64>,
32    total_cost_microdollars: Option<i64>,
33    avg_latency_ms: Option<i64>,
34}
35
36impl TraceRepository {
37    pub async fn get_ai_request_stats(
38        &self,
39        since: Option<DateTime<Utc>>,
40    ) -> Result<AiRequestStats> {
41        let totals = self.fetch_request_totals(since).await?;
42        let provider_rows = self.fetch_provider_stats(since).await?;
43        let model_rows = self.fetch_model_stats(since).await?;
44
45        Ok(AiRequestStats {
46            total_requests: totals.request_count.unwrap_or(0),
47            total_input_tokens: totals.total_input_tokens.unwrap_or(0),
48            total_output_tokens: totals.total_output_tokens.unwrap_or(0),
49            total_cost_microdollars: totals.total_cost_microdollars.unwrap_or(0),
50            avg_latency_ms: totals.avg_latency_ms.unwrap_or(0),
51            by_provider: provider_rows
52                .into_iter()
53                .map(|r| ProviderStatsRow {
54                    provider: r.provider,
55                    request_count: r.request_count.unwrap_or(0),
56                    total_tokens: r.total_tokens.unwrap_or(0),
57                    total_cost_microdollars: r.total_cost_microdollars.unwrap_or(0),
58                    avg_latency_ms: r.avg_latency_ms.unwrap_or(0),
59                })
60                .collect(),
61            by_model: model_rows
62                .into_iter()
63                .map(|r| ModelStatsRow {
64                    model: r.model,
65                    provider: r.provider,
66                    request_count: r.request_count.unwrap_or(0),
67                    total_tokens: r.total_tokens.unwrap_or(0),
68                    total_cost_microdollars: r.total_cost_microdollars.unwrap_or(0),
69                    avg_latency_ms: r.avg_latency_ms.unwrap_or(0),
70                })
71                .collect(),
72        })
73    }
74
75    async fn fetch_request_totals(&self, since: Option<DateTime<Utc>>) -> Result<TotalRow> {
76        sqlx::query_as!(
77            TotalRow,
78            r#"
79        SELECT
80            COUNT(*) as "request_count",
81            COALESCE(SUM(input_tokens), 0) as "total_input_tokens",
82            COALESCE(SUM(output_tokens), 0) as "total_output_tokens",
83            COALESCE(SUM(cost_microdollars), 0)::bigint as "total_cost_microdollars",
84            COALESCE(AVG(latency_ms), 0)::bigint as "avg_latency_ms"
85        FROM ai_requests
86        WHERE ($1::timestamptz IS NULL OR created_at >= $1)
87        "#,
88            since
89        )
90        .fetch_one(&*self.pool)
91        .await
92        .map_err(Into::into)
93    }
94
95    async fn fetch_provider_stats(&self, since: Option<DateTime<Utc>>) -> Result<Vec<ProviderRow>> {
96        sqlx::query_as!(
97            ProviderRow,
98            r#"
99        SELECT
100            provider as "provider!",
101            COUNT(*) as "request_count",
102            COALESCE(SUM(input_tokens), 0) + COALESCE(SUM(output_tokens), 0) as "total_tokens",
103            COALESCE(SUM(cost_microdollars), 0)::bigint as "total_cost_microdollars",
104            COALESCE(AVG(latency_ms), 0)::bigint as "avg_latency_ms"
105        FROM ai_requests
106        WHERE ($1::timestamptz IS NULL OR created_at >= $1)
107          AND provider IS NOT NULL
108        GROUP BY provider
109        ORDER BY request_count DESC
110        "#,
111            since
112        )
113        .fetch_all(&*self.pool)
114        .await
115        .map_err(Into::into)
116    }
117
118    async fn fetch_model_stats(&self, since: Option<DateTime<Utc>>) -> Result<Vec<ModelRow>> {
119        sqlx::query_as!(
120            ModelRow,
121            r#"
122        SELECT
123            model as "model!",
124            provider as "provider!",
125            COUNT(*) as "request_count",
126            COALESCE(SUM(input_tokens), 0) + COALESCE(SUM(output_tokens), 0) as "total_tokens",
127            COALESCE(SUM(cost_microdollars), 0)::bigint as "total_cost_microdollars",
128            COALESCE(AVG(latency_ms), 0)::bigint as "avg_latency_ms"
129        FROM ai_requests
130        WHERE ($1::timestamptz IS NULL OR created_at >= $1)
131          AND model IS NOT NULL AND provider IS NOT NULL
132        GROUP BY model, provider
133        ORDER BY request_count DESC
134        LIMIT 10
135        "#,
136            since
137        )
138        .fetch_all(&*self.pool)
139        .await
140        .map_err(Into::into)
141    }
142}
systemprompt_runtime/trace/repository/request_stats.rs

systemprompt_runtime/trace/repository/
request_stats.rs
