car-inference 0.60.0

//! A thinking budget for vLLM-MLX models whose prompt opens a `<think>` block.
//!
//! Qwen3.5 and later (and GLM-4.7, DeepSeek-R1) end the rendered prompt with
//! `<think>\n`, so generation starts inside the reasoning block and the model
//! emits only the closing tag. Qwen3.8 regularly thinks for longer than a
//! caller's whole `max_tokens`. Two vllm-mlx 0.6.0 defects then meet:
//!
//! - Its non-streaming reasoning parser treats output with no `</think>` as
//!   plain content (the streaming parser knows about the open block and does
//!   not), so a response cut off mid-thought comes back as the *answer*.
//! - Its own `thinking_token_budget` is a logits processor, and the MTP
//!   drafter's decode path never applies it.
//!
//! So CAR spends the budget itself, the way Qwen documents for its models:
//! the first request gets a share of `max_tokens`; if the model is still
//! thinking when that runs out, CAR closes the thought and asks for the answer
//! with the rest, through `/v1/completions` over the same prompt the server
//! rendered. A response that finishes inside the budget is returned untouched,
//! and no second request is made.
//!
//! The prompt is rendered from the *wire* request (the messages and tools the
//! server was actually sent), normalized the way vllm-mlx normalizes them
//! before its own render. A mismatch would only shift the continuation's
//! prefix; the caller compares the server's prompt-token count against the
//! expected one and logs a difference.

use std::path::Path;

use serde_json::Value;

use crate::tasks::chat_template::ChatTemplate;
use crate::tasks::generate::{ThinkingBlock, ThinkingMode, ToolCall};

/// Below this, half the budget would leave too little to think or to answer.
const MIN_SPLIT_MAX_TOKENS: usize = 512;

/// Qwen's own wording for ending a thought at a budget.
const CLOSE_THOUGHT: &str = "\n\nConsidering the limited time by the user, I have to give the \
                             solution based on the thinking directly now.\n</think>\n\n";

/// How a request to a `<think>`-opening model is split into a thinking pass
/// and, when the model runs out of budget mid-thought, an answer pass.
#[derive(Debug)]
pub(crate) struct ThinkingSplit {
    /// The prompt as the server rendered it, ending in `<think>\n`.
    prompt: String,
    /// `max_tokens` for the first request.
    pub(crate) budget: usize,
    /// The caller's `max_tokens`, which both requests together stay within.
    max_tokens: usize,
}

/// What the second request is: a prompt that continues where the first one
/// stopped, and the part of the answer already written.
#[derive(Debug, PartialEq)]
pub(crate) struct Continuation {
    pub(crate) prompt: String,
    pub(crate) max_tokens: usize,
    /// The thought, when the first request never closed it.
    pub(crate) thought: Option<String>,
    /// The answer text the first request already produced, which the second
    /// continues.
    pub(crate) answer_so_far: String,
}

impl ThinkingSplit {
    /// A split for this request, or `None` when it does not apply: the model
    /// directory has no chat template, the template does not open `<think>`
    /// for this conversation, the budget is too small to split, or the caller
    /// asked for structured output (a constrained answer cannot be resumed
    /// from raw text), or the conversation carries an image or other non-text
    /// part (a text completion would continue without it).
    ///
    /// `budget_tokens` is the caller's thinking budget, used when it fits
    /// inside `max_tokens`; otherwise half of `max_tokens` goes to thinking.
    /// `thinking` is the mode the server was sent, so the prompt is rendered
    /// as the server rendered it (with thinking off, a template closes the
    /// block itself and there is nothing to split).
    pub(crate) fn plan(
        model_dir: &Path,
        body: &Value,
        max_tokens: usize,
        budget_tokens: usize,
        thinking: ThinkingMode,
        structured_output: bool,
    ) -> Option<Self> {
        if structured_output || max_tokens < MIN_SPLIT_MAX_TOKENS || !model_dir.is_dir() {
            return None;
        }
        let template = match ChatTemplate::load(model_dir) {
            Ok(Some(template)) => template,
            Ok(None) => return None,
            Err(error) => {
                tracing::debug!(%error, dir = %model_dir.display(), "no thinking budget: chat template unusable");
                return None;
            }
        };
        let wire_messages = body.get("messages")?.as_array()?;
        let has_media = wire_messages.iter().any(|message| {
            message
                .get("content")
                .and_then(Value::as_array)
                .is_some_and(|parts| {
                    parts
                        .iter()
                        .any(|part| part.get("type").and_then(Value::as_str) != Some("text"))
                })
        });
        if has_media {
            return None;
        }
        let messages = served_messages(wire_messages);
        let tools = body
            .get("tools")
            .and_then(Value::as_array)
            .map(Vec::as_slice);
        let prompt = match template.render_values(&messages, tools, thinking) {
            Ok(prompt) => prompt,
            Err(error) => {
                tracing::debug!(%error, "no thinking budget: chat template did not render");
                return None;
            }
        };
        if !prompt.ends_with("<think>\n") {
            return None;
        }
        let budget = if budget_tokens > 0 && budget_tokens < max_tokens {
            budget_tokens
        } else {
            max_tokens / 2
        };
        Some(Self {
            prompt,
            budget,
            max_tokens,
        })
    }

