diff --git a/crates/tokenizer/src/chat_template.rs b/crates/tokenizer/src/chat_template.rs index ef0cc9aab..d3d9da2d5 100644 --- a/crates/tokenizer/src/chat_template.rs +++ b/crates/tokenizer/src/chat_template.rs @@ -246,13 +246,21 @@ impl<'a> Detector<'a> { } /// Check if a list of statements contains `` in EmitRaw or string constants. + /// `` not closed by a later `` in the same literal — an + /// open tag leaves the completion mid-reasoning, a closed pair (e.g. + /// Qwen3's thinking-off `\n\n` filler) does not. + fn str_has_open_think_tag(s: &str) -> bool { + s.rfind("") + .is_some_and(|idx| !s[idx..].contains("")) + } + fn body_has_think_tag(stmts: &[Stmt]) -> bool { for stmt in stmts { match stmt { - Stmt::EmitRaw(raw) if raw.raw.contains("") => return true, + Stmt::EmitRaw(raw) if Self::str_has_open_think_tag(raw.raw) => return true, Stmt::EmitExpr(e) => { if let Expr::Const(c) = &e.expr { - if c.value.as_str().is_some_and(|s| s.contains("")) { + if c.value.as_str().is_some_and(Self::str_has_open_think_tag) { return true; } } @@ -1144,6 +1152,37 @@ impl ChatTemplateState { mod tests { use super::*; + #[test] + fn think_in_prefill_requires_open_think_tag() { + // Qwen3-style: the generation prompt injects a CLOSED empty think + // block only when thinking is disabled — completions never start + // mid-reasoning, so this must not count as think-in-prefill. + let qwen3_style = r" +{%- if add_generation_prompt %} + {{- '<|im_start|>assistant\n' }} + {%- if enable_thinking is defined and enable_thinking is false %} + {{- '\n\n\n\n' }} + {%- endif %} +{%- endif %}"; + let (_, think_in_prefill) = detect_all_with_ast(qwen3_style); + assert!( + !think_in_prefill, + "closed pair is not a prefill think" + ); + + // Thinking-only style: the generation prompt ends with an OPEN + // , so completions genuinely start mid-reasoning. + let thinking_style = r" +{%- if add_generation_prompt %} + {{- '<|im_start|>assistant\n\n' }} +{%- endif %}"; + let (_, think_in_prefill) = detect_all_with_ast(thinking_style); + assert!( + think_in_prefill, + "open in the prefill must be detected" + ); + } + #[test] fn test_chat_template_state_no_template() { let state = ChatTemplateState::new(None).unwrap(); diff --git a/e2e_test/chat_completions/test_function_calling.py b/e2e_test/chat_completions/test_function_calling.py index 35b9d37bd..e26c4e716 100644 --- a/e2e_test/chat_completions/test_function_calling.py +++ b/e2e_test/chat_completions/test_function_calling.py @@ -1565,6 +1565,134 @@ def test_conflicting_defs_required_tool_choice(self, model, api_client): """Skip: Mistral uses structural tags which don't consolidate $defs.""" +# ============================================================================= +# Tool Choice Required + Thinking Tests (reasoning parser interaction) +# Regression for: with a JSON-schema tool constraint (tool_choice "required" / +# named function) and thinking effectively ON, the pre-armed reasoning parser +# classified the grammar-forced JSON payload as reasoning_content, returning +# no tool_calls and finish_reason "stop". +# ============================================================================= + +THINKING_WEATHER_TOOLS = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather for a city", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "Name of the city"}, + }, + "required": ["city"], + }, + }, + } +] + + +# TokenSpeed excluded: its cold-start kernel compilation for Qwen3-30B-A3B +# exceeds the worker launch timeout on the 1-GPU CI rig. +@pytest.mark.engine("sglang", "vllm", "trtllm") +@pytest.mark.gpu(1) +@pytest.mark.model("Qwen/Qwen3-30B-A3B") +@pytest.mark.gateway( + extra_args=[ + "--tool-call-parser", + "qwen", + "--reasoning-parser", + "qwen3", + "--history-backend", + "memory", + ] +) +@pytest.mark.parametrize("setup_backend", ["grpc"], indirect=True) +class TestToolChoiceRequiredThinking: + """tool_choice "required" on a thinking model (enable_thinking defaults ON).""" + + def _assert_weather_tool_call(self, tool_calls) -> None: + assert tool_calls, "tool_choice='required' must produce tool_calls" + call = tool_calls[0] + assert call.function.name == "get_weather" + args = json.loads(call.function.arguments) + assert isinstance(args.get("city"), str), f"expected string 'city' arg, got: {args}" + + def test_required_non_streaming(self, model, api_client): + response = api_client.chat.completions.create( + model=model, + max_tokens=512, + messages=[{"role": "user", "content": "What is the weather in Paris right now?"}], + stream=False, + tools=THINKING_WEATHER_TOOLS, + tool_choice="required", + ) + + choice = response.choices[0] + self._assert_weather_tool_call(choice.message.tool_calls) + assert choice.finish_reason == "tool_calls", ( + f"expected finish_reason 'tool_calls', got: {choice.finish_reason}" + ) + # The grammar-forced payload must not be misclassified as reasoning. + reasoning = getattr(choice.message, "reasoning_content", None) + if reasoning: + try: + parsed = json.loads(reasoning) + except ValueError: + parsed = None + assert not isinstance(parsed, list), ( + f"tool payload leaked into reasoning_content: {reasoning}" + ) + + def test_required_streaming(self, model, api_client): + stream = api_client.chat.completions.create( + model=model, + max_tokens=512, + messages=[{"role": "user", "content": "What is the weather in Paris right now?"}], + stream=True, + tools=THINKING_WEATHER_TOOLS, + tool_choice="required", + ) + + finish_reason = None + tool_name = None + arguments = "" + reasoning_parts: list[str] = [] + for chunk in stream: + if not chunk.choices: + continue + choice = chunk.choices[0] + if choice.finish_reason: + finish_reason = choice.finish_reason + delta = choice.delta + if delta is None: + continue + for tool_call in delta.tool_calls or []: + if tool_call.function and tool_call.function.name: + tool_name = tool_call.function.name + if tool_call.function and tool_call.function.arguments: + arguments += tool_call.function.arguments + reasoning = getattr(delta, "reasoning_content", None) + if reasoning: + reasoning_parts.append(reasoning) + + assert tool_name == "get_weather", f"expected streamed tool call, got: {tool_name}" + args = json.loads(arguments) + assert isinstance(args.get("city"), str), f"expected string 'city' arg, got: {args}" + assert finish_reason == "tool_calls", ( + f"expected finish_reason 'tool_calls', got: {finish_reason}" + ) + # The grammar-forced payload must not stream out as reasoning deltas. + reasoning_text = "".join(reasoning_parts) + if reasoning_text: + try: + parsed = json.loads(reasoning_text) + except ValueError: + parsed = None + assert not isinstance(parsed, list), ( + f"tool payload leaked into reasoning deltas: {reasoning_text}" + ) + + # ============================================================================= # Multi-Turn Tool Call Tests # Regression for: assistant messages with tool_calls serialized without diff --git a/model_gateway/src/routers/grpc/regular/processor.rs b/model_gateway/src/routers/grpc/regular/processor.rs index 2513d1e6c..e7fd1cb05 100644 --- a/model_gateway/src/routers/grpc/regular/processor.rs +++ b/model_gateway/src/routers/grpc/regular/processor.rs @@ -103,6 +103,29 @@ impl ResponseProcessor { final_text.push_str(&t); } + // Check if a JSON schema constraint was used (specific function or + // required mode): it changes both reasoning arming and tool parsing. + let tool_choice_enabled = !matches!( + &original_request.tool_choice, + Some(ToolChoice::Value(ToolChoiceValue::None)) + ); + let has_structural_tag = self + .tool_parser_factory + .registry() + .has_structural_tag_for_parser(tool_parser_name); + let used_json_schema = if has_structural_tag { + false + } else { + match &original_request.tool_choice { + Some(ToolChoice::Function { .. }) => true, + Some(ToolChoice::Value(ToolChoiceValue::Required)) => true, + Some(ToolChoice::AllowedTools { mode, .. }) => mode == "required", + _ => false, + } + }; + let tool_constraint_active = + tool_choice_enabled && original_request.tools.is_some() && used_json_schema; + // Step 1: Handle reasoning content parsing let mut reasoning_text: Option = None; let mut processed_text = final_text; @@ -115,25 +138,25 @@ impl ResponseProcessor { reasoning_parser_name, &original_request.model, ) { - // If the template injected `` in the prefill (thinking toggle - // is supported and effectively ON), start in reasoning mode. - if utils::should_mark_reasoning_started( + if utils::should_start_in_reasoning( utils::resolve_user_thinking( original_request.chat_template_kwargs.as_ref(), original_request.reasoning_effort.as_deref(), tokenizer.as_ref(), ), tokenizer.as_ref(), + tool_constraint_active, ) { parser.mark_reasoning_started(); } match parser.detect_and_parse_reasoning(&processed_text) { Ok(result) => { - if !result.reasoning_text.is_empty() { - reasoning_text = Some(result.reasoning_text); - } - processed_text = result.normal_text; + (reasoning_text, processed_text) = utils::split_reasoning_result( + result, + processed_text, + tool_constraint_active, + ); } Err(e) => { warn!("Reasoning parsing error, skipping parsing: {e}"); @@ -144,28 +167,8 @@ impl ResponseProcessor { // Step 2: Handle tool call parsing let mut tool_calls: Option> = None; - let tool_choice_enabled = !matches!( - &original_request.tool_choice, - Some(ToolChoice::Value(ToolChoiceValue::None)) - ); if tool_choice_enabled && original_request.tools.is_some() { - // Check if JSON schema constraint was used (specific function or required mode) - let has_structural_tag = self - .tool_parser_factory - .registry() - .has_structural_tag_for_parser(tool_parser_name); - let used_json_schema = if has_structural_tag { - false - } else { - match &original_request.tool_choice { - Some(ToolChoice::Function { .. }) => true, - Some(ToolChoice::Value(ToolChoiceValue::Required)) => true, - Some(ToolChoice::AllowedTools { mode, .. }) => mode == "required", - _ => false, - } - }; - if used_json_schema { (tool_calls, processed_text) = utils::parse_json_schema_response( &processed_text, @@ -579,6 +582,20 @@ impl ResponseProcessor { &messages_request.model, ); + // Check if a JSON schema constraint was used (specific tool or any/required + // mode): it changes both reasoning arming and tool parsing. + let has_structural_tag = self + .tool_parser_factory + .registry() + .has_structural_tag_for_parser(tool_parser_name.as_deref()); + let used_json_schema = !has_structural_tag + && matches!( + &messages_request.tool_choice, + Some(messages::ToolChoice::Tool { .. } | messages::ToolChoice::Any { .. }) + ); + let tool_constraint_active = + tool_choice_enabled && messages_request.tools.is_some() && used_json_schema; + if separate_reasoning && !reasoning_parser_available { tracing::debug!( "No reasoning parser found for model '{}', reasoning content will not be separated", @@ -648,17 +665,22 @@ impl ResponseProcessor { Some(messages::ThinkingConfig::Disabled) => Some(false), None => None, }; - if utils::should_mark_reasoning_started(user_thinking, tokenizer.as_ref()) { + if utils::should_start_in_reasoning( + user_thinking, + tokenizer.as_ref(), + tool_constraint_active, + ) { parser.mark_reasoning_started(); } } match parser.detect_and_parse_reasoning(&processed_text) { Ok(result) => { - if !result.reasoning_text.is_empty() { - reasoning_text = Some(result.reasoning_text); - } - processed_text = result.normal_text; + (reasoning_text, processed_text) = utils::split_reasoning_result( + result, + processed_text, + tool_constraint_active, + ); } Err(e) => { warn!("Reasoning parsing error, skipping parsing: {e}"); @@ -671,17 +693,6 @@ impl ResponseProcessor { let mut tool_calls: Option> = None; if tool_choice_enabled && messages_request.tools.is_some() { - // Check if JSON schema constraint was used (specific tool or any/required mode) - let has_structural_tag = self - .tool_parser_factory - .registry() - .has_structural_tag_for_parser(tool_parser_name.as_deref()); - let used_json_schema = !has_structural_tag - && matches!( - &messages_request.tool_choice, - Some(messages::ToolChoice::Tool { .. } | messages::ToolChoice::Any { .. }) - ); - if used_json_schema { // Bridge Messages ToolChoice to Chat ToolChoice for reuse let chat_tool_choice = messages_request diff --git a/model_gateway/src/routers/grpc/regular/streaming.rs b/model_gateway/src/routers/grpc/regular/streaming.rs index 43bd04eae..20f51272e 100644 --- a/model_gateway/src/routers/grpc/regular/streaming.rs +++ b/model_gateway/src/routers/grpc/regular/streaming.rs @@ -328,17 +328,6 @@ impl StreamingProcessor { model, ); - // If the template supports a thinking toggle and the user enabled it, - // the template injected `` in the prefill — parsers should start - // in reasoning mode. - let thinking_override = utils::should_mark_reasoning_started( - utils::resolve_user_thinking( - original_request.chat_template_kwargs.as_ref(), - original_request.reasoning_effort.as_deref(), - tokenizer.as_ref(), - ), - tokenizer.as_ref(), - ); let think_in_prefill = tokenizer.think_in_prefill(); // Check if JSON schema constraint was used (specific function or required mode) @@ -356,6 +345,19 @@ impl StreamingProcessor { _ => false, } }; + let tool_constraint_active = used_json_schema && tools.is_some(); + + // Thinking effectively ON: parsers start in reasoning mode — unless a + // JSON-schema tool constraint is active (see should_start_in_reasoning). + let thinking_override = utils::should_start_in_reasoning( + utils::resolve_user_thinking( + original_request.chat_template_kwargs.as_ref(), + original_request.reasoning_effort.as_deref(), + tokenizer.as_ref(), + ), + tokenizer.as_ref(), + tool_constraint_active, + ); // Check if this is the specific function case (LLM generates parameters only, no name field). // Only applies when json_schema is used — structural tags include framing tokens @@ -668,6 +670,59 @@ impl StreamingProcessor { } } + // Phase 3.5: JSON-schema constraint recovery. An engine that enforces + // the tool grammar from the first token emits pure JSON; a pre-armed + // (think-in-prefill) parser then classifies the whole payload as + // reasoning and the tool stage never runs. The buffered raw text is + // the tool payload — parse it so tool calls are still delivered. + if tool_constraint_active { + for (index, buffer) in &stream_buffers { + if buffer.is_empty() || has_tool_calls.get(index).copied().unwrap_or(false) { + continue; + } + let still_in_reasoning = match reasoning_parsers.get(index) { + Some(parser) => parser.lock().await.is_in_reasoning(), + None => false, + }; + if !still_in_reasoning { + continue; + } + let (calls, _) = utils::parse_json_schema_response( + buffer, + tool_choice.as_ref(), + model, + history_tool_calls_count, + ); + let Some(calls) = calls else { continue }; + if calls.is_empty() { + continue; + } + for (call_index, call) in calls.into_iter().enumerate() { + let tool_call_delta = ToolCallDelta { + index: call_index as u32, + id: Some(call.id), + tool_type: Some("function".to_string()), + function: Some(FunctionCallDelta { + name: Some(call.function.name), + arguments: call.function.arguments, + }), + }; + let tool_chunk = ChatCompletionStreamResponse::builder(request_id, model) + .created(created) + .add_choice_tool_call_delta(*index, tool_call_delta) + .maybe_system_fingerprint(system_fingerprint) + .build(); + let sse_chunk = sse_encoder + .encode_data(&tool_chunk) + .map_err(|e| format!("Failed to serialize recovered tool chunk: {e}"))?; + tx.send(Ok(sse_chunk)) + .await + .map_err(|_| "Failed to send recovered tool chunk".to_string())?; + } + has_tool_calls.insert(*index, true); + } + } + // Phase 4: Finish reason chunks for (index, finish_reason) in &finish_reasons { let final_finish_reason = @@ -1878,6 +1933,10 @@ impl StreamingProcessor { // Parser state (simple variables — Messages is always n=1) let mut reasoning_parser: Option>>> = None; + // Raw text accumulated under a JSON-schema tool constraint, for + // end-of-stream recovery when a pre-armed parser ate the payload. + let mut constrained_raw_text = String::new(); + // Stop decoder let mut stop_decoder = { let (ref stop, ref stop_token_ids, skip_special_tokens, no_stop_trim, ignore_eos) = @@ -1932,19 +1991,6 @@ impl StreamingProcessor { model, ); - // Determine if thinking is effectively ON (for mark_reasoning_started). - let user_thinking = match &original_request.thinking { - Some( - messages::ThinkingConfig::Enabled { .. } - | messages::ThinkingConfig::Adaptive { .. }, - ) => Some(true), - Some(messages::ThinkingConfig::Disabled) => Some(false), - None => None, - }; - let thinking_override = - utils::should_mark_reasoning_started(user_thinking, tokenizer.as_ref()); - let think_in_prefill = tokenizer.think_in_prefill(); - let tool_choice_enabled = !matches!( &original_request.tool_choice, Some(messages::ToolChoice::None) @@ -1967,6 +2013,25 @@ impl StreamingProcessor { &original_request.tool_choice, Some(messages::ToolChoice::Tool { .. } | messages::ToolChoice::Any { .. }) ); + let tool_constraint_active = has_tools && used_json_schema; + + // Determine if thinking is effectively ON (for mark_reasoning_started). + // Under a JSON-schema tool constraint the parser must not be pre-armed + // (see should_start_in_reasoning). + let user_thinking = match &original_request.thinking { + Some( + messages::ThinkingConfig::Enabled { .. } + | messages::ThinkingConfig::Adaptive { .. }, + ) => Some(true), + Some(messages::ThinkingConfig::Disabled) => Some(false), + None => None, + }; + let thinking_override = utils::should_start_in_reasoning( + user_thinking, + tokenizer.as_ref(), + tool_constraint_active, + ); + let think_in_prefill = tokenizer.think_in_prefill(); // Check if model output is arguments-only for a specific function (ToolChoice::Tool). // Only applies when json_schema is used — structural tags include framing tokens. @@ -2064,6 +2129,10 @@ impl StreamingProcessor { continue; } + if tool_constraint_active { + constrained_raw_text.push_str(&chunk_text); + } + // Apply reasoning parser let (normal_text, reasoning_chunk_text, in_reasoning) = if reasoning_parser_available { @@ -2368,6 +2437,95 @@ impl StreamingProcessor { } } + // JSON-schema constraint recovery: a pre-armed (think-in-prefill) + // parser that never saw a think-end token streamed the grammar-forced + // payload as thinking. Parse the raw text so tool_use is still delivered. + if tool_constraint_active && !has_tool_calls && !constrained_raw_text.is_empty() { + let still_in_reasoning = match &reasoning_parser { + Some(parser) => parser.lock().await.is_in_reasoning(), + None => false, + }; + if still_in_reasoning { + let chat_tool_choice = original_request + .tool_choice + .as_ref() + .map(message_utils::convert_message_tool_choice); + let (calls, _) = utils::parse_json_schema_response( + &constrained_raw_text, + chat_tool_choice.as_ref(), + model, + history_tool_calls_count, + ); + if let Some(calls) = calls.filter(|calls| !calls.is_empty()) { + // Close any open block before emitting tool_use blocks. + if thinking_block_open { + Self::send_messages_event( + tx, + &mut sse_buffer, + &MessageStreamEvent::ContentBlockStop { + index: current_block_index, + }, + ) + .await?; + thinking_block_open = false; + current_block_index += 1; + } + if text_block_open { + Self::send_messages_event( + tx, + &mut sse_buffer, + &MessageStreamEvent::ContentBlockStop { + index: current_block_index, + }, + ) + .await?; + text_block_open = false; + current_block_index += 1; + } + for call in calls { + has_tool_calls = true; + Self::send_messages_event( + tx, + &mut sse_buffer, + &MessageStreamEvent::ContentBlockStart { + index: current_block_index, + content_block: ContentBlock::ToolUse { + id: message_utils::anthropic_tool_use_id(&call.id), + name: call.function.name.clone(), + input: Value::Object(serde_json::Map::new()), + }, + }, + ) + .await?; + if let Some(args) = call.function.arguments { + if !args.is_empty() { + Self::send_messages_event( + tx, + &mut sse_buffer, + &MessageStreamEvent::ContentBlockDelta { + index: current_block_index, + delta: ContentBlockDelta::InputJsonDelta { + partial_json: args, + }, + }, + ) + .await?; + } + } + Self::send_messages_event( + tx, + &mut sse_buffer, + &MessageStreamEvent::ContentBlockStop { + index: current_block_index, + }, + ) + .await?; + current_block_index += 1; + } + } + } + } + // Phase 3: Flush unstreamed tool args from the incremental parser if let Some(ref parser) = streaming_tool_parser { if let Some(unstreamed_items) = parser.get_unstreamed_tool_args() { diff --git a/model_gateway/src/routers/grpc/utils/mod.rs b/model_gateway/src/routers/grpc/utils/mod.rs index 517e25fec..0b0ff69cb 100644 --- a/model_gateway/src/routers/grpc/utils/mod.rs +++ b/model_gateway/src/routers/grpc/utils/mod.rs @@ -22,7 +22,8 @@ pub(crate) use logprobs::{ pub(crate) use metrics::{error_type_from_status, route_to_endpoint}; pub(crate) use parsers::{ check_reasoning_parser_availability, check_tool_parser_availability, create_reasoning_parser, - create_tool_parser, get_tool_parser, reasoning_parser_requires_special_tokens, ParserResolver, + create_tool_parser, get_tool_parser, reasoning_parser_requires_special_tokens, + should_start_in_reasoning, split_reasoning_result, ParserResolver, }; // `pub` (not `pub(crate)`) so the Go bindings can reuse the gateway's reasoning // detection instead of duplicating it. diff --git a/model_gateway/src/routers/grpc/utils/parsers.rs b/model_gateway/src/routers/grpc/utils/parsers.rs index 51ea413ed..5ffeed2cd 100644 --- a/model_gateway/src/routers/grpc/utils/parsers.rs +++ b/model_gateway/src/routers/grpc/utils/parsers.rs @@ -125,6 +125,42 @@ pub fn should_mark_reasoning_started( } } +/// Whether the reasoning parser should start pre-armed (in reasoning mode). +/// +/// Under a JSON-schema tool constraint (`tool_choice: required` / named +/// function without a structural tag) engines may grammar-force pure JSON +/// from the first token; a pre-armed parser would classify that payload as +/// truncated reasoning. Templates without `` in the prefill emit an +/// explicit start token whenever reasoning does happen, so they stay +/// un-armed under the constraint. Think-in-prefill templates keep the +/// pre-arm: their completions genuinely begin mid-reasoning. +pub(crate) fn should_start_in_reasoning( + user_thinking: Option, + tokenizer: &dyn Tokenizer, + used_json_schema: bool, +) -> bool { + should_mark_reasoning_started(user_thinking, tokenizer) + && (tokenizer.think_in_prefill() || !used_json_schema) +} + +/// Split a completed generation into `(reasoning_content, normal_text)`. +/// +/// Under a JSON-schema tool constraint an all-reasoning result is impossible +/// (the grammar forces a JSON payload): it means a pre-armed parser never saw +/// a think-end token, so the original text is returned for tool parsing. +pub(crate) fn split_reasoning_result( + result: reasoning_parser::ParserResult, + original_text: String, + used_json_schema: bool, +) -> (Option, String) { + if used_json_schema && result.normal_text.is_empty() && !result.reasoning_text.is_empty() { + (None, original_text) + } else { + let reasoning = (!result.reasoning_text.is_empty()).then_some(result.reasoning_text); + (reasoning, result.normal_text) + } +} + /// Extract the user's thinking preference from chat_template_kwargs. /// /// Only checks the key that the template actually uses (e.g. `enable_thinking` @@ -316,6 +352,128 @@ pub(crate) fn create_tool_parser( mod tests { use super::*; + /// MockTokenizer wrapper with a configurable thinking toggle and + /// think-in-prefill flag (trait defaults are `None`/`false`). + struct ToggleTok { + inner: llm_tokenizer::MockTokenizer, + toggle: ThinkingToggle, + prefill: bool, + } + + impl ToggleTok { + fn new(toggle: ThinkingToggle, prefill: bool) -> Self { + Self { + inner: llm_tokenizer::MockTokenizer::new(), + toggle, + prefill, + } + } + } + + impl llm_tokenizer::traits::Encoder for ToggleTok { + fn encode(&self, i: &str, s: bool) -> anyhow::Result { + self.inner.encode(i, s) + } + fn encode_batch( + &self, + i: &[&str], + s: bool, + ) -> anyhow::Result> { + self.inner.encode_batch(i, s) + } + } + + impl llm_tokenizer::traits::Decoder for ToggleTok { + fn decode(&self, ids: &[u32], s: bool) -> anyhow::Result { + self.inner.decode(ids, s) + } + } + + impl Tokenizer for ToggleTok { + fn vocab_size(&self) -> usize { + self.inner.vocab_size() + } + fn get_special_tokens(&self) -> &llm_tokenizer::traits::SpecialTokens { + self.inner.get_special_tokens() + } + fn token_to_id(&self, t: &str) -> Option { + self.inner.token_to_id(t) + } + fn id_to_token(&self, id: u32) -> Option { + self.inner.id_to_token(id) + } + fn as_any(&self) -> &dyn std::any::Any { + self + } + fn thinking_toggle(&self) -> ThinkingToggle { + self.toggle + } + fn think_in_prefill(&self) -> bool { + self.prefill + } + } + + #[test] + fn start_in_reasoning_unarmed_under_json_schema_for_toggle_templates() { + // Qwen3-style template: thinking toggle DefaultOn, no in prefill. + let tok = ToggleTok::new(ThinkingToggle::DefaultOn, false); + + // Unconstrained: thinking ON arms the parser. + assert!(should_start_in_reasoning(None, &tok, false)); + // JSON-schema constraint: must not pre-arm (payload has no think tokens). + assert!(!should_start_in_reasoning(None, &tok, true)); + // Thinking explicitly off never arms. + assert!(!should_start_in_reasoning(Some(false), &tok, false)); + } + + #[test] + fn start_in_reasoning_keeps_prearm_for_think_in_prefill_templates() { + // DeepSeek-style: `` in the prefill, completions start mid-reasoning. + let tok = ToggleTok::new(ThinkingToggle::DefaultOn, true); + + assert!(should_start_in_reasoning(None, &tok, true)); + assert!(should_start_in_reasoning(None, &tok, false)); + assert!(!should_start_in_reasoning(Some(false), &tok, true)); + + // No thinking toggle at all: never armed, constraint or not. + let plain = ToggleTok::new(ThinkingToggle::None, false); + assert!(!should_start_in_reasoning(None, &plain, true)); + assert!(!should_start_in_reasoning(None, &plain, false)); + } + + #[test] + fn split_reasoning_result_recovers_constrained_payload() { + use reasoning_parser::ParserResult; + + let payload = r#"[{"name":"get_weather","parameters":{"city":"Paris"}}]"#; + + // All-reasoning under the constraint: the payload is recovered. + let all_reasoning = ParserResult::reasoning(payload.to_string()); + let (reasoning, normal) = split_reasoning_result(all_reasoning, payload.to_string(), true); + assert_eq!(reasoning, None); + assert_eq!(normal, payload); + + // All-reasoning without the constraint: kept as truncated reasoning. + let all_reasoning = ParserResult::reasoning("partial thought".to_string()); + let (reasoning, normal) = + split_reasoning_result(all_reasoning, "partial thought".to_string(), false); + assert_eq!(reasoning.as_deref(), Some("partial thought")); + assert_eq!(normal, ""); + + // A proper reasoning/payload split passes through unchanged. + let split = ParserResult::new(payload.to_string(), "thought".to_string()); + let (reasoning, normal) = + split_reasoning_result(split, format!("thought{payload}"), true); + assert_eq!(reasoning.as_deref(), Some("thought")); + assert_eq!(normal, payload); + + // Empty reasoning stays None. + let plain = ParserResult::normal("hi".to_string()); + let (reasoning, normal) = split_reasoning_result(plain, "hi".to_string(), true); + assert_eq!(reasoning, None); + assert_eq!(normal, "hi"); + } + #[test] fn resolve_thinking_pref_precedence() { // Explicit toggle > native template effort > reasoning_effort mapping.