feat: add streaming

src/routes/v1/chat.rs

@@ -1,4 +1,11 @@
-use axum::{Json, extract::State};
+use axum::{
+    Json,
+    extract::State,
+    response::{
+        IntoResponse, Response,
+        sse::{KeepAlive, Sse},
+    },
+};
 use serde_json::Value;
 
 use crate::errors::OllamaError;
@@ -7,38 +14,77 @@ use crate::state::app_state::AppState;
 pub async fn completions(
     State(state): State<AppState>,
     Json(body): Json<Value>,
-) -> Result<Json<Value>, (axum::http::StatusCode, String)> {
-    match state.ollama.completions(body).await {
-        Ok(response) => Ok(Json(response)),
-        Err(OllamaError::MissingPrompt) => Err((
-            axum::http::StatusCode::BAD_REQUEST,
-            "prompt is required and cannot be empty".to_string(),
-        )),
-        Err(OllamaError::ModelNotFound(m)) => Err((
-            axum::http::StatusCode::UNPROCESSABLE_ENTITY,
-            format!("model '{m}' is not available — run `ollama pull {m}` first"),
-        )),
-        Err(e) => Err((axum::http::StatusCode::INTERNAL_SERVER_ERROR, e.to_string())),
+) -> Result<Response, (axum::http::StatusCode, String)> {
+    let wants_stream = body
+        .get("stream")
+        .and_then(|v| v.as_bool())
+        .unwrap_or(false);
+
+    if wants_stream {
+        let stream = state
+            .ollama
+            .completions_stream(body)
+            .await
+            .map_err(ollama_err)?;
+
+        Ok(Sse::new(stream)
+            .keep_alive(KeepAlive::default())
+            .into_response())
+    } else {
+        let response = state.ollama.completions(body).await.map_err(ollama_err)?;
+
+        Ok(Json(response).into_response())
     }
 }
 
 pub async fn chat_completions(
     State(state): State<AppState>,
     Json(body): Json<Value>,
-) -> Result<Json<Value>, (axum::http::StatusCode, String)> {
-    match state.ollama.chat_completions(body).await {
-        Ok(response) => Ok(Json(response)),
-        Err(OllamaError::MissingMessages) => Err((
+) -> Result<Response, (axum::http::StatusCode, String)> {
+    let wants_stream = body
+        .get("stream")
+        .and_then(|v| v.as_bool())
+        .unwrap_or(false);
+
+    if wants_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,
-            "messages must be a non-empty array with at least one user message".to_string(),
-        )),
-        Err(OllamaError::ModelNotFound(m)) => Err((
+            "prompt 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"),
-        )),
-        Err(OllamaError::Http(e)) => {
-            Err((axum::http::StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))
-        }
-        Err(e) => Err((axum::http::StatusCode::BAD_REQUEST, e.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()),
     }
 }