    /// The second request, or `None` when the first one needs none: it
    /// finished inside its budget, or it called a tool. `text` and `thinking`
    /// are the first response's content and reasoning as the server split
    /// them.
    pub(crate) fn continuation(
        &self,
        text: &str,
        thinking: &[ThinkingBlock],
        tool_calls: &[ToolCall],
        completion_tokens: u64,
    ) -> Option<Continuation> {
        // `finish_reason` cannot tell a natural stop from the limit (the
        // drafter reports "stop" at it), so spending the budget is the
        // trigger. An answer that happened to end exactly there gets a
        // continuation whose prompt ends in a finished turn; the model's next
        // token is its end of turn. A response with no usage reads as zero
        // tokens and is returned as it is; vllm-mlx always reports usage.
        let spent = completion_tokens as usize;
        if !tool_calls.is_empty() || spent < self.budget {
            return None;
        }
        // The drafter can overshoot its limit by a token or two.
        let max_tokens = self.max_tokens.checked_sub(spent).filter(|n| *n > 0)?;
        let reasoning: String = thinking.iter().map(|block| block.text.as_str()).collect();
        if reasoning.is_empty() {
            // Still thinking: with no `</think>` in its output, the server
            // returned the whole thought as content.
            if text.contains("</think>") {
                // The server left the tags in, so its split cannot be read.
                // Callers plan a split only for families served with a
                // `<think>` parser, so this is the server misbehaving.
                tracing::warn!("vllm-mlx returned `</think>` unsplit; no thinking continuation");
                return None;
            }
            Some(Continuation {
                prompt: format!("{}{}{}", self.prompt, text, CLOSE_THOUGHT),
                max_tokens,
                thought: Some(text.to_string()),
                answer_so_far: String::new(),
            })
        } else {
            // The thought closed and the answer ran out of budget: continue
            // the answer, rendered the way the template renders a turn.
            Some(Continuation {
                prompt: format!("{}{}\n</think>\n\n{}", self.prompt, reasoning, text),
                max_tokens,
                thought: None,
                answer_so_far: text.to_string(),
            })
        }
    }
}

/// `messages` as vllm-mlx hands them to the chat template: a tool call's
/// JSON-string `arguments` parsed into the object templates iterate, content
/// parts joined to text, and missing content as `""`.
fn served_messages(messages: &[Value]) -> Vec<Value> {
    messages
        .iter()
        .map(|message| {
            let mut message = message.clone();
            if let Some(object) = message.as_object_mut() {
                let content = match object.get("content") {
                    None | Some(Value::Null) => Value::String(String::new()),
                    Some(Value::Array(parts)) => Value::String(
                        parts
                            .iter()
                            .filter_map(|part| part.get("text").and_then(Value::as_str))
                            .collect::<Vec<_>>()
                            .join("\n"),
                    ),
                    Some(other) => other.clone(),
                };
                object.insert("content".into(), content);
                if let Some(Value::Array(calls)) = object.get_mut("tool_calls") {
                    for call in calls {
                        if let Some(arguments) = call.pointer_mut("/function/arguments") {
                            if let Some(parsed) = arguments
                                .as_str()
                                .and_then(|raw| serde_json::from_str::<Value>(raw).ok())
                            {
                                *arguments = parsed;
                            }
                        }
                    }
                }
            }
            message
        })
        .collect()
}

