systemprompt_logging/trace/
service.rs

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

systemprompt_logging/trace/
service.rs
