From aad18f3d4747739c0484356c2e4711c2a77f3b29 Mon Sep 17 00:00:00 2001 From: ConnorLi96 Date: Tue, 7 Apr 2026 13:46:25 -0700 Subject: [PATCH 1/2] fix(tool_parser): fix func call parsing for native tool-call token models - Prioritize explicitly configured tool parser over JSON schema parsing - Support alternative delimiters (<|func_start|>/<|func_end|>) in KimiK2 parser - Prevent reasoning parser from consuming tool call markers when is skipped Signed-off-by: ConnorLi96 Made-with: Cursor --- crates/reasoning_parser/src/parsers/base.rs | 16 ++++++++-- crates/tool_parser/src/parsers/kimik2.rs | 32 +++++++++---------- .../src/routers/grpc/regular/processor.rs | 23 +++++++++++-- .../src/routers/grpc/regular/streaming.rs | 12 ++++--- 4 files changed, 57 insertions(+), 26 deletions(-) diff --git a/crates/reasoning_parser/src/parsers/base.rs b/crates/reasoning_parser/src/parsers/base.rs index 6214acd695..61420b6179 100644 --- a/crates/reasoning_parser/src/parsers/base.rs +++ b/crates/reasoning_parser/src/parsers/base.rs @@ -63,7 +63,12 @@ impl ReasoningParser for BaseReasoningParser { .to_string(); if !processed_text.contains(&self.config.think_end_token) { - // Assume reasoning was truncated before end token + // Don't consume tool call markers as reasoning content + if let Some(tool_pos) = processed_text.find("<|tool_calls_section_begin|>") { + let reasoning_text = processed_text[..tool_pos].trim().to_string(); + let normal_text = processed_text[tool_pos..].to_string(); + return Ok(ParserResult::new(normal_text, reasoning_text)); + } return Ok(ParserResult::reasoning(processed_text)); } @@ -133,7 +138,14 @@ impl ReasoningParser for BaseReasoningParser { // Continue with reasoning content if self.in_reasoning && self.config.stream_reasoning { - // Stream the content immediately + // Some models skip and go straight to tool calls + if let Some(tool_pos) = current_text.find("<|tool_calls_section_begin|>") { + let reasoning_text = current_text[..tool_pos].trim().to_string(); + let normal_text = current_text[tool_pos..].to_string(); + self.buffer.clear(); + self.in_reasoning = false; + return Ok(ParserResult::new(normal_text, reasoning_text)); + } let reasoning_text = current_text; self.buffer.clear(); Ok(ParserResult::reasoning(reasoning_text)) diff --git a/crates/tool_parser/src/parsers/kimik2.rs b/crates/tool_parser/src/parsers/kimik2.rs index 18e2f3e50a..4f089dfa09 100644 --- a/crates/tool_parser/src/parsers/kimik2.rs +++ b/crates/tool_parser/src/parsers/kimik2.rs @@ -57,16 +57,14 @@ impl KimiK2Parser { reason = "regex patterns are compile-time string literals" )] pub fn new() -> Self { - // Pattern for complete tool calls - let tool_call_pattern = r"<\|tool_call_begin\|>\s*(?P[\w\.]+:\d+)\s*<\|tool_call_argument_begin\|>\s*(?P\{.*?\})\s*<\|tool_call_end\|>"; + // Supports alternative delimiters: <|func_start|>/<|func_end|>; (?s) for multi-line JSON + let tool_call_pattern = r"(?s)<\|tool_call_begin\|>\s*(?P[\w\.]+:\d+)\s*(?:<\|tool_call_argument_begin\|>\s*|<\|func_start\|>\s*)?(?P\{.*?\})\s*(?:<\|tool_call_end\|>|<\|func_end\|>)"; let tool_call_extractor = Regex::new(tool_call_pattern).expect("Valid regex pattern"); - // Pattern for streaming (partial) tool calls - let stream_pattern = r"<\|tool_call_begin\|>\s*(?P[\w\.]+:\d+)\s*<\|tool_call_argument_begin\|>\s*(?P\{.*)"; + let stream_pattern = r"(?s)<\|tool_call_begin\|>\s*(?P[\w\.]+:\d+)\s*(?:<\|tool_call_argument_begin\|>\s*|<\|func_start\|>\s*)?(?P\{.*)"; let stream_tool_call_extractor = Regex::new(stream_pattern).expect("Valid regex pattern"); - // Pattern for removing completed tool calls - let end_pattern = r"<\|tool_call_begin\|>.*?<\|tool_call_end\|>"; + let end_pattern = r"<\|tool_call_begin\|>.*?(?:<\|tool_call_end\|>|<\|func_end\|>)"; let tool_call_end_pattern = Regex::new(end_pattern).expect("Valid regex pattern"); // Robust parser for ids like "functions.search:0" or fallback "search:0" @@ -181,7 +179,7 @@ impl ToolParser for KimiK2Parser { // No tool markers detected - return all buffered content as normal text let mut normal_text = std::mem::take(&mut self.buffer); // Remove end tokens if present - for e_token in ["<|tool_calls_section_end|>", "<|tool_call_end|>"] { + for e_token in ["<|tool_calls_section_end|>", "<|tool_call_end|>", "<|func_end|>"] { normal_text = normal_text.replace(e_token, ""); } return Ok(StreamingParseResult { @@ -242,13 +240,13 @@ impl ToolParser for KimiK2Parser { function_args }; - // Split by end token before sending (like Python does) - let parsed_args_diff = - if let Some(pos) = argument_diff.find("<|tool_call_end|>") { - &argument_diff[..pos] - } else { - argument_diff - }; + let end_pos = argument_diff.find("<|tool_call_end|>") + .or_else(|| argument_diff.find("<|func_end|>")); + let parsed_args_diff = if let Some(pos) = end_pos { + &argument_diff[..pos] + } else { + argument_diff + }; if !parsed_args_diff.is_empty() { calls.push(ToolCallItem { @@ -264,9 +262,9 @@ impl ToolParser for KimiK2Parser { } } - // Check completeness - split by end token first - let parsed_args = if let Some(pos) = function_args.find("<|tool_call_end|>") - { + let end_pos2 = function_args.find("<|tool_call_end|>") + .or_else(|| function_args.find("<|func_end|>")); + let parsed_args = if let Some(pos) = end_pos2 { &function_args[..pos] } else { function_args diff --git a/model_gateway/src/routers/grpc/regular/processor.rs b/model_gateway/src/routers/grpc/regular/processor.rs index a9f927dde0..362cfed344 100644 --- a/model_gateway/src/routers/grpc/regular/processor.rs +++ b/model_gateway/src/routers/grpc/regular/processor.rs @@ -148,7 +148,16 @@ impl ResponseProcessor { _ => false, }; - if used_json_schema { + if self.configured_tool_parser.is_some() && tool_parser_available { + // Explicitly configured parser takes priority (models may emit native tokens regardless of tool_choice) + (tool_calls, processed_text) = self + .parse_tool_calls( + &processed_text, + &original_request.model, + history_tool_calls_count, + ) + .await; + } else if used_json_schema { (tool_calls, processed_text) = utils::parse_json_schema_response( &processed_text, original_request.tool_choice.as_ref(), @@ -621,7 +630,17 @@ impl ResponseProcessor { Some(messages::ToolChoice::Tool { .. } | messages::ToolChoice::Any { .. }) ); - if used_json_schema { + if self.configured_tool_parser.is_some() && tool_parser_available { + (tool_calls, processed_text) = self + .parse_tool_calls( + &processed_text, + &messages_request.model, + utils::message_utils::get_history_tool_calls_count_messages( + &messages_request, + ), + ) + .await; + } else if used_json_schema { // Bridge Messages ToolChoice to Chat ToolChoice for reuse let chat_tool_choice = messages_request .tool_choice diff --git a/model_gateway/src/routers/grpc/regular/streaming.rs b/model_gateway/src/routers/grpc/regular/streaming.rs index a0e391cd1a..a0d1b2f3d8 100644 --- a/model_gateway/src/routers/grpc/regular/streaming.rs +++ b/model_gateway/src/routers/grpc/regular/streaming.rs @@ -377,8 +377,10 @@ impl StreamingProcessor { && tool_choice_enabled && (tool_parser_available || used_json_schema) { - let tool_chunks = if is_specific_function { - // Handle specific function case - emit tool call deltas with arguments + let tool_chunks = if is_specific_function + && !(self.configured_tool_parser.is_some() + && tool_parser_available) + { Self::process_specific_function_stream( &delta, index, @@ -391,7 +393,6 @@ impl StreamingProcessor { history_tool_calls_count, ) } else { - // Use incremental parser for regular/required modes self.process_tool_calls_stream( &delta, index, @@ -1753,8 +1754,9 @@ impl StreamingProcessor { // Tool call handling: incremental streaming parser if !in_reasoning && streaming_tool_parser.is_some() { - if is_specific_function { - // Specific function: entire output is arguments for one tool + if is_specific_function + && !(self.configured_tool_parser.is_some() && tool_parser_available) + { if !has_tool_calls { has_tool_calls = true; // Close text block if open before starting tool block From 468f1dfa1b0cdd17d4798528a26e859d8fe81455 Mon Sep 17 00:00:00 2001 From: ConnorLi96 Date: Wed, 8 Apr 2026 00:38:58 -0700 Subject: [PATCH 2/2] fix(tool_parser): address review feedback on func call parsing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Replace hardcoded `<|tool_calls_section_begin|>` in BaseReasoningParser with configurable `tool_section_start_markers` in ParserConfig, injected only for Kimi models that need it - Fix `used_json_schema` flag leaking into `process_tool_calls_stream` when a configured native parser should take priority — compute `force_native_parser` and clear the JSON parser flag accordingly - Same fix for Messages API streaming path: `streaming_tool_parser` now respects configured parser over JSON schema parser - Add missing `(?s)` dotall flag to KimiK2 `end_pattern` regex so multi-line JSON arguments are correctly cleared from the buffer Signed-off-by: ConnorLi96 Signed-off-by: ConnorLi96 Made-with: Cursor Signed-off-by: ConnorLi96 Made-with: Cursor Signed-off-by: ConnorLi96 Made-with: Cursor --- crates/reasoning_parser/src/factory.rs | 60 +++++++++-------- crates/reasoning_parser/src/parsers/base.rs | 67 +++++++++++++++++-- .../src/parsers/cohere_cmd.rs | 6 +- .../src/parsers/deepseek_r1.rs | 5 +- crates/reasoning_parser/src/parsers/glm45.rs | 6 +- crates/reasoning_parser/src/parsers/kimi.rs | 6 +- .../reasoning_parser/src/parsers/minimax.rs | 5 +- .../reasoning_parser/src/parsers/nano_v3.rs | 6 +- crates/reasoning_parser/src/parsers/qwen3.rs | 9 +-- crates/reasoning_parser/src/parsers/step3.rs | 5 +- crates/reasoning_parser/src/traits.rs | 10 ++- crates/tool_parser/src/parsers/kimik2.rs | 29 ++++++-- .../src/routers/grpc/regular/processor.rs | 15 ++++- .../src/routers/grpc/regular/streaming.rs | 24 ++++--- 14 files changed, 171 insertions(+), 82 deletions(-) diff --git a/crates/reasoning_parser/src/factory.rs b/crates/reasoning_parser/src/factory.rs index eeda6cbd13..598df10da4 100644 --- a/crates/reasoning_parser/src/factory.rs +++ b/crates/reasoning_parser/src/factory.rs @@ -208,35 +208,45 @@ impl ParserFactory { let config = ParserConfig { think_start_token: "".to_string(), think_end_token: "".to_string(), - stream_reasoning: true, - max_buffer_size: DEFAULT_MAX_BUFFER_SIZE, - always_in_reasoning: false, + ..Default::default() }; Box::new(BaseReasoningParser::new(config).with_model_type("deepseek_v31".to_string())) }); + let kimi_tool_markers = vec!["<|tool_calls_section_begin|>".to_string()]; + // Register Kimi-K2.5 parser (standard think tokens, always_in_reasoning=false) - registry.register_parser("kimi_k25", || { - let config = ParserConfig { - think_start_token: "".to_string(), - think_end_token: "".to_string(), - stream_reasoning: true, - max_buffer_size: DEFAULT_MAX_BUFFER_SIZE, - always_in_reasoning: false, - }; - Box::new(BaseReasoningParser::new(config).with_model_type("kimi_k25".to_string())) + registry.register_parser("kimi_k25", { + let markers = kimi_tool_markers.clone(); + move || { + let config = ParserConfig { + think_start_token: "".to_string(), + think_end_token: "".to_string(), + stream_reasoning: true, + max_buffer_size: DEFAULT_MAX_BUFFER_SIZE, + always_in_reasoning: false, + tool_section_start_markers: markers.clone(), + }; + Box::new(BaseReasoningParser::new(config).with_model_type("kimi_k25".to_string())) + } }); // Register Kimi-K2-Thinking parser (standard think tokens, always_in_reasoning=true) - registry.register_parser("kimi_thinking", || { - let config = ParserConfig { - think_start_token: "".to_string(), - think_end_token: "".to_string(), - stream_reasoning: true, - max_buffer_size: DEFAULT_MAX_BUFFER_SIZE, - always_in_reasoning: true, - }; - Box::new(BaseReasoningParser::new(config).with_model_type("kimi_thinking".to_string())) + registry.register_parser("kimi_thinking", { + let markers = kimi_tool_markers; + move || { + let config = ParserConfig { + think_start_token: "".to_string(), + think_end_token: "".to_string(), + stream_reasoning: true, + max_buffer_size: DEFAULT_MAX_BUFFER_SIZE, + always_in_reasoning: true, + tool_section_start_markers: markers.clone(), + }; + Box::new( + BaseReasoningParser::new(config).with_model_type("kimi_thinking".to_string()), + ) + } }); // Register model patterns @@ -293,9 +303,7 @@ impl ParserFactory { let config = ParserConfig { think_start_token: String::new(), think_end_token: String::new(), - stream_reasoning: true, - max_buffer_size: DEFAULT_MAX_BUFFER_SIZE, - always_in_reasoning: false, + ..Default::default() }; Box::new( BaseReasoningParser::new(config).with_model_type("passthrough".to_string()), @@ -320,9 +328,7 @@ impl ParserFactory { let config = ParserConfig { think_start_token: String::new(), think_end_token: String::new(), - stream_reasoning: true, - max_buffer_size: DEFAULT_MAX_BUFFER_SIZE, - always_in_reasoning: false, + ..Default::default() }; Box::new(BaseReasoningParser::new(config).with_model_type("passthrough".to_string())) } diff --git a/crates/reasoning_parser/src/parsers/base.rs b/crates/reasoning_parser/src/parsers/base.rs index 61420b6179..0db48c6267 100644 --- a/crates/reasoning_parser/src/parsers/base.rs +++ b/crates/reasoning_parser/src/parsers/base.rs @@ -41,6 +41,15 @@ impl BaseReasoningParser { || (self.config.think_end_token.starts_with(text) && self.config.think_end_token != text) } + + /// Find the earliest tool-section start marker in `text`. + fn find_tool_section_start(&self, text: &str) -> Option { + self.config + .tool_section_start_markers + .iter() + .filter_map(|marker| text.find(marker.as_str())) + .min() + } } impl ReasoningParser for BaseReasoningParser { @@ -63,8 +72,7 @@ impl ReasoningParser for BaseReasoningParser { .to_string(); if !processed_text.contains(&self.config.think_end_token) { - // Don't consume tool call markers as reasoning content - if let Some(tool_pos) = processed_text.find("<|tool_calls_section_begin|>") { + if let Some(tool_pos) = self.find_tool_section_start(&processed_text) { let reasoning_text = processed_text[..tool_pos].trim().to_string(); let normal_text = processed_text[tool_pos..].to_string(); return Ok(ParserResult::new(normal_text, reasoning_text)); @@ -138,8 +146,7 @@ impl ReasoningParser for BaseReasoningParser { // Continue with reasoning content if self.in_reasoning && self.config.stream_reasoning { - // Some models skip and go straight to tool calls - if let Some(tool_pos) = current_text.find("<|tool_calls_section_begin|>") { + if let Some(tool_pos) = self.find_tool_section_start(¤t_text) { let reasoning_text = current_text[..tool_pos].trim().to_string(); let normal_text = current_text[tool_pos..].to_string(); self.buffer.clear(); @@ -200,6 +207,7 @@ mod tests { stream_reasoning, max_buffer_size: DEFAULT_MAX_BUFFER_SIZE, always_in_reasoning, + ..Default::default() }; BaseReasoningParser::new(config) } @@ -380,4 +388,55 @@ mod tests { _ => panic!("Expected BufferOverflow error"), } } + + fn create_parser_with_markers() -> BaseReasoningParser { + let config = ParserConfig { + think_start_token: "".to_string(), + think_end_token: "".to_string(), + tool_section_start_markers: vec!["<|tool_calls_section_begin|>".to_string()], + ..Default::default() + }; + BaseReasoningParser::new(config) + } + + #[test] + fn test_tool_marker_stops_reasoning_non_streaming() { + let mut parser = create_parser_with_markers(); + let input = "thinking here<|tool_calls_section_begin|>tool call data"; + let result = parser.detect_and_parse_reasoning(input).unwrap(); + assert_eq!(result.reasoning_text, "thinking here"); + assert_eq!( + result.normal_text, + "<|tool_calls_section_begin|>tool call data" + ); + } + + #[test] + fn test_tool_marker_stops_reasoning_streaming() { + let mut parser = create_parser_with_markers(); + let r1 = parser + .parse_reasoning_streaming_incremental("reasoning ") + .unwrap(); + assert_eq!(r1.reasoning_text, "reasoning "); + assert!(parser.is_in_reasoning()); + + let r2 = parser + .parse_reasoning_streaming_incremental("more<|tool_calls_section_begin|>tool data") + .unwrap(); + assert_eq!(r2.reasoning_text, "more"); + assert_eq!(r2.normal_text, "<|tool_calls_section_begin|>tool data"); + assert!(!parser.is_in_reasoning()); + } + + #[test] + fn test_no_markers_does_not_stop_reasoning() { + let mut parser = create_test_parser(false, true); + let input = "thinking<|tool_calls_section_begin|>stuff"; + let result = parser.detect_and_parse_reasoning(input).unwrap(); + assert_eq!( + result.reasoning_text, + "thinking<|tool_calls_section_begin|>stuff" + ); + assert_eq!(result.normal_text, ""); + } } diff --git a/crates/reasoning_parser/src/parsers/cohere_cmd.rs b/crates/reasoning_parser/src/parsers/cohere_cmd.rs index 1d04acde22..f2b738eb12 100644 --- a/crates/reasoning_parser/src/parsers/cohere_cmd.rs +++ b/crates/reasoning_parser/src/parsers/cohere_cmd.rs @@ -5,7 +5,7 @@ use crate::{ parsers::BaseReasoningParser, - traits::{ParseError, ParserConfig, ParserResult, ReasoningParser, DEFAULT_MAX_BUFFER_SIZE}, + traits::{ParseError, ParserConfig, ParserResult, ReasoningParser}, }; /// Cohere Command model reasoning parser. @@ -22,9 +22,7 @@ impl CohereCmdParser { let config = ParserConfig { think_start_token: "<|START_THINKING|>".to_string(), think_end_token: "<|END_THINKING|>".to_string(), - stream_reasoning: true, - max_buffer_size: DEFAULT_MAX_BUFFER_SIZE, - always_in_reasoning: false, + ..Default::default() }; Self { diff --git a/crates/reasoning_parser/src/parsers/deepseek_r1.rs b/crates/reasoning_parser/src/parsers/deepseek_r1.rs index 60036ed7d5..d937446f85 100644 --- a/crates/reasoning_parser/src/parsers/deepseek_r1.rs +++ b/crates/reasoning_parser/src/parsers/deepseek_r1.rs @@ -4,7 +4,7 @@ use crate::{ parsers::BaseReasoningParser, - traits::{ParseError, ParserConfig, ParserResult, ReasoningParser, DEFAULT_MAX_BUFFER_SIZE}, + traits::{ParseError, ParserConfig, ParserResult, ReasoningParser}, }; /// DeepSeek-R1 reasoning parser. @@ -21,9 +21,8 @@ impl DeepSeekR1Parser { let config = ParserConfig { think_start_token: "".to_string(), think_end_token: "".to_string(), - stream_reasoning: true, - max_buffer_size: DEFAULT_MAX_BUFFER_SIZE, always_in_reasoning: true, + ..Default::default() }; Self { diff --git a/crates/reasoning_parser/src/parsers/glm45.rs b/crates/reasoning_parser/src/parsers/glm45.rs index 450e7b208e..dcc4128c64 100644 --- a/crates/reasoning_parser/src/parsers/glm45.rs +++ b/crates/reasoning_parser/src/parsers/glm45.rs @@ -3,7 +3,7 @@ use crate::{ parsers::BaseReasoningParser, - traits::{ParseError, ParserConfig, ParserResult, ReasoningParser, DEFAULT_MAX_BUFFER_SIZE}, + traits::{ParseError, ParserConfig, ParserResult, ReasoningParser}, }; /// GLM45 reasoning parser. @@ -20,9 +20,7 @@ impl Glm45Parser { let config = ParserConfig { think_start_token: "".to_string(), think_end_token: "".to_string(), - stream_reasoning: true, - max_buffer_size: DEFAULT_MAX_BUFFER_SIZE, - always_in_reasoning: false, + ..Default::default() }; Self { diff --git a/crates/reasoning_parser/src/parsers/kimi.rs b/crates/reasoning_parser/src/parsers/kimi.rs index 0894ed675a..441f697f4e 100644 --- a/crates/reasoning_parser/src/parsers/kimi.rs +++ b/crates/reasoning_parser/src/parsers/kimi.rs @@ -3,7 +3,7 @@ use crate::{ parsers::BaseReasoningParser, - traits::{ParseError, ParserConfig, ParserResult, ReasoningParser, DEFAULT_MAX_BUFFER_SIZE}, + traits::{ParseError, ParserConfig, ParserResult, ReasoningParser}, }; /// Kimi reasoning parser. @@ -20,9 +20,7 @@ impl KimiParser { let config = ParserConfig { think_start_token: "◁think▷".to_string(), think_end_token: "◁/think▷".to_string(), - stream_reasoning: true, - max_buffer_size: DEFAULT_MAX_BUFFER_SIZE, - always_in_reasoning: false, + ..Default::default() }; Self { diff --git a/crates/reasoning_parser/src/parsers/minimax.rs b/crates/reasoning_parser/src/parsers/minimax.rs index 7687d0bdfd..d4cd88e51e 100644 --- a/crates/reasoning_parser/src/parsers/minimax.rs +++ b/crates/reasoning_parser/src/parsers/minimax.rs @@ -4,7 +4,7 @@ use crate::{ parsers::BaseReasoningParser, - traits::{ParseError, ParserConfig, ParserResult, ReasoningParser, DEFAULT_MAX_BUFFER_SIZE}, + traits::{ParseError, ParserConfig, ParserResult, ReasoningParser}, }; /// MiniMax M2 reasoning parser. @@ -21,9 +21,8 @@ impl MiniMaxParser { let config = ParserConfig { think_start_token: "".to_string(), think_end_token: "".to_string(), - stream_reasoning: true, - max_buffer_size: DEFAULT_MAX_BUFFER_SIZE, always_in_reasoning: true, + ..Default::default() }; Self { diff --git a/crates/reasoning_parser/src/parsers/nano_v3.rs b/crates/reasoning_parser/src/parsers/nano_v3.rs index 40ce06b753..1dd2e03266 100644 --- a/crates/reasoning_parser/src/parsers/nano_v3.rs +++ b/crates/reasoning_parser/src/parsers/nano_v3.rs @@ -10,7 +10,7 @@ use crate::{ parsers::BaseReasoningParser, - traits::{ParseError, ParserConfig, ParserResult, ReasoningParser, DEFAULT_MAX_BUFFER_SIZE}, + traits::{ParseError, ParserConfig, ParserResult, ReasoningParser}, }; /// NanoV3 / Nemotron reasoning parser. @@ -24,9 +24,7 @@ impl NanoV3Parser { let config = ParserConfig { think_start_token: "".to_string(), think_end_token: "".to_string(), - stream_reasoning: true, - max_buffer_size: DEFAULT_MAX_BUFFER_SIZE, - always_in_reasoning: false, + ..Default::default() }; Self { diff --git a/crates/reasoning_parser/src/parsers/qwen3.rs b/crates/reasoning_parser/src/parsers/qwen3.rs index bdd283c024..3ed0273098 100644 --- a/crates/reasoning_parser/src/parsers/qwen3.rs +++ b/crates/reasoning_parser/src/parsers/qwen3.rs @@ -4,7 +4,7 @@ use crate::{ parsers::BaseReasoningParser, - traits::{ParseError, ParserConfig, ParserResult, ReasoningParser, DEFAULT_MAX_BUFFER_SIZE}, + traits::{ParseError, ParserConfig, ParserResult, ReasoningParser}, }; /// Qwen3 reasoning parser. @@ -21,9 +21,7 @@ impl Qwen3Parser { let config = ParserConfig { think_start_token: "".to_string(), think_end_token: "".to_string(), - stream_reasoning: true, - max_buffer_size: DEFAULT_MAX_BUFFER_SIZE, - always_in_reasoning: false, + ..Default::default() }; Self { @@ -84,9 +82,8 @@ impl QwenThinkingParser { let config = ParserConfig { think_start_token: "".to_string(), think_end_token: "".to_string(), - stream_reasoning: true, - max_buffer_size: DEFAULT_MAX_BUFFER_SIZE, always_in_reasoning: true, + ..Default::default() }; Self { diff --git a/crates/reasoning_parser/src/parsers/step3.rs b/crates/reasoning_parser/src/parsers/step3.rs index 009b740dc3..b3b153c27e 100644 --- a/crates/reasoning_parser/src/parsers/step3.rs +++ b/crates/reasoning_parser/src/parsers/step3.rs @@ -3,7 +3,7 @@ use crate::{ parsers::BaseReasoningParser, - traits::{ParseError, ParserConfig, ParserResult, ReasoningParser, DEFAULT_MAX_BUFFER_SIZE}, + traits::{ParseError, ParserConfig, ParserResult, ReasoningParser}, }; /// Step3 reasoning parser. @@ -20,9 +20,8 @@ impl Step3Parser { let config = ParserConfig { think_start_token: "".to_string(), think_end_token: "".to_string(), - stream_reasoning: true, - max_buffer_size: DEFAULT_MAX_BUFFER_SIZE, always_in_reasoning: true, + ..Default::default() }; Self { diff --git a/crates/reasoning_parser/src/traits.rs b/crates/reasoning_parser/src/traits.rs index 47689936c6..11aa21c840 100644 --- a/crates/reasoning_parser/src/traits.rs +++ b/crates/reasoning_parser/src/traits.rs @@ -127,6 +127,13 @@ pub struct ParserConfig { /// For models with a template thinking toggle, this should be `false` — /// the runtime will call `mark_reasoning_started()` when appropriate. pub always_in_reasoning: bool, + + /// Tokens that mark the start of a tool-call section in the model's output. + /// + /// When the model skips the reasoning end token and jumps straight to + /// tool calls, the reasoning parser uses these markers to stop consuming + /// content as reasoning and return it as normal text for the tool parser. + pub tool_section_start_markers: Vec, } impl Default for ParserConfig { @@ -136,7 +143,8 @@ impl Default for ParserConfig { think_end_token: "".to_string(), stream_reasoning: true, max_buffer_size: DEFAULT_MAX_BUFFER_SIZE, - always_in_reasoning: false, // Default to false (explicit reasoning) + always_in_reasoning: false, + tool_section_start_markers: Vec::new(), } } } diff --git a/crates/tool_parser/src/parsers/kimik2.rs b/crates/tool_parser/src/parsers/kimik2.rs index 4f089dfa09..491b526c7b 100644 --- a/crates/tool_parser/src/parsers/kimik2.rs +++ b/crates/tool_parser/src/parsers/kimik2.rs @@ -57,14 +57,15 @@ impl KimiK2Parser { reason = "regex patterns are compile-time string literals" )] pub fn new() -> Self { - // Supports alternative delimiters: <|func_start|>/<|func_end|>; (?s) for multi-line JSON + // (?s) for multi-line JSON args. Delimiter group is optional (trailing ?) to + // be lenient with models that omit it between the ID and JSON body. let tool_call_pattern = r"(?s)<\|tool_call_begin\|>\s*(?P[\w\.]+:\d+)\s*(?:<\|tool_call_argument_begin\|>\s*|<\|func_start\|>\s*)?(?P\{.*?\})\s*(?:<\|tool_call_end\|>|<\|func_end\|>)"; let tool_call_extractor = Regex::new(tool_call_pattern).expect("Valid regex pattern"); let stream_pattern = r"(?s)<\|tool_call_begin\|>\s*(?P[\w\.]+:\d+)\s*(?:<\|tool_call_argument_begin\|>\s*|<\|func_start\|>\s*)?(?P\{.*)"; let stream_tool_call_extractor = Regex::new(stream_pattern).expect("Valid regex pattern"); - let end_pattern = r"<\|tool_call_begin\|>.*?(?:<\|tool_call_end\|>|<\|func_end\|>)"; + let end_pattern = r"(?s)<\|tool_call_begin\|>.*?(?:<\|tool_call_end\|>|<\|func_end\|>)"; let tool_call_end_pattern = Regex::new(end_pattern).expect("Valid regex pattern"); // Robust parser for ids like "functions.search:0" or fallback "search:0" @@ -179,7 +180,11 @@ impl ToolParser for KimiK2Parser { // No tool markers detected - return all buffered content as normal text let mut normal_text = std::mem::take(&mut self.buffer); // Remove end tokens if present - for e_token in ["<|tool_calls_section_end|>", "<|tool_call_end|>", "<|func_end|>"] { + for e_token in [ + "<|tool_calls_section_end|>", + "<|tool_call_end|>", + "<|func_end|>", + ] { normal_text = normal_text.replace(e_token, ""); } return Ok(StreamingParseResult { @@ -240,8 +245,13 @@ impl ToolParser for KimiK2Parser { function_args }; - let end_pos = argument_diff.find("<|tool_call_end|>") - .or_else(|| argument_diff.find("<|func_end|>")); + let end_pos = [ + argument_diff.find("<|tool_call_end|>"), + argument_diff.find("<|func_end|>"), + ] + .into_iter() + .flatten() + .min(); let parsed_args_diff = if let Some(pos) = end_pos { &argument_diff[..pos] } else { @@ -262,8 +272,13 @@ impl ToolParser for KimiK2Parser { } } - let end_pos2 = function_args.find("<|tool_call_end|>") - .or_else(|| function_args.find("<|func_end|>")); + let end_pos2 = [ + function_args.find("<|tool_call_end|>"), + function_args.find("<|func_end|>"), + ] + .into_iter() + .flatten() + .min(); let parsed_args = if let Some(pos) = end_pos2 { &function_args[..pos] } else { diff --git a/model_gateway/src/routers/grpc/regular/processor.rs b/model_gateway/src/routers/grpc/regular/processor.rs index 362cfed344..8d6239df88 100644 --- a/model_gateway/src/routers/grpc/regular/processor.rs +++ b/model_gateway/src/routers/grpc/regular/processor.rs @@ -148,8 +148,12 @@ impl ResponseProcessor { _ => false, }; - if self.configured_tool_parser.is_some() && tool_parser_available { - // Explicitly configured parser takes priority (models may emit native tokens regardless of tool_choice) + let has_native_parser = self + .configured_tool_parser + .as_deref() + .is_some_and(|p| p != "json"); + + if has_native_parser && tool_parser_available { (tool_calls, processed_text) = self .parse_tool_calls( &processed_text, @@ -630,7 +634,12 @@ impl ResponseProcessor { Some(messages::ToolChoice::Tool { .. } | messages::ToolChoice::Any { .. }) ); - if self.configured_tool_parser.is_some() && tool_parser_available { + let has_native_parser = self + .configured_tool_parser + .as_deref() + .is_some_and(|p| p != "json"); + + if has_native_parser && tool_parser_available { (tool_calls, processed_text) = self .parse_tool_calls( &processed_text, diff --git a/model_gateway/src/routers/grpc/regular/streaming.rs b/model_gateway/src/routers/grpc/regular/streaming.rs index a0d1b2f3d8..27f958d348 100644 --- a/model_gateway/src/routers/grpc/regular/streaming.rs +++ b/model_gateway/src/routers/grpc/regular/streaming.rs @@ -377,10 +377,12 @@ impl StreamingProcessor { && tool_choice_enabled && (tool_parser_available || used_json_schema) { - let tool_chunks = if is_specific_function - && !(self.configured_tool_parser.is_some() - && tool_parser_available) - { + let force_native_parser = self + .configured_tool_parser + .as_deref() + .is_some_and(|p| p != "json") + && tool_parser_available; + let tool_chunks = if is_specific_function && !force_native_parser { Self::process_specific_function_stream( &delta, index, @@ -404,7 +406,7 @@ impl StreamingProcessor { created, system_fingerprint, history_tool_calls_count, - used_json_schema, + used_json_schema && !force_native_parser, ) .await }; @@ -1637,10 +1639,16 @@ impl StreamingProcessor { .map(message_utils::extract_chat_tools) .unwrap_or_default(); + let force_native_parser = self + .configured_tool_parser + .as_deref() + .is_some_and(|p| p != "json") + && tool_parser_available; + // Create fresh streaming tool parser (not pooled — streaming parsers maintain state) let mut streaming_tool_parser: Option> = if has_tools && tool_choice_enabled && (tool_parser_available || used_json_schema) { - let parser_name = if used_json_schema { + let parser_name = if used_json_schema && !force_native_parser { Some("json") } else { self.configured_tool_parser.as_deref() @@ -1754,9 +1762,7 @@ impl StreamingProcessor { // Tool call handling: incremental streaming parser if !in_reasoning && streaming_tool_parser.is_some() { - if is_specific_function - && !(self.configured_tool_parser.is_some() && tool_parser_available) - { + if is_specific_function && !force_native_parser { if !has_tool_calls { has_tool_calls = true; // Close text block if open before starting tool block