use serde::{Deserialize, Serialize}; use crate::config::{LABELS, SERVICE_NAME, SERVICE_VERSION}; use crate::model::{ModelService, Prediction}; use crate::error::ServiceError; use crate::image::preprocess_image; use crate::telemetry; use crate::telemetry::RequestMetricsGuard; use axum::{ extract::{State, Multipart}, routing::{get, post}, Json, Router, }; use std::sync::Arc; #[derive(Debug, Serialize, Deserialize)] pub struct HealthResponse { pub status: String, pub model_loaded: bool, } #[derive(Debug, Serialize, Deserialize)] pub struct MetadataResponse { pub service_name: String, pub service_version: String, pub model_path: String, pub model_loaded: bool, pub input_size: u32, pub labels: Vec, } #[derive(Debug, Serialize, Deserialize)] pub struct PredictionResponse { pub status: String, pub label: String, pub confidence: f32, pub probabilities: std::collections::BTreeMap, } #[derive(Clone)] pub struct AppState { pub model: Arc, } pub fn health_response(model_loaded: bool) -> HealthResponse { HealthResponse { status: "ok".to_string(), model_loaded, } } pub fn metadata_response(model_path: String, model_loaded: bool, input_size: u32) -> MetadataResponse { MetadataResponse { service_name: SERVICE_NAME.to_string(), service_version: SERVICE_VERSION.to_string(), model_path, model_loaded, input_size, labels: LABELS.iter().map(|s| s.to_string()).collect(), } } pub fn prediction_response(prediction: Prediction) -> PredictionResponse { PredictionResponse { status: prediction.status, label: prediction.label, confidence: prediction.confidence, probabilities: prediction.probabilities, } } pub async fn metrics() -> (axum::http::StatusCode, String) { (axum::http::StatusCode::OK, telemetry::encode_metrics()) } pub async fn health(State(state): State) -> Json { let _guard = RequestMetricsGuard::new(); let res = health_response(state.model.is_loaded()); _guard.finish(); Json(res) } pub async fn metadata(State(state): State) -> Json { let _guard = RequestMetricsGuard::new(); let res = metadata_response( state.model.model_path().to_string_lossy().to_string(), state.model.is_loaded(), state.model.input_size(), ); _guard.finish(); Json(res) } pub async fn predict( State(state): State, mut multipart: Multipart, ) -> Result, ServiceError> { let _guard = RequestMetricsGuard::new(); // Extract the file field from multipart let mut file_data = None; while let Ok(Some(field)) = multipart.next_field().await { if field.name() == Some("file") { if let Some(content_type) = field.content_type() { if !content_type.starts_with("image/") { return Err(ServiceError::BadRequest( "Uploaded file must be an image".to_string(), )); } } else { return Err(ServiceError::BadRequest( "Uploaded file must be an image".to_string(), )); } file_data = Some(field.bytes().await); break; } } // Handle missing or multipart read errors let bytes = match file_data { Some(Ok(b)) => b, Some(Err(_)) => { return Err(ServiceError::BadRequest( "Uploaded file must be an image".to_string(), )) } None => { return Err(ServiceError::BadRequest( "Uploaded file must be an image".to_string(), )) } }; // Preprocess the image let _preprocess_start = std::time::Instant::now(); let input = preprocess_image(&bytes, state.model.input_size())?; // Record image size metric telemetry::image_size_bytes().observe(bytes.len() as f64); // Run prediction with timing let inference_start = std::time::Instant::now(); let prediction = state.model.predict(input)?; telemetry::inference_duration_seconds().observe(inference_start.elapsed().as_secs_f64()); // Record business telemetry telemetry::predictions_total().inc(); telemetry::predictions_by_class() .with_label_values(&[&prediction.label]) .inc(); telemetry::predictions_confidence().observe(prediction.confidence as f64); _guard.finish(); Ok(Json(prediction_response(prediction))) } /// Helper to record errors from route handlers pub fn record_error(kind: &str) { telemetry::errors_total().with_label_values(&[kind]).inc(); } pub fn router(state: AppState) -> Router { Router::new() .route("/health", get(health)) .route("/metadata", get(metadata)) .route("/predict", post(predict)) .route("/metrics", get(metrics)) .with_state(state) } #[cfg(test)] mod tests { use super::*; #[test] fn health_response_with_model_loaded() { let response = health_response(true); assert_eq!(response.status, "ok"); assert_eq!(response.model_loaded, true); } #[test] fn health_response_without_model_loaded() { let response = health_response(false); assert_eq!(response.status, "ok"); assert_eq!(response.model_loaded, false); } #[test] fn metadata_response_includes_service_info() { let response = metadata_response("/path/to/model.onnx".to_string(), true, 224); assert_eq!(response.service_name, "zeavis-ml-service"); assert_eq!(response.service_version, "0.1.0"); assert_eq!(response.model_path, "/path/to/model.onnx"); assert_eq!(response.model_loaded, true); assert_eq!(response.input_size, 224); } #[test] fn metadata_response_includes_all_labels() { let response = metadata_response("/path/to/model.onnx".to_string(), true, 224); assert_eq!(response.labels.len(), 4); assert_eq!( response.labels, vec!["Bercak Daun", "Daun Sehat", "Hawar Daun", "Karat Daun"] ); } #[test] fn prediction_response_matches_prediction_contract() { let prediction = Prediction { status: "confident".to_string(), label: "Hawar Daun".to_string(), confidence: 0.6, probabilities: { let mut map = std::collections::BTreeMap::new(); map.insert("Bercak Daun".to_string(), 0.1); map.insert("Daun Sehat".to_string(), 0.2); map.insert("Hawar Daun".to_string(), 0.6); map.insert("Karat Daun".to_string(), 0.1); map }, }; let response = prediction_response(prediction); assert_eq!(response.status, "confident"); assert_eq!(response.label, "Hawar Daun"); assert_eq!(response.confidence, 0.6); assert_eq!(response.probabilities.len(), 4); assert_eq!(response.probabilities.get("Bercak Daun"), Some(&0.1)); assert_eq!(response.probabilities.get("Daun Sehat"), Some(&0.2)); assert_eq!(response.probabilities.get("Hawar Daun"), Some(&0.6)); assert_eq!(response.probabilities.get("Karat Daun"), Some(&0.1)); } }