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, Json(body): Json, ) -> Result)> { 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, Json(body): Json, ) -> Result { 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()), } }