synwire_core/observability/
metrics_collector.rs1use crate::BoxFuture;
4use crate::observability::gen_ai_metrics;
5use std::collections::HashMap;
6
7pub trait MetricsCollector: Send + Sync {
9 fn record_token_usage(
11 &self,
12 input_tokens: u64,
13 output_tokens: u64,
14 attributes: &HashMap<String, String>,
15 ) -> BoxFuture<'_, ()>;
16
17 fn record_operation_duration(
19 &self,
20 duration_secs: f64,
21 attributes: &HashMap<String, String>,
22 ) -> BoxFuture<'_, ()>;
23
24 fn record_time_to_first_chunk(
26 &self,
27 duration_secs: f64,
28 attributes: &HashMap<String, String>,
29 ) -> BoxFuture<'_, ()>;
30}
31
32pub struct OTelMetricsCollector {
36 token_usage: opentelemetry::metrics::Histogram<u64>,
37 operation_duration: opentelemetry::metrics::Histogram<f64>,
38 time_to_first_chunk: opentelemetry::metrics::Histogram<f64>,
39}
40
41impl std::fmt::Debug for OTelMetricsCollector {
42 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
43 f.debug_struct("OTelMetricsCollector").finish()
44 }
45}
46
47impl OTelMetricsCollector {
48 pub fn new(meter: &opentelemetry::metrics::Meter) -> Self {
50 let token_usage = meter
51 .u64_histogram(gen_ai_metrics::CLIENT_TOKEN_USAGE)
52 .with_description("GenAI client token usage")
53 .build();
54 let operation_duration = meter
55 .f64_histogram(gen_ai_metrics::CLIENT_OPERATION_DURATION)
56 .with_description("GenAI client operation duration in seconds")
57 .build();
58 let time_to_first_chunk = meter
59 .f64_histogram(gen_ai_metrics::CLIENT_OPERATION_TIME_TO_FIRST_CHUNK)
60 .with_description("GenAI client time-to-first-chunk in seconds")
61 .build();
62
63 Self {
64 token_usage,
65 operation_duration,
66 time_to_first_chunk,
67 }
68 }
69}
70
71fn to_key_values(attrs: &HashMap<String, String>) -> Vec<opentelemetry::KeyValue> {
73 attrs
74 .iter()
75 .map(|(k, v)| opentelemetry::KeyValue::new(k.clone(), v.clone()))
76 .collect()
77}
78
79impl MetricsCollector for OTelMetricsCollector {
80 fn record_token_usage(
81 &self,
82 input_tokens: u64,
83 output_tokens: u64,
84 attributes: &HashMap<String, String>,
85 ) -> BoxFuture<'_, ()> {
86 let kvs = to_key_values(attributes);
87 self.token_usage.record(input_tokens + output_tokens, &kvs);
88 Box::pin(async {})
89 }
90
91 fn record_operation_duration(
92 &self,
93 duration_secs: f64,
94 attributes: &HashMap<String, String>,
95 ) -> BoxFuture<'_, ()> {
96 let kvs = to_key_values(attributes);
97 self.operation_duration.record(duration_secs, &kvs);
98 Box::pin(async {})
99 }
100
101 fn record_time_to_first_chunk(
102 &self,
103 duration_secs: f64,
104 attributes: &HashMap<String, String>,
105 ) -> BoxFuture<'_, ()> {
106 let kvs = to_key_values(attributes);
107 self.time_to_first_chunk.record(duration_secs, &kvs);
108 Box::pin(async {})
109 }
110}
111
112#[cfg(test)]
113mod tests {
114 use super::*;
115
116 #[test]
117 fn to_key_values_converts_correctly() {
118 let mut attrs = HashMap::new();
119 let _ = attrs.insert("key1".to_owned(), "val1".to_owned());
120 let _ = attrs.insert("key2".to_owned(), "val2".to_owned());
121
122 let kvs = to_key_values(&attrs);
123 assert_eq!(kvs.len(), 2);
124 }
125
126 #[tokio::test]
127 async fn otel_metrics_collector_records() {
128 let meter = opentelemetry::global::meter("test");
129 let collector = OTelMetricsCollector::new(&meter);
130 let attrs = HashMap::new();
131
132 collector.record_token_usage(100, 50, &attrs).await;
133 collector.record_operation_duration(1.5, &attrs).await;
134 collector.record_time_to_first_chunk(0.2, &attrs).await;
135 }
136}