Skip to main content

systemprompt_logging/trace/
service.rs

1use anyhow::Result;
2use sqlx::PgPool;
3use std::sync::Arc;
4
5use super::models::{AiRequestSummary, ExecutionStepSummary, McpExecutionSummary, TraceEvent};
6use super::queries;
7
8#[derive(Debug, Clone)]
9pub struct TraceQueryService {
10    pool: Arc<PgPool>,
11}
12
13impl TraceQueryService {
14    pub const fn new(pool: Arc<PgPool>) -> Self {
15        Self { pool }
16    }
17
18    pub async fn get_log_events(&self, trace_id: &str) -> Result<Vec<TraceEvent>> {
19        queries::fetch_log_events(&self.pool, trace_id).await
20    }
21
22    pub async fn get_ai_request_summary(&self, trace_id: &str) -> Result<AiRequestSummary> {
23        queries::fetch_ai_request_summary(&self.pool, trace_id).await
24    }
25
26    pub async fn get_ai_request_events(&self, trace_id: &str) -> Result<Vec<TraceEvent>> {
27        queries::fetch_ai_request_events(&self.pool, trace_id).await
28    }
29
30    pub async fn get_mcp_execution_summary(&self, trace_id: &str) -> Result<McpExecutionSummary> {
31        queries::fetch_mcp_execution_summary(&self.pool, trace_id).await
32    }
33
34    pub async fn get_mcp_execution_events(&self, trace_id: &str) -> Result<Vec<TraceEvent>> {
35        queries::fetch_mcp_execution_events(&self.pool, trace_id).await
36    }
37
38    pub async fn get_task_id(&self, trace_id: &str) -> Result<Option<String>> {
39        queries::fetch_task_id_for_trace(&self.pool, trace_id).await
40    }
41
42    pub async fn get_execution_step_summary(&self, trace_id: &str) -> Result<ExecutionStepSummary> {
43        queries::fetch_execution_step_summary(&self.pool, trace_id).await
44    }
45
46    pub async fn get_execution_step_events(&self, trace_id: &str) -> Result<Vec<TraceEvent>> {
47        queries::fetch_execution_step_events(&self.pool, trace_id).await
48    }
49
50    pub async fn get_all_trace_data(
51        &self,
52        trace_id: &str,
53    ) -> Result<(
54        Vec<TraceEvent>,
55        Vec<TraceEvent>,
56        Vec<TraceEvent>,
57        Vec<TraceEvent>,
58        AiRequestSummary,
59        McpExecutionSummary,
60        ExecutionStepSummary,
61        Option<String>,
62    )> {
63        tokio::try_join!(
64            self.get_log_events(trace_id),
65            self.get_ai_request_events(trace_id),
66            self.get_mcp_execution_events(trace_id),
67            self.get_execution_step_events(trace_id),
68            self.get_ai_request_summary(trace_id),
69            self.get_mcp_execution_summary(trace_id),
70            self.get_execution_step_summary(trace_id),
71            self.get_task_id(trace_id),
72        )
73    }
74}