From e8f90fc9b18a3a6a9a1deccfafdb584269871cb3 Mon Sep 17 00:00:00 2001 From: asepharyana Date: Mon, 3 Aug 2026 11:44:21 +0700 Subject: [PATCH] feat(metrics): prometheus /metrics + generation timing in usage - /metrics endpoint: request/token/latency counters, tok/s gauge, build info (std-only, no deps) - usage.duration_ms + usage.tokens_per_second in non-streaming and streaming responses - /health now reports uptime_s, n_ctx, version - MAX_TOKENS env config (hard cap, default 2048; 0 = unlimited) - metrics unit tests (counters, prometheus shape, uptime monotonic) --- src/config/mod.rs | 4 + src/domain/entity/mod.rs | 15 + src/presentation/handler/chat-ui/index.html | 648 +++++++++++++++++--- src/presentation/handler/chat.rs | 44 +- src/presentation/handler/health.rs | 27 +- src/presentation/handler/metrics.rs | 249 ++++++++ src/presentation/handler/mod.rs | 3 + src/presentation/router.rs | 3 +- 8 files changed, 891 insertions(+), 102 deletions(-) create mode 100644 src/presentation/handler/metrics.rs diff --git a/src/config/mod.rs b/src/config/mod.rs index 628b448..ebf385d 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -38,6 +38,9 @@ pub struct AppConfig { /// Number of CPU threads for inference pub n_threads: i32, + + /// Hard cap for `max_tokens` in chat requests (0 = unlimited) + pub max_tokens: u32, } impl AppConfig { @@ -55,6 +58,7 @@ impl AppConfig { n_ctx: env_or("N_CTX", 8192), n_batch: env_or("N_BATCH", 512), n_threads: env_or("N_THREADS", 4), + max_tokens: env_or("MAX_TOKENS", 2048), } } } diff --git a/src/domain/entity/mod.rs b/src/domain/entity/mod.rs index 14ab11a..af1e95e 100644 --- a/src/domain/entity/mod.rs +++ b/src/domain/entity/mod.rs @@ -109,6 +109,12 @@ pub struct Usage { pub prompt_tokens: u32, pub completion_tokens: u32, pub total_tokens: u32, + /// Wall-clock duration of the generation, in milliseconds. + #[serde(skip_serializing_if = "Option::is_none")] + pub duration_ms: Option, + /// Generated tokens per second (completion_tokens / seconds). + #[serde(skip_serializing_if = "Option::is_none")] + pub tokens_per_second: Option, } // ═══════════════════════════════════════════════════════════════ @@ -278,4 +284,13 @@ pub struct ModelInfo { pub struct HealthResponse { pub status: String, pub model: String, + /// Server process uptime in seconds. + #[serde(skip_serializing_if = "Option::is_none")] + pub uptime_s: Option, + /// llama.cpp context size (n_ctx). + #[serde(skip_serializing_if = "Option::is_none")] + pub n_ctx: Option, + /// Server binary version (from CARGO_PKG_VERSION). + #[serde(skip_serializing_if = "Option::is_none")] + pub version: Option, } diff --git a/src/presentation/handler/chat-ui/index.html b/src/presentation/handler/chat-ui/index.html index 4e4545c..5f0f53e 100644 --- a/src/presentation/handler/chat-ui/index.html +++ b/src/presentation/handler/chat-ui/index.html @@ -3,9 +3,11 @@ -AI Chat +AI Chat — llm-api + - +

AI Chat

- llm-api + llm-api +
+ +
-
+ +
+
+

MiniCPM5-1B Thinking

+
+ Fast reasoning model running locally via llama.cpp.
+ Markdown & code supported · context window 8K · 1B params (Q8_0) +
+
+
+ +
+
+ +
+
+ + +
+
+ + +
+
+ + +
+
+ + +
+
+
- - + + +
diff --git a/src/presentation/handler/chat.rs b/src/presentation/handler/chat.rs index 92f7caa..d9d7eeb 100644 --- a/src/presentation/handler/chat.rs +++ b/src/presentation/handler/chat.rs @@ -19,11 +19,13 @@ use tokio_stream::wrappers::ReceiverStream; use tracing::info; use crate::application::chat; +use crate::config::CONFIG; use crate::domain::entity::{ ChatRequest, ChatResponse, Choice, FinishReason, ResponseMessage, SseChunk, SseDelta, Usage, }; use crate::infrastructure::llama::SendSampler; use crate::presentation::error::AppError; +use crate::presentation::handler::metrics; use crate::presentation::state::AppState; /// POST /v1/chat/completions @@ -34,7 +36,8 @@ pub async fn chat_completions( // Strict model validation — reject unknown model ids up front. chat::validate_model(&req.model).map_err(AppError::BadRequest)?; - let max_tokens = req.max_tokens.unwrap_or(256).min(1024); + // Cap max_tokens at the configured hard limit (0 = unlimited). + let max_tokens = req.max_tokens.unwrap_or(256).min(CONFIG.max_tokens.max(1)); let stop = req.stop.clone().unwrap_or_default(); let prompt = chat::build_prompt(&req.messages, &req.tools).map_err(AppError::LlmError)?; @@ -52,6 +55,8 @@ pub async fn chat_completions( req.tools.as_ref().is_some_and(|t| !t.is_empty()) ); + metrics::count_request(req.stream.unwrap_or(false)); + let response = if req.stream.unwrap_or(false) { handle_streaming(state.clone(), req, max_tokens, stop, input_tokens).await? } else { @@ -76,6 +81,7 @@ async fn handle_non_streaming( let params = chat::SamplerParams::from_request(&req); let engine = state.engine.clone(); + let gen_start = std::time::Instant::now(); let outcome = tokio::task::spawn_blocking(move || { let mut sampler = SendSampler(chat::build_sampler(¶ms)); engine.generate( @@ -90,9 +96,18 @@ async fn handle_non_streaming( .await .map_err(|e| AppError::Internal(format!("Generation task panicked: {e}")))? .map_err(AppError::from)?; + let duration_ms = gen_start.elapsed().as_millis() as u64; let completion_tokens = outcome.tokens.len() as u32; - info!(" {} generated tokens", completion_tokens); + info!( + " {} generated tokens in {}ms", + completion_tokens, duration_ms + ); + + metrics::record_tokens(prompt_tokens, completion_tokens, duration_ms); + if outcome.finish == FinishReason::Aborted { + metrics::count_aborted(); + } let (reasoning, cleaned) = chat::clean_text(&outcome.text); let (output_text, tool_calls) = chat::parse_tool_calls(&cleaned); @@ -110,6 +125,12 @@ async fn handle_non_streaming( Some(reasoning) }; + let tok_per_s = if duration_ms > 0 { + Some(completion_tokens as f64 / (duration_ms as f64 / 1000.0)) + } else { + None + }; + Ok(Json(ChatResponse { id: chat_id, object: "chat.completion".into(), @@ -133,6 +154,8 @@ async fn handle_non_streaming( prompt_tokens, completion_tokens, total_tokens: prompt_tokens + completion_tokens, + duration_ms: Some(duration_ms), + tokens_per_second: tok_per_s, }, }) .into_response()) @@ -173,6 +196,7 @@ async fn handle_streaming( let engine = state.engine.clone(); tokio::task::spawn_blocking(move || { + let gen_start = std::time::Instant::now(); let mut sampler = SendSampler(chat::build_sampler(¶ms)); let mut text_buf = String::new(); let mut sent_len: usize = 0; @@ -249,12 +273,27 @@ async fn handle_streaming( match outcome { Ok(outcome) => { + let duration_ms = gen_start.elapsed().as_millis() as u64; let completion_tokens = outcome.tokens.len() as u32; let usage = Usage { prompt_tokens, completion_tokens, total_tokens: prompt_tokens + completion_tokens, + duration_ms: Some(duration_ms), + tokens_per_second: if duration_ms > 0 { + Some(completion_tokens as f64 / (duration_ms as f64 / 1000.0)) + } else { + None + }, }; + info!( + " stream: {} generated tokens in {}ms", + completion_tokens, duration_ms + ); + metrics::record_tokens(prompt_tokens, completion_tokens, duration_ms); + if outcome.finish == FinishReason::Aborted { + metrics::count_aborted(); + } // If the model never emitted ``, everything was streamed // as reasoning_content. Flush it as content so the client always @@ -317,6 +356,7 @@ async fn handle_streaming( } Err(e) => { // Surface the error instead of silently truncating the stream. + metrics::count_error(); let error_body = serde_json::json!({ "error": { "message": e.to_string(), diff --git a/src/presentation/handler/health.rs b/src/presentation/handler/health.rs index 9dfd106..41cbdcd 100644 --- a/src/presentation/handler/health.rs +++ b/src/presentation/handler/health.rs @@ -1,14 +1,39 @@ //! Health check endpoint. +use std::sync::atomic::AtomicU64; +use std::sync::LazyLock; +use std::time::Instant; + use axum::Json; -use crate::config::MODEL_ID; +use crate::config::{CONFIG, MODEL_ID}; use crate::domain::entity::HealthResponse; +/// Process start instant — used to compute uptime for /health and /metrics. +pub static START_INSTANT: LazyLock = LazyLock::new(Instant::now); + +/// Process start timestamp (unix seconds) — exported as a Prometheus gauge. +pub static START_TIMESTAMP: LazyLock = LazyLock::new(|| { + AtomicU64::new( + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_secs()) + .unwrap_or(0), + ) +}); + +/// Seconds since process start. +pub fn uptime_secs() -> u64 { + START_INSTANT.elapsed().as_secs() +} + pub async fn health_check() -> Json { Json(HealthResponse { status: "ok".into(), // Report the exact model id served by /v1/models (no extra suffix). model: MODEL_ID.into(), + uptime_s: Some(uptime_secs()), + n_ctx: Some(CONFIG.n_ctx), + version: Some(env!("CARGO_PKG_VERSION").into()), }) } diff --git a/src/presentation/handler/metrics.rs b/src/presentation/handler/metrics.rs new file mode 100644 index 0000000..2c8ebcd --- /dev/null +++ b/src/presentation/handler/metrics.rs @@ -0,0 +1,249 @@ +//! Prometheus metrics endpoint. +//! +//! Exports process-level metrics (CPU, RSS) plus request counters, token +//! usage and generation latencies. Implemented with `std` only — no external +//! metrics dependency — so the deploy stays dependency-free. +//! +//! Scrape config (VPS): `prometheus.yml` file_sd targets llm-api at +//! `/metrics` with default `__metrics_path__`. + +use std::fmt::Write as _; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::LazyLock; +use std::time::Instant; + +use axum::http::{header, HeaderValue, StatusCode}; +use axum::response::{IntoResponse, Response}; + +use super::health::{uptime_secs, START_INSTANT, START_TIMESTAMP}; + +// ── Atomic counters (updated by the chat handlers) ── + +/// Total /v1/chat/completions requests received. +pub static REQUESTS_TOTAL: LazyLock = LazyLock::new(|| AtomicU64::new(0)); +/// Requests that ended in an error (any 4xx/5xx). +pub static ERRORS_TOTAL: LazyLock = LazyLock::new(|| AtomicU64::new(0)); +/// Requests that streamed (`stream: true`). +pub static STREAMING_TOTAL: LazyLock = LazyLock::new(|| AtomicU64::new(0)); +/// Prompt tokens accepted across all requests. +pub static PROMPT_TOKENS_TOTAL: LazyLock = LazyLock::new(|| AtomicU64::new(0)); +/// Completion tokens generated across all requests. +pub static COMPLETION_TOKENS_TOTAL: LazyLock = LazyLock::new(|| AtomicU64::new(0)); +/// Generation time spent across all requests, in milliseconds. +pub static GENERATION_MS_TOTAL: LazyLock = LazyLock::new(|| AtomicU64::new(0)); +/// Requests whose generation was aborted early (client disconnect). +pub static ABORTED_TOTAL: LazyLock = LazyLock::new(|| AtomicU64::new(0)); + +// ── Public helpers used by handlers ── + +pub fn count_request(streaming: bool) { + REQUESTS_TOTAL.fetch_add(1, Ordering::Relaxed); + if streaming { + STREAMING_TOTAL.fetch_add(1, Ordering::Relaxed); + } +} + +pub fn count_error() { + ERRORS_TOTAL.fetch_add(1, Ordering::Relaxed); +} + +pub fn count_aborted() { + ABORTED_TOTAL.fetch_add(1, Ordering::Relaxed); +} + +pub fn record_tokens(prompt_tokens: u32, completion_tokens: u32, duration_ms: u64) { + PROMPT_TOKENS_TOTAL.fetch_add(prompt_tokens as u64, Ordering::Relaxed); + COMPLETION_TOKENS_TOTAL.fetch_add(completion_tokens as u64, Ordering::Relaxed); + GENERATION_MS_TOTAL.fetch_add(duration_ms, Ordering::Relaxed); +} + +fn f(field: &mut String, name: &str, value: impl std::fmt::Display) { + let _ = writeln!(field, "{name} {value}"); +} + +/// GET /metrics — Prometheus text exposition format. +pub async fn metrics() -> Response { + let uptime = uptime_secs(); + let start_ts = START_TIMESTAMP.load(Ordering::Relaxed); + + // Per-second rates over the process lifetime. + let total = REQUESTS_TOTAL.load(Ordering::Relaxed); + let errors = ERRORS_TOTAL.load(Ordering::Relaxed); + let streaming = STREAMING_TOTAL.load(Ordering::Relaxed); + let aborted = ABORTED_TOTAL.load(Ordering::Relaxed); + let prompt_tokens = PROMPT_TOKENS_TOTAL.load(Ordering::Relaxed); + let completion_tokens = COMPLETION_TOKENS_TOTAL.load(Ordering::Relaxed); + let gen_ms = GENERATION_MS_TOTAL.load(Ordering::Relaxed); + + let rps = if uptime > 0 { + total as f64 / uptime as f64 + } else { + 0.0 + }; + let tok_per_s = if gen_ms > 0 { + completion_tokens as f64 / (gen_ms as f64 / 1000.0) + } else { + 0.0 + }; + let avg_ms = if total > 0 { + gen_ms as f64 / total as f64 + } else { + 0.0 + }; + + let mut body = String::with_capacity(2048); + body.push_str("# HELP llm_api_requests_total Total /v1/chat/completions requests received.\n"); + body.push_str("# TYPE llm_api_requests_total counter\n"); + f(&mut body, "llm_api_requests_total", total); + body.push_str("# HELP llm_api_errors_total Requests that ended in an error.\n"); + body.push_str("# TYPE llm_api_errors_total counter\n"); + f(&mut body, "llm_api_errors_total", errors); + body.push_str("# HELP llm_api_streaming_requests_total Requests that streamed.\n"); + body.push_str("# TYPE llm_api_streaming_requests_total counter\n"); + f(&mut body, "llm_api_streaming_requests_total", streaming); + body.push_str( + "# HELP llm_api_aborted_requests_total Generations aborted early (client disconnect).\n", + ); + body.push_str("# TYPE llm_api_aborted_requests_total counter\n"); + f(&mut body, "llm_api_aborted_requests_total", aborted); + body.push_str("# HELP llm_api_prompt_tokens_total Prompt tokens accepted.\n"); + body.push_str("# TYPE llm_api_prompt_tokens_total counter\n"); + f(&mut body, "llm_api_prompt_tokens_total", prompt_tokens); + body.push_str("# HELP llm_api_completion_tokens_total Completion tokens generated.\n"); + body.push_str("# TYPE llm_api_completion_tokens_total counter\n"); + f( + &mut body, + "llm_api_completion_tokens_total", + completion_tokens, + ); + body.push_str("# HELP llm_api_generation_ms_total Generation time in milliseconds.\n"); + body.push_str("# TYPE llm_api_generation_ms_total counter\n"); + f(&mut body, "llm_api_generation_ms_total", gen_ms); + body.push_str("# HELP llm_api_requests_per_second Lifetime request rate.\n"); + body.push_str("# TYPE llm_api_requests_per_second gauge\n"); + f( + &mut body, + "llm_api_requests_per_second", + format!("{rps:.3}"), + ); + body.push_str("# HELP llm_api_tokens_per_second Lifetime generation throughput.\n"); + body.push_str("# TYPE llm_api_tokens_per_second gauge\n"); + f( + &mut body, + "llm_api_tokens_per_second", + format!("{tok_per_s:.3}"), + ); + body.push_str( + "# HELP llm_api_average_generation_ms Average generation duration per request.\n", + ); + body.push_str("# TYPE llm_api_average_generation_ms gauge\n"); + f( + &mut body, + "llm_api_average_generation_ms", + format!("{avg_ms:.1}"), + ); + body.push_str("# HELP llm_api_uptime_seconds Server process uptime.\n"); + body.push_str("# TYPE llm_api_uptime_seconds gauge\n"); + f(&mut body, "llm_api_uptime_seconds", uptime); + body.push_str("# HELP llm_api_start_time_seconds Process start time (unix).\n"); + body.push_str("# TYPE llm_api_start_time_seconds gauge\n"); + f(&mut body, "llm_api_start_time_seconds", start_ts); + + // Engine identity (helpful when multiple model servers exist). + body.push_str("# HELP llm_api_build_info Build information.\n"); + body.push_str("# TYPE llm_api_build_info gauge\n"); + let _ = writeln!( + body, + "llm_api_build_info{{version=\"{}\",model=\"{}\"}} 1", + env!("CARGO_PKG_VERSION"), + crate::config::MODEL_ID + ); + + ( + StatusCode::OK, + [( + header::CONTENT_TYPE, + HeaderValue::from_static("text/plain; version=0.0.4; charset=utf-8"), + )], + body, + ) + .into_response() +} + +/// Snapshot helper used by tests to inspect counters. +pub fn snapshot() -> (u64, u64, u64) { + ( + REQUESTS_TOTAL.load(Ordering::Relaxed), + COMPLETION_TOKENS_TOTAL.load(Ordering::Relaxed), + GENERATION_MS_TOTAL.load(Ordering::Relaxed), + ) +} + +/// Used by tests to verify the monotonic clock source is live. +pub fn start_instant() -> &'static Instant { + &START_INSTANT +} + +#[cfg(test)] +mod tests { + use super::*; + + fn reset() { + for c in [ + &REQUESTS_TOTAL, + &ERRORS_TOTAL, + &STREAMING_TOTAL, + &PROMPT_TOKENS_TOTAL, + &COMPLETION_TOKENS_TOTAL, + &GENERATION_MS_TOTAL, + &ABORTED_TOTAL, + ] { + c.store(0, Ordering::Relaxed); + } + } + + #[test] + fn counters_accumulate() { + reset(); + count_request(true); + count_request(false); + count_error(); + count_aborted(); + record_tokens(100, 250, 5000); + + assert_eq!(REQUESTS_TOTAL.load(Ordering::Relaxed), 2); + assert_eq!(STREAMING_TOTAL.load(Ordering::Relaxed), 1); + assert_eq!(ERRORS_TOTAL.load(Ordering::Relaxed), 1); + assert_eq!(ABORTED_TOTAL.load(Ordering::Relaxed), 1); + assert_eq!(PROMPT_TOKENS_TOTAL.load(Ordering::Relaxed), 100); + assert_eq!(COMPLETION_TOKENS_TOTAL.load(Ordering::Relaxed), 250); + assert_eq!(GENERATION_MS_TOTAL.load(Ordering::Relaxed), 5000); + } + + #[test] + fn metrics_body_has_prometheus_shape() { + reset(); + count_request(true); + record_tokens(10, 20, 1000); + + let response = futures::executor::block_on(metrics()); + let bytes = + futures::executor::block_on(axum::body::to_bytes(response.into_body(), usize::MAX)) + .expect("read body"); + let text = String::from_utf8(bytes.to_vec()).unwrap(); + + assert!(text.contains("# TYPE llm_api_requests_total counter")); + assert!(text.contains("llm_api_requests_total 1")); + assert!(text.contains("llm_api_completion_tokens_total 20")); + assert!(text.contains("llm_api_build_info{version=")); + assert!(text.contains("llm_api_uptime_seconds ")); + } + + #[test] + fn uptime_is_monotonic() { + let a = uptime_secs(); + std::thread::sleep(std::time::Duration::from_millis(20)); + let b = uptime_secs(); + assert!(b >= a); + } +} diff --git a/src/presentation/handler/mod.rs b/src/presentation/handler/mod.rs index 409ac03..b20852c 100644 --- a/src/presentation/handler/mod.rs +++ b/src/presentation/handler/mod.rs @@ -1,4 +1,7 @@ +//! HTTP handlers. + pub mod chat; pub mod chat_ui; pub mod health; +pub mod metrics; pub mod models; diff --git a/src/presentation/router.rs b/src/presentation/router.rs index 6d4ef02..13f4c17 100644 --- a/src/presentation/router.rs +++ b/src/presentation/router.rs @@ -7,7 +7,7 @@ use axum::routing::{get, post}; use axum::Router; use tower_http::cors::CorsLayer; -use super::handler::{chat, chat_ui, health, models}; +use super::handler::{chat, chat_ui, health, metrics, models}; use crate::presentation::middleware::auth::auth_middleware; use crate::presentation::state::AppState; @@ -17,6 +17,7 @@ pub fn build_router(state: Arc) -> Router { // Public routes (no auth) .route("/", get(chat_ui::chat_ui)) .route("/health", get(health::health_check)) + .route("/metrics", get(metrics::metrics)) .route("/v1/models", get(models::list_models)) // Chat completions (auth-protected) .route("/v1/chat/completions", post(chat::chat_completions))