xserv/crates/xserv-server/src/api.rs

use axum::Extension;
use axum::Json;
use axum::http::StatusCode;
use axum::response::sse::{Event, KeepAlive, Sse};
use axum::response::{IntoResponse, Response};
use serde::{Deserialize, Serialize};
use std::convert::Infallible;
use std::sync::Arc;
use tokio_stream::StreamExt;
use tokio_stream::wrappers::ReceiverStream;
use uuid::Uuid;

use crate::AppState;
use crate::engine::{GenerateEvent, GenerateRequest};
use xserv_model::SamplingParams;

#[derive(Deserialize)]
pub struct ChatRequest {
    #[serde(default)]
    pub model: Option<String>,
    pub messages: Vec<Message>,
    #[serde(default = "default_max_tokens")]
    pub max_tokens: usize,
    #[serde(default)]
    pub stream: Option<bool>,
    #[serde(default)]
    pub temperature: Option<f32>,
    #[serde(default)]
    pub top_k: Option<usize>,
    #[serde(default)]
    pub top_p: Option<f32>,
}

#[derive(Deserialize)]
pub struct Message {
    pub role: String,
    pub content: String,
}

fn default_max_tokens() -> usize {
    256
}

#[derive(Serialize)]
pub struct ModelsResponse {
    object: &'static str,
    data: Vec<ModelInfo>,
}

#[derive(Serialize)]
pub struct ModelInfo {
    id: String,
    object: &'static str,
    owned_by: &'static str,
}

pub async fn health() -> &'static str {
    "ok"
}

pub async fn list_models(Extension(state): Extension<Arc<AppState>>) -> Json<ModelsResponse> {
    Json(ModelsResponse {
        object: "list",
        data: vec![ModelInfo {
            id: state.model_name.clone(),
            object: "model",
            owned_by: "xserv",
        }],
    })
}

pub async fn chat_completions(
    Extension(state): Extension<Arc<AppState>>,
    Json(req): Json<ChatRequest>,
) -> Response {
    if req.stream == Some(true) {
        chat_stream(state, req)
    } else {
        chat_non_stream(state, req).await
    }
}

async fn chat_non_stream(state: Arc<AppState>, req: ChatRequest) -> Response {
    let id = format!("chatcmpl-{}", Uuid::new_v4());
    let model_name = state.model_name.clone();
    let created = unix_timestamp();

    if let Some(response) = validate_request(&req, &model_name) {
        return response;
    }

    let prompt = build_prompt(&req.messages);
    let prompt_tokens = state.engine_tokenizer.lock().unwrap().encode(&prompt);
    let prompt_token_count = prompt_tokens.len();

    let max_seq_len = state.max_seq_len;
    if prompt_token_count >= max_seq_len {
        return bad_request(format!(
            "prompt is {} tokens, exceeds max_seq_len {}",
            prompt_token_count, max_seq_len
        ));
    }
    let max_tokens = req.max_tokens.min(max_seq_len - prompt_token_count);

    let (tx, mut rx) = tokio::sync::mpsc::channel::<GenerateEvent>(64);
    let gen_req = GenerateRequest {
        prompt_tokens,
        max_tokens,
        sampling: sampling_params(&req),
        sender: tx,
    };
    if let Err(resp) = submit_to_engine(&state, gen_req) {
        return resp;
    }

    let mut content = String::new();
    let mut completion_token_count: usize = 0;
    let mut finish_reason = "length".to_string();
    while let Some(event) = rx.recv().await {
        match event {
            GenerateEvent::Token { text, .. } => {
                completion_token_count += 1;
                content.push_str(&text);
            }
            GenerateEvent::Done { finish_reason: fr } => {
                finish_reason = fr;
                break;
            }
        }
    }

    Json(serde_json::json!({
        "id": id,
        "object": "chat.completion",
        "created": created,
        "model": model_name,
        "choices": [{
            "index": 0,
            "message": { "role": "assistant", "content": content },
            "finish_reason": finish_reason,
        }],
        "usage": {
            "prompt_tokens": prompt_token_count,
            "completion_tokens": completion_token_count,
            "total_tokens": prompt_token_count + completion_token_count
        }
    })).into_response()
}

