Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions bindings/python/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -526,6 +526,7 @@ struct Router {
worker_overload_token_usage: Option<f64>,
worker_overload_protection: bool,
disable_load_monitoring: bool,
tool_choice_none_ban: bool,
}

impl Router {
Expand Down Expand Up @@ -900,6 +901,7 @@ impl Router {
.maybe_redis(redis_config)
.maybe_reasoning_parser(self.reasoning_parser.as_ref())
.maybe_tool_call_parser(self.tool_call_parser.as_ref())
.tool_choice_none_ban(self.tool_choice_none_ban)
.maybe_mcp_config_path(self.mcp_config_path.as_ref())
.maybe_storage_hook_wasm_path(self.storage_hook_wasm_path.as_deref())
.enable_wasm(self.enable_wasm)
Expand Down Expand Up @@ -1083,6 +1085,7 @@ impl Router {
worker_overload_token_usage = None,
worker_overload_protection = false,
disable_load_monitoring = false,
tool_choice_none_ban = false,
))]
#[expect(clippy::too_many_arguments)]
#[expect(
Expand Down Expand Up @@ -1233,6 +1236,7 @@ impl Router {
worker_overload_token_usage: Option<f64>,
worker_overload_protection: bool,
disable_load_monitoring: bool,
tool_choice_none_ban: bool,
) -> PyResult<Self> {
let mut all_urls = worker_urls.clone();

Expand Down Expand Up @@ -1397,6 +1401,7 @@ impl Router {
worker_overload_token_usage,
worker_overload_protection,
disable_load_monitoring,
tool_choice_none_ban,
})
}

Expand Down
11 changes: 11 additions & 0 deletions bindings/python/src/smg/router_args.py
Original file line number Diff line number Diff line change
Expand Up @@ -242,6 +242,7 @@ class RouterArgs:
worker_overload_protection: bool = False
# Restore the conditional load-monitor poll gate (default: poll always)
disable_load_monitoring: bool = False
tool_choice_none_ban: bool = False

@staticmethod
def add_cli_args(
Expand Down Expand Up @@ -1280,6 +1281,16 @@ def add_cli_args(
choices=tool_call_parser_choices,
help="Specify the parser for tool-call interactions (e.g., json, qwen)",
)
parser_group.add_argument(
f"--{prefix}tool-choice-none-ban",
action="store_true",
help=(
"With tools present but tool_choice 'none', ban the resolved"
" parser's tool-call opener strings at decode time (requires"
" engine support for the any_text/excludes structural-tag"
" format)"
),
)
Comment thread
coderabbitai[bot] marked this conversation as resolved.
parser_group.add_argument(
f"--{prefix}mcp-config-path",
type=str,
Expand Down
1 change: 1 addition & 0 deletions bindings/python/tests/test_arg_parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -1401,6 +1401,7 @@ class TestRouterArgsFieldOrder:
"worker_overload_token_usage",
"worker_overload_protection",
"disable_load_monitoring",
"tool_choice_none_ban",
]

def test_complete_field_sequence_is_frozen(self):
Expand Down
40 changes: 40 additions & 0 deletions crates/tool_parser/src/factory.rs
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,12 @@ impl ToolConstraint {
struct ParserEntry {
creator: ParserCreator,
build_structural_tag: Option<BuildStructuralTagFn>,
/// Strings that exclusively open this parser's native tool-call syntax,
/// used to build the `tool_choice: "none"` suppression constraint. Empty
/// for parsers without one (no suppression possible). Curated per parser:
/// this is NOT the structural-tag trigger list — triggers may include
/// markers ordinary generation must emit.
tool_call_ban_strings: &'static [&'static str],
}

/// Registry for model-specific tool parsers with pooling support.
Expand Down Expand Up @@ -97,6 +103,7 @@ impl ParserRegistry {
Arc::new(ParserEntry {
creator: Arc::new(creator),
build_structural_tag: None,
tool_call_ban_strings: &[],
}),
);
}
Expand All @@ -105,11 +112,16 @@ impl ParserRegistry {
///
/// The `build_structural_tag` function takes `(&[Tool], at_least_one)` and returns
/// the full xgrammar structural tag JSON for this parser's native tool-call format.
/// `tool_call_ban_strings` is the parser's curated `tool_choice: "none"`
/// suppression inventory: strings that exclusively open its native
/// tool-call syntax (NOT the trigger list — triggers may include markers
/// ordinary generation must emit). Pass an empty slice to opt out.
pub fn register_parser_with_structural_tag<F>(
&self,
name: &str,
creator: F,
build_structural_tag: fn(&[Tool], bool) -> serde_json::Value,
tool_call_ban_strings: &'static [&'static str],
) where
F: Fn() -> Box<dyn ToolParser> + Send + Sync + 'static,
{
Expand All @@ -119,6 +131,7 @@ impl ParserRegistry {
Arc::new(ParserEntry {
creator: Arc::new(creator),
build_structural_tag: Some(Arc::new(build_structural_tag)),
tool_call_ban_strings,
}),
);
}
Expand Down Expand Up @@ -183,6 +196,29 @@ impl ParserRegistry {
configured.is_some_and(|p| self.has_structural_tag(p))
}

