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