diff options
Diffstat (limited to 'makima/src')
| -rw-r--r-- | makima/src/bin/server.rs | 32 | ||||
| -rw-r--r-- | makima/src/server/mod.rs | 23 |
2 files changed, 46 insertions, 9 deletions
diff --git a/makima/src/bin/server.rs b/makima/src/bin/server.rs index 06b6585..470e295 100644 --- a/makima/src/bin/server.rs +++ b/makima/src/bin/server.rs @@ -8,10 +8,10 @@ use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; use makima::server::{run_server, state::AppState}; -/// Default model paths relative to the working directory. -const PARAKEET_MODEL_DIR: &str = "models/parakeet-tdt-0.6b-v3"; -const PARAKEET_EOU_DIR: &str = "models/realtime_eou_120m-v1-onnx"; -const SORTFORMER_MODEL_PATH: &str = "models/diarization/diar_streaming_sortformer_4spk-v2.onnx"; +/// Default model paths (can be overridden via environment variables). +const DEFAULT_PARAKEET_MODEL_DIR: &str = "models/parakeet-tdt-0.6b-v3"; +const DEFAULT_PARAKEET_EOU_DIR: &str = "models/realtime_eou_120m-v1-onnx"; +const DEFAULT_SORTFORMER_MODEL_PATH: &str = "models/diarization/diar_streaming_sortformer_4spk-v2.onnx"; #[tokio::main] async fn main() -> anyhow::Result<()> { @@ -19,22 +19,38 @@ async fn main() -> anyhow::Result<()> { tracing_subscriber::registry() .with( tracing_subscriber::EnvFilter::try_from_default_env() - .unwrap_or_else(|_| "makima=debug,tower_http=debug".into()), + .unwrap_or_else(|_| "makima=info,tower_http=info".into()), ) .with(tracing_subscriber::fmt::layer()) .init(); tracing::info!("Starting Makima Listening API Server"); - tracing::info!("Loading ML models..."); + + // Read configuration from environment + let port = std::env::var("PORT").unwrap_or_else(|_| "8080".to_string()); + let parakeet_dir = std::env::var("PARAKEET_MODEL_DIR") + .unwrap_or_else(|_| DEFAULT_PARAKEET_MODEL_DIR.to_string()); + let parakeet_eou_dir = std::env::var("PARAKEET_EOU_DIR") + .unwrap_or_else(|_| DEFAULT_PARAKEET_EOU_DIR.to_string()); + let sortformer_path = std::env::var("SORTFORMER_MODEL_PATH") + .unwrap_or_else(|_| DEFAULT_SORTFORMER_MODEL_PATH.to_string()); + + tracing::info!( + parakeet = %parakeet_dir, + eou = %parakeet_eou_dir, + sortformer = %sortformer_path, + "Loading ML models..." + ); // Load ML models let state = Arc::new( - AppState::new(PARAKEET_MODEL_DIR, PARAKEET_EOU_DIR, SORTFORMER_MODEL_PATH) + AppState::new(¶keet_dir, ¶keet_eou_dir, &sortformer_path) .map_err(|e| anyhow::anyhow!("Failed to load models: {}", e))?, ); tracing::info!("Models loaded successfully"); // Run the server - run_server(state, "0.0.0.0:8080").await + let addr = format!("0.0.0.0:{}", port); + run_server(state, &addr).await } diff --git a/makima/src/server/mod.rs b/makima/src/server/mod.rs index a6e0525..c509afa 100644 --- a/makima/src/server/mod.rs +++ b/makima/src/server/mod.rs @@ -6,9 +6,12 @@ pub mod openapi; pub mod state; use axum::{ + http::StatusCode, + response::IntoResponse, routing::get, - Router, + Json, Router, }; +use serde::Serialize; use tower_http::cors::{Any, CorsLayer}; use tower_http::trace::TraceLayer; use utoipa::OpenApi; @@ -18,6 +21,23 @@ use crate::server::handlers::listen; use crate::server::openapi::ApiDoc; use crate::server::state::SharedState; +#[derive(Serialize)] +struct HealthResponse { + status: &'static str, + version: &'static str, +} + +/// Health check endpoint for load balancers and orchestrators. +async fn health_check() -> impl IntoResponse { + ( + StatusCode::OK, + Json(HealthResponse { + status: "healthy", + version: env!("CARGO_PKG_VERSION"), + }), + ) +} + /// Create the axum Router with all routes configured. pub fn make_router(state: SharedState) -> Router { // API v1 routes @@ -29,6 +49,7 @@ pub fn make_router(state: SharedState) -> Router { .url("/api-docs/openapi.json", ApiDoc::openapi()); Router::new() + .route("/api/v1/healthcheck", get(health_check)) .nest("/api/v1", api_v1) .merge(swagger) .layer( |
