systemprompt_logging/trace/
service.rs

1use anyhow::Result;
2use chrono::{DateTime, Utc};
3use sqlx::PgPool;
4use std::sync::Arc;
5use systemprompt_identifiers::TaskId;
6
7use crate::models::LogEntry;
8
9use super::models::{
10    AiRequestDetail, AiRequestListItem, AiRequestStats, AiRequestSummary, AuditLookupResult,
11    AuditToolCallRow, ConversationMessage, ExecutionStepSummary, LevelCount, LinkedMcpCall,
12    LogSearchItem, LogTimeRange, McpExecutionSummary, ModuleCount, ToolExecutionFilter,
13    ToolExecutionItem, TraceEvent, TraceListFilter, TraceListItem,
14};
15use super::{
16    audit_queries, list_queries, log_lookup_queries, log_search_queries, log_summary_queries,
17    queries, request_queries, tool_queries,
18};
19
20#[derive(Debug, Clone)]
21pub struct TraceQueryService {
22    pool: Arc<PgPool>,
23}
24
25impl TraceQueryService {
26    pub const fn new(pool: Arc<PgPool>) -> Self {
27        Self { pool }
28    }
29
30    pub async fn get_log_events(&self, trace_id: &str) -> Result<Vec<TraceEvent>> {
31        queries::fetch_log_events(&self.pool, trace_id).await
32    }
33
34    pub async fn get_ai_request_summary(&self, trace_id: &str) -> Result<AiRequestSummary> {
35        queries::fetch_ai_request_summary(&self.pool, trace_id).await
36    }
37
38    pub async fn get_ai_request_events(&self, trace_id: &str) -> Result<Vec<TraceEvent>> {
39        queries::fetch_ai_request_events(&self.pool, trace_id).await
40    }
41
42    pub async fn get_mcp_execution_summary(&self, trace_id: &str) -> Result<McpExecutionSummary> {
43        queries::fetch_mcp_execution_summary(&self.pool, trace_id).await
44    }
45
46    pub async fn get_mcp_execution_events(&self, trace_id: &str) -> Result<Vec<TraceEvent>> {
47        queries::fetch_mcp_execution_events(&self.pool, trace_id).await
48    }
49
50    pub async fn get_task_id(&self, trace_id: &str) -> Result<Option<TaskId>> {
51        Ok(queries::fetch_task_id_for_trace(&self.pool, trace_id)
52            .await?
53            .map(TaskId::from))
54    }
55
56    pub async fn get_execution_step_summary(&self, trace_id: &str) -> Result<ExecutionStepSummary> {
57        queries::fetch_execution_step_summary(&self.pool, trace_id).await
58    }
59
60    pub async fn get_execution_step_events(&self, trace_id: &str) -> Result<Vec<TraceEvent>> {
61        queries::fetch_execution_step_events(&self.pool, trace_id).await
62    }
63
64    pub async fn get_all_trace_data(
65        &self,
66        trace_id: &str,
67    ) -> Result<(
68        Vec<TraceEvent>,
69        Vec<TraceEvent>,
70        Vec<TraceEvent>,
71        Vec<TraceEvent>,
72        AiRequestSummary,
73        McpExecutionSummary,
74        ExecutionStepSummary,
75        Option<TaskId>,
76    )> {
77        tokio::try_join!(
78            self.get_log_events(trace_id),
79            self.get_ai_request_events(trace_id),
80            self.get_mcp_execution_events(trace_id),
81            self.get_execution_step_events(trace_id),
82            self.get_ai_request_summary(trace_id),
83            self.get_mcp_execution_summary(trace_id),
84            self.get_execution_step_summary(trace_id),
85            self.get_task_id(trace_id),
86        )
87    }
88
89    pub async fn list_traces(&self, filter: &TraceListFilter) -> Result<Vec<TraceListItem>> {
90        list_queries::list_traces(&self.pool, filter).await
91    }
92
93    pub async fn list_tool_executions(
94        &self,
95        filter: &ToolExecutionFilter,
96    ) -> Result<Vec<ToolExecutionItem>> {
97        tool_queries::list_tool_executions(&self.pool, filter).await
98    }
99
100    pub async fn search_logs(
101        &self,
102        pattern: &str,
103        since: Option<DateTime<Utc>>,
104        level: Option<&str>,
105        limit: i64,
106    ) -> Result<Vec<LogSearchItem>> {
107        log_search_queries::search_logs(&self.pool, pattern, since, level, limit).await
108    }
109
110    pub async fn search_tool_executions(
111        &self,
112        pattern: &str,
113        since: Option<DateTime<Utc>>,
114        limit: i64,
115    ) -> Result<Vec<ToolExecutionItem>> {
116        log_search_queries::search_tool_executions(&self.pool, pattern, since, limit).await
117    }
118
119    pub async fn list_ai_requests(
120        &self,
121        since: Option<DateTime<Utc>>,
122        model: Option<&str>,
123        provider: Option<&str>,
124        limit: i64,
125    ) -> Result<Vec<AiRequestListItem>> {
126        request_queries::list_ai_requests(&self.pool, since, model, provider, limit).await
127    }
128
129    pub async fn get_ai_request_stats(
130        &self,
131        since: Option<DateTime<Utc>>,
132    ) -> Result<AiRequestStats> {
133        request_queries::get_ai_request_stats(&self.pool, since).await
134    }
135
136    pub async fn find_ai_request_detail(&self, id: &str) -> Result<Option<AiRequestDetail>> {
137        request_queries::find_ai_request_detail(&self.pool, id).await
138    }
139
140    pub async fn find_ai_request_for_audit(&self, id: &str) -> Result<Option<AuditLookupResult>> {
141        audit_queries::find_ai_request_for_audit(&self.pool, id).await
142    }
143
144    pub async fn list_audit_messages(&self, request_id: &str) -> Result<Vec<ConversationMessage>> {
145        audit_queries::list_audit_messages(&self.pool, request_id).await
146    }
147
148    pub async fn list_audit_tool_calls(&self, request_id: &str) -> Result<Vec<AuditToolCallRow>> {
149        audit_queries::list_audit_tool_calls(&self.pool, request_id).await
150    }
151
152    pub async fn list_linked_mcp_calls(&self, request_id: &str) -> Result<Vec<LinkedMcpCall>> {
153        audit_queries::list_linked_mcp_calls(&self.pool, request_id).await
154    }
155
156    pub async fn find_log_by_id(&self, id: &str) -> Result<Option<LogEntry>> {
157        log_lookup_queries::find_log_by_id(&self.pool, id).await
158    }
159
160    pub async fn find_log_by_partial_id(&self, id_prefix: &str) -> Result<Option<LogEntry>> {
161        log_lookup_queries::find_log_by_partial_id(&self.pool, id_prefix).await
162    }
163
164    pub async fn find_logs_by_trace_id(&self, trace_id: &str) -> Result<Vec<LogEntry>> {
165        log_lookup_queries::find_logs_by_trace_id(&self.pool, trace_id).await
166    }
167
168    pub async fn list_logs_filtered(
169        &self,
170        since: Option<DateTime<Utc>>,
171        level: Option<&str>,
172        limit: i64,
173    ) -> Result<Vec<LogEntry>> {
174        log_lookup_queries::list_logs_filtered(&self.pool, since, level, limit).await
175    }
176
177    pub async fn count_logs_by_level(
178        &self,
179        since: Option<DateTime<Utc>>,
180    ) -> Result<Vec<LevelCount>> {
181        log_summary_queries::count_logs_by_level(&self.pool, since).await
182    }
183
184    pub async fn top_modules(
185        &self,
186        since: Option<DateTime<Utc>>,
187        limit: i64,
188    ) -> Result<Vec<ModuleCount>> {
189        log_summary_queries::top_modules(&self.pool, since, limit).await
190    }
191
192    pub async fn log_time_range(&self, since: Option<DateTime<Utc>>) -> Result<LogTimeRange> {
193        log_summary_queries::log_time_range(&self.pool, since).await
194    }
195
196    pub async fn total_log_count(&self) -> Result<i64> {
197        log_summary_queries::total_log_count(&self.pool).await
198    }
199}
systemprompt_logging/trace/service.rs

systemprompt_logging/trace/
service.rs
