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

use axum::{
    Json,
    extract::State,
    response::{
        IntoResponse, Response,
        sse::{KeepAlive, Sse},
    },
};

use crate::dto::api;
use crate::errors::OllamaError;
use crate::state::app_state::AppState;

#[utoipa::path(
    post,
    path = "/completions",
    tag = "chat",
    request_body(
        content = api::CompletionRequest,
        description = "Text completion request",
        content_type = "application/json"
    ),
    responses(
        (
            status = 200,
            description = "Text completion response. If stream=true, response is SSE stream of chunks ending in [DONE].",
            body = api::CompletionResponse,
            content_type = "application/json"
        ),
        (
            status = 400,
            description = "Invalid request: missing prompt, model, or invalid format",
            body = api::ErrorResponse,
            example = json!({ "error": "prompt is required and cannot be empty" })
        ),
        (
            status = 422,
            description = "Model not found or not available locally",
            body = api::ErrorResponse,
            example = json!({ "error": "model 'llama3' is not available — run `ollama pull llama3` first" })
        ),
        (
            status = 500,
            description = "Internal server error (Ollama or network failure)",
            body = api::ErrorResponse,
            example = json!({ "error": "connection refused" })
        )
    )
)]
pub async fn completions(
    State(state): State<AppState>,
    Json(body): Json<api::CompletionRequest>,
) -> Result<Response, (axum::http::StatusCode, Json<api::ErrorResponse>)> {
    tracing::debug!("Received /completion with body {:?}", body);
    if body.base.stream {
        let stream = state.ollama.completions_stream(&body).await.map_err(|e| {
            let (code, msg) = ollama_err(e);
            (code, Json(api::ErrorResponse::new(msg)))
        })?;

        Ok(Sse::new(stream)
            .keep_alive(KeepAlive::default())
            .into_response())
    } else {
        let response = state.ollama.completions(&body).await.map_err(|e| {
            let (code, msg) = ollama_err(e);
            (code, Json(api::ErrorResponse::new(msg)))
        })?;

        Ok(Json(response).into_response())
    }
}

#[utoipa::path(
    post,
    path = "/chat/completions",
    tag = "chat",
    request_body(
        content = api::ChatRequest,
        description = "Chat completion request with message history",
        content_type = "application/json"
    ),
    responses(
        (
            status = 200,
            description = "Chat completion response. If stream=false returns JSON. If stream=true returns SSE stream of chunks ending with [DONE].",
            body = api::ChatCompletionResponse,
            content_type = "application/json"
        ),
        (
            status = 400,
            description = "Invalid request",
            body = api::ErrorResponse,
            example = json!({ "error": "messages array with at least one user message is required" })
        ),
        (
            status = 401,
            description = "Unauthorized",
            body = api::ErrorResponse,
            example = json!({ "error": "missing or invalid token" })
        ),
        (
            status = 422,
            description = "Model not found or unavailable",
            body = api::ErrorResponse,
            example = json!({ "error": "model 'llama3' is not available — run `ollama pull llama3` first" })
        ),
        (
            status = 500,
            description = "Internal server error",
            body = api::ErrorResponse,
            example = json!({ "error": "connection refused" })
        )
    )
)]
pub async fn chat_completions(
    State(state): State<AppState>,
    Json(body): Json<api::ChatRequest>,
) -> Result<Response, (axum::http::StatusCode, String)> {
    tracing::debug!("Received /completion with body {:?}", body);
    if body.base.stream {
        let stream = state
            .ollama
            .chat_completions_stream(&body)
            .await
            .map_err(ollama_err)?;

        Ok(Sse::new(stream)
            .keep_alive(KeepAlive::default())
            .into_response())
    } else {
        let response = state
            .ollama
            .chat_completions(&body)
            .await
            .map_err(ollama_err)?;

        Ok(Json(response).into_response())
    }
}

fn ollama_err(e: OllamaError) -> (axum::http::StatusCode, String) {
    match e {
        OllamaError::MissingPrompt => (
            axum::http::StatusCode::BAD_REQUEST,
            "prompt is required and cannot be empty".to_string(),
        ),
        OllamaError::MissingModel => (
            axum::http::StatusCode::BAD_REQUEST,
            "model is required and cannot be empty".to_string(),
        ),
        OllamaError::ModelNotFound(m) => (
            axum::http::StatusCode::UNPROCESSABLE_ENTITY,
            format!("model '{m}' is not available — run `ollama pull {m}` first"),
        ),
        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}'"),
        ),
        OllamaError::MissingMessages => (
            axum::http::StatusCode::BAD_REQUEST,
            "messages array with at least one user message is required".to_string(),
        ),
        OllamaError::Http(e) => (axum::http::StatusCode::INTERNAL_SERVER_ERROR, e.to_string()),
    }
}