Skip to main content

algocline_core/
metrics.rs

1use std::sync::{Arc, Mutex};
2use std::time::Instant;
3
4use crate::observer::ExecutionObserver;
5use crate::{CustomMetrics, LlmQuery};
6
7/// Metrics automatically derived from the execution lifecycle.
8pub(crate) struct AutoMetrics {
9    started_at: Instant,
10    ended_at: Option<Instant>,
11    llm_calls: u64,
12    pauses: u64,
13}
14
15impl AutoMetrics {
16    fn new() -> Self {
17        Self {
18            started_at: Instant::now(),
19            ended_at: None,
20            llm_calls: 0,
21            pauses: 0,
22        }
23    }
24
25    fn to_json(&self) -> serde_json::Value {
26        let elapsed_ms = self
27            .ended_at
28            .map(|end| end.duration_since(self.started_at).as_millis() as u64)
29            .unwrap_or_else(|| self.started_at.elapsed().as_millis() as u64);
30
31        serde_json::json!({
32            "elapsed_ms": elapsed_ms,
33            "llm_calls": self.llm_calls,
34            "pauses": self.pauses,
35        })
36    }
37}
38
39/// Measurement data for a single execution.
40pub struct ExecutionMetrics {
41    auto: Arc<Mutex<AutoMetrics>>,
42    custom: Arc<Mutex<CustomMetrics>>,
43}
44
45impl ExecutionMetrics {
46    pub fn new() -> Self {
47        Self {
48            auto: Arc::new(Mutex::new(AutoMetrics::new())),
49            custom: Arc::new(Mutex::new(CustomMetrics::new())),
50        }
51    }
52
53    /// JSON snapshot combining auto and custom metrics.
54    pub fn to_json(&self) -> serde_json::Value {
55        let auto_json = self
56            .auto
57            .lock()
58            .map(|m| m.to_json())
59            .unwrap_or(serde_json::Value::Null);
60
61        let custom_json = self
62            .custom
63            .lock()
64            .map(|m| m.to_json())
65            .unwrap_or(serde_json::Value::Null);
66
67        serde_json::json!({
68            "auto": auto_json,
69            "custom": custom_json,
70        })
71    }
72
73    /// Handle for custom metrics, passed to the Lua bridge.
74    pub fn custom_handle(&self) -> Arc<Mutex<CustomMetrics>> {
75        Arc::clone(&self.custom)
76    }
77
78    pub fn create_observer(&self) -> MetricsObserver {
79        MetricsObserver::new(Arc::clone(&self.auto))
80    }
81}
82
83impl Default for ExecutionMetrics {
84    fn default() -> Self {
85        Self::new()
86    }
87}
88
89/// Updates AutoMetrics via the ExecutionObserver trait.
90pub struct MetricsObserver {
91    auto: Arc<Mutex<AutoMetrics>>,
92}
93
94impl MetricsObserver {
95    pub(crate) fn new(auto: Arc<Mutex<AutoMetrics>>) -> Self {
96        Self { auto }
97    }
98}
99
100impl ExecutionObserver for MetricsObserver {
101    fn on_paused(&self, queries: &[LlmQuery]) {
102        if let Ok(mut m) = self.auto.lock() {
103            m.pauses += 1;
104            m.llm_calls += queries.len() as u64;
105        }
106    }
107
108    fn on_completed(&self, _result: &serde_json::Value) {
109        if let Ok(mut m) = self.auto.lock() {
110            m.ended_at = Some(Instant::now());
111        }
112    }
113
114    fn on_failed(&self, _error: &str) {
115        if let Ok(mut m) = self.auto.lock() {
116            m.ended_at = Some(Instant::now());
117        }
118    }
119
120    fn on_cancelled(&self) {
121        if let Ok(mut m) = self.auto.lock() {
122            m.ended_at = Some(Instant::now());
123        }
124    }
125}
126
127#[cfg(test)]
128mod tests {
129    use super::*;
130    use crate::{LlmQuery, QueryId};
131
132    #[test]
133    fn metrics_to_json_has_auto_and_custom() {
134        let metrics = ExecutionMetrics::new();
135        let json = metrics.to_json();
136        assert!(json.get("auto").is_some());
137        assert!(json.get("custom").is_some());
138    }
139
140    #[test]
141    fn custom_handle_shares_state() {
142        let metrics = ExecutionMetrics::new();
143        let handle = metrics.custom_handle();
144
145        handle
146            .lock()
147            .unwrap()
148            .record("key".into(), serde_json::json!("value"));
149
150        let json = metrics.to_json();
151        let custom = json.get("custom").unwrap();
152        assert_eq!(custom.get("key").unwrap(), "value");
153    }
154
155    #[test]
156    fn observer_updates_auto_metrics() {
157        let metrics = ExecutionMetrics::new();
158        let observer = metrics.create_observer();
159
160        let queries = vec![LlmQuery {
161            id: QueryId::batch(0),
162            prompt: "test".into(),
163            system: None,
164            max_tokens: 100,
165        }];
166
167        observer.on_paused(&queries);
168        observer.on_completed(&serde_json::json!(null));
169
170        let json = metrics.to_json();
171        let auto = json.get("auto").unwrap();
172        assert_eq!(auto.get("llm_calls").unwrap(), 1);
173        assert_eq!(auto.get("pauses").unwrap(), 1);
174    }
175}