Gahow Wang
d96ee0766c
Three related hardening changes for the API surface: - validate_request rejects NaN/negative temperature, out-of-range top_p, and absurd top_k before those values reach the CUDA sampling paths. Prevents NaN logits from downstream sampling and matches typical OpenAI-compatible server behavior (400 instead of 500). - normalize_finish_reason maps engine strings to the OpenAI-standard subset. Currently only "error" (from tp/pp engine client-stall) needs normalization — it collapses to null so SDK clients see a clean stream close instead of an unknown finish_reason value. Applied to both streaming (SSE) and non-streaming JSON responses. - Replace the unbounded std::sync::mpsc engine channel with a bounded sync_channel(256) and switch submit_to_engine to try_send. A saturated engine now returns 503 "engine is busy" instead of letting requests pile up in RAM. Also add axum DefaultBodyLimit(4 MiB) so a malicious or misbehaving client cannot exhaust memory with an arbitrary JSON POST.
574 lines
19 KiB
Rust
574 lines
19 KiB
Rust
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::path::Path;
|
|
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, Serialize, Clone)]
|
|
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,
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Chat Template: Jinja2 rendering via minijinja
|
|
// ---------------------------------------------------------------------------
|
|
|
|
pub struct ChatTemplate {
|
|
source: String,
|
|
model_type: String,
|
|
}
|
|
|
|
impl ChatTemplate {
|
|
pub fn load(model_dir: &Path, model_type: &str) -> Self {
|
|
// 1. Try standalone chat_template.jinja file
|
|
let jinja_path = model_dir.join("chat_template.jinja");
|
|
if jinja_path.exists() {
|
|
let source = std::fs::read_to_string(&jinja_path)
|
|
.unwrap_or_else(|e| panic!("failed to read {}: {e}", jinja_path.display()));
|
|
eprintln!("[chat-template] loaded from {}", jinja_path.display());
|
|
return Self {
|
|
source,
|
|
model_type: model_type.to_string(),
|
|
};
|
|
}
|
|
|
|
// 2. Try tokenizer_config.json → chat_template field
|
|
let tok_cfg_path = model_dir.join("tokenizer_config.json");
|
|
if tok_cfg_path.exists() {
|
|
if let Ok(data) = std::fs::read_to_string(&tok_cfg_path) {
|
|
if let Ok(v) = serde_json::from_str::<serde_json::Value>(&data) {
|
|
if let Some(ct) = v.get("chat_template").and_then(|v| v.as_str()) {
|
|
eprintln!("[chat-template] loaded from tokenizer_config.json");
|
|
return Self {
|
|
source: ct.to_string(),
|
|
model_type: model_type.to_string(),
|
|
};
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// 3. No template found — use empty source, will fall back to hardcoded
|
|
eprintln!("[chat-template] no Jinja template found, using hardcoded fallback");
|
|
Self {
|
|
source: String::new(),
|
|
model_type: model_type.to_string(),
|
|
}
|
|
}
|
|
|
|
pub fn render(&self, messages: &[Message]) -> String {
|
|
if self.source.is_empty() {
|
|
return build_prompt_hardcoded(messages, &self.model_type);
|
|
}
|
|
|
|
match self.render_jinja(messages) {
|
|
Ok(prompt) => prompt,
|
|
Err(e) => {
|
|
eprintln!("[chat-template] Jinja render error: {e}, falling back to hardcoded");
|
|
build_prompt_hardcoded(messages, &self.model_type)
|
|
}
|
|
}
|
|
}
|
|
|
|
fn render_jinja(&self, messages: &[Message]) -> Result<String, minijinja::Error> {
|
|
let mut env = minijinja::Environment::new();
|
|
|
|
// Register custom functions the template may call.
|
|
env.add_function("strftime_now", strftime_now);
|
|
env.add_function("raise_exception", raise_exception);
|
|
|
|
// Python str methods used by harmony/gpt-oss templates.
|
|
env.add_filter("startswith", |s: String, prefix: String| -> bool {
|
|
s.starts_with(&prefix)
|
|
});
|
|
|
|
env.add_template("chat", &self.source)?;
|
|
let tmpl = env.get_template("chat")?;
|
|
|
|
let ctx = minijinja::context! {
|
|
messages => minijinja::Value::from_serialize(messages),
|
|
add_generation_prompt => true,
|
|
bos_token => "",
|
|
eos_token => "",
|
|
};
|
|
|
|
tmpl.render(ctx)
|
|
}
|
|
}
|
|
|
|
fn strftime_now(fmt: String) -> String {
|
|
use std::time::SystemTime;
|
|
let now = SystemTime::now()
|
|
.duration_since(SystemTime::UNIX_EPOCH)
|
|
.unwrap()
|
|
.as_secs();
|
|
// Only support %Y-%m-%d (the only format used by known templates)
|
|
let days = now / 86400;
|
|
let (y, m, d) = days_to_ymd(days);
|
|
fmt.replace("%Y", &format!("{y:04}"))
|
|
.replace("%m", &format!("{m:02}"))
|
|
.replace("%d", &format!("{d:02}"))
|
|
}
|
|
|
|
fn days_to_ymd(days_since_epoch: u64) -> (u32, u32, u32) {
|
|
// Civil calendar from days since 1970-01-01 (Rata Die algorithm)
|
|
let z = days_since_epoch as i64 + 719468;
|
|
let era = (if z >= 0 { z } else { z - 146096 }) / 146097;
|
|
let doe = (z - era * 146097) as u32;
|
|
let yoe = (doe - doe / 1460 + doe / 36524 - doe / 146096) / 365;
|
|
let y = yoe as i64 + era * 400;
|
|
let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
|
|
let mp = (5 * doy + 2) / 153;
|
|
let d = doy - (153 * mp + 2) / 5 + 1;
|
|
let m = if mp < 10 { mp + 3 } else { mp - 9 };
|
|
let y = if m <= 2 { y + 1 } else { y };
|
|
(y as u32, m, d)
|
|
}
|
|
|
|
fn raise_exception(msg: String) -> Result<String, minijinja::Error> {
|
|
Err(minijinja::Error::new(
|
|
minijinja::ErrorKind::InvalidOperation,
|
|
msg,
|
|
))
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Hardcoded fallback templates (for models without a Jinja template)
|
|
// ---------------------------------------------------------------------------
|
|
|
|
fn build_prompt_hardcoded(messages: &[Message], model_type: &str) -> String {
|
|
if model_type == "gpt_oss" {
|
|
return build_prompt_gpt_oss(messages);
|
|
}
|
|
// Default: Qwen3 ChatML format
|
|
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
|
|
}
|
|
|
|
fn build_prompt_gpt_oss(messages: &[Message]) -> String {
|
|
let mut prompt = String::new();
|
|
// Canonical harmony system message (mirrors the model's chat_template.jinja
|
|
// build_system_message macro). A hand-rolled substitute puts gpt-oss out of
|
|
// distribution and destabilizes channel selection. This hardcoded builder is
|
|
// only a fallback for gpt-oss models that ship no Jinja template; the
|
|
// gpt-oss-20b release does ship one, so the template path is normally used.
|
|
prompt.push_str("<|start|>system<|message|>");
|
|
prompt.push_str("You are ChatGPT, a large language model trained by OpenAI.\n");
|
|
prompt.push_str("Knowledge cutoff: 2024-06\n");
|
|
prompt.push_str(&format!(
|
|
"Current date: {}\n\n",
|
|
strftime_now("%Y-%m-%d".to_string())
|
|
));
|
|
prompt.push_str("Reasoning: low\n\n");
|
|
prompt.push_str("# Valid channels: analysis, commentary, final. Channel must be included for every message.");
|
|
prompt.push_str("<|end|>");
|
|
let dev_instructions: String = messages
|
|
.iter()
|
|
.filter(|m| m.role == "system")
|
|
.map(|m| m.content.as_str())
|
|
.collect::<Vec<_>>()
|
|
.join("\n\n");
|
|
if !dev_instructions.is_empty() {
|
|
prompt.push_str("<|start|>developer<|message|># Instructions\n\n");
|
|
prompt.push_str(&dev_instructions);
|
|
prompt.push_str("<|end|>");
|
|
}
|
|
for msg in messages {
|
|
match msg.role.as_str() {
|
|
"user" => {
|
|
prompt.push_str("<|start|>user<|message|>");
|
|
prompt.push_str(&msg.content);
|
|
prompt.push_str("<|end|>");
|
|
}
|
|
"assistant" => {
|
|
prompt.push_str("<|start|>assistant<|channel|>final<|message|>");
|
|
prompt.push_str(&msg.content);
|
|
prompt.push_str("<|end|>");
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
prompt.push_str("<|start|>assistant<|channel|>final<|message|>");
|
|
prompt
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// HTTP handlers
|
|
// ---------------------------------------------------------------------------
|
|
|
|
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 = state.chat_template.render(&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;
|
|
}
|
|
}
|
|
}
|
|
|
|
let fr_value = match normalize_finish_reason(&finish_reason) {
|
|
Some(s) => serde_json::Value::String(s.to_string()),
|
|
None => serde_json::Value::Null,
|
|
};
|
|
Json(serde_json::json!({
|
|
"id": id,
|
|
"object": "chat.completion",
|
|
"created": created,
|
|
"model": model_name,
|
|
"choices": [{
|
|
"index": 0,
|
|
"message": { "role": "assistant", "content": content },
|
|
"finish_reason": fr_value,
|
|
}],
|
|
"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 = state.chat_template.render(&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;
|
|
}
|
|
// Only "stop" and "length" are OpenAI-standard values. Internal
|
|
// codes like "error" (client-stalled from tp/pp engine) map to
|
|
// null so SDK clients see a clean stream close.
|
|
let fr = normalize_finish_reason(&finish_reason);
|
|
let chunk = make_chunk(&id, &model_name, created, None, None, fr);
|
|
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"));
|
|
}
|
|
|
|
if let Some(t) = req.temperature {
|
|
if !t.is_finite() || t < 0.0 {
|
|
return Some(bad_request("temperature must be a finite value >= 0"));
|
|
}
|
|
}
|
|
if let Some(p) = req.top_p {
|
|
if !p.is_finite() || !(0.0..=1.0).contains(&p) {
|
|
return Some(bad_request("top_p must be in [0, 1]"));
|
|
}
|
|
}
|
|
if let Some(k) = req.top_k {
|
|
if k > 1_000_000 {
|
|
return Some(bad_request("top_k must be <= 1_000_000"));
|
|
}
|
|
}
|
|
|
|
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.try_send(req).map_err(|err| match err {
|
|
std::sync::mpsc::TrySendError::Full(_) => {
|
|
service_unavailable("inference engine is busy, retry later")
|
|
}
|
|
std::sync::mpsc::TrySendError::Disconnected(_) => {
|
|
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),
|
|
}
|
|
}
|
|
|
|
/// Map engine finish_reason strings to OpenAI-standard values. Any engine-internal
|
|
/// code (e.g. "error" from tp/pp client-stall) collapses to None so SDK clients see
|
|
/// a clean null instead of an unknown value.
|
|
fn normalize_finish_reason(fr: &str) -> Option<&'static str> {
|
|
match fr {
|
|
"stop" => Some("stop"),
|
|
"length" => Some("length"),
|
|
_ => None,
|
|
}
|
|
}
|