1use metrics::{counter, histogram};
7use metrics_exporter_prometheus::{PrometheusBuilder, PrometheusHandle};
8
9pub const METRIC_REQUESTS_TOTAL: &str = "sokuto_requests_total";
11
12pub const METRIC_INFERENCE_DURATION_SECONDS: &str = "sokuto_inference_duration_seconds";
14
15pub const METRIC_HTTP_DURATION_SECONDS: &str = "sokuto_http_duration_seconds";
17
18pub const METRIC_BATCH_SIZE: &str = "sokuto_batch_size";
20
21pub const METRIC_QUESTION_TYPE_TOTAL: &str = "sokuto_question_type_total";
23
24pub const METRIC_CONFIDENCE_SCORE: &str = "sokuto_confidence_score";
26
27pub const METRIC_GATING_ROUTES_TOTAL: &str = "sokuto_gating_routes_total";
29
30pub const METRIC_DAG_EXECUTION_DURATION_SECONDS: &str = "sokuto_dag_execution_duration_seconds";
32
33pub const METRIC_DAG_NODE_DURATION_SECONDS: &str = "sokuto_dag_node_duration_seconds";
35
36pub const METRIC_DAG_NODE_EXECUTIONS_TOTAL: &str = "sokuto_dag_node_executions_total";
38
39pub const METRIC_HIERARCHICAL_SOFT_BEAM_TOTAL: &str = "sokuto_hierarchical_soft_beam_total";
41
42pub const METRIC_HIERARCHICAL_EARLY_EXIT_TOTAL: &str = "sokuto_hierarchical_early_exit_total";
44
45pub fn setup_metrics_recorder() -> Result<PrometheusHandle, Box<dyn std::error::Error + Send + Sync>>
47{
48 let handle = PrometheusBuilder::new().install_recorder()?;
49 Ok(handle)
50}
51
52pub fn record_http_request(endpoint: &'static str, status: u16) {
54 counter!(
55 METRIC_REQUESTS_TOTAL,
56 "endpoint" => endpoint,
57 "status" => status.to_string(),
58 )
59 .increment(1);
60}
61
62pub fn record_http_duration(endpoint: &'static str, duration_secs: f64) {
64 histogram!(
65 METRIC_HTTP_DURATION_SECONDS,
66 "endpoint" => endpoint,
67 )
68 .record(duration_secs);
69}
70
71pub fn record_inference_metrics(duration_secs: f64, batch_size: usize) {
73 histogram!(METRIC_INFERENCE_DURATION_SECONDS).record(duration_secs);
74 histogram!(METRIC_BATCH_SIZE).record(batch_size as f64);
75}
76
77pub fn record_question_type(question_type: &'static str) {
79 counter!(
80 METRIC_QUESTION_TYPE_TOTAL,
81 "type" => question_type,
82 )
83 .increment(1);
84}
85
86pub fn record_confidence(confidence: f64) {
88 histogram!(METRIC_CONFIDENCE_SCORE).record(confidence);
89}
90
91pub fn record_gating_route(route: &str, question_type: &str) {
93 counter!(
94 METRIC_GATING_ROUTES_TOTAL,
95 "route" => route.to_string(),
96 "type" => question_type.to_string(),
97 )
98 .increment(1);
99}
100
101pub fn record_dag_execution(dag_id: &str, status: &str, duration_secs: f64) {
103 histogram!(
104 METRIC_DAG_EXECUTION_DURATION_SECONDS,
105 "dag_id" => dag_id.to_string(),
106 "status" => status.to_string(),
107 )
108 .record(duration_secs);
109}
110
111pub fn record_dag_node_execution(dag_id: &str, node_id: &str, node_type: &str, duration_secs: f64) {
113 histogram!(
114 METRIC_DAG_NODE_DURATION_SECONDS,
115 "dag_id" => dag_id.to_string(),
116 "node_id" => node_id.to_string(),
117 "node_type" => node_type.to_string(),
118 )
119 .record(duration_secs);
120
121 counter!(
122 METRIC_DAG_NODE_EXECUTIONS_TOTAL,
123 "dag_id" => dag_id.to_string(),
124 "node_id" => node_id.to_string(),
125 "node_type" => node_type.to_string(),
126 )
127 .increment(1);
128}
129
130pub fn record_hierarchical_soft_beam() {
132 counter!(METRIC_HIERARCHICAL_SOFT_BEAM_TOTAL).increment(1);
133}
134
135pub fn record_hierarchical_early_exit() {
137 counter!(METRIC_HIERARCHICAL_EARLY_EXIT_TOTAL).increment(1);
138}