diff --git a/Cargo.lock b/Cargo.lock index 0cf0890..cf3cd69 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -177,6 +177,7 @@ dependencies = [ "thiserror 2.0.18", "tokio", "tokio-stream", + "utoipa", "wiremock", ] @@ -692,6 +693,8 @@ checksum = "45a8a2b9cb3e0b0c1803dbb0758ffac5de2f425b23c28f518faabd9d805342ff" dependencies = [ "equivalent", "hashbrown", + "serde", + "serde_core", ] [[package]] @@ -1801,6 +1804,30 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" +[[package]] +name = "utoipa" +version = "5.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2fcc29c80c21c31608227e0912b2d7fddba57ad76b606890627ba8ee7964e993" +dependencies = [ + "indexmap", + "serde", + "serde_json", + "utoipa-gen", +] + +[[package]] +name = "utoipa-gen" +version = "5.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d79d08d92ab8af4c5e8a6da20c47ae3f61a0f1dabc1997cdf2d082b757ca08b" +dependencies = [ + "proc-macro2", + "quote", + "regex", + "syn", +] + [[package]] name = "walkdir" version = "2.5.0" diff --git a/Cargo.toml b/Cargo.toml index e301768..02e9182 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,6 +9,7 @@ tokio = { version = "1", features = ["macros", "rt-multi-thread"] } [dependencies] axum = "0.8.8" +utoipa = { version = "5.4.0", features = ["axum_extras"] } tokio = { version = "1", features = ["full"] } serde = { version = "1", features = ["derive"] } serde_json = "1" diff --git a/src/errors.rs b/src/errors.rs index e36c286..1de226c 100644 --- a/src/errors.rs +++ b/src/errors.rs @@ -6,6 +6,9 @@ pub enum OllamaError { #[error("prompt is required and cannot be empty")] MissingPrompt, + #[error("model is required and cannot be empty")] + MissingModel, + #[error("messages must be a non-empty array containing at least one user message")] MissingMessages, diff --git a/src/main.rs b/src/main.rs index 9f8d122..135725d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,5 +1,6 @@ mod auth; mod errors; +mod openapi; mod providers; mod routes; mod state; diff --git a/src/openapi.rs b/src/openapi.rs new file mode 100644 index 0000000..ed06820 --- /dev/null +++ b/src/openapi.rs @@ -0,0 +1,18 @@ +use utoipa::OpenApi; + +#[derive(OpenApi)] +#[openapi( + paths( + crate::routes::v1::chat::completions + ), + components( + schemas( + // add your request/response structs here later + ) + ), + tags( + (name = "chat", description = "Chat endpoints"), + (name = "models", description = "Model management") + ) +)] +pub struct V1ApiDoc; diff --git a/src/providers/ollama.rs b/src/providers/ollama.rs index cbd5f44..cd79d19 100644 --- a/src/providers/ollama.rs +++ b/src/providers/ollama.rs @@ -82,7 +82,7 @@ impl OllamaProvider { let model = body .get("model") .and_then(|v| v.as_str()) - .unwrap_or("llama3"); + .ok_or(OllamaError::MissingModel)?; Ok((prompt, model)) } diff --git a/src/routes/v1/chat.rs b/src/routes/v1/chat.rs index c3108af..213e168 100644 --- a/src/routes/v1/chat.rs +++ b/src/routes/v1/chat.rs @@ -11,6 +11,13 @@ use serde_json::Value; use crate::errors::OllamaError; use crate::state::app_state::AppState; +#[utoipa::path( + post, + path = "/chat/completions", + responses( + (status = 200, description = "Chat completion", body = Value), + ) +)] pub async fn completions( State(state): State, Json(body): Json, @@ -73,6 +80,10 @@ fn ollama_err(e: OllamaError) -> (axum::http::StatusCode, String) { 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"), diff --git a/src/routes/v1/mod.rs b/src/routes/v1/mod.rs index 30059e7..78112f9 100644 --- a/src/routes/v1/mod.rs +++ b/src/routes/v1/mod.rs @@ -1,16 +1,26 @@ pub mod chat; pub mod models; +mod openapi; use crate::auth::middleware::auth_middleware; use crate::state::app_state::AppState; use axum::{Router, middleware, routing::get, routing::post}; -pub fn router() -> Router { +fn public_router() -> Router { + Router::new().route("/openapi.json", get(openapi::openapi_json)) +} + +pub fn protected_router() -> Router { Router::new() .route("/models", get(models::list_models)) .route("/completions", post(chat::completions)) .route("/chat/completions", post(chat::chat_completions)) .route("/models/{model}/load", post(models::load_model)) .route("/models/{model}/unload", post(models::unload_model)) - .layer(middleware::from_fn(auth_middleware)) +} + +pub fn router() -> Router { + Router::new() + .merge(public_router()) + .merge(protected_router().layer(middleware::from_fn(auth_middleware))) } diff --git a/src/routes/v1/models.rs b/src/routes/v1/models.rs index 2f9d629..19aa281 100644 --- a/src/routes/v1/models.rs +++ b/src/routes/v1/models.rs @@ -53,6 +53,10 @@ fn ollama_err(e: OllamaError) -> (axum::http::StatusCode, String) { 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(), diff --git a/src/routes/v1/openapi.rs b/src/routes/v1/openapi.rs new file mode 100644 index 0000000..bb53a58 --- /dev/null +++ b/src/routes/v1/openapi.rs @@ -0,0 +1,8 @@ +use axum::Json; +use utoipa::OpenApi; + +use crate::openapi::V1ApiDoc; + +pub async fn openapi_json() -> Json { + Json(V1ApiDoc::openapi()) +}