Skip to main content

systemprompt_analytics/repository/
conversations.rs

1use anyhow::Result;
2use chrono::{DateTime, Utc};
3use sqlx::PgPool;
4use std::sync::Arc;
5use systemprompt_database::DbPool;
6
7use crate::models::cli::{ConversationListRow, TimestampRow};
8
9#[derive(Debug)]
10pub struct ConversationAnalyticsRepository {
11    pool: Arc<PgPool>,
12}
13
14impl ConversationAnalyticsRepository {
15    pub fn new(db: &DbPool) -> Result<Self> {
16        let pool = db.pool_arc()?;
17        Ok(Self { pool })
18    }
19
20    pub async fn list_conversations(
21        &self,
22        start: DateTime<Utc>,
23        end: DateTime<Utc>,
24        limit: i64,
25    ) -> Result<Vec<ConversationListRow>> {
26        sqlx::query_as!(
27            ConversationListRow,
28            r#"
29            SELECT
30                uc.context_id as "context_id!",
31                uc.name,
32                (SELECT COUNT(*) FROM agent_tasks at WHERE at.context_id = uc.context_id)::bigint as "task_count!",
33                (SELECT COUNT(*) FROM task_messages tm
34                 JOIN agent_tasks at ON at.task_id = tm.task_id
35                 WHERE at.context_id = uc.context_id)::bigint as "message_count!",
36                uc.created_at as "created_at!",
37                uc.updated_at as "updated_at!"
38            FROM user_contexts uc
39            WHERE uc.created_at >= $1 AND uc.created_at < $2
40            ORDER BY uc.updated_at DESC
41            LIMIT $3
42            "#,
43            start,
44            end,
45            limit
46        )
47        .fetch_all(&*self.pool)
48        .await
49        .map_err(Into::into)
50    }
51
52    pub async fn get_context_count(&self, start: DateTime<Utc>, end: DateTime<Utc>) -> Result<i64> {
53        let count = sqlx::query_scalar!(
54            r#"SELECT COUNT(*)::bigint as "count!" FROM user_contexts WHERE created_at >= $1 AND created_at < $2"#,
55            start,
56            end
57        )
58        .fetch_one(&*self.pool)
59        .await?;
60        Ok(count)
61    }
62
63    pub async fn get_task_stats(
64        &self,
65        start: DateTime<Utc>,
66        end: DateTime<Utc>,
67    ) -> Result<(i64, Option<f64>)> {
68        let row = sqlx::query!(
69            r#"
70            SELECT COUNT(*)::bigint as "count!", AVG(execution_time_ms)::float8 as avg_time
71            FROM agent_tasks
72            WHERE started_at >= $1 AND started_at < $2
73            "#,
74            start,
75            end
76        )
77        .fetch_one(&*self.pool)
78        .await?;
79        Ok((row.count, row.avg_time))
80    }
81
82    pub async fn get_message_count(&self, start: DateTime<Utc>, end: DateTime<Utc>) -> Result<i64> {
83        let count = sqlx::query_scalar!(
84            r#"SELECT COUNT(*)::bigint as "count!" FROM task_messages WHERE created_at >= $1 AND created_at < $2"#,
85            start,
86            end
87        )
88        .fetch_one(&*self.pool)
89        .await?;
90        Ok(count)
91    }
92
93    pub async fn get_context_timestamps(
94        &self,
95        start: DateTime<Utc>,
96        end: DateTime<Utc>,
97    ) -> Result<Vec<TimestampRow>> {
98        sqlx::query_as!(
99            TimestampRow,
100            r#"
101            SELECT created_at as "timestamp!"
102            FROM user_contexts
103            WHERE created_at >= $1 AND created_at < $2
104            "#,
105            start,
106            end
107        )
108        .fetch_all(&*self.pool)
109        .await
110        .map_err(Into::into)
111    }
112
113    pub async fn get_task_timestamps(
114        &self,
115        start: DateTime<Utc>,
116        end: DateTime<Utc>,
117    ) -> Result<Vec<TimestampRow>> {
118        sqlx::query_as!(
119            TimestampRow,
120            r#"
121            SELECT started_at as "timestamp!"
122            FROM agent_tasks
123            WHERE started_at >= $1 AND started_at < $2
124            "#,
125            start,
126            end
127        )
128        .fetch_all(&*self.pool)
129        .await
130        .map_err(Into::into)
131    }
132
133    pub async fn get_message_timestamps(
134        &self,
135        start: DateTime<Utc>,
136        end: DateTime<Utc>,
137    ) -> Result<Vec<TimestampRow>> {
138        sqlx::query_as!(
139            TimestampRow,
140            r#"
141            SELECT created_at as "timestamp!"
142            FROM task_messages
143            WHERE created_at >= $1 AND created_at < $2
144            "#,
145            start,
146            end
147        )
148        .fetch_all(&*self.pool)
149        .await
150        .map_err(Into::into)
151    }
152}