use axum::{ Json, extract::{Path, State}, }; use crate::dto::api; use crate::{errors::OllamaError, state::app_state::AppState}; #[utoipa::path( get, path = "/models", tag = "models", responses( ( status = 200, description = "List of locally available Ollama models", body = api::ModelsResponse, content_type = "application/json", ), ( status = 500, description = "Internal server error (Ollama or network failure)", body = api::ErrorResponse, example = json!({ "error": "connection refused" }) ) ) )] pub async fn list_models( State(state): State, ) -> Result, (axum::http::StatusCode, String)> { match state.ollama.list_models().await { Ok(models) => Ok(Json(models)), Err(e) => Err(ollama_err(e)), } } #[utoipa::path( post, path = "/models/{model}/load", tag = "models", params( ("model" = String, Path, description = "Name of the model to load into memory (e.g. 'llama3')") ), request_body( content = api::LoadModelBody, description = "Load model request", content_type = "application/json", example = json!({ "keep_alive": "10m" }) ), responses( ( status = 200, description = "Model successfully loaded into memory", body = api::LoadModelResponse, content_type = "application/json", ), ( status = 400, description = "Invalid or missing keep_alive format", body = api::ErrorResponse, examples( ("Missing" = (value = json!({ "error": "keep alive is required and cannot be empty" }))), ("Invalid" = (value = json!({ "error": "invalid keep_alive '10x' — use 30s / 10m / 2h, a plain integer, or -1" }))) ) ), ( status = 404, description = "Model not found locally", body = api::ErrorResponse, example = json!({ "error": "model 'llama3' not found — run `ollama pull llama3`" }) ), ( status = 500, description = "Internal server error (Ollama or network failure)", body = api::ErrorResponse, example = json!({ "error": "connection refused" }) ) ) )] pub async fn load_model( State(state): State, Path(model): Path, Json(body): Json, ) -> Result, (axum::http::StatusCode, String)> { let response = state .ollama .load_model(&model, body.keep_alive.as_deref()) .await .map_err(ollama_err)?; Ok(Json(response)) } #[utoipa::path( delete, path = "/models/{model}/load", tag = "models", params( ("model" = String, Path, description = "Name of the model to unload from memory (e.g. 'llama3')") ), responses( ( status = 200, description = "Model successfully unloaded from memory", body = api::UnloadModelResponse, content_type = "application/json", ), ( status = 404, description = "Model not found locally", body = api::ErrorResponse, example = json!({ "error": "model 'llama3' not found — run `ollama pull llama3`" }) ), ( status = 500, description = "Internal server error (Ollama or network failure)", body = api::ErrorResponse, example = json!({ "error": "connection refused" }) ) ) )] pub async fn unload_model( State(state): State, Path(model): Path, ) -> Result, (axum::http::StatusCode, String)> { let response = state .ollama .unload_model(&model) .await .map_err(ollama_err)?; Ok(Json(response)) } fn ollama_err(e: OllamaError) -> (axum::http::StatusCode, String) { match e { OllamaError::ModelNotFound(m) => ( axum::http::StatusCode::NOT_FOUND, format!("model '{m}' not found — run `ollama pull {m}`"), ), OllamaError::MissingModel => ( axum::http::StatusCode::BAD_REQUEST, "model is required and cannot be empty".to_string(), ), OllamaError::MissingPrompt => ( axum::http::StatusCode::BAD_REQUEST, "prompt is required".to_string(), ), OllamaError::MissingMessages => ( axum::http::StatusCode::BAD_REQUEST, "messages array with at least one user message is required".to_string(), ), OllamaError::MissingKeepAlive => ( axum::http::StatusCode::BAD_REQUEST, "keep alive is required and cannot be empty".to_string(), ), OllamaError::InvalidKeepAlive(v) => ( axum::http::StatusCode::BAD_REQUEST, format!("invalid keep_alive '{v}' — use 30s / 10m / 2h, a plain integer, or -1"), ), OllamaError::Http(e) => (axum::http::StatusCode::INTERNAL_SERVER_ERROR, e.to_string()), } }