systemprompt_runtime/trace/repository/
events.rs

1//! Log-event and AI-request timeline queries for one trace.
2//!
3//! Copyright (c) systemprompt.io — Business Source License 1.1.
4//! See <https://systemprompt.io> for licensing details.
5
6use serde_json::json;
7
8use systemprompt_identifiers::{ContextId, SessionId, TaskId, TraceId, UserId};
9
10use super::{Result, TraceRepository};
11use crate::trace::models::{AiRequestSummary, TraceEvent};
12
13impl TraceRepository {
14    pub async fn fetch_log_events(&self, trace_id: &TraceId) -> Result<Vec<TraceEvent>> {
15        let rows = sqlx::query!(
16            r#"
17        SELECT
18            timestamp,
19            level as type,
20            CONCAT(module, ': ', message) as details,
21            user_id,
22            session_id,
23            task_id,
24            context_id,
25            metadata
26        FROM logs
27        WHERE trace_id = $1
28        ORDER BY timestamp ASC
29        "#,
30            trace_id.as_str()
31        )
32        .fetch_all(&*self.pool)
33        .await?;
34
35        Ok(rows
36        .into_iter()
37        .map(|row| TraceEvent {
38            event_type: row.r#type,
39            timestamp: row.timestamp,
40            details: row.details.unwrap_or_default(),
41            user_id: row.user_id.map(UserId::new),
42            session_id: row.session_id.map(SessionId::new),
43            task_id: row.task_id.map(TaskId::new),
44            context_id: row.context_id.and_then(|s| {
45                ContextId::try_new(&s)
46                    .map_err(|e| {
47                        tracing::warn!(error = %e, raw = %s, "Skipping non-UUID context_id from logs row");
48                        e
49                    })
50                    .ok()
51            }),
52            metadata: row.metadata,
53        })
54        .collect())
55    }
56
57    pub async fn fetch_ai_request_summary(&self, trace_id: &TraceId) -> Result<AiRequestSummary> {
58        let row = sqlx::query!(
59        r#"
60        SELECT
61            COALESCE(SUM(cost_microdollars), 0)::bigint as total_cost_microdollars,
62            COALESCE(SUM(COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)), 0)::bigint as total_tokens,
63            COALESCE(SUM(input_tokens), 0)::bigint as total_input_tokens,
64            COALESCE(SUM(output_tokens), 0)::bigint as total_output_tokens,
65            COUNT(*)::bigint as request_count,
66            COALESCE(SUM(latency_ms), 0)::bigint as total_latency_ms
67        FROM ai_requests
68        WHERE trace_id = $1
69        "#,
70        trace_id.as_str()
71    )
72    .fetch_one(&*self.pool)
73    .await?;
74
75        Ok(AiRequestSummary {
76            total_cost_microdollars: row.total_cost_microdollars.unwrap_or(0),
77            total_tokens: row.total_tokens.unwrap_or(0),
78            total_input_tokens: row.total_input_tokens.unwrap_or(0),
79            total_output_tokens: row.total_output_tokens.unwrap_or(0),
80            request_count: row.request_count.unwrap_or(0),
81            total_latency_ms: row.total_latency_ms.unwrap_or(0),
82        })
83    }
84
85    pub async fn fetch_ai_request_events(&self, trace_id: &TraceId) -> Result<Vec<TraceEvent>> {
86        let rows = sqlx::query!(
87            r#"
88        SELECT
89            created_at as timestamp,
90            provider,
91            model,
92            input_tokens,
93            output_tokens,
94            cost_microdollars,
95            latency_ms,
96            status,
97            user_id,
98            session_id,
99            task_id,
100            context_id
101        FROM ai_requests
102        WHERE trace_id = $1
103        ORDER BY created_at ASC
104        "#,
105            trace_id.as_str()
106        )
107        .fetch_all(&*self.pool)
108        .await?;
109
110        Ok(rows
111        .into_iter()
112        .map(|row| {
113            let provider = row.provider.clone().unwrap_or_else(|| "-".to_owned());
114            let model = row.model.clone().unwrap_or_else(|| "-".to_owned());
115            let details = format!(
116                "{provider}/{model}: {} (in:{}, out:{}, {}ms)",
117                row.status,
118                row.input_tokens.unwrap_or(0),
119                row.output_tokens.unwrap_or(0),
120                row.latency_ms.unwrap_or(0)
121            );
122
123            let metadata = json!({
124                "cost_microdollars": row.cost_microdollars,
125                "latency_ms": row.latency_ms,
126                "input_tokens": row.input_tokens,
127                "output_tokens": row.output_tokens,
128                "tokens_used": row.input_tokens.unwrap_or(0) + row.output_tokens.unwrap_or(0),
129                "provider": row.provider,
130                "model": row.model
131            });
132
133            TraceEvent {
134                event_type: "AI".to_owned(),
135                timestamp: row.timestamp,
136                details,
137                user_id: Some(UserId::new(row.user_id)),
138                session_id: row.session_id.map(SessionId::new),
139                task_id: row.task_id.map(TaskId::new),
140                context_id: ContextId::try_new(&row.context_id)
141                    .map_err(|e| {
142                        tracing::warn!(error = %e, raw = %row.context_id, "Skipping non-UUID context_id from ai_requests row");
143                        e
144                    })
145                    .ok(),
146                metadata: Some(metadata.to_string()),
147            }
148        })
149        .collect())
150    }
151}
systemprompt_runtime/trace/repository/events.rs

systemprompt_runtime/trace/repository/
events.rs