/// Build the decode-time suppression constraint for `tool_choice: "none"`:
/// free-form output that may never contain the strings opening the
/// configured parser's native tool-call syntax, expressed as an xgrammar
/// structural tag (`any_text` with `excludes`).
///
/// Returns `None` when no parser is configured, the name is unknown, or
/// the parser has no curated ban inventory — callers fall back to today's
/// behavior (no constraint; parsing is separately disabled for `"none"`).
pub fn tool_call_ban_constraint(&self, configured: Option<&str>) -> Option<ToolConstraint> {
let entries = self.entries.read();
let entry = entries.get(configured?)?;
if entry.tool_call_ban_strings.is_empty() {
return None;
}
let tag = serde_json::json!({
"format": {
"type": "any_text",
"excludes": entry.tool_call_ban_strings,
}
});
Some(ToolConstraint::StructuralTag(tag.to_string()))
}

/// Generate tool call constraint.
///
/// If `configured_parser` supports structural tags → `StructuralTag(json)`.
Expand Down Expand Up @@ -314,6 +350,7 @@ impl ParserFactory {
"mistral",
|| Box::new(MistralParser::new()),
MistralParser::build_structural_tag,
MistralParser::TOOL_CALL_BAN_STRINGS,
);
registry.register_parser("qwen", || Box::new(QwenParser::new()));
registry.register_parser("qwen_xml", || Box::new(QwenXmlParser::new()));
Expand All @@ -335,16 +372,19 @@ impl ParserFactory {
"kimik2",
|| Box::new(KimiK2Parser::new()),
KimiK2Parser::build_structural_tag,
KimiK2Parser::TOOL_CALL_BAN_STRINGS,
);
registry.register_parser_with_structural_tag(
"kimi_k3",
|| Box::new(KimiK3Parser::new()),
KimiK3Parser::build_structural_tag,
KimiK3Parser::TOOL_CALL_BAN_STRINGS,
);
registry.register_parser_with_structural_tag(
"inkling",
|| Box::new(InklingParser::new()),
InklingParser::build_structural_tag,
InklingParser::TOOL_CALL_BAN_STRINGS,
);
registry.register_parser("minimax_m2", || Box::new(MinimaxM2Parser::new()));
registry.register_parser("cohere", || Box::new(CohereParser::new()));
Expand Down
6 changes: 6 additions & 0 deletions crates/tool_parser/src/parsers/inkling.rs
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,12 @@ impl InklingParser {
}
}

/// Strings that only ever open Inkling's native tool-call syntax; banning
/// them makes tool calls unreachable when `tool_choice` is `"none"`. Both
/// invocation modes are covered: JSON-arguments and TML text-mode.
pub const TOOL_CALL_BAN_STRINGS: &'static [&'static str] =
&[TOOL_CALL_JSON_START, TOOL_CALL_TEXT_START];

