use async_trait::async_trait;
#[async_trait]
pub trait Metrics: Send + Sync + 'static {
async fn incr_counter(&self, name: &str, value: u64);
async fn record_gauge(&self, name: &str, value: f64);
async fn start_timer(&self, name: &str) -> Option<Box<dyn Timer + Send>>;
async fn record_histogram(&self, name: &str, value: f64);
async fn record_histogram_with_tags(&self, name: &str, value: f64, tags: &[(&str, &str)]);
async fn incr_counter_with_tags(&self, name: &str, value: u64, tags: &[(&str, &str)]);
async fn record_gauge_with_tags(&self, name: &str, value: f64, tags: &[(&str, &str)]);
async fn record_error(&self, name: &str, error_type: &str);
async fn record_success(&self, name: &str, success: bool);
}
pub trait Timer: Send {
fn stop(self: Box<Self>);
}
pub struct NoopMetrics;
#[async_trait]
impl Metrics for NoopMetrics {
async fn incr_counter(&self, _name: &str, _value: u64) {}
async fn record_gauge(&self, _name: &str, _value: f64) {}
async fn start_timer(&self, _name: &str) -> Option<Box<dyn Timer + Send>> {
None
}
async fn record_histogram(&self, _name: &str, _value: f64) {}
async fn record_histogram_with_tags(&self, _name: &str, _value: f64, _tags: &[(&str, &str)]) {}
async fn incr_counter_with_tags(&self, _name: &str, _value: u64, _tags: &[(&str, &str)]) {}
async fn record_gauge_with_tags(&self, _name: &str, _value: f64, _tags: &[(&str, &str)]) {}
async fn record_error(&self, _name: &str, _error_type: &str) {}
async fn record_success(&self, _name: &str, _success: bool) {}
}
pub struct NoopTimer;
impl Timer for NoopTimer {
fn stop(self: Box<Self>) {}
}
impl NoopMetrics {
pub fn new() -> Self {
NoopMetrics
}
}
impl Default for NoopMetrics {
fn default() -> Self {
Self::new()
}
}
#[allow(async_fn_in_trait)]
pub trait MetricsExt: Metrics {
async fn record_request(
&self,
name: &str,
timer: Option<Box<dyn Timer + Send>>,
success: bool,
) {
if let Some(t) = timer {
t.stop();
}
self.record_success(name, success).await;
}
async fn record_request_with_tags(
&self,
name: &str,
timer: Option<Box<dyn Timer + Send>>,
success: bool,
tags: &[(&str, &str)],
) {
if let Some(t) = timer {
t.stop();
}
self.record_success(name, success).await;
self.incr_counter_with_tags(&format!("{}.total", name), 1, tags)
.await;
if success {
self.incr_counter_with_tags(&format!("{}.success", name), 1, tags)
.await;
} else {
self.incr_counter_with_tags(&format!("{}.failure", name), 1, tags)
.await;
}
}
async fn record_error_with_context(&self, name: &str, error_type: &str, context: &str) {
self.record_error(name, error_type).await;
self.incr_counter_with_tags(name, 1, &[("error_type", error_type), ("context", context)])
.await;
}
async fn record_complete_request(
&self,
name: &str,
timer: Option<Box<dyn Timer + Send>>,
status_code: u16,
success: bool,
tags: &[(&str, &str)],
) {
if let Some(t) = timer {
t.stop();
}
self.record_success(name, success).await;
self.record_gauge_with_tags(&format!("{}.status_code", name), status_code as f64, tags)
.await;
self.incr_counter_with_tags(&format!("{}.total", name), 1, tags)
.await;
if success {
self.incr_counter_with_tags(&format!("{}.success", name), 1, tags)
.await;
} else {
self.incr_counter_with_tags(&format!("{}.failure", name), 1, tags)
.await;
}
}
async fn record_batch_latency_percentiles(
&self,
name: &str,
measurements: &[f64],
tags: &[(&str, &str)],
) {
if measurements.is_empty() {
return;
}
let mut sorted = measurements.to_vec();
sorted.sort_by(|a, b| a.partial_cmp(b).unwrap());
let len = sorted.len();
let p50 = sorted[(len * 50 / 100).min(len - 1)];
let p95 = sorted[(len * 95 / 100).min(len - 1)];
let p99 = sorted[(len * 99 / 100).min(len - 1)];
self.record_gauge_with_tags(&format!("{}.latency.p50", name), p50, tags)
.await;
self.record_gauge_with_tags(&format!("{}.latency.p95", name), p95, tags)
.await;
self.record_gauge_with_tags(&format!("{}.latency.p99", name), p99, tags)
.await;
}
}
impl<T: Metrics> MetricsExt for T {}
pub mod keys {
pub fn requests(provider: &str) -> String {
format!("{}.requests", provider)
}
pub fn request_duration_ms(provider: &str) -> String {
format!("{}.request_duration_ms", provider)
}
pub fn success(provider: &str) -> String {
format!("{}.success", provider)
}
pub fn failure(provider: &str) -> String {
format!("{}.failure", provider)
}
pub fn latency_p50(provider: &str) -> String {
format!("{}.latency_p50", provider)
}
pub fn latency_p95(provider: &str) -> String {
format!("{}.latency_p95", provider)
}
pub fn latency_p99(provider: &str) -> String {
format!("{}.latency_p99", provider)
}
pub fn status_codes(provider: &str) -> String {
format!("{}.status_codes", provider)
}
pub fn error_rate(provider: &str) -> String {
format!("{}.error_rate", provider)
}
pub fn throughput(provider: &str) -> String {
format!("{}.throughput", provider)
}
pub fn cost_usd(provider: &str) -> String {
format!("{}.cost_usd", provider)
}
pub fn cost_per_request(provider: &str) -> String {
format!("{}.cost_per_request", provider)
}
pub fn tokens_input(provider: &str) -> String {
format!("{}.tokens_input", provider)
}
pub fn tokens_output(provider: &str) -> String {
format!("{}.tokens_output", provider)
}
pub fn routing_requests(route: &str) -> String {
format!("routing.{}.requests", route)
}
pub fn routing_selected(route: &str) -> String {
format!("routing.{}.selected", route)
}
pub fn routing_health_fail(route: &str) -> String {
format!("routing.{}.health_fail", route)
}
}
#[cfg(feature = "cost_metrics")]
pub mod cost {
use crate::metrics::Metrics;
pub fn estimate_usd(input_tokens: u32, output_tokens: u32) -> f64 {
let in_rate = std::env::var("COST_INPUT_PER_1K")
.ok()
.and_then(|s| s.parse::<f64>().ok())
.unwrap_or(0.0);
let out_rate = std::env::var("COST_OUTPUT_PER_1K")
.ok()
.and_then(|s| s.parse::<f64>().ok())
.unwrap_or(0.0);
(input_tokens as f64 / 1000.0) * in_rate + (output_tokens as f64 / 1000.0) * out_rate
}
pub async fn record_cost<M: Metrics + ?Sized>(m: &M, provider: &str, model: &str, usd: f64) {
m.record_histogram_with_tags("cost.usd", usd, &[("provider", provider), ("model", model)])
.await;
}
}