fn chat_stream(
    state: Arc<AppState>,
    req: ChatRequest,
) -> Response {
    let id = format!("chatcmpl-{}", Uuid::new_v4());
    let model_name = state.model_name.clone();
    let created = unix_timestamp();

    if let Some(response) = validate_request(&req, &model_name) {
        return response;
    }

    let prompt = build_prompt(&req.messages);
    let prompt_tokens = state.engine_tokenizer.lock().unwrap().encode(&prompt);

    let max_seq_len = state.max_seq_len;
    if prompt_tokens.len() >= max_seq_len {
        return bad_request(format!(
            "prompt is {} tokens, exceeds max_seq_len {}",
            prompt_tokens.len(), max_seq_len
        ));
    }
    let max_tokens = req.max_tokens.min(max_seq_len - prompt_tokens.len());

    let (engine_tx, engine_rx) = tokio::sync::mpsc::channel::<GenerateEvent>(64);
    let gen_req = GenerateRequest {
        prompt_tokens,
        max_tokens,
        sampling: sampling_params(&req),
        sender: engine_tx,
    };
    if let Err(resp) = submit_to_engine(&state, gen_req) {
        return resp;
    }

    // SSE event channel: engine events -> SSE events
    let (sse_tx, sse_rx) = tokio::sync::mpsc::channel::<Result<Event, Infallible>>(64);

    tokio::spawn(async move {
        let mut engine_stream = ReceiverStream::new(engine_rx);
        let mut first = true;

        while let Some(event) = engine_stream.next().await {
            match event {
                GenerateEvent::Token { text, .. } => {
                    if first {
                        // First chunk: role announcement
                        let chunk =
                            make_chunk(&id, &model_name, created, None, Some("assistant"), None);
                        let _ = sse_tx.send(Ok(Event::default().data(chunk))).await;
                        first = false;
                    }
                    let chunk = make_chunk(&id, &model_name, created, Some(&text), None, None);
                    if sse_tx.send(Ok(Event::default().data(chunk))).await.is_err() {
                        return; // client disconnected
                    }
                }
                GenerateEvent::Done { finish_reason } => {
                    if first {
                        // Edge case: Done arrived with no tokens
                        let chunk =
                            make_chunk(&id, &model_name, created, None, Some("assistant"), None);
                        let _ = sse_tx.send(Ok(Event::default().data(chunk))).await;
                    }
                    let chunk =
                        make_chunk(&id, &model_name, created, None, None, Some(&finish_reason));
                    let _ = sse_tx.send(Ok(Event::default().data(chunk))).await;
                    let _ = sse_tx
                        .send(Ok(Event::default().data("[DONE]".to_string())))
                        .await;
                    return;
                }
            }
        }
    });

    Sse::new(ReceiverStream::new(sse_rx)).keep_alive(KeepAlive::default()).into_response()
}

fn validate_request(req: &ChatRequest, model_name: &str) -> Option<Response> {
    if let Some(model) = &req.model {
        if model != model_name {
            return Some(bad_request(format!(
                "model '{model}' is not loaded; available model is '{model_name}'"
            )));
        }
    }

    if req.max_tokens == 0 {
        return Some(bad_request("max_tokens must be greater than 0"));
    }

    None
}

/// Hand a request to the engine thread. Poison-tolerant (recovers the lock if a
/// prior handler panicked) and returns a clean 503 instead of panicking when the
/// engine thread is gone, so one dead engine doesn't cascade into every request.
fn submit_to_engine(state: &AppState, req: GenerateRequest) -> Result<(), Response> {
    let sender = state.engine_sender.lock().unwrap_or_else(|e| e.into_inner());
    sender.send(req).map_err(|_| service_unavailable("inference engine is not available"))
}

fn service_unavailable(message: impl Into<String>) -> Response {
    (
        StatusCode::SERVICE_UNAVAILABLE,
        Json(serde_json::json!({
            "error": { "message": message.into(), "type": "server_error" }
        })),
    )
        .into_response()
}

fn bad_request(message: impl Into<String>) -> Response {
    (
        StatusCode::BAD_REQUEST,
        Json(serde_json::json!({
            "error": {
                "message": message.into(),
                "type": "invalid_request_error"
            }
        })),
    )
        .into_response()
}

fn make_chunk(
    id: &str,
    model: &str,
    created: u64,
    content: Option<&str>,
    role: Option<&str>,
    finish_reason: Option<&str>,
) -> String {
    let mut delta = serde_json::Map::new();
    if let Some(r) = role {
        delta.insert("role".into(), serde_json::Value::String(r.into()));
        // Role chunk also includes empty content per OpenAI spec
        delta.insert("content".into(), serde_json::Value::String(String::new()));
    }
    if let Some(c) = content {
        delta.insert("content".into(), serde_json::Value::String(c.into()));
    }

    let fr = match finish_reason {
        Some(r) => serde_json::Value::String(r.into()),
        None => serde_json::Value::Null,
    };

    serde_json::json!({
        "id": id,
        "object": "chat.completion.chunk",
        "created": created,
        "model": model,
        "choices": [{
            "index": 0,
            "delta": delta,
            "finish_reason": fr,
        }]
    })
    .to_string()
}

fn unix_timestamp() -> u64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap()
        .as_secs()
}

fn sampling_params(req: &ChatRequest) -> SamplingParams {
    SamplingParams {
        temperature: req.temperature.unwrap_or(0.0),
        top_k: req.top_k.unwrap_or(0),
        top_p: req.top_p.unwrap_or(1.0),
    }
}

fn build_prompt(messages: &[Message]) -> String {
    let mut prompt = String::new();
    for msg in messages {
        match msg.role.as_str() {
            "system" | "user" | "assistant" => {
                prompt.push_str("<|im_start|>");
                prompt.push_str(&msg.role);
                prompt.push('\n');
                prompt.push_str(&msg.content);
                prompt.push_str("<|im_end|>\n");
            }
            _ => {}
        }
    }
    prompt.push_str("<|im_start|>assistant\n");
    prompt.push_str("<think>\n\n</think>\n\n");
    prompt
}