systemprompt_logging/trace/
service.rs

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

systemprompt_logging/trace/
service.rs
