Skip to main content

ai_agent/utils/
query_profiler.rs

1//! Query profiler utilities.
2
3use std::time::{Duration, Instant};
4
5/// Query profiler for tracking query performance
6pub struct QueryProfiler {
7    start_time: Instant,
8    queries: Vec<QueryProfileEntry>,
9}
10
11#[derive(Debug, Clone)]
12pub struct QueryProfileEntry {
13    query: String,
14    duration: Duration,
15    success: bool,
16}
17
18impl QueryProfiler {
19    pub fn new() -> Self {
20        Self {
21            start_time: Instant::now(),
22            queries: Vec::new(),
23        }
24    }
25
26    pub fn record(&mut self, query: String, duration: Duration, success: bool) {
27        self.queries.push(QueryProfileEntry {
28            query,
29            duration,
30            success,
31        });
32    }
33
34    pub fn get_total_queries(&self) -> usize {
35        self.queries.len()
36    }
37
38    pub fn get_successful_queries(&self) -> usize {
39        self.queries.iter().filter(|q| q.success).count()
40    }
41
42    pub fn get_total_duration(&self) -> Duration {
43        self.queries.iter().map(|q| q.duration).sum()
44    }
45
46    pub fn get_average_duration(&self) -> Duration {
47        let count = self.queries.len();
48        if count == 0 {
49            Duration::ZERO
50        } else {
51            self.get_total_duration() / count as u32
52        }
53    }
54
55    pub fn get_slowest_query(&self) -> Option<&QueryProfileEntry> {
56        self.queries.iter().max_by_key(|q| q.duration)
57    }
58}
59
60impl Default for QueryProfiler {
61    fn default() -> Self {
62        Self::new()
63    }
64}