Skip to main content

systemprompt_analytics/repository/agents/
stats_queries.rs

1//! Per-agent aggregate statistics queries.
2//!
3//! Copyright (c) systemprompt.io — Business Source License 1.1.
4//! See <https://systemprompt.io> for licensing details.
5
6use crate::Result;
7use chrono::{DateTime, Utc};
8
9use super::AgentAnalyticsRepository;
10use crate::models::cli::{AgentAiStatsRow, AgentStatsRow, AgentTaskRow};
11
12impl AgentAnalyticsRepository {
13    pub async fn get_stats(
14        &self,
15        start: DateTime<Utc>,
16        end: DateTime<Utc>,
17        agent_filter: Option<&str>,
18    ) -> Result<AgentStatsRow> {
19        if let Some(agent) = agent_filter {
20            let pattern = format!("%{}%", agent);
21            sqlx::query_as!(
22                AgentStatsRow,
23                r#"
24                SELECT
25                    COUNT(DISTINCT agent_name)::bigint as "total_agents!",
26                    COUNT(*)::bigint as "total_tasks!",
27                    -- agent_tasks.status stores the canonical A2A TaskState value
28                    -- written by task_state_to_db_string (TASK_STATE_*), which the
29                    -- agent_tasks_status_check CHECK constraint also enforces.
30                    COUNT(*) FILTER (WHERE status = 'TASK_STATE_COMPLETED')::bigint as "completed_tasks!",
31                    COUNT(*) FILTER (WHERE status = 'TASK_STATE_FAILED')::bigint as "failed_tasks!",
32                    COALESCE(AVG(execution_time_ms)::float8, 0) as "avg_execution_time_ms!"
33                FROM agent_tasks
34                WHERE started_at >= $1 AND started_at < $2
35                  AND agent_name ILIKE $3
36                "#,
37                start,
38                end,
39                pattern
40            )
41            .fetch_one(&*self.pool)
42            .await
43            .map_err(Into::into)
44        } else {
45            sqlx::query_as!(
46                AgentStatsRow,
47                r#"
48                SELECT
49                    COUNT(DISTINCT agent_name)::bigint as "total_agents!",
50                    COUNT(*)::bigint as "total_tasks!",
51                    COUNT(*) FILTER (WHERE status = 'TASK_STATE_COMPLETED')::bigint as "completed_tasks!",
52                    COUNT(*) FILTER (WHERE status = 'TASK_STATE_FAILED')::bigint as "failed_tasks!",
53                    COALESCE(AVG(execution_time_ms)::float8, 0) as "avg_execution_time_ms!"
54                FROM agent_tasks
55                WHERE started_at >= $1 AND started_at < $2
56                "#,
57                start,
58                end
59            )
60            .fetch_one(&*self.pool)
61            .await
62            .map_err(Into::into)
63        }
64    }
65
66    pub async fn get_ai_stats(
67        &self,
68        start: DateTime<Utc>,
69        end: DateTime<Utc>,
70    ) -> Result<AgentAiStatsRow> {
71        sqlx::query_as!(
72            AgentAiStatsRow,
73            r#"
74            SELECT
75                COUNT(*)::bigint as "total_ai_requests!",
76                COALESCE(SUM(cost_microdollars), 0)::bigint as "total_cost_microdollars!"
77            FROM ai_requests
78            WHERE created_at >= $1 AND created_at < $2
79            "#,
80            start,
81            end
82        )
83        .fetch_one(&*self.pool)
84        .await
85        .map_err(Into::into)
86    }
87
88    pub async fn get_tasks_for_trends(
89        &self,
90        start: DateTime<Utc>,
91        end: DateTime<Utc>,
92        agent_filter: Option<&str>,
93    ) -> Result<Vec<AgentTaskRow>> {
94        if let Some(agent) = agent_filter {
95            let pattern = format!("%{}%", agent);
96            sqlx::query_as!(
97                AgentTaskRow,
98                r#"
99                SELECT
100                    started_at as "started_at!",
101                    status,
102                    execution_time_ms
103                FROM agent_tasks
104                WHERE started_at >= $1 AND started_at < $2
105                  AND agent_name ILIKE $3
106                ORDER BY started_at
107                "#,
108                start,
109                end,
110                pattern
111            )
112            .fetch_all(&*self.pool)
113            .await
114            .map_err(Into::into)
115        } else {
116            sqlx::query_as!(
117                AgentTaskRow,
118                r#"
119                SELECT
120                    started_at as "started_at!",
121                    status,
122                    execution_time_ms
123                FROM agent_tasks
124                WHERE started_at >= $1 AND started_at < $2
125                ORDER BY started_at
126                "#,
127                start,
128                end
129            )
130            .fetch_all(&*self.pool)
131            .await
132            .map_err(Into::into)
133        }
134    }
135}