@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 1c23c42030

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

[Rust Frontend] Support Kimi K2 tool call IDs (#44901)

Maria Guevara committed Jun 9, 2026 at 13:31 UTC 1c23c4203007d16d96b1f07847f5bf1b663ad0d7

5 files changed +147 -8

rust/src/chat/src/output/default/tool.rs

+103 -1

@@ -121,7 +121,11 @@ impl ToolState {

121 None => true,

122 };

123 if is_new_tool {

124 - let id = generate_tool_call_id();

124 + let id = self

125 + .parser

126 + .tool_call_id(item.tool_index)

127 + .map(str::to_string)

128 + .unwrap_or_else(generate_tool_call_id);

129 self.open_call_index = Some(item.tool_index);

130 events.push(AssistantEvent::ToolCallStart { id, name });

131 }

@@ -291,6 +295,11 @@ mod tests {

295 buffered: String,

296 }

297

298 + struct IdScriptedParser {

299 + output: ToolParserOutput,

300 + tool_call_id: Option<String>,

301 + }

302 +

303 impl ToolParser for FailingParser {

304 fn create(_tools: &[ChatTool]) -> vllm_tool_parser::Result<Box<dyn ToolParser>>

305 where

@@ -351,6 +360,35 @@ mod tests {

360 }

361 }

362

363 + impl ToolParser for IdScriptedParser {

364 + fn create(_tools: &[ChatTool]) -> vllm_tool_parser::Result<Box<dyn ToolParser>>

365 + where

366 + Self: Sized + 'static,

367 + {

368 + Ok(Box::new(Self {

369 + output: ToolParserOutput::default(),

370 + tool_call_id: None,

371 + }))

372 + }

373 +

374 + fn tool_call_id(&self, tool_index: usize) -> Option<&str> {

375 + (tool_index == 0).then_some(self.tool_call_id.as_deref()).flatten()

376 + }

377 +

378 + fn parse_into(&mut self, _chunk: &str, output: &mut ToolParserOutput) -> Result<()> {

379 + output.append(std::mem::take(&mut self.output));

380 + Ok(())

381 + }

382 +

383 + fn finish(&mut self) -> Result<ToolParserOutput> {

384 + Ok(ToolParserOutput::default())

385 + }

386 +

387 + fn reset(&mut self) -> String {

388 + String::new()

389 + }

390 + }

391 +

392 impl ToolParser for PartialThenFailParser {

393 fn create(_tools: &[ChatTool]) -> vllm_tool_parser::Result<Box<dyn ToolParser>>

394 where

@@ -505,6 +543,70 @@ mod tests {

543 assert!(matches!(events[3], AssistantEvent::Done { .. }));

544 }

545

546 + #[tokio::test]

547 + async fn tool_stream_preserves_parser_provided_tool_call_id() {

548 + let events = stream::iter(vec![Ok(ContentEvent::TextDelta {

549 + kind: AssistantBlockKind::Text,

550 + delta: "ignored".to_string(),

551 + })]);

552 + let parser = IdScriptedParser {

553 + output: ToolParserOutput {

554 + normal_text: String::new(),

555 + calls: vec![crate::parser::tool::ToolCallDelta {

556 + tool_index: 0,

557 + name: Some("get_weather".to_string()),

558 + arguments: "{}".to_string(),

559 + }],

560 + },

561 + tool_call_id: Some("functions.get_weather:0".to_string()),

562 + };

563 +

564 + let events = tool_event_stream(events, Some(Box::new(parser)))

565 + .collect::<Vec<_>>()

566 + .await

567 + .into_iter()

568 + .collect::<crate::Result<Vec<_>>>()

569 + .unwrap();

570 +

571 + assert!(matches!(

572 + &events[0],

573 + AssistantEvent::ToolCallStart { id, name }

574 + if id == "functions.get_weather:0" && name == "get_weather"

575 + ));

576 + }

577 +

578 + #[tokio::test]

579 + async fn tool_stream_generates_tool_call_id_when_parser_omits_one() {

580 + let events = stream::iter(vec![Ok(ContentEvent::TextDelta {

581 + kind: AssistantBlockKind::Text,

582 + delta: "ignored".to_string(),

583 + })]);

584 + let parser = IdScriptedParser {

585 + output: ToolParserOutput {

586 + normal_text: String::new(),

587 + calls: vec![crate::parser::tool::ToolCallDelta {

588 + tool_index: 0,

589 + name: Some("get_weather".to_string()),

590 + arguments: "{}".to_string(),

591 + }],

592 + },

593 + tool_call_id: None,

594 + };

595 +

596 + let events = tool_event_stream(events, Some(Box::new(parser)))

597 + .collect::<Vec<_>>()

598 + .await

599 + .into_iter()

600 + .collect::<crate::Result<Vec<_>>>()

601 + .unwrap();

602 +

603 + assert!(matches!(

604 + &events[0],

605 + AssistantEvent::ToolCallStart { id, name }

606 + if id.starts_with("call_") && name == "get_weather"

607 + ));

608 + }

609 +

610 #[tokio::test]

611 async fn real_buffered_parser_error_matches_streaming_and_non_streaming() {

612 let prefix = "I will check both.\n";

rust/src/chat/src/output/mod.rs

-2

@@ -128,8 +128,6 @@ trait_set! {

128

129 /// Generate the northbound tool-call ID using the OpenAI-style `call_<id>`

130 /// format.

131 -// TODO: support other ID scheme like Kimi-K2's

132 -// `functions.{name}:{global_index}`.

131 pub(crate) fn generate_tool_call_id() -> String {

132 format!("call_{}", &Uuid::new_v4().simple().to_string()[..24])

133 }

rust/src/chat/tests/roundtrip.rs

+1 -2

@@ -170,8 +170,7 @@ roundtrip_tests! {

170 glm47 => [reasoning_and_content, tool_call_mix],

171

172 // Note: Kimi K2.5 strips the reasoning content in history.

173 - // TODO: we don't respect model-generated tool call id now so `tool_call_mix` cannot pass.

174 - // kimi_k25 => [tool_call_mix],

173 + kimi_k25 => [tool_call_mix],

174 }

175

176 /// Run the fixed reasoning+content fixture for one model/parser case.

rust/src/tool-parser/src/kimi_k2.rs

+37 -3

@@ -1,3 +1,5 @@

1 +use std::collections::BTreeMap;

2 +

3 use winnow::ascii::{digit1, multispace0 as ws0};

4 use winnow::combinator::{alt, eof, repeat, seq};

5 use winnow::prelude::*;

@@ -33,6 +35,7 @@ enum KimiK2Event {

35 ToolCallsStart,

36 ToolCallStart,

37 ToolCallHeader {

38 + tool_call_id: String,

39 function_name: String,

40 function_index: usize,

41 },

@@ -60,6 +63,7 @@ pub struct KimiK2ToolParser {

63 buffer: String,

64 mode: KimiK2Mode,

65 active_tool_index: Option<usize>,

66 + call_ids: BTreeMap<usize, String>,

67 }

69 impl KimiK2ToolParser {

@@ -69,6 +73,7 @@ impl KimiK2ToolParser {

73 buffer: String::new(),

74 mode: KimiK2Mode::Text,

75 active_tool_index: None,

76 + call_ids: BTreeMap::new(),

77 }

78 }

@@ -81,6 +86,7 @@ impl KimiK2ToolParser {

86 KimiK2Event::ToolCallsStart => self.mode = KimiK2Mode::ToolBlock,

87 KimiK2Event::ToolCallStart => self.mode = KimiK2Mode::Header,

88 KimiK2Event::ToolCallHeader {

89 + tool_call_id,

90 function_name,

91 function_index,

92 } => {

@@ -89,6 +95,7 @@ impl KimiK2ToolParser {

95 self.mode = KimiK2Mode::Arguments {

96 json_scan: JsonObjectScanState::default(),

97 };

98 + self.call_ids.insert(tool_index, tool_call_id);

99 output.calls.push(ToolCallDelta {

100 tool_index,

101 name: Some(function_name),

@@ -123,6 +130,7 @@ impl KimiK2ToolParser {

130 fn reset(&mut self) -> String {

131 self.mode = KimiK2Mode::Text;

132 self.active_tool_index = None;

133 + self.call_ids.clear();

134 std::mem::take(&mut self.buffer)

135 }

136 }

@@ -139,6 +147,10 @@ impl ToolParser for KimiK2ToolParser {

147 true

148 }

149

150 + fn tool_call_id(&self, tool_index: usize) -> Option<&str> {

151 + self.call_ids.get(&tool_index).map(String::as_str)

152 + }

153 +

154 fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {

155 self.buffer.push_str(chunk);

156

@@ -232,16 +244,18 @@ fn tool_call_end_event(input: &mut KimiK2Input<'_>) -> ModalResult<KimiK2Event>

244

245 /// Parse a Kimi K2 tool-call header before the argument marker.

246 fn tool_call_header_event(input: &mut KimiK2Input<'_>) -> ModalResult<KimiK2Event> {

235 - let (header, _) = (

247 + let (raw_header, _) = (

248 take_until(1.., TOOL_CALL_ARGUMENT_START),

249 literal(TOOL_CALL_ARGUMENT_START),

250 )

251 .parse_next(input)?;

252

241 - let mut header_input = header;

253 + let tool_call_id = raw_header.trim().to_string();

254 + let mut header_input = raw_header;

255 let (header, _, _) = (tool_header, ws0, eof).parse_next(&mut header_input)?;

256

257 Ok(KimiK2Event::ToolCallHeader {

258 + tool_call_id,

259 function_name: header.function_name,

260 function_index: header.function_index,

261 })

@@ -502,6 +516,25 @@ mod tests {

516 .assert_debug_eq(&output);

517 }

518

519 + #[test]

520 + fn kimi_k2_preserves_model_generated_tool_call_ids() {

521 + let mut parser = KimiK2ToolParser::new(&test_tools());

522 + let input = build_tool_section(&[

523 + build_tool_call("get_weather", 0, r#"{"location":"Shanghai"}"#),

524 + build_tool_call("add", 1, r#"{"x":1,"y":2}"#),

525 + ]);

526 +

527 + for chunk in split_by_chars(&input, 7) {

528 + parser.parse_chunk(chunk).unwrap();

529 + }

530 +

531 + // IDs are available after parsing but before finish(), which calls reset().

532 + assert_eq!(parser.tool_call_id(0), Some("functions.get_weather:0"));

533 + assert_eq!(parser.tool_call_id(1), Some("functions.add:1"));

534 + parser.finish().unwrap();

535 + assert_eq!(parser.tool_call_id(0), None);

536 + }

537 +

538 #[test]

539 fn kimi_k2_accepts_non_functions_header_prefix() {

540 let mut parser = KimiK2ToolParser::new(&test_tools());

@@ -509,9 +542,10 @@ mod tests {

542 "{TOOL_CALLS_START}{TOOL_CALL_START}api.tools.search:42{TOOL_CALL_ARGUMENT_START}{{}}{TOOL_CALL_END}{TOOL_CALLS_END}"

543 );

544

512 - let output = parser.parse_complete(&input).unwrap();

545 + let output = parser.parse_chunk(&input).unwrap().coalesce_calls();

546

547 assert_eq!(output.calls[0].tool_index, 42);

548 + assert_eq!(parser.tool_call_id(42), Some("api.tools.search:42"));

549 assert_eq!(output.calls[0].name.as_deref(), Some("search"));

550 assert_eq!(output.calls[0].arguments, "{}");

551 }

rust/src/tool-parser/src/lib.rs

@@ -121,6 +121,12 @@ pub trait ToolParser: Send {

121 false

122 }

123

124 + /// Return the parser-provided ID for a tool call by index, if the model

125 + /// emitted one.

126 + fn tool_call_id(&self, _tool_index: usize) -> Option<&str> {

127 + None

128 + }

129 +

130 /// Feed one decoded text delta into the parser, appending committed output

131 /// into `output`.

132 ///