offline_intelligence/
metrics.rs

1// _Aud.CLI/_Server/metrics.rs
2
3use prometheus::{Encoder, TextEncoder, Registry, IntCounterVec, IntGauge, Histogram, HistogramVec};
4use lazy_static::lazy_static;
5use std::sync::Mutex;
6use axum::response::IntoResponse;
7use axum::http::StatusCode;
8
9lazy_static! {
10    static ref REGISTRY: Registry = Registry::new();
11    static ref REQ_COUNTER: Mutex<IntCounterVec> = Mutex::new(
12        IntCounterVec::new(
13            prometheus::opts!("requests_total", "Total requests per route"),
14            &["route", "status"]
15        ).unwrap()
16    );
17    static ref ACTIVE_SESSIONS: Mutex<IntGauge> = Mutex::new(
18        IntGauge::new("active_sessions", "Active streaming sessions").unwrap()
19    );
20    static ref QUEUE_DEPTH: Mutex<IntGauge> = Mutex::new(
21        IntGauge::new("queue_depth", "Number of requests waiting in queue").unwrap()
22    );
23    static ref RESPONSE_TIME: Mutex<HistogramVec> = Mutex::new(
24        HistogramVec::new(
25            prometheus::HistogramOpts::new(
26                "response_time_seconds",
27                "Response time by endpoint"
28            ),
29            &["endpoint"]
30        ).unwrap()
31    );
32    static ref CONTEXT_OPTIMIZATION_TIME: Mutex<Histogram> = Mutex::new(
33        Histogram::with_opts(prometheus::HistogramOpts::new(
34            "context_optimization_time_seconds",
35            "Time spent in context optimization"
36        )).unwrap()
37    );
38    static ref BACKEND_LATENCY: Mutex<Histogram> = Mutex::new(
39        Histogram::with_opts(prometheus::HistogramOpts::new(
40            "backend_latency_seconds",
41            "Latency to backend service"
42        )).unwrap()
43    );
44}
45
46pub fn init_metrics() {
47    REGISTRY.register(Box::new(REQ_COUNTER.lock().unwrap().clone())).ok();
48    REGISTRY.register(Box::new(ACTIVE_SESSIONS.lock().unwrap().clone())).ok();
49    REGISTRY.register(Box::new(QUEUE_DEPTH.lock().unwrap().clone())).ok();
50    REGISTRY.register(Box::new(RESPONSE_TIME.lock().unwrap().clone())).ok();
51    REGISTRY.register(Box::new(CONTEXT_OPTIMIZATION_TIME.lock().unwrap().clone())).ok();
52    REGISTRY.register(Box::new(BACKEND_LATENCY.lock().unwrap().clone())).ok();
53}
54
55// Add this missing function
56pub fn inc_request(route: &str, status: &str) {
57    REQ_COUNTER.lock().unwrap().with_label_values(&[route, status]).inc();
58}
59
60pub fn inc_sessions() {
61    ACTIVE_SESSIONS.lock().unwrap().inc();
62}
63
64pub fn dec_sessions() {
65    ACTIVE_SESSIONS.lock().unwrap().dec();
66}
67
68pub fn inc_queue() {
69    QUEUE_DEPTH.lock().unwrap().inc();
70}
71
72pub fn dec_queue() {
73    QUEUE_DEPTH.lock().unwrap().dec();
74}
75
76pub fn observe_queue_wait(duration: f64) {
77    // Placeholder for queue wait time
78}
79
80pub fn observe_response_time(endpoint: &str, duration: f64) {
81    RESPONSE_TIME.lock().unwrap().with_label_values(&[endpoint]).observe(duration);
82}
83
84pub fn observe_context_optimization(duration: f64) {
85    CONTEXT_OPTIMIZATION_TIME.lock().unwrap().observe(duration);
86}
87
88pub fn observe_backend_latency(duration: f64) {
89    BACKEND_LATENCY.lock().unwrap().observe(duration);
90}
91
92pub async fn get_metrics() -> impl IntoResponse {
93    let encoder = TextEncoder::new();
94    let metric_families = REGISTRY.gather();
95    let mut buffer = vec![];
96    encoder.encode(&metric_families, &mut buffer).unwrap();
97    
98    (
99        StatusCode::OK,
100        [("content-type", "text/plain; version=0.0.4")],
101        buffer,
102    )
103}