Dynamic output-token cap + determinism invariants (api/vLLM mode)

ltengine/src/llm_api.rs

@@ -16,30 +16,46 @@ struct ChatMessage {
     content: String,
 }
 
+#[derive(Serialize)]
+struct ChatTemplateKwargs {
+    enable_thinking: bool,
+}
+
 #[derive(Serialize)]
 struct ChatCompletionRequest {
     model: String,
     messages: Vec<ChatMessage>,
     temperature: f32,
     #[serde(skip_serializing_if = "Option::is_none")]
     max_tokens: Option<u32>,
+    chat_template_kwargs: ChatTemplateKwargs,
 }
 
 #[derive(Deserialize)]
 struct ChatCompletionResponse {
     choices: Vec<Choice>,
+    #[serde(default)]
+    usage: Option<Usage>,
 }
 
 #[derive(Deserialize)]
 struct Choice {
     message: ResponseMessage,
+    #[serde(default)]
+    finish_reason: Option<String>,
 }
 
 #[derive(Deserialize)]
 struct ResponseMessage {
     content: String,
 }
 
+#[derive(Deserialize)]
+struct Usage {
+    #[serde(default)]
+    completion_tokens: Option<u32>,
+}
+
 #[derive(Deserialize)]
 struct ModelsResponse {
     data: Vec<ModelEntry>,
@@ -52,8 +68,14 @@ struct ModelEntry {
 
 impl LLM {
     pub fn new(url: String, api_key: String, model: String, max_tokens: Option<u32>, timeout_secs: u64) -> Result<Self> {
-        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<String> {
+    pub async fn run_prompt(
+        &self,
+        system: String,
+        user: String,
+        max_output_tokens: Option<u32>,
+    ) -> Result<String> {
         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());
     }
 }

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<Arc<
 
     let translated_text = if source != target {
         #[cfg(feature = "api")]
-        { llm.run_prompt(prompt.system, prompt.user).await.unwrap_or(q.clone()) }
+        {
+            let budget_cfg = token_budget::TokenBudgetConfig {
+                chars_per_token: args.llm_chars_per_token,
+                output_mult: args.llm_max_tokens_mult,
+                floor: args.llm_max_tokens_floor,
+                ceiling: if args.llm_max_tokens > 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{

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<u32>,
+}
+
+/// 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<u32> {
+    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));
+    }
+}