/// Build an xgrammar structural tag that constrains the JSON arguments for
/// each declared tool while retaining Inkling's native TML framing.
pub fn build_structural_tag(tools: &[Tool], at_least_one: bool) -> Value {
Expand Down
8 changes: 8 additions & 0 deletions crates/tool_parser/src/parsers/kimi_k3.rs
Original file line number Diff line number Diff line change
Expand Up @@ -123,6 +123,14 @@ pub struct KimiK3Parser {
}

impl KimiK3Parser {
/// Strings that only ever open K3's native tool-call syntax; banning them
/// makes tool calls unreachable when `tool_choice` is `"none"`. ONLY the
/// tools-section opener qualifies: the structural-tag triggers also list
/// `<|close|>think<|sep|>` / `<|close|>response<|sep|>` (as tag-begin
/// prefixes), but those close sections ordinary generation must emit, so
/// excluding them would corrupt normal output.
pub const TOOL_CALL_BAN_STRINGS: &'static [&'static str] = &[TOOLS_OPEN];

/// Build an xgrammar structural tag that constrains Kimi-K3 XTML tool
/// calls to the declared `tools`.
///
Expand Down
7 changes: 7 additions & 0 deletions crates/tool_parser/src/parsers/kimik2.rs
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,13 @@ pub struct KimiK2Parser {
}

impl KimiK2Parser {
/// Strings that only ever open K2's native tool-call syntax; banning them
/// makes tool calls unreachable when `tool_choice` is `"none"`. Both the
/// section opener and the per-call opener appear exclusively in tool-call
/// output, so both are safe to exclude.
pub const TOOL_CALL_BAN_STRINGS: &'static [&'static str] =
&["<|tool_calls_section_begin|>", "<|tool_call_begin|>"];

/// Build structural tag for Kimi K2 tool call format.
///
/// Uses dual triggers following sglang's approach:
Expand Down
4 changes: 4 additions & 0 deletions crates/tool_parser/src/parsers/mistral.rs
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,10 @@ pub struct MistralParser {
}

impl MistralParser {
/// Strings that only ever open Mistral's native tool-call syntax; banning
/// them makes tool calls unreachable when `tool_choice` is `"none"`.
pub const TOOL_CALL_BAN_STRINGS: &'static [&'static str] = &["[TOOL_CALLS]"];

/// Build structural tag for Mistral tool call format.
///
/// Mistral outputs tool calls as a JSON array after `[TOOL_CALLS]`:
Expand Down
126 changes: 126 additions & 0 deletions crates/tool_parser/tests/tool_constraint_ban.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,126 @@
#![expect(
clippy::expect_used,
reason = "test-only helpers outside #[test] fns; failures are test failures"
)]

//! Suppression-constraint coverage: the `tool_choice: "none"` ban tag.
//!
//! The ban is a structural tag whose format is free text excluding the
//! strings that open a parser's native tool-call syntax. Only parsers with
//! model-native framing (and a curated opener inventory) produce one.

use serde_json::Value;
use tool_parser::ParserFactory;

fn ban_tag_json(factory: &ParserFactory, parser: &str) -> Value {
let constraint = factory
.registry()
.tool_call_ban_constraint(Some(parser))
.expect("parser should produce a ban constraint");
let (kind, json) = constraint.to_tuple();
assert_eq!(kind, "structural_tag", "ban constraint for {parser}");
serde_json::from_str(&json).expect("ban tag must be valid JSON")
}

fn excludes(tag: &Value) -> Vec<&str> {
tag["format"]["excludes"]
.as_array()
.expect("excludes must be an array")
.iter()
.map(|v| v.as_str().expect("excludes entries must be strings"))
.collect()
}

