algocline_core/
metrics.rs1use std::sync::{Arc, Mutex};
2use std::time::Instant;
3
4use crate::observer::ExecutionObserver;
5use crate::{CustomMetrics, LlmQuery};
6
7pub(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
39pub 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 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 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
89pub 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}