diff --git a/crates/apr-cli/src/code_output_format_help_tests.rs b/crates/apr-cli/src/code_output_format_help_tests.rs new file mode 100644 index 0000000000..348f5caf54 --- /dev/null +++ b/crates/apr-cli/src/code_output_format_help_tests.rs @@ -0,0 +1,74 @@ +//! #5211: `apr code --output-format json` help names the envelope that +//! `code_envelope::envelope` actually writes. It used to list only +//! `{type, subtype:"success", result, session_id, duration_ms}`, so a harness +//! reading `--help` could not learn that `backend` (the device that ran) or the +//! error object exist. + +/// Every top-level key `envelope()` writes, the `backend` and `error` sub-keys, +/// and the two `subtype` values. +const ENVELOPE_FIELDS: [&str; 20] = [ + "type", + "subtype", + "\"success\"", + "\"error\"", + "is_error", + "status", + "duration_ms", + "result", + "session_id", + "num_turns", + "tokens_in", + "tokens_out", + "total_cost_usd", + "backend", + "requested", + "ran", + "fell_back", + "kind", + "message", + "exit_code", +]; + +/// Clap's recursive build of the full `Commands` enum overflows the default +/// test-thread stack in debug builds, as in run_tests_top_k_default_3754.rs. +#[test] +fn code_output_format_help_names_every_envelope_field() { + std::thread::Builder::new() + .stack_size(16 * 1024 * 1024) + .spawn(code_output_format_help_body) + .expect("spawn") + .join() + .expect("--output-format help row panicked"); +} + +fn code_output_format_help_body() { + use clap::CommandFactory; + let cmd = crate::Cli::command(); + let code = cmd.find_subcommand("code").expect("code subcommand"); + let arg = code + .get_arguments() + .find(|a| a.get_id() == "output_format") + .expect("--output-format"); + let arg_help = arg + .get_long_help() + .or_else(|| arg.get_help()) + .expect("--output-format has help") + .to_string(); + let json_value = arg + .get_possible_values() + .into_iter() + .find(|v| v.get_name() == "json") + .expect("json value"); + let value_help = json_value + .get_help() + .expect("json value has help") + .to_string(); + for (what, help) in [("--output-format", &arg_help), ("json value", &value_help)] { + for field in ENVELOPE_FIELDS { + assert!( + help.contains(field), + "{what} help omits envelope field {field}: {help}" + ); + } + } +} diff --git a/crates/apr-cli/src/commands_enum.rs b/crates/apr-cli/src/commands_enum.rs index b3705bea24..170b6dcd58 100644 --- a/crates/apr-cli/src/commands_enum.rs +++ b/crates/apr-cli/src/commands_enum.rs @@ -132,7 +132,10 @@ pub enum CodeOutputFormat { /// Plain assistant text on stdout (default; existing behavior). #[default] Text, - /// Structured JSON envelope: `{type:"result", subtype:"success", result, session_id, duration_ms}`. + /// Structured JSON envelope: `{type:"result", subtype:"success"|"error", is_error, status, + /// duration_ms, result, session_id, num_turns, tokens_in, tokens_out, total_cost_usd, + /// backend:{requested, ran, fell_back}|null}`, plus `error:{kind, message, exit_code}` + /// when `subtype` is `"error"`. Json, } @@ -911,8 +914,11 @@ pub enum Commands { /// Output format for non-interactive (`-p`) mode (PMAT-CODE-OUTPUT-FORMAT-001). /// `text` (default): plain assistant text. - /// `json`: structured `{type:"result", subtype:"success", result, session_id, duration_ms}` - /// envelope matching Claude Code's `claude -p --output-format json` shape. + /// `json`: structured `{type:"result", subtype:"success"|"error", is_error, status, + /// duration_ms, result, session_id, num_turns, tokens_in, tokens_out, total_cost_usd, + /// backend:{requested, ran, fell_back}|null}` envelope, plus + /// `error:{kind, message, exit_code}` when `subtype` is `"error"`. The shape matches + /// Claude Code's `claude -p --output-format json`; `backend` names the device that ran. #[arg(long, value_enum, default_value_t = CodeOutputFormat::Text)] output_format: CodeOutputFormat, @@ -1042,3 +1048,7 @@ pub enum DebugCommands { force: bool, }, } + +#[cfg(all(test, feature = "training"))] +#[path = "code_output_format_help_tests.rs"] +mod code_output_format_help_tests; diff --git a/crates/aprender-serve/src/api/constrained_chat_backend.rs b/crates/aprender-serve/src/api/constrained_chat_backend.rs index ea4ffd3d53..90d1470bca 100644 --- a/crates/aprender-serve/src/api/constrained_chat_backend.rs +++ b/crates/aprender-serve/src/api/constrained_chat_backend.rs @@ -358,6 +358,11 @@ fn constrained_turn( } } +/// 499, "client closed request": the status a turn cancelled by a disconnect is logged with. +fn client_closed_status() -> StatusCode { + StatusCode::from_u16(499).expect("499 is a valid status code") +} + /// The verdict on a finished constrained reply: `Ok` when it ships. A reply the budget cut is /// `Truncated` (no partial document is returned); an empty one, or one the second reader /// refuses, is `SchemaViolation`. @@ -496,6 +501,15 @@ async fn constrained_chat( )) }, }; + if cancel.peek_cancelled() { + // The client went away mid-turn: no budget cut it, so it is not `Truncated`, and no + // one reads this reply + return Err(fail_response( + state, + client_closed_status(), + "the client closed the request before the constrained turn ended", + )); + } let reply = &turn.tokens[prompt_tokens..]; let text = mapped.model.decode(reply); constrained_verdict(&constraint, stop, &text, reply.len()) diff --git a/crates/aprender-serve/src/api/constrained_chat_backend_tests.rs b/crates/aprender-serve/src/api/constrained_chat_backend_tests.rs index ef998e6389..bf5d9cd619 100644 --- a/crates/aprender-serve/src/api/constrained_chat_backend_tests.rs +++ b/crates/aprender-serve/src/api/constrained_chat_backend_tests.rs @@ -502,7 +502,36 @@ async fn a_budget_too_small_for_the_document_is_422_with_no_partial_text() { let (status, body) = post(state, constrained_body(schema, 1)).await; assert_eq!(status, StatusCode::UNPROCESSABLE_ENTITY, "{body}"); assert_eq!(body["type"], "constraint_truncated"); + assert!( + body["error"] + .as_str() + .is_some_and(|e| e.contains("budget of 1 tokens")), + "the refusal names the request's budget: {body}" + ); + assert!(body.get("choices").is_none(), "{body}"); +} + +/// The disconnect path end to end: a turn whose token is already cancelled is answered 499, +/// never as the document or as a budget cut. +#[cfg(feature = "structured-output")] +#[tokio::test(flavor = "multi_thread")] +async fn a_turn_the_client_cancelled_is_answered_499_not_truncated() { + let (_f, state) = cpu_state(true); + let schema = serde_json::json!({"type": "boolean"}); + let request: ChatCompletionRequest = + serde_json::from_value(constrained_body(schema, 8)).expect("a chat request"); + let cancel = CancelToken::new(); + cancel.cancel(); + let response = try_constrained_chat(&state, &request, "id", None, Instant::now(), &cancel) + .await + .expect("a constrained request is answered here"); + assert_eq!(response.status().as_u16(), 499); + let bytes = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .expect("a body"); + let body: serde_json::Value = serde_json::from_slice(&bytes).expect("a JSON body"); assert!(body.get("choices").is_none(), "{body}"); + assert_ne!(body["type"], "constraint_truncated", "{body}"); } /// A server that retains no GGUF has no vocabulary to compile against: refused by name, never @@ -525,3 +554,38 @@ async fn a_server_without_the_gguf_refuses_by_name() { "{body}" ); } + +/// #3568: a `response_format.type` apr does not serve is a 422 that names the field and the +/// types it serves, not the sanitizer's "Invalid request body"; a known type still parses. +#[tokio::test(flavor = "multi_thread")] +async fn an_unknown_response_format_type_is_refused_naming_the_field() { + for (format, got) in [ + (serde_json::json!({"type": "xml"}), "got xml"), + (serde_json::json!({"type": 12}), "got 12"), + (serde_json::json!({}), "got no type"), + ] { + let (_f, state) = cpu_state(false); + let body = serde_json::json!({ + "model": "default", + "messages": [{"role": "user", "content": "hi"}], + "response_format": format, + }); + let (status, body) = post(state, body).await; + assert_eq!(status, StatusCode::UNPROCESSABLE_ENTITY, "{body}"); + let error = body["error"].as_str().expect("an error message"); + assert!( + error.contains("response_format.type must be one of text, json_object, json_schema"), + "{error}" + ); + assert!(error.ends_with(got), "{error}"); + } + assert_eq!( + format(serde_json::json!({"type": "json_object"})), + ResponseFormat::JsonObject + ); + // A known type goes through the request's own deserializer, the hook under test + let request = chat_request(serde_json::json!({"response_format": {"type": "json_object"}})); + assert_eq!(request.response_format, Some(ResponseFormat::JsonObject)); + let request = chat_request(serde_json::json!({"response_format": null})); + assert_eq!(request.response_format, None); +} diff --git a/crates/aprender-serve/src/api/mod_create_demo.rs b/crates/aprender-serve/src/api/mod_create_demo.rs index 87fd478649..e591a5f0d6 100644 --- a/crates/aprender-serve/src/api/mod_create_demo.rs +++ b/crates/aprender-serve/src/api/mod_create_demo.rs @@ -163,8 +163,13 @@ pub struct ChatCompletionRequest { pub think: Option, /// #3568 PR 4: OpenAI's structured output. `json_schema` and `json_object` run the turn /// under a constraint and a second reader; `text` (or absent) is today's turn. A `type` - /// apr does not know is refused at deserialization, never dropped. - #[serde(default, skip_serializing_if = "Option::is_none")] + /// apr does not know is refused at deserialization, never dropped, with a message that + /// names the field and the types apr serves. + #[serde( + default, + deserialize_with = "response_format_of_known_type", + skip_serializing_if = "Option::is_none" + )] pub response_format: Option, /// #4971: OpenAI's `logprobs`. When true the reply carries /// `choices[0].logprobs`, one entry per generated token, or the backend @@ -192,6 +197,43 @@ pub enum ResponseFormat { }, } +/// The `type`s [`ResponseFormat`] serves, in the order a refusal lists them. +const RESPONSE_FORMAT_TYPES: [&str; 3] = ["text", "json_object", "json_schema"]; + +/// `deserialize_with` for `response_format` (#3568): a `type` apr does not serve is refused +/// with a client-visible reason naming the field, where serde's own "unknown variant" text is +/// hidden by the rejection sanitizer and the caller read only "Invalid request body". +fn response_format_of_known_type<'de, D>( + deserializer: D, +) -> Result, D::Error> +where + D: serde::Deserializer<'de>, +{ + use serde::de::Error; + let value = serde_json::Value::deserialize(deserializer)?; + if value.is_null() { + return Ok(None); + } + let kind = value.get("type").and_then(serde_json::Value::as_str); + if !kind.is_some_and(|k| RESPONSE_FORMAT_TYPES.contains(&k)) { + return Err(D::Error::custom(format!( + "{}response_format.type must be one of {}, got {}", + crate::api::CLIENT_VISIBLE_MARKER, + RESPONSE_FORMAT_TYPES.join(", "), + // Unquoted: the sanitizer trims a trailing `"` off a client-visible reason + match kind { + Some(k) => k.to_string(), + None => value + .get("type") + .map_or_else(|| "no type".to_string(), ToString::to_string), + } + ))); + } + ResponseFormat::deserialize(value) + .map(Some) + .map_err(D::Error::custom) +} + /// The `json_schema` object of a `response_format` (#3568 PR 4). #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct JsonSchemaFormat { diff --git a/crates/aprender-serve/src/api/ollama_handlers.rs b/crates/aprender-serve/src/api/ollama_handlers.rs index 0305b78174..9fd1444de6 100644 --- a/crates/aprender-serve/src/api/ollama_handlers.rs +++ b/crates/aprender-serve/src/api/ollama_handlers.rs @@ -106,6 +106,9 @@ pub struct OllamaChatResponse { pub prompt_eval_count: usize, /// Generated token count. pub eval_count: usize, + /// Why generation ended: `"length"` or `"stop"`; absent when unknown. + #[serde(skip_serializing_if = "Option::is_none")] + pub done_reason: Option, } /// Ollama `/api/generate` request (single prompt, non-chat). @@ -142,6 +145,9 @@ pub struct OllamaGenerateResponse { pub prompt_eval_count: usize, /// Generated token count. pub eval_count: usize, + /// Why generation ended: `"length"` or `"stop"`; absent when unknown. + #[serde(skip_serializing_if = "Option::is_none")] + pub done_reason: Option, } /// One newline-delimited object of an Ollama `/api/chat` stream. @@ -436,6 +442,7 @@ fn chat_stream_objects( content: &str, prompt_eval_count: usize, eval_count: usize, + done_reason: &str, ) -> Vec { let created_at = created_at_now(); let mut out: Vec = content_fragments(content) @@ -461,7 +468,7 @@ fn chat_stream_objects( content: String::new(), }, done: true, - done_reason: Some("stop".to_string()), + done_reason: Some(done_reason.to_string()), prompt_eval_count: Some(prompt_eval_count), eval_count: Some(eval_count), }); @@ -474,6 +481,7 @@ fn generate_stream_objects( content: &str, prompt_eval_count: usize, eval_count: usize, + done_reason: &str, ) -> Vec { let created_at = created_at_now(); let mut out: Vec = content_fragments(content) @@ -493,7 +501,7 @@ fn generate_stream_objects( created_at, response: String::new(), done: true, - done_reason: Some("stop".to_string()), + done_reason: Some(done_reason.to_string()), prompt_eval_count: Some(prompt_eval_count), eval_count: Some(eval_count), }); @@ -552,24 +560,31 @@ fn modified_at(source: Option<&ModelSourceInfo>) -> Option { ) } -/// Extract `(content, prompt_tokens, completion_tokens)` from the OpenAI chat -/// response, or a fallback `(error_text, 0, 0)` when generation failed. +/// Extract `(content, prompt_tokens, completion_tokens, done_reason)` from the +/// OpenAI chat response, or a fallback `(error_text, 0, 0, "stop")` when +/// generation failed. `done_reason` comes from [`ollama_done_reason`]. /// /// Crucially this ALWAYS yields an Ollama-shaped body — even on a backend error /// or a missing model — so a wired route is observably distinct from the axum /// `not_found` fallback (which has no `done` field). -fn chat_response_to_parts(status: StatusCode, body: &[u8]) -> (String, usize, usize) { +fn chat_response_to_parts(status: StatusCode, body: &[u8]) -> (String, usize, usize, &'static str) { if status.is_success() { if let Ok(resp) = serde_json::from_slice::(body) { - let content = resp + let (content, done_reason) = resp .choices .first() - .map(|c| c.message.content.clone()) - .unwrap_or_default(); + .map(|c| { + ( + c.message.content.clone(), + ollama_done_reason(&c.finish_reason), + ) + }) + .unwrap_or((String::new(), "stop")); return ( content, resp.usage.prompt_tokens, resp.usage.completion_tokens, + done_reason, ); } } @@ -579,7 +594,19 @@ fn chat_response_to_parts(status: StatusCode, body: &[u8]) -> (String, usize, us .ok() .and_then(|v| v.get("error").and_then(|e| e.as_str().map(str::to_string))) .unwrap_or_else(|| "generation unavailable".to_string()); - (msg, 0, 0) + (msg, 0, 0, "stop") +} + +/// Map an OpenAI `finish_reason` to Ollama's `done_reason` (aprender#5357). +/// +/// Ollama has two terminal reasons: `"length"` when the turn ran out of token +/// budget (`num_predict`), and `"stop"` for everything else. +fn ollama_done_reason(finish_reason: &str) -> &'static str { + if finish_reason == "length" { + "length" + } else { + "stop" + } } /// Stamp the upstream generation status onto an Ollama-shaped response. @@ -639,7 +666,7 @@ pub async fn ollama_chat_handler( openai_chat_completions_handler(State(state), headers, Extension(cancel), Json(chat_req)) .await; let (status, body) = split_response(inner).await; - let (content, prompt_tokens, eval_count) = chat_response_to_parts(status, &body); + let (content, prompt_tokens, eval_count, done_reason) = chat_response_to_parts(status, &body); if stream { return with_upstream_status( @@ -649,27 +676,65 @@ pub async fn ollama_chat_handler( &content, prompt_tokens, eval_count, + done_reason, )), ); } with_upstream_status( status, - Json(OllamaChatResponse { + Json(chat_final_response( model, - created_at: created_at_now(), - message: OllamaMessage { - role: "assistant".to_string(), - content, - }, - done: true, - prompt_eval_count: prompt_tokens, + content, + prompt_tokens, eval_count, - }) + done_reason, + )) .into_response(), ) } +/// The coalesced `/api/chat` reply (`stream: false`), carrying `done_reason` (aprender#5357). +fn chat_final_response( + model: String, + content: String, + prompt_eval_count: usize, + eval_count: usize, + done_reason: &str, +) -> OllamaChatResponse { + OllamaChatResponse { + model, + created_at: created_at_now(), + message: OllamaMessage { + role: "assistant".to_string(), + content, + }, + done: true, + prompt_eval_count, + eval_count, + done_reason: Some(done_reason.to_string()), + } +} + +/// The coalesced `/api/generate` reply (`stream: false`), carrying `done_reason` (aprender#5357). +fn generate_final_response( + model: String, + response: String, + prompt_eval_count: usize, + eval_count: usize, + done_reason: &str, +) -> OllamaGenerateResponse { + OllamaGenerateResponse { + model, + created_at: created_at_now(), + response, + done: true, + prompt_eval_count, + eval_count, + done_reason: Some(done_reason.to_string()), + } +} + /// `POST /api/generate` — Ollama single-prompt generate endpoint. /// /// Folds `system` + `prompt` into a chat request and reuses the same generation @@ -700,7 +765,7 @@ pub async fn ollama_generate_handler( openai_chat_completions_handler(State(state), headers, Extension(cancel), Json(chat_req)) .await; let (status, body) = split_response(inner).await; - let (content, prompt_tokens, eval_count) = chat_response_to_parts(status, &body); + let (content, prompt_tokens, eval_count, done_reason) = chat_response_to_parts(status, &body); if stream { return with_upstream_status( @@ -710,20 +775,20 @@ pub async fn ollama_generate_handler( &content, prompt_tokens, eval_count, + done_reason, )), ); } with_upstream_status( status, - Json(OllamaGenerateResponse { + Json(generate_final_response( model, - created_at: created_at_now(), - response: content, - done: true, - prompt_eval_count: prompt_tokens, + content, + prompt_tokens, eval_count, - }) + done_reason, + )) .into_response(), ) } @@ -896,7 +961,7 @@ mod tests { "choices":[{"index":0,"message":{"role":"assistant","content":"4"},"finish_reason":"stop"}], "usage":{"prompt_tokens":3,"completion_tokens":1,"total_tokens":4} }"#; - let (content, p, c) = chat_response_to_parts(StatusCode::OK, body); + let (content, p, c, _) = chat_response_to_parts(StatusCode::OK, body); assert_eq!(content, "4"); assert_eq!(p, 3); assert_eq!(c, 1); @@ -907,7 +972,7 @@ mod tests { // On a backend error (e.g. no model), the Ollama body must still be // well-formed: the error text becomes the assistant content, tokens 0. let body = br#"{"error":"model not found"}"#; - let (content, p, c) = chat_response_to_parts(StatusCode::NOT_FOUND, body); + let (content, p, c, _) = chat_response_to_parts(StatusCode::NOT_FOUND, body); assert_eq!(content, "model not found"); assert_eq!(p, 0); assert_eq!(c, 0); @@ -925,6 +990,7 @@ mod tests { done: true, prompt_eval_count: 1, eval_count: 2, + done_reason: None, }; let json = serde_json::to_value(&resp).expect("serialize"); assert_eq!(json["message"]["role"], "assistant"); @@ -941,6 +1007,7 @@ mod tests { done: true, prompt_eval_count: 0, eval_count: 1, + done_reason: None, }; let json = serde_json::to_value(&resp).expect("serialize"); assert_eq!(json["response"], "hi"); @@ -992,6 +1059,7 @@ mod tests { done: true, prompt_eval_count: 1, eval_count: 2, + done_reason: None, }; let json = serde_json::to_value(&resp).expect("serialize"); let created_at = json["created_at"].as_str().expect("created_at is a string"); @@ -1010,6 +1078,7 @@ mod tests { done: true, prompt_eval_count: 0, eval_count: 1, + done_reason: None, }; let json = serde_json::to_value(&resp).expect("serialize"); let created_at = json["created_at"].as_str().expect("created_at is a string"); @@ -1062,7 +1131,7 @@ mod stream_and_discovery_tests { /// and is indistinguishable from a non-streaming reply. #[test] fn chat_stream_is_a_sequence_terminated_by_done_true() { - let objs = chat_stream_objects("apr", "The capital of France is Paris.", 21, 7); + let objs = chat_stream_objects("apr", "The capital of France is Paris.", 21, 7, "stop"); assert!( objs.len() > 2, "a 6-word answer must arrive as several chunks, got {}", @@ -1091,7 +1160,7 @@ mod stream_and_discovery_tests { /// Same contract on `/api/generate`, whose chunks carry a flat `response`. #[test] fn generate_stream_is_a_sequence_terminated_by_done_true() { - let objs = generate_stream_objects("apr", "1 2 3 4 5 6 7 8", 16, 16); + let objs = generate_stream_objects("apr", "1 2 3 4 5 6 7 8", 16, 16, "stop"); let (last, rest) = objs.split_last().expect("non-empty"); assert!(last.done); assert_eq!(last.done_reason.as_deref(), Some("stop")); @@ -1103,11 +1172,91 @@ mod stream_and_discovery_tests { assert_eq!(assembled, "1 2 3 4 5 6 7 8"); } + /// aprender#5357: a turn cut off by `max_tokens` / `num_predict` must say + /// so. The OpenAI `finish_reason: "length"` maps to Ollama's + /// `done_reason: "length"`; every other reason maps to `"stop"`. Before the + /// fix the terminal chunk said `"stop"` unconditionally, so an Ollama client + /// could not tell a truncated answer from a finished one. + fn openai_body(finish_reason: &str) -> Vec { + format!( + r#"{{ + "id":"x","object":"chat.completion","created":0,"model":"m", + "choices":[{{"index":0,"message":{{"role":"assistant","content":"one two"}},"finish_reason":"{finish_reason}"}}], + "usage":{{"prompt_tokens":4,"completion_tokens":2,"total_tokens":6}} + }}"# + ) + .into_bytes() + } + + #[test] + fn length_finish_maps_to_length_done_reason_on_both_routes() { + let (content, p, c, reason) = + chat_response_to_parts(StatusCode::OK, &openai_body("length")); + assert_eq!(reason, "length"); + let chat = chat_stream_objects("apr", &content, p, c, reason); + let last = chat.last().expect("non-empty"); + assert!(last.done); + assert_eq!(last.done_reason.as_deref(), Some("length")); + let generate = generate_stream_objects("apr", &content, p, c, reason); + let last = generate.last().expect("non-empty"); + assert!(last.done); + assert_eq!(last.done_reason.as_deref(), Some("length")); + } + + /// The `stream: false` bodies (the default for several clients) carry the reason too. + #[test] + fn non_stream_replies_carry_the_done_reason() { + for (finish, want) in [("length", "length"), ("stop", "stop")] { + let (content, p, c, reason) = + chat_response_to_parts(StatusCode::OK, &openai_body(finish)); + let chat = serde_json::to_value(chat_final_response( + "apr".to_string(), + content.clone(), + p, + c, + reason, + )) + .expect("serialize"); + assert_eq!(chat["done_reason"], want, "chat, finish {finish:?}"); + let generate = serde_json::to_value(generate_final_response( + "apr".to_string(), + content, + p, + c, + reason, + )) + .expect("serialize"); + assert_eq!(generate["done_reason"], want, "generate, finish {finish:?}"); + } + } + + #[test] + fn non_length_finish_maps_to_stop_done_reason_on_both_routes() { + for finish in ["stop", "eos", "content_filter", "tool_calls", ""] { + let (content, p, c, reason) = + chat_response_to_parts(StatusCode::OK, &openai_body(finish)); + assert_eq!(reason, "stop", "finish_reason {finish:?}"); + let chat = chat_stream_objects("apr", &content, p, c, reason); + assert_eq!( + chat.last().and_then(|o| o.done_reason.as_deref()), + Some("stop") + ); + let generate = generate_stream_objects("apr", &content, p, c, reason); + assert_eq!( + generate.last().and_then(|o| o.done_reason.as_deref()), + Some("stop") + ); + } + let (_, _, _, reason) = + chat_response_to_parts(StatusCode::NOT_FOUND, br#"{"error":"model not found"}"#); + assert_eq!(reason, "stop", "an upstream error is a terminal stop"); + } + /// Even an empty generation must terminate the stream, or a client's read /// loop hangs waiting for `done:true`. #[test] fn empty_generation_still_terminates_the_stream() { - let objs = chat_stream_objects("apr", "", 3, 0); + let objs = chat_stream_objects("apr", "", 3, 0, "stop"); assert_eq!(objs.len(), 1); assert!(objs[0].done); } @@ -1117,7 +1266,7 @@ mod stream_and_discovery_tests { /// halves in the client's line reader. #[test] fn each_object_serializes_to_exactly_one_json_line() { - for chunk in chat_stream_objects("apr", "a b\nc", 1, 3) { + for chunk in chat_stream_objects("apr", "a b\nc", 1, 3, "stop") { let line = serde_json::to_string(&chunk).expect("serialize"); assert!( !line.contains('\n'), @@ -1132,7 +1281,7 @@ mod stream_and_discovery_tests { /// told the model generated nothing. #[test] fn non_terminal_chunks_omit_counts_and_done_reason() { - let objs = chat_stream_objects("apr", "two words", 5, 2); + let objs = chat_stream_objects("apr", "two words", 5, 2, "stop"); let first = serde_json::to_value(&objs[0]).expect("serialize"); assert!(first.get("eval_count").is_none()); assert!(first.get("prompt_eval_count").is_none()); @@ -1144,7 +1293,7 @@ mod stream_and_discovery_tests { /// decodes every chunk into the same `api.ChatResponse`. #[test] fn stream_chunk_created_at_is_rfc3339() { - for chunk in chat_stream_objects("apr", "hi there", 1, 2) { + for chunk in chat_stream_objects("apr", "hi there", 1, 2, "stop") { let json = serde_json::to_value(&chunk).expect("serialize"); let created_at = json["created_at"].as_str().expect("string"); chrono::DateTime::parse_from_rfc3339(created_at) diff --git a/crates/aprender-serve/src/infer/constrained_generate.rs b/crates/aprender-serve/src/infer/constrained_generate.rs index 56fc5c3ef5..bf7fa5e5f3 100644 --- a/crates/aprender-serve/src/infer/constrained_generate.rs +++ b/crates/aprender-serve/src/infer/constrained_generate.rs @@ -99,6 +99,23 @@ fn unconstrained_gguf_path( gen_config.trace.then_some(("gguf-trace", TRACE_REMOVED_BY)) } +/// `gen_config` with the constraint vocabulary's end-of-sequence id among its stop tokens, as +/// `apr serve`'s constrained route adds it (#3568). Once the document is whole the mask admits +/// only that id; a stop list built from the model config's EOS alone, when the two ids differ, +/// would take it as text and keep going past the end of the document. +fn with_constraint_eos( + gen_config: &crate::gguf::QuantizedGenerateConfig, + eos: Option, +) -> crate::gguf::QuantizedGenerateConfig { + let mut config = gen_config.clone(); + if let Some(eos) = eos { + if !config.stop_tokens.contains(&eos) { + config.stop_tokens.push(eos); + } + } + config +} + /// Generate under `request` through the one engine (#3793): `(tokens, used_gpu, /// gpu_attempted, stop)`, the GPU pair as the unconstrained run reports it (#3826). The schema /// or grammar is compiled against the model's vocabulary here, after load and before the first @@ -133,6 +150,7 @@ fn generate_gguf_constrained( } let env = mapped.constraint_env().map_err(RealizarError::Constraint)?; let mut constraint = request.compile(&env).map_err(RealizarError::Constraint)?; + let gen_config = &with_constraint_eos(gen_config, mapped.model.eos_token_id()); let context_length = model.config.context_length; let (turn, stop, gpu_attempted) = if let Some(qwen) = qwen35_host { crate::gguf::forward_qwen35::qwen35_check_context(input_tokens.len(), context_length)?; diff --git a/crates/aprender-serve/src/infer/constrained_generate_tests.rs b/crates/aprender-serve/src/infer/constrained_generate_tests.rs index 0ca103c3c5..43e7e0107b 100644 --- a/crates/aprender-serve/src/infer/constrained_generate_tests.rs +++ b/crates/aprender-serve/src/infer/constrained_generate_tests.rs @@ -244,3 +244,52 @@ fn a_dense_constrained_run_reports_the_attempt_the_dispatch_made() { assert!(matches!(tokens.get(1), Some(2 | 3)), "{tokens:?}"); } } + +#[test] +fn the_constraint_eos_joins_the_stop_tokens_once() { + let gen_config = crate::gguf::QuantizedGenerateConfig { + stop_tokens: vec![7], + ..Default::default() + }; + assert_eq!( + with_constraint_eos(&gen_config, Some(1)).stop_tokens, + vec![7, 1] + ); + assert_eq!( + with_constraint_eos(&gen_config, Some(7)).stop_tokens, + vec![7] + ); + assert_eq!(with_constraint_eos(&gen_config, None).stop_tokens, vec![7]); +} + +/// #3568: a caller whose stop list lacks the vocabulary's EOS (the model config names another +/// id, or none) still ends a whole document at it. The mask admits only EOS once "a" or "b" is +/// out; without it among the stops the turn would take it as a token of the reply. +#[cfg(feature = "structured-output")] +#[test] +fn a_constrained_run_stops_at_the_vocabulary_eos_the_caller_did_not_list() { + let (_f, mapped) = q4k_model_with_vocab(); + let request = ConstraintRequest::Lark(r#"start: "a" | "b""#.to_string()); + let gen_config = crate::gguf::QuantizedGenerateConfig { + max_tokens: 4, + temperature: 0.0, + top_k: 1, + stop_tokens: vec![], + ..Default::default() + }; + let model = crate::gguf::OwnedQuantizedModel::from_mapped(&mapped).expect("load"); + let config = InferenceConfig::new("m.gguf").without_gpu(); + let (tokens, _, _, stop) = generate_gguf_constrained( + &request, + &config, + &mapped, + Some(model), + None, + &[2], + &gen_config, + ) + .expect("a constrained turn"); + assert_eq!(stop, ConstrainedStop::Complete); + assert_eq!(tokens.len(), 2, "prompt plus one token, no EOS: {tokens:?}"); + assert!(matches!(tokens[1], 2 | 3), "{tokens:?}"); +} diff --git a/docs/audits/surface_audit.csv b/docs/audits/surface_audit.csv index 8fe7d73b82..55d87c90b2 100644 --- a/docs/audits/surface_audit.csv +++ b/docs/audits/surface_audit.csv @@ -54,10 +54,10 @@ apr,GET /v1/metrics,6,x86_64-linux,llama.cpp,yes,3,http-apr-serve,crates/aprende apr,GET /v1/models,9,x86_64-linux,llama.cpp,yes,3,http-apr-serve,crates/aprender-serve/src/api/router.rs:130,medium apr,HTTP 404 fallback route index,6,x86_64-linux,llama.cpp,no,3,http-apr-serve,crates/aprender-serve/src/api/router.rs:329,high apr,POST /api/chat,3,x86_64-linux,ollama,yes,3,http-apr-serve,crates/aprender-serve/src/api/router.rs:174,high -apr,"POST /api/chat (stream:true, NDJSON)",3,x86_64-linux,ollama,no,3,http-apr-serve,crates/aprender-serve/src/api/ollama_handlers.rs:614 (framing :387-425),high +apr,"POST /api/chat (stream:true, NDJSON)",3,x86_64-linux,ollama,no,3,http-apr-serve,crates/aprender-serve/src/api/ollama_handlers.rs:654 (framing :397-437),high apr,POST /api/embeddings,6,x86_64-linux,ollama,yes,3,http-apr-serve,crates/aprender-serve/src/api/router.rs:184,high apr,POST /api/generate,3,x86_64-linux,ollama,yes,3,http-apr-serve,crates/aprender-serve/src/api/router.rs:175,high -apr,"POST /api/generate (stream:true, NDJSON)",3,x86_64-linux,ollama,no,3,http-apr-serve,crates/aprender-serve/src/api/ollama_handlers.rs:669 (framing :387-425),high +apr,"POST /api/generate (stream:true, NDJSON)",3,x86_64-linux,ollama,no,3,http-apr-serve,crates/aprender-serve/src/api/ollama_handlers.rs:742 (framing :397-437),high apr,POST /api/pull,6,x86_64-linux,ollama,no,10,apr-core-commands,crates/apr-cli/src/commands/serve/ollama.rs:650,medium apr,POST /api/show,6,x86_64-linux,ollama,yes,3,http-apr-serve,crates/aprender-serve/src/api/router.rs:180,high apr,POST /batch/generate,2,x86_64-linux,llama.cpp,yes,3,http-apr-serve,crates/aprender-serve/src/api/router.rs:99,high