#[test]
fn ban_tag_shape_is_any_text_with_excludes() {
let factory = ParserFactory::new();
let tag = ban_tag_json(&factory, "mistral");
assert_eq!(tag["format"]["type"], "any_text");
assert!(tag["format"]["excludes"].is_array());
// Top-level envelope matches the positive structural tags: one "format" key.
assert!(tag.get("format").is_some());
assert_eq!(tag.as_object().map(serde_json::Map::len), Some(1));
}

#[test]
fn curated_ban_inventories_per_parser() {
let factory = ParserFactory::new();
assert_eq!(
excludes(&ban_tag_json(&factory, "mistral")),
vec!["[TOOL_CALLS]"]
);
assert_eq!(
excludes(&ban_tag_json(&factory, "kimik2")),
vec!["<|tool_calls_section_begin|>", "<|tool_call_begin|>"]
);
// K3 bans ONLY the tools-section opener: its structural-tag triggers also
// include think/response section closers that ordinary generation must
// emit, and excluding those would corrupt normal output.
assert_eq!(
excludes(&ban_tag_json(&factory, "kimi_k3")),
vec!["<|open|>tools<|sep|>"]
);
// Both Inkling invocation modes: JSON-arguments and TML text-mode.
assert_eq!(
excludes(&ban_tag_json(&factory, "inkling")),
vec![
"<|content_invoke_tool_json|>",
"<|content_invoke_tool_text|>"
]
);
}

#[test]
fn parsers_without_native_framing_produce_no_ban() {
let factory = ParserFactory::new();
for parser in ["json", "qwen", "qwen_xml", "pythonic", "llama", "deepseek"] {
assert!(
factory
.registry()
.tool_call_ban_constraint(Some(parser))
.is_none(),
"{parser} has no curated ban inventory"
);
}
}

#[test]
fn unknown_or_absent_parser_produces_no_ban() {
let factory = ParserFactory::new();
assert!(factory
.registry()
.tool_call_ban_constraint(Some("no-such-parser"))
.is_none());
assert!(factory.registry().tool_call_ban_constraint(None).is_none());
}

#[test]
fn generate_tool_constraint_still_skips_none_choice() {
use openai_protocol::common::{Function, Tool, ToolChoice, ToolChoiceValue};

let factory = ParserFactory::new();
let tools = vec![Tool {
tool_type: "function".to_string(),
function: Function {
name: "get_weather".to_string(),
description: None,
parameters: serde_json::json!({"type": "object"}),
strict: None,
},
}];
// The ban is a separate opt-in surface; the standard constraint generator
// keeps returning no constraint for tool_choice "none" and "auto".
for choice in [
ToolChoice::Value(ToolChoiceValue::None),
ToolChoice::Value(ToolChoiceValue::Auto),
] {
let constraint = factory
.registry()
.generate_tool_constraint(Some("mistral"), &tools, &choice)
.expect("constraint generation must not error");
assert!(
constraint.is_none(),
"no constraint expected for {choice:?}"
);
}
}
5 changes: 5 additions & 0 deletions model_gateway/src/config/builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -613,6 +613,11 @@ impl RouterConfigBuilder {
self
}

pub fn tool_choice_none_ban(mut self, enable: bool) -> Self {
self.config.tool_choice_none_ban = enable;
self

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Nit: The Python bindings at bindings/python/src/lib.rs build RouterConfig via this builder but don't set tool_choice_none_ban — so Python-binding users have no way to enable the feature. The default (false) is safe, but REVIEW.md flags config changes as the #1 bug source precisely because of these cross-surface gaps. Consider adding the field to the Python RouterArgs dataclass and the PyO3 struct to keep the surfaces in sync, even if it's a follow-up.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done — the flag is now exposed end to end in the Python bindings: RouterArgs.tool_choice_none_ban + --tool-choice-none-ban (argparse), and the PyO3 Router constructor passes it through .tool_choice_none_ban(...) in the config conversion. All additions appended at list tails per the positional-compat rule.

}

// ==================== Tokenizer Cache ====================

pub fn tokenizer_cache(mut self, cache: TokenizerCacheConfig) -> Self {
Expand Down
Loading
Loading