/// The generation-prompt tail of Qwen3.8's `chat_template.jinja`, with its
/// message loop reduced to user and assistant turns.
#[cfg(test)]
pub(crate) const THINKING_TEMPLATE: &str = r#"{%- for m in messages %}{{- '<|im_start|>' + m.role + '\n' + m.content + '<|im_end|>\n' }}{%- endfor %}{%- if add_generation_prompt %}{{- '<|im_start|>assistant\n' }}{%- if enable_thinking is defined and enable_thinking is false %}{{- '<think>\n\n</think>\n\n' }}{%- else %}{{- '<think>\n' }}{%- endif %}{%- endif %}"#;

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    fn model_dir(template: &str) -> tempfile::TempDir {
        let dir = tempfile::tempdir().unwrap();
        std::fs::write(dir.path().join("chat_template.jinja"), template).unwrap();
        dir
    }

    fn body() -> Value {
        json!({"model": "m", "messages": [{"role": "user", "content": "hi"}]})
    }

    fn block(text: &str) -> ThinkingBlock {
        ThinkingBlock {
            text: text.into(),
            signature: None,
            redacted_data: None,
        }
    }

    #[test]
    fn a_think_opening_template_splits_the_budget() {
        let dir = model_dir(THINKING_TEMPLATE);
        let split =
            ThinkingSplit::plan(dir.path(), &body(), 1024, 0, ThinkingMode::Auto, false).unwrap();
        assert_eq!(split.budget, 512);
        assert!(split.prompt.ends_with("<|im_start|>assistant\n<think>\n"));
        // A caller's thinking budget is honored when it fits.
        let split =
            ThinkingSplit::plan(dir.path(), &body(), 1024, 300, ThinkingMode::Auto, false).unwrap();
        assert_eq!(split.budget, 300);
    }

    #[test]
    fn thinking_off_closes_the_block_and_is_not_split() {
        let dir = model_dir(THINKING_TEMPLATE);
        assert!(
            ThinkingSplit::plan(dir.path(), &body(), 1024, 0, ThinkingMode::Off, false).is_none()
        );
        assert!(
            ThinkingSplit::plan(dir.path(), &body(), 1024, 0, ThinkingMode::On, false).is_some()
        );
    }

    #[test]
    fn no_split_where_it_cannot_apply() {
        let plain = model_dir(
            r#"{%- for m in messages %}{{ m.content }}{%- endfor %}{{- '<|im_start|>assistant\n' }}"#,
        );
        assert!(
            ThinkingSplit::plan(plain.path(), &body(), 1024, 0, ThinkingMode::Auto, false)
                .is_none()
        );
        let dir = model_dir(THINKING_TEMPLATE);
        assert!(
            ThinkingSplit::plan(dir.path(), &body(), 256, 0, ThinkingMode::Auto, false).is_none()
        );
        assert!(
            ThinkingSplit::plan(dir.path(), &body(), 1024, 0, ThinkingMode::Auto, true).is_none()
        );
        // A model id rather than a local directory (a user-run server).
        assert!(ThinkingSplit::plan(
            Path::new("org/model"),
            &body(),
            1024,
            0,
            ThinkingMode::Auto,
            false
        )
        .is_none());
    }

    #[test]
    fn a_conversation_with_an_image_is_not_split() {
        let dir = model_dir(THINKING_TEMPLATE);
        let with_image = json!({"messages": [{"role": "user", "content": [
            {"type": "text", "text": "Describe this image"},
            {"type": "image_url", "image_url": {"url": "data:image/png;base64,aGVsbG8="}}
        ]}]});
        assert!(
            ThinkingSplit::plan(dir.path(), &with_image, 1024, 0, ThinkingMode::Auto, false)
                .is_none()
        );
        let text_parts = json!({"messages": [{"role": "user", "content": [
            {"type": "text", "text": "first"},
            {"type": "text", "text": "second"}
        ]}]});
        let split =
            ThinkingSplit::plan(dir.path(), &text_parts, 1024, 0, ThinkingMode::Auto, false)
                .unwrap();
        assert!(split.prompt.contains("first\nsecond"));
    }

    fn split() -> (tempfile::TempDir, ThinkingSplit) {
        let dir = model_dir(THINKING_TEMPLATE);
        let split =
            ThinkingSplit::plan(dir.path(), &body(), 1000, 0, ThinkingMode::Auto, false).unwrap();
        (dir, split)
    }

    #[test]
    fn a_thought_cut_off_at_the_budget_is_closed_and_answered() {
        let (_dir, split) = split();
        let next = split
            .continuation("We need to add 2 and 2", &[], &[], 502)
            .unwrap();
        assert_eq!(next.thought.as_deref(), Some("We need to add 2 and 2"));
        assert_eq!(next.max_tokens, 498);
        assert!(next
            .prompt
            .ends_with(&format!("<think>\nWe need to add 2 and 2{CLOSE_THOUGHT}")));
        assert!(next.answer_so_far.is_empty());
    }

    #[test]
    fn an_answer_cut_off_at_the_budget_is_continued() {
        let (_dir, split) = split();
        let next = split
            .continuation("The answer is", &[block("2 + 2 = 4")], &[], 500)
            .unwrap();
        assert_eq!(next.thought, None);
        assert_eq!(next.answer_so_far, "The answer is");
        assert!(next
            .prompt
            .ends_with("<think>\n2 + 2 = 4\n</think>\n\nThe answer is"));
    }

    #[test]
    fn a_response_inside_its_budget_or_calling_a_tool_needs_no_second_request() {
        let (_dir, split) = split();
        assert_eq!(split.continuation("4", &[block("2 + 2")], &[], 499), None);
        let call = ToolCall {
            id: None,
            name: "f".into(),
            arguments: Default::default(),
        };
        assert_eq!(split.continuation("x", &[], &[call], 500), None);
        // Tags the server failed to strip mean its split cannot be trusted.
        assert_eq!(split.continuation("a</think>b", &[], &[], 500), None);
        // Nothing left of max_tokens.
        assert_eq!(split.continuation("thinking", &[], &[], 1000), None);
    }

    #[test]
    fn served_messages_match_what_vllm_mlx_renders() {
        let messages = served_messages(&[
            json!({"role": "assistant", "content": null, "tool_calls": [
                {"id": "c", "type": "function", "function": {"name": "f", "arguments": "{\"k\": 1}"}}
            ]}),
            json!({"role": "user", "content": [{"type": "text", "text": "a"}, {"type": "text", "text": "b"}]}),
        ]);
        assert_eq!(messages[0]["content"], json!(""));
        assert_eq!(
            messages[0]["tool_calls"][0]["function"]["arguments"],
            json!({"k": 1})
        );
        assert_eq!(messages[1]["content"], json!("a\nb"));
    }
}