Skip to main content

systemprompt_analytics/repository/agents/
stats_queries.rs

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