From c62a0bc33eb66855a829b7e78dc3f6e568d85968 Mon Sep 17 00:00:00 2001 From: Dmitri Afanasjev Date: Mon, 8 Jun 2026 08:11:55 +0300 Subject: [PATCH] feat(ltengine): dynamic output-token cap + determinism invariants (api mode); fix --llm-timeout 0 to mean no timeout --- ltengine/src/llm_api.rs | 90 +++++++++++++++++++++++++++++++++--- ltengine/src/main.rs | 28 ++++++++++- ltengine/src/token_budget.rs | 89 +++++++++++++++++++++++++++++++++++ 3 files changed, 199 insertions(+), 8 deletions(-) create mode 100644 ltengine/src/token_budget.rs diff --git a/ltengine/src/llm_api.rs b/ltengine/src/llm_api.rs index cf6d94f..8ab486e 100644 --- a/ltengine/src/llm_api.rs +++ b/ltengine/src/llm_api.rs @@ -16,6 +16,11 @@ struct ChatMessage { content: String, } +#[derive(Serialize)] +struct ChatTemplateKwargs { + enable_thinking: bool, +} + #[derive(Serialize)] struct ChatCompletionRequest { model: String, @@ -23,16 +28,21 @@ struct ChatCompletionRequest { temperature: f32, #[serde(skip_serializing_if = "Option::is_none")] max_tokens: Option, + chat_template_kwargs: ChatTemplateKwargs, } #[derive(Deserialize)] struct ChatCompletionResponse { choices: Vec, + #[serde(default)] + usage: Option, } #[derive(Deserialize)] struct Choice { message: ResponseMessage, + #[serde(default)] + finish_reason: Option, } #[derive(Deserialize)] @@ -40,6 +50,12 @@ struct ResponseMessage { content: String, } +#[derive(Deserialize)] +struct Usage { + #[serde(default)] + completion_tokens: Option, +} + #[derive(Deserialize)] struct ModelsResponse { data: Vec, @@ -52,8 +68,14 @@ struct ModelEntry { impl LLM { pub fn new(url: String, api_key: String, model: String, max_tokens: Option, timeout_secs: u64) -> Result { - let client = reqwest::Client::builder() - .timeout(Duration::from_secs(timeout_secs)) + // timeout_secs == 0 means "no timeout": reqwest has no timeout by default, + // so we only apply one when a positive value is configured. Passing + // Duration::from_secs(0) would instead make every request time out instantly. + let mut builder = reqwest::Client::builder(); + if timeout_secs > 0 { + builder = builder.timeout(Duration::from_secs(timeout_secs)); + } + let client = builder .build() .with_context(|| "Failed to build HTTP client")?; Ok(LLM { url, api_key, model, max_tokens, client }) @@ -88,17 +110,28 @@ impl LLM { Ok(model_id) } - pub async fn run_prompt(&self, system: String, user: String) -> Result { + pub async fn run_prompt( + &self, + system: String, + user: String, + max_output_tokens: Option, + ) -> Result { let endpoint = format!("{}/v1/chat/completions", self.url.trim_end_matches('/')); + // Dynamic per-request cap when provided, otherwise the static ceiling. + let max_tokens = max_output_tokens.or(self.max_tokens); + let request = ChatCompletionRequest { model: self.model.clone(), messages: vec![ ChatMessage { role: "system".to_string(), content: system }, ChatMessage { role: "user".to_string(), content: user }, ], + // Translation pipeline invariants — enforced here, not relied on from + // the vLLM serve flags or model defaults. temperature: 0.0, - max_tokens: self.max_tokens, + max_tokens, + chat_template_kwargs: ChatTemplateKwargs { enable_thinking: false }, }; let mut req = self.client.post(&endpoint) @@ -120,9 +153,52 @@ impl LLM { let completion: ChatCompletionResponse = response.json().await .with_context(|| "Failed to parse LLM API response")?; - completion.choices.into_iter() + let completion_tokens = completion.usage.and_then(|u| u.completion_tokens); + let choice = completion.choices.into_iter() .next() - .map(|c| c.message.content) - .ok_or_else(|| anyhow::anyhow!("LLM API returned no choices")) + .ok_or_else(|| anyhow::anyhow!("LLM API returned no choices"))?; + + // Truncation safety signal: the cap may have cut a real translation. + if choice.finish_reason.as_deref() == Some("length") { + eprintln!( + "⚠️ LLM stopped at max_tokens cap (max_tokens={:?}, completion_tokens={:?}) — output may be truncated", + max_tokens, completion_tokens + ); + } + + Ok(choice.message.content) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn request_enforces_determinism_invariants() { + let req = ChatCompletionRequest { + model: "m".to_string(), + messages: vec![], + temperature: 0.0, + max_tokens: Some(123), + chat_template_kwargs: ChatTemplateKwargs { enable_thinking: false }, + }; + let v = serde_json::to_value(&req).unwrap(); + assert_eq!(v["temperature"], serde_json::json!(0.0)); + assert_eq!(v["chat_template_kwargs"]["enable_thinking"], serde_json::json!(false)); + assert_eq!(v["max_tokens"], serde_json::json!(123)); + } + + #[test] + fn max_tokens_omitted_when_none() { + let req = ChatCompletionRequest { + model: "m".to_string(), + messages: vec![], + temperature: 0.0, + max_tokens: None, + chat_template_kwargs: ChatTemplateKwargs { enable_thinking: false }, + }; + let v = serde_json::to_value(&req).unwrap(); + assert!(v.get("max_tokens").is_none()); } } diff --git a/ltengine/src/main.rs b/ltengine/src/main.rs index 4869093..6b0c56a 100644 --- a/ltengine/src/main.rs +++ b/ltengine/src/main.rs @@ -20,6 +20,8 @@ mod llm; mod llm; mod banner; mod prompt; +#[cfg(feature = "api")] +mod token_budget; use languages::{detect_lang, get_language_from_code, LANGUAGES}; use error_response::ErrorResponse; @@ -84,6 +86,21 @@ struct Args { #[arg(long, default_value_t = 0)] llm_max_tokens: u32, + /// Dynamic cap: conservative characters-per-token divisor + #[cfg(feature = "api")] + #[arg(long, default_value_t = 2.0)] + llm_chars_per_token: f32, + + /// Dynamic cap: output safety multiple (0 disables the dynamic cap) + #[cfg(feature = "api")] + #[arg(long, default_value_t = 3.0)] + llm_max_tokens_mult: f32, + + /// Dynamic cap: minimum output tokens for tiny inputs + #[cfg(feature = "api")] + #[arg(long, default_value_t = 64)] + llm_max_tokens_floor: u32, + /// HTTP timeout in seconds for LLM API requests (0 = no timeout) #[cfg(feature = "api")] #[arg(long, default_value_t = 120)] @@ -290,7 +307,16 @@ async fn translate(req: HttpRequest, payload: web::Payload, args: web::Data 0 { Some(args.llm_max_tokens) } else { None }, + }; + let cap = token_budget::dynamic_output_cap(q.chars().count(), &budget_cfg); + llm.run_prompt(prompt.system, prompt.user, cap).await.unwrap_or(q.clone()) + } #[cfg(not(feature = "api"))] { llm.run_prompt(prompt.system, prompt.user).unwrap_or(q.clone()) } }else{ diff --git a/ltengine/src/token_budget.rs b/ltengine/src/token_budget.rs new file mode 100644 index 0000000..81f32b3 --- /dev/null +++ b/ltengine/src/token_budget.rs @@ -0,0 +1,89 @@ +//! Dynamic per-request output-token budget for api/vLLM translation requests. +//! +//! Turns the input text size into a `max_tokens` cap so a hallucinating +//! generation is cut early instead of running to the static ceiling. + +/// Configuration for the dynamic output-token cap. +#[derive(Debug, Clone, Copy)] +pub struct TokenBudgetConfig { + /// Conservative characters-per-token divisor used to estimate input tokens. + pub chars_per_token: f32, + /// Output safety multiple applied to estimated input tokens (the "x3"). + pub output_mult: f32, + /// Minimum cap (tokens) so tiny inputs are never starved. + pub floor: u32, + /// Hard ceiling (tokens); `None` means no ceiling. + pub ceiling: Option, +} + +/// Compute the dynamic output-token cap for an input of `input_chars` characters. +/// +/// Returns `None` when the dynamic policy is disabled (non-positive +/// `output_mult` or `chars_per_token`), signalling the caller to fall back to +/// the static cap. +#[must_use] +pub fn dynamic_output_cap(input_chars: usize, cfg: &TokenBudgetConfig) -> Option { + if cfg.output_mult <= 0.0 || cfg.chars_per_token <= 0.0 { + return None; + } + let est_input_tokens = (input_chars as f32 / cfg.chars_per_token).ceil(); + let budget = (est_input_tokens * cfg.output_mult).ceil(); + // Guard NaN/negative before the saturating f32 -> u32 cast. + let budget = if budget.is_finite() && budget >= 0.0 { budget } else { 0.0 }; + let mut cap = budget as u32; // saturates to u32::MAX on overflow + cap = cap.max(cfg.floor); + if let Some(c) = cfg.ceiling { + // Ceiling is a hard limit (model/server max) and takes priority over the + // floor if misconfigured with ceiling < floor. + cap = cap.min(c); + } + Some(cap) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn cfg() -> TokenBudgetConfig { + TokenBudgetConfig { chars_per_token: 2.0, output_mult: 3.0, floor: 64, ceiling: Some(16384) } + } + + #[test] + fn disabled_when_mult_non_positive() { + let mut c = cfg(); + c.output_mult = 0.0; + assert_eq!(dynamic_output_cap(1000, &c), None); + } + + #[test] + fn disabled_when_chars_per_token_non_positive() { + let mut c = cfg(); + c.chars_per_token = 0.0; + assert_eq!(dynamic_output_cap(1000, &c), None); + } + + #[test] + fn tiny_input_uses_floor() { + // 12 chars / 2.0 = 6 -> *3 = 18 -> raised to floor 64 + assert_eq!(dynamic_output_cap(12, &cfg()), Some(64)); + } + + #[test] + fn mid_input_uses_formula() { + // 1000 / 2.0 = 500 -> *3 = 1500 (above floor, below ceiling) + assert_eq!(dynamic_output_cap(1000, &cfg()), Some(1500)); + } + + #[test] + fn huge_input_clamped_to_ceiling() { + // 100000 / 2.0 = 50000 -> *3 = 150000 -> clamped to 16384 + assert_eq!(dynamic_output_cap(100_000, &cfg()), Some(16384)); + } + + #[test] + fn no_ceiling_allows_large_cap() { + let mut c = cfg(); + c.ceiling = None; + assert_eq!(dynamic_output_cap(100_000, &c), Some(150_000)); + } +}