From 7acfc161ec0d068e2856624dccf7cede9c07de2c Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Mon, 6 Jul 2026 16:59:48 +0800 Subject: [PATCH] add tool schema resolve Signed-off-by: Bugen Zhao --- rust/src/parser/src/tool/mod.rs | 1 + rust/src/parser/src/tool/parameters.rs | 19 ++++++++---- .../src/unified/config_driven/gemma4.rs | 3 +- .../parser/src/unified/config_driven/mod.rs | 31 +++++++++++++++---- 4 files changed, 41 insertions(+), 13 deletions(-) diff --git a/rust/src/parser/src/tool/mod.rs b/rust/src/parser/src/tool/mod.rs index 5a06f2311ed..9ba4433da92 100644 --- a/rust/src/parser/src/tool/mod.rs +++ b/rust/src/parser/src/tool/mod.rs @@ -28,6 +28,7 @@ pub use json::{ pub use kimi_k2::KimiK2ToolParser; pub use minimax_m2::MinimaxM2ToolParser; pub use minimax_m3::MinimaxM3ToolParser; +pub use parameters::{ToolSchema, ToolSchemas}; pub use qwen_coder::Qwen3CoderToolParser; use serde::{Deserialize, Serialize}; use serde_json::Value; diff --git a/rust/src/parser/src/tool/parameters.rs b/rust/src/parser/src/tool/parameters.rs index f9abb50f8b1..1cb3901d52b 100644 --- a/rust/src/parser/src/tool/parameters.rs +++ b/rust/src/parser/src/tool/parameters.rs @@ -6,7 +6,7 @@ use crate::tool::Tool; /// Normalized parameter schemas for all tools in one request. #[derive(Debug, Clone, Default, PartialEq, Eq)] -pub(super) struct ToolSchemas { +pub struct ToolSchemas { tools: BTreeMap, } @@ -17,7 +17,7 @@ pub(super) struct ToolSchemas { /// coercing raw string parameter values into more specific JSON types for /// downstream tool call execution. #[derive(Debug, Clone, Default, PartialEq, Eq)] -pub(super) struct ToolSchema { +pub struct ToolSchema { params: BTreeMap, } @@ -64,7 +64,7 @@ pub(super) enum JsonParamType { impl ToolSchemas { /// Normalize OpenAI-style tool parameter JSON schemas for one request. - pub(super) fn from_tools(tools: &[Tool]) -> Self { + pub(crate) fn from_tools(tools: &[Tool]) -> Self { let tools = tools .iter() .map(|tool| (tool.name.clone(), ToolSchema::from_schema(&tool.parameters))) @@ -73,6 +73,14 @@ impl ToolSchemas { Self { tools } } + /// Resolve the parameter schema for one named tool. + /// + /// Unknown tool names resolve to the empty schema, so all parameters fall + /// back to strings or object-like JSON for structured inputs. + pub(crate) fn resolve(&self, function_name: &str) -> &ToolSchema { + self.tools.get(function_name).unwrap_or(ToolSchema::empty()) + } + /// Convert parameter values for one named tool. /// /// Unknown tool names use an empty schema, so all parameters fall back to @@ -85,7 +93,7 @@ impl ToolSchemas { where P: Into, { - let tool_schema = self.tools.get(function_name).unwrap_or(ToolSchema::empty()); + let tool_schema = self.resolve(function_name); let mut converted = Map::with_capacity(params.len()); for (name, value) in params { let value = tool_schema.convert(&name, value.into()); @@ -104,8 +112,7 @@ impl ToolSchemas { where P: Into, { - let tool_schema = self.tools.get(function_name).unwrap_or(ToolSchema::empty()); - tool_schema.convert(name, value.into()) + self.resolve(function_name).convert(name, value.into()) } } diff --git a/rust/src/parser/src/unified/config_driven/gemma4.rs b/rust/src/parser/src/unified/config_driven/gemma4.rs index 817296b8071..618cc41556d 100644 --- a/rust/src/parser/src/unified/config_driven/gemma4.rs +++ b/rust/src/parser/src/unified/config_driven/gemma4.rs @@ -18,6 +18,7 @@ use winnow::stream::Stream; use winnow::token::{literal, take_till, take_until}; use super::{ArgsEndScan, ConfigDrivenParser, Input, ParserFormat}; +use crate::tool::ToolSchema; use crate::utils::{incomplete, partial_prefix_len}; const CHANNEL_START: &str = "<|channel>"; @@ -52,7 +53,7 @@ impl ParserFormat for Gemma4Format { Ok(name) } - fn tool_args(body: &str) -> ModalResult> { + fn tool_args(_schema: &ToolSchema, body: &str) -> ModalResult> { let Some(args_input) = body.strip_suffix('}') else { return Err(ErrMode::Cut(ContextError::new())); }; diff --git a/rust/src/parser/src/unified/config_driven/mod.rs b/rust/src/parser/src/unified/config_driven/mod.rs index 8d611a145e6..d28aaf8d74c 100644 --- a/rust/src/parser/src/unified/config_driven/mod.rs +++ b/rust/src/parser/src/unified/config_driven/mod.rs @@ -42,7 +42,7 @@ use winnow::token::{literal, rest}; use super::{Result, UnifiedParser, UnifiedParserError, UnifiedParserOutput}; use crate::reasoning::last_reasoning_boundary; -use crate::tool::{Tool, ToolCallDelta}; +use crate::tool::{Tool, ToolCallDelta, ToolSchema, ToolSchemas}; use crate::unified::parsing_failed; use crate::utils::{MarkerScanState, parse_buffered_event, safe_text_len_mul, take_until_marker}; @@ -87,6 +87,12 @@ pub trait ParserFormat: 'static { /// Whether decoded output must keep tokenizer special tokens. const PRESERVE_SPECIAL_TOKENS: bool = true; + /// Whether to normalize the request tools' parameter JSON schemas at + /// creation for schema-aware argument conversion. When `false`, the + /// engine skips schema construction and `tool_args` resolves to the + /// empty schema (all raw values keep their string-level interpretation). + const USES_TOOL_SCHEMAS: bool = false; + /// Resumable scanner locating `CALL_END` for the raw argument body. type ArgsScan: ArgsEndScan; @@ -96,7 +102,10 @@ pub trait ParserFormat: 'static { /// Parse one complete raw argument body (end marker already stripped) /// into a JSON object. - fn tool_args(body: &str) -> ModalResult>; + /// + /// `schema` is the parameter schema resolved for this call's tool name — + /// the empty schema unless `USES_TOOL_SCHEMAS` is enabled. + fn tool_args(schema: &ToolSchema, body: &str) -> ModalResult>; } /// Resumable scanner that locates a tool-call end marker in buffered input. @@ -229,6 +238,7 @@ pub struct ConfigDrivenParser { buffer: String, mode: Mode, markers: ModeMarkerSet, + tool_schemas: ToolSchemas, emitted_tool_count: usize, tokenizer: DynTokenizer, /// Resolved reasoning boundary token IDs for prompt-state derivation. @@ -238,7 +248,13 @@ pub struct ConfigDrivenParser { impl ConfigDrivenParser { /// Create a parser for one request stream. - pub fn new(_tools: &[Tool], tokenizer: DynTokenizer) -> Result { + pub fn new(tools: &[Tool], tokenizer: DynTokenizer) -> Result { + let tool_schemas = if F::USES_TOOL_SCHEMAS { + ToolSchemas::from_tools(tools) + } else { + ToolSchemas::default() + }; + let boundary_ids = match F::REASONING_BOUNDARY_TOKENS { Some((start, end)) => { let start_id = tokenizer.token_to_id(start).ok_or_else(|| { @@ -259,6 +275,7 @@ impl ConfigDrivenParser { buffer: String::new(), mode: Mode::Text, markers: ModeMarkerSet::of::(), + tool_schemas, emitted_tool_count: 0, tokenizer, boundary_ids, @@ -389,7 +406,7 @@ impl UnifiedParser for ConfigDrivenParser { while let Some((event, consumed_len)) = { parse_buffered_event(&self.buffer, |input| { - parse_next_event::(input, &mut self.mode, &self.markers) + parse_next_event::(input, &mut self.mode, &self.markers, &self.tool_schemas) })? } { self.apply_event(event, consumed_len, output)?; @@ -430,13 +447,14 @@ fn parse_next_event( input: &mut Input<'_>, mode: &mut Mode, markers: &ModeMarkerSet, + schemas: &ToolSchemas, ) -> ModalResult { match mode { Mode::Text => boundary_or_content(input, &markers.text, Event::Text), Mode::Reasoning => boundary_or_content(input, &markers.reasoning, Event::Reasoning), Mode::ToolBetween => boundary_or_content(input, &markers.between, Event::Ignored), Mode::ToolHeader => F::tool_header(input).map(|name| Event::CallHeader { name }), - Mode::ToolArgs { scan, .. } => args_event::(input, scan), + Mode::ToolArgs { name, scan, .. } => args_event::(input, scan, schemas.resolve(name)), Mode::Done => rest.value(Event::Ignored).parse_next(input), } } @@ -445,9 +463,10 @@ fn parse_next_event( fn args_event( input: &mut Input<'_>, scan: &mut F::ArgsScan, + schema: &ToolSchema, ) -> ModalResult { let body = scan.scan(input, F::CALL_END)?; - let args = F::tool_args(body)?; + let args = F::tool_args(schema, body)?; Ok(Event::CallComplete { args }) }