Skip to main content

synwire_core/observability/
metrics_collector.rs

1//! Metrics collector trait and OpenTelemetry implementation.
2
3use crate::BoxFuture;
4use crate::observability::gen_ai_metrics;
5use std::collections::HashMap;
6
7/// Trait for collecting observability metrics.
8pub trait MetricsCollector: Send + Sync {
9    /// Records token usage (input and output tokens).
10    fn record_token_usage(
11        &self,
12        input_tokens: u64,
13        output_tokens: u64,
14        attributes: &HashMap<String, String>,
15    ) -> BoxFuture<'_, ()>;
16
17    /// Records operation duration in seconds.
18    fn record_operation_duration(
19        &self,
20        duration_secs: f64,
21        attributes: &HashMap<String, String>,
22    ) -> BoxFuture<'_, ()>;
23
24    /// Records time-to-first-chunk for streaming operations, in seconds.
25    fn record_time_to_first_chunk(
26        &self,
27        duration_secs: f64,
28        attributes: &HashMap<String, String>,
29    ) -> BoxFuture<'_, ()>;
30}
31
32/// OpenTelemetry-based metrics collector using histogram instruments.
33///
34/// Creates `OTel` histograms for the `GenAI` semantic convention metrics.
35pub 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    /// Creates a new `OTelMetricsCollector` using the given meter.
49    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
71/// Converts a `HashMap<String, String>` into `OTel` `KeyValue` pairs.
72fn 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}