chat-api/src/routes/v1/models.rs

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<AppState>,
) -> Result<Json<api::ModelsResponse>, (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<AppState>,
    Path(model): Path<String>,
    Json(body): Json<api::LoadModelBody>,
) -> Result<Json<api::LoadModelResponse>, (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<AppState>,
    Path(model): Path<String>,
) -> Result<Json<api::UnloadModelResponse>, (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()),
    }
}