@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 03ce1c6ed9

Fork 0

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

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 4, 23:45 UTC

[Bugfix] Kimi-K2 tool parser streaming - fix token leakage, argument truncation, and content dropping (#38579)

Signed-off-by: sfeng33 <4florafeng@gmail.com>

Flora Feng committed Apr 19, 2026 at 04:30 UTC 03ce1c6ed908788cf508219b666720783a723123

2 files changed +684 -1405

tests/tool_parsers/test_kimi_k2_tool_parser.py

+525 -921

@@ -3,14 +3,20 @@

3 # ruff: noqa: E501

5 import json

6 +from unittest.mock import MagicMock

8 import pytest

9 -from vllm.entrypoints.openai.engine.protocol import FunctionCall, ToolCall

10 +from tests.tool_parsers.utils import (

11 + run_tool_extraction,

12 + run_tool_extraction_streaming,

13 +)

14 +from vllm.entrypoints.openai.chat_completion.protocol import (

15 + ChatCompletionRequest,

16 +)

17 from vllm.tokenizers import get_tokenizer

18 from vllm.tool_parsers.kimi_k2_tool_parser import KimiK2ToolParser

13 -# Use a common model that is likely to be available

20 MODEL = "moonshotai/Kimi-K2-Instruct"

@@ -20,959 +26,557 @@ def kimi_k2_tokenizer():

28 @pytest.fixture

23 -def kimi_k2_tool_parser(kimi_k2_tokenizer):

29 +def parser(kimi_k2_tokenizer):

30 return KimiK2ToolParser(kimi_k2_tokenizer)

27 -def assert_tool_calls(

28 - actual_tool_calls: list[ToolCall], expected_tool_calls: list[ToolCall]

29 -):

30 - assert len(actual_tool_calls) == len(expected_tool_calls)

33 +SECTION_BEGIN = "<|tool_calls_section_begin|>"

34 +SECTION_END = "<|tool_calls_section_end|>"

35 +TOOL_BEGIN = "<|tool_call_begin|>"

36 +TOOL_END = "<|tool_call_end|>"

37 +ARG_BEGIN = "<|tool_call_argument_begin|>"

38 +

39 +

40 +def _tool(tool_id: str, args: str) -> str:

41 + return f"{TOOL_BEGIN}{tool_id} {ARG_BEGIN}{args}{TOOL_END}"

32 - for actual_tool_call, expected_tool_call in zip(

33 - actual_tool_calls, expected_tool_calls

34 - ):

35 - assert actual_tool_call.type == "function"

36 - assert actual_tool_call.function == expected_tool_call.function

37 -

38 - # assert tool call id format: should contain function name and numeric index

39 - # Format can be either "functions.func_name:0" or "func_name:0"

40 - assert actual_tool_call.id.split(":")[-1].isdigit()

41 - assert (

42 - actual_tool_call.id.split(":")[0].split(".")[-1]

43 - == expected_tool_call.function.name

44 - )

44 +def _wrap(*tool_strs: str) -> str:

45 + return SECTION_BEGIN + "".join(tool_strs) + SECTION_END

47 -def run_streaming_sequence(parser, deltas):

48 - """Helper to simulate a streaming sequence and return results."""

49 - previous_text = ""

50 - previous_token_ids: list[int] = []

51 - results = []

52 -

53 - for delta_text, delta_token_ids in deltas:

54 - current_text = previous_text + delta_text

55 - current_token_ids = previous_token_ids + delta_token_ids

56 -

57 - result = parser.extract_tool_calls_streaming(

58 - previous_text=previous_text,

59 - current_text=current_text,

60 - delta_text=delta_text,

61 - previous_token_ids=previous_token_ids,

62 - current_token_ids=current_token_ids,

63 - delta_token_ids=delta_token_ids,

64 - request=None,

47 +

48 +class TestExtractToolCalls:

49 + def test_no_tools(self, parser):

50 + content, tool_calls = run_tool_extraction(

51 + parser, "This is a test", streaming=False

52 )

66 - results.append(result)

67 -

68 - previous_text = current_text

69 - previous_token_ids = current_token_ids

70 -

71 - return results

72 -

73 -

74 -def test_extract_tool_calls_no_tools(kimi_k2_tool_parser):

75 - model_output = "This is a test"

76 - extracted_tool_calls = kimi_k2_tool_parser.extract_tool_calls(

77 - model_output, request=None

78 - ) # type: ignore[arg-type]

79 - assert not extracted_tool_calls.tools_called

80 - assert extracted_tool_calls.tool_calls == []

81 - assert extracted_tool_calls.content == model_output

82 -

83 -

84 -@pytest.mark.parametrize(

85 - ids=[

86 - "tool_call_with_content_before",

87 - "multi_tool_call_with_content_before",

88 - "concatenated_tool_calls_bug_fix",

89 - "three_concatenated_tool_calls",

90 - "mixed_spacing_tool_calls",

91 - "angle_brackets_in_json",

92 - "newlines_in_json",

93 - ],

94 - argnames=["model_output", "expected_tool_calls", "expected_content"],

95 - argvalues=[

96 - (

97 - """I'll help you check the weather. <|tool_calls_section_begin|> <|tool_call_begin|>

98 -functions.get_weather:0 <|tool_call_argument_begin|> {"city": "Beijing"} <|tool_call_end|> <|tool_calls_section_end|>""",

99 - [

100 - ToolCall(

101 - id="functions.get_weather:0",

102 - function=FunctionCall(

103 - name="get_weather",

104 - arguments=json.dumps(

105 - {

106 - "city": "Beijing",

107 - },

108 - ),

109 - ),

110 - type="function",

111 - )

112 - ],

113 - "I'll help you check the weather. ",

114 - ),

115 - (

116 - """I'll help you check the weather. <|tool_calls_section_begin|> <|tool_call_begin|>

117 -functions.get_weather:0 <|tool_call_argument_begin|> {"city": "Beijing"} <|tool_call_end|> <|tool_call_begin|>

118 -functions.get_weather:1 <|tool_call_argument_begin|> {"city": "Shanghai"} <|tool_call_end|> <|tool_calls_section_end|>""",

119 - [

120 - ToolCall(

121 - id="functions.get_weather:0",

122 - function=FunctionCall(

123 - name="get_weather",

124 - arguments=json.dumps(

125 - {

126 - "city": "Beijing",

127 - },

128 - ),

129 - ),

130 - type="function",

131 - ),

132 - ToolCall(

133 - id="functions.get_weather:1",

134 - function=FunctionCall(

135 - name="get_weather",

136 - arguments=json.dumps(

137 - {

138 - "city": "Shanghai",

139 - },

140 - ),

141 - ),

142 - type="function",

143 - ),

144 - ],

145 - "I'll help you check the weather. ",

146 - ),

147 - (

148 - """I'll get the weather and news for LA today. First, let me get the weather using Los Angeles coordinates, and then get the latest news. <|tool_calls_section_begin|><|tool_call_begin|>functions.get_weather:0<|tool_call_argument_begin|>{"latitude": 34.0522, "longitude": -118.2437}<|tool_call_end|><|tool_call_begin|>functions.get_news:1<|tool_call_argument_begin|>{"content": "Los Angeles today"}<|tool_call_end|><|tool_calls_section_end|>""",

149 - [

150 - ToolCall(

151 - id="functions.get_weather:0",

152 - function=FunctionCall(

153 - name="get_weather",

154 - arguments=json.dumps(

155 - {"latitude": 34.0522, "longitude": -118.2437}

156 - ),

157 - ),

158 - type="function",

159 - ),

160 - ToolCall(

161 - id="functions.get_news:1",

162 - function=FunctionCall(

163 - name="get_news",

164 - arguments=json.dumps({"content": "Los Angeles today"}),

165 - ),

166 - type="function",

167 - ),

168 - ],

169 - "I'll get the weather and news for LA today. First, let me get the weather using Los Angeles coordinates, and then get the latest news. ",

170 - ),

171 - (

172 - """I'll help you with multiple tasks. <|tool_calls_section_begin|><|tool_call_begin|>functions.get_weather:0<|tool_call_argument_begin|>{"city": "New York"}<|tool_call_end|><|tool_call_begin|>functions.get_news:1<|tool_call_argument_begin|>{"topic": "technology"}<|tool_call_end|><|tool_call_begin|>functions.send_email:2<|tool_call_argument_begin|>{"to": "user@example.com", "subject": "Daily Update"}<|tool_call_end|><|tool_calls_section_end|>""",

173 - [

174 - ToolCall(

175 - id="functions.get_weather:0",

176 - function=FunctionCall(

177 - name="get_weather",

178 - arguments=json.dumps({"city": "New York"}),

179 - ),

180 - type="function",

53 + assert content == "This is a test"

54 + assert tool_calls == []

55 +

56 + @pytest.mark.parametrize(

57 + "model_output, expected_names, expected_args_list, expected_content",

58 + [

59 + pytest.param(

60 + "I'll check. "

61 + + _wrap(_tool("functions.get_weather:0", '{"city": "Beijing"}')),

62 + ["get_weather"],

63 + [{"city": "Beijing"}],

64 + "I'll check. ",

65 + id="single_tool_call",

66 + ),

67 + pytest.param(

68 + "Compare weather. "

69 + + _wrap(

70 + _tool("functions.get_weather:0", '{"city": "Beijing"}'),

71 + _tool("functions.get_weather:1", '{"city": "Shanghai"}'),

72 ),

182 - ToolCall(

183 - id="functions.get_news:1",

184 - function=FunctionCall(

185 - name="get_news",

186 - arguments=json.dumps({"topic": "technology"}),

73 + ["get_weather", "get_weather"],

74 + [{"city": "Beijing"}, {"city": "Shanghai"}],

75 + "Compare weather. ",

76 + id="parallel_tool_calls",

77 + ),

78 + pytest.param(

79 + "Multiple tasks. "

80 + + _wrap(

81 + _tool("functions.get_weather:0", '{"city": "New York"}'),

82 + _tool("functions.get_news:1", '{"topic": "technology"}'),

83 + _tool(

84 + "functions.send_email:2",

85 + '{"to": "user@example.com", "subject": "Daily Update"}',

86 ),

188 - type="function",

87 ),

190 - ToolCall(

191 - id="functions.send_email:2",

192 - function=FunctionCall(

193 - name="send_email",

194 - arguments=json.dumps(

195 - {"to": "user@example.com", "subject": "Daily Update"}

196 - ),

197 - ),

198 - type="function",

88 + ["get_weather", "get_news", "send_email"],

89 + [

90 + {"city": "New York"},

91 + {"topic": "technology"},

92 + {"to": "user@example.com", "subject": "Daily Update"},

93 + ],

94 + "Multiple tasks. ",

95 + id="three_tool_calls",

96 + ),

97 + pytest.param(

98 + "Process HTML. "

99 + + _wrap(

100 + _tool("functions.process_html:0", '{"html": "<div>content</div>"}')

101 ),

200 - ],

201 - "I'll help you with multiple tasks. ",

202 - ),

203 - (

204 - """Mixed spacing test. <|tool_calls_section_begin|> <|tool_call_begin|> functions.test:0 <|tool_call_argument_begin|> {} <|tool_call_end|><|tool_call_begin|>functions.test2:1<|tool_call_argument_begin|>{}<|tool_call_end|> <|tool_calls_section_end|>""",

205 - [

206 - ToolCall(

207 - id="functions.test:0",

208 - function=FunctionCall(

209 - name="test",

210 - arguments=json.dumps({}),

211 - ),

212 - type="function",

102 + ["process_html"],

103 + [{"html": "<div>content</div>"}],

104 + "Process HTML. ",

105 + id="angle_brackets_in_json",

106 + ),

107 + pytest.param(

108 + "Formatted. "

109 + + _wrap(

110 + _tool(

111 + "functions.process_data:0",

112 + '{\n "name": "test",\n "value": 123\n}',

113 + )

114 ),

214 - ToolCall(

215 - id="functions.test2:1",

216 - function=FunctionCall(

217 - name="test2",

218 - arguments=json.dumps({}),

219 - ),

220 - type="function",

221 - ),

222 - ],

223 - "Mixed spacing test. ",

224 - ),

225 - (

226 - """I need to process HTML content. <|tool_calls_section_begin|><|tool_call_begin|>functions.process_html:0<|tool_call_argument_begin|>{"html": "<div>content</div>", "text": "normal text"}<|tool_call_end|><|tool_calls_section_end|>""",

227 - [

228 - ToolCall(

229 - id="functions.process_html:0",

230 - function=FunctionCall(

231 - name="process_html",

232 - arguments=json.dumps(

233 - {"html": "<div>content</div>", "text": "normal text"}

234 - ),

235 - ),

236 - type="function",

237 - )

238 - ],

239 - "I need to process HTML content. ",

240 - ),

241 - (

242 - """I need to process formatted JSON. <|tool_calls_section_begin|><|tool_call_begin|>functions.process_data:0<|tool_call_argument_begin|>{

243 - "name": "test",

244 - "value": 123,

245 - "nested": {

246 - "key": "value"

247 - }

248 -}<|tool_call_end|><|tool_calls_section_end|>""",

249 - [

250 - ToolCall(

251 - id="functions.process_data:0",

252 - function=FunctionCall(

253 - name="process_data",

254 - arguments=json.dumps(

255 - {"name": "test", "value": 123, "nested": {"key": "value"}},

256 - indent=2,

257 - ),

258 - ),

259 - type="function",

260 - )

261 - ],

262 - "I need to process formatted JSON. ",

263 - ),

264 - ],

265 -)

266 -def test_extract_tool_calls(

267 - kimi_k2_tool_parser, model_output, expected_tool_calls, expected_content

268 -):

269 - extracted_tool_calls = kimi_k2_tool_parser.extract_tool_calls(

270 - model_output, request=None

271 - ) # type: ignore[arg-type]

272 - assert extracted_tool_calls.tools_called

273 -

274 - assert_tool_calls(extracted_tool_calls.tool_calls, expected_tool_calls)

275 -

276 - assert extracted_tool_calls.content == expected_content

277 -

278 -

279 -def test_extract_tool_calls_invalid_json(kimi_k2_tool_parser):

280 - """we'll return every funcall result"""

281 - model_output = """I'll help you check the weather. <|tool_calls_section_begin|> <|tool_call_begin|>

282 -functions.invalid_get_weather:0 <|tool_call_argument_begin|> {"city": "Beijing" <|tool_call_end|> <|tool_call_begin|>

283 -functions.valid_get_weather:1 <|tool_call_argument_begin|> {"city": "Shanghai"} <|tool_call_end|> <|tool_calls_section_end|>"""

284 -

285 - extracted_tool_calls = kimi_k2_tool_parser.extract_tool_calls(

286 - model_output, request=None

287 - ) # type: ignore[arg-type]

288 -

289 - assert extracted_tool_calls.tools_called

290 - # Should extract only the valid JSON tool calls

291 - assert len(extracted_tool_calls.tool_calls) == 2

292 - assert extracted_tool_calls.tool_calls[0].function.name == "invalid_get_weather"

293 - assert extracted_tool_calls.tool_calls[1].function.name == "valid_get_weather"

294 -

295 -

296 -def test_extract_tool_calls_invalid_funcall(kimi_k2_tool_parser):

297 - """we'll return every funcall result"""

298 - model_output = """I'll help you check the weather. <|tool_calls_section_begin|> <|tool_call_begin|>

299 -functions.invalid_get_weather.0 <|tool_call_argument_begin|> {"city": "Beijing"} <|tool_call_end|> <|tool_call_begin|>

300 -functions.valid_get_weather:1 <|tool_call_argument_begin|> {"city": "Shanghai"} <|tool_call_end|> <|tool_calls_section_end|>"""

301 -

302 - extracted_tool_calls = kimi_k2_tool_parser.extract_tool_calls(

303 - model_output, request=None

304 - ) # type: ignore[arg-type]

305 -

306 - assert extracted_tool_calls.tools_called

307 - # Should extract only the valid JSON tool calls

308 - assert len(extracted_tool_calls.tool_calls) == 1

309 - assert extracted_tool_calls.tool_calls[0].function.name == "valid_get_weather"

310 -

311 -

312 -def test_streaming_basic_functionality(kimi_k2_tool_parser):

313 - """Test basic streaming functionality."""

314 - # Reset streaming state

315 - kimi_k2_tool_parser.current_tool_name_sent = False

316 - kimi_k2_tool_parser.prev_tool_call_arr = []

317 - kimi_k2_tool_parser.current_tool_id = -1

318 - kimi_k2_tool_parser.streamed_args_for_tool = []

319 -

320 - # Test with a simple tool call

321 - current_text = """ check the weather. <|tool_calls_section_begin|> <|tool_call_begin|>

322 -functions.get_weather:0 <|tool_call_argument_begin|> {"city": "Beijing"} <|tool_call_end|> <|tool_calls_section_end|>"""

323 -

324 - # First call should handle the initial setup

325 - result = kimi_k2_tool_parser.extract_tool_calls_streaming(

326 - previous_text="I'll help you",

327 - current_text=current_text,

328 - delta_text="<|tool_calls_section_end|>",

329 - previous_token_ids=[],

330 - current_token_ids=[],

331 - delta_token_ids=[],

332 - request=None,

333 - )

334 -

335 - # The result might be None or contain tool call information

336 - # This depends on the internal state management

337 - if result is not None and hasattr(result, "tool_calls") and result.tool_calls:

338 - assert len(result.tool_calls) >= 0

339 -

340 -

341 -def test_streaming_no_tool_calls(kimi_k2_tool_parser):

342 - """Test streaming when there are no tool calls."""

343 - current_text = "This is just regular text without any tool calls."

344 -

345 - result = kimi_k2_tool_parser.extract_tool_calls_streaming(

346 - previous_text="This is just regular text",

347 - current_text=current_text,

348 - delta_text=" without any tool calls.",

349 - previous_token_ids=[],

350 - current_token_ids=[],

351 - delta_token_ids=[],

352 - request=None,

353 - )

354 -

355 - # Should return the delta text as content

356 - assert result is not None

357 - assert hasattr(result, "content")

358 - assert result.content == " without any tool calls."

359 -

360 -

361 -def test_token_leak_between_section_and_tool_begin(kimi_k2_tool_parser):

362 - """

363 - Test that text between <|tool_calls_section_begin|> and <|tool_call_begin|>

364 - is suppressed and does not leak into reasoning_delta.

365 - This is the main vulnerability being fixed.

366 - """

367 - kimi_k2_tool_parser.reset_streaming_state()

368 -

369 - # Get token IDs for the markers

370 - section_begin_token_id = kimi_k2_tool_parser.vocab.get(

371 - "<|tool_calls_section_begin|>"

372 - )

373 - tool_call_begin_token_id = kimi_k2_tool_parser.vocab.get("<|tool_call_begin|>")

374 -

375 - # Simulate streaming sequence:

376 - deltas = [

377 - ("I'll help you with that. ", [1, 2, 3]),

378 - ("<|tool_calls_section_begin|>", [section_begin_token_id]),

379 - (" spurious text ", [4, 5]),

380 - ("<|tool_call_begin|>", [tool_call_begin_token_id]),

381 - ]

382 -

383 - results = run_streaming_sequence(kimi_k2_tool_parser, deltas)

384 -

385 - # Delta 1: "I'll help you with that. "

386 - assert results[0] is not None

387 - assert results[0].content == "I'll help you with that. "

388 -

389 - # Delta 2: "<|tool_calls_section_begin|>"

390 - # Section marker should be stripped and suppressed

391 - assert results[1] is None or (

392 - results[1].content is None or results[1].content == ""

393 - )

394 -

395 - # Delta 3: " spurious text or tokens " (THE LEAK SCENARIO)

396 - # CRITICAL: This text should be suppressed, NOT returned as reasoning_delta

397 - assert results[2] is None or (

398 - results[2].content is None or results[2].content == ""

399 - )

400 -

401 - # Delta 4: "<|tool_call_begin|>..."

402 - # Now we're in tool call mode, result depends on internal state

403 - # The key is that the spurious text from Delta 3 was not leaked

404 -

405 -

406 -def test_split_markers_across_deltas(kimi_k2_tool_parser):

407 - """

408 - Test that markers split across delta chunks are correctly detected

409 - via the rolling buffer mechanism.

410 - """

411 - kimi_k2_tool_parser.reset_streaming_state()

412 -

413 - section_begin_token_id = kimi_k2_tool_parser.vocab.get(

414 - "<|tool_calls_section_begin|>"

115 + ["process_data"],

116 + [{"name": "test", "value": 123}],

117 + "Formatted. ",

118 + id="multiline_json",

119 + ),

120 + pytest.param(

121 + "No prefix. " + _wrap(_tool("get_weather:0", '{"city": "Tokyo"}')),

122 + ["get_weather"],

123 + [{"city": "Tokyo"}],

124 + "No prefix. ",

125 + id="no_functions_prefix",

126 + ),

127 + pytest.param(

128 + "Empty args. " + _wrap(_tool("functions.test:0", "{}")),

129 + ["test"],

130 + [{}],

131 + "Empty args. ",

132 + id="empty_arguments",

133 + ),

134 + ],

135 )

416 -

417 - # Delta 1: partial token, Delta 2: complete marker

418 - deltas = [

419 - ("<|tool_calls_sec", [3]),

420 - ("tion_begin|> ", [section_begin_token_id, 4]),

421 - ]

422 -

423 - _results = run_streaming_sequence(kimi_k2_tool_parser, deltas)

424 -

425 - # Now the complete marker should be detected via buffer

426 - assert kimi_k2_tool_parser.in_tool_section is True

427 -

428 -

429 -def test_marker_variants(kimi_k2_tool_parser):

430 - """Test that both singular and plural marker variants are recognized."""

431 - kimi_k2_tool_parser.reset_streaming_state()

432 -

433 - # Test singular variant: <|tool_call_section_begin|> (note: singular "call")

434 - singular_token_id = kimi_k2_tool_parser.vocab.get("<|tool_call_section_begin|>")

435 -

436 - if singular_token_id is not None: # Only test if tokenizer supports it

437 - _result = kimi_k2_tool_parser.extract_tool_calls_streaming(

438 - previous_text="Reasoning ",

439 - current_text="Reasoning <|tool_call_section_begin|>",

440 - delta_text="<|tool_call_section_begin|>",

441 - previous_token_ids=[1, 2],

442 - current_token_ids=[1, 2, singular_token_id],

443 - delta_token_ids=[singular_token_id],

444 - request=None,

136 + def test_extract_tool_calls(

137 + self, parser, model_output, expected_names, expected_args_list, expected_content

138 + ):

139 + content, tool_calls = run_tool_extraction(parser, model_output, streaming=False)

140 + assert content == expected_content

141 + assert len(tool_calls) == len(expected_names)

142 + for tc, name, expected_args in zip(

143 + tool_calls, expected_names, expected_args_list

144 + ):

145 + assert tc.type == "function"

146 + assert tc.function.name == name

147 + assert json.loads(tc.function.arguments) == expected_args

148 + # id format: "something:digit"

149 + assert tc.id.split(":")[-1].isdigit()

150 +

151 + def test_invalid_json_still_extracted(self, parser):

152 + """Tool calls with invalid JSON are still returned (arguments as-is)."""

153 + model_output = (

154 + "Help. "

155 + + SECTION_BEGIN

156 + + _tool("functions.bad:0", '{"city": "Beijing"')

157 + + _tool("functions.good:1", '{"city": "Shanghai"}')

158 + + SECTION_END

159 )

446 - # Should enter tool section mode with singular variant too

447 - assert kimi_k2_tool_parser.in_tool_section is True

448 -

449 -

450 -def test_reentry_to_reasoning_after_tool_section(kimi_k2_tool_parser):

451 - """

452 - Test that after exiting a tool section with <|tool_calls_section_end|>,

453 - subsequent text is correctly returned as reasoning content.

454 - """

455 - kimi_k2_tool_parser.reset_streaming_state()

456 -

457 - section_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_begin|>")

458 - section_end_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_end|>")

459 -

460 - deltas = [

461 - ("<|tool_calls_section_begin|>", [section_begin_id]),

462 - ("<|tool_calls_section_end|>", [section_end_id]),

463 - (" More reasoning", [10, 11]),

464 - ]

465 -

466 - results = run_streaming_sequence(kimi_k2_tool_parser, deltas)

467 -

468 - assert kimi_k2_tool_parser.in_tool_section is False

469 - assert results[2] is not None

470 - assert results[2].content == " More reasoning"

471 -

472 -

473 -def test_empty_tool_section(kimi_k2_tool_parser):

474 - """Test an empty tool section (begin immediately followed by end)."""

475 - kimi_k2_tool_parser.reset_streaming_state()

476 -

477 - section_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_begin|>")

478 - section_end_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_end|>")

479 -

480 - # Section begin

481 - _result1 = kimi_k2_tool_parser.extract_tool_calls_streaming(

482 - previous_text="Reasoning ",

483 - current_text="Reasoning <|tool_calls_section_begin|>",

484 - delta_text="<|tool_calls_section_begin|>",

485 - previous_token_ids=[1],

486 - current_token_ids=[1, section_begin_id],

487 - delta_token_ids=[section_begin_id],

488 - request=None,

489 - )

490 -

491 - # Immediate section end

492 - _result2 = kimi_k2_tool_parser.extract_tool_calls_streaming(

493 - previous_text="Reasoning <|tool_calls_section_begin|>",

494 - current_text="Reasoning <|tool_calls_section_begin|><|tool_calls_section_end|>",

495 - delta_text="<|tool_calls_section_end|>",

496 - previous_token_ids=[1, section_begin_id],

497 - current_token_ids=[1, section_begin_id, section_end_id],

498 - delta_token_ids=[section_end_id],

499 - request=None,

500 - )

501 - # Should exit cleanly without errors

502 - assert kimi_k2_tool_parser.in_tool_section is False

503 -

504 -

505 -def test_malformed_tool_section_recovery(kimi_k2_tool_parser):

506 - """

507 - Test that the parser recovers from a malformed tool section

508 - that never closes properly.

509 - """

510 - kimi_k2_tool_parser.reset_streaming_state()

511 -

512 - section_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_begin|>")

513 -

514 - # Enter tool section

515 - _result1 = kimi_k2_tool_parser.extract_tool_calls_streaming(

516 - previous_text="",

517 - current_text="<|tool_calls_section_begin|>",

518 - delta_text="<|tool_calls_section_begin|>",

519 - previous_token_ids=[],

520 - current_token_ids=[section_begin_id],

521 - delta_token_ids=[section_begin_id],

522 - request=None,

523 - )

524 - assert kimi_k2_tool_parser.in_tool_section is True

525 -

526 - # Simulate a lot of text without proper tool calls or section end

527 - # This should trigger the error recovery mechanism

528 - large_text = "x" * 10000 # Exceeds max_section_chars

529 -

530 - result2 = kimi_k2_tool_parser.extract_tool_calls_streaming(

531 - previous_text="<|tool_calls_section_begin|>",

532 - current_text="<|tool_calls_section_begin|>" + large_text,

533 - delta_text=large_text,

534 - previous_token_ids=[section_begin_id],

535 - current_token_ids=[section_begin_id] + list(range(100, 100 + len(large_text))),

536 - delta_token_ids=list(range(100, 100 + len(large_text))),

537 - request=None,

538 - )

539 -

540 - # Parser should have force-exited the tool section

541 - assert kimi_k2_tool_parser.in_tool_section is False

542 - # And returned the content as reasoning

543 - assert result2 is not None

544 - assert result2.content == large_text

545 -

546 -

547 -def test_state_reset(kimi_k2_tool_parser):

548 - """Test that reset_streaming_state() properly clears all state."""

549 - # Put parser in a complex state

550 - kimi_k2_tool_parser.in_tool_section = True

551 - kimi_k2_tool_parser.token_buffer = "some buffer"

552 - kimi_k2_tool_parser.current_tool_id = 5

553 - kimi_k2_tool_parser.prev_tool_call_arr = [{"id": "test"}]

554 - kimi_k2_tool_parser.section_char_count = 1000

555 -

556 - # Reset

557 - kimi_k2_tool_parser.reset_streaming_state()

558 -

559 - # Verify all state is cleared

560 - assert kimi_k2_tool_parser.in_tool_section is False

561 - assert kimi_k2_tool_parser.token_buffer == ""

562 - assert kimi_k2_tool_parser.current_tool_id == -1

563 - assert kimi_k2_tool_parser.prev_tool_call_arr == []

564 - assert kimi_k2_tool_parser.section_char_count == 0

565 - assert kimi_k2_tool_parser.current_tool_name_sent is False

566 - assert kimi_k2_tool_parser.streamed_args_for_tool == []

567 -

568 -

569 -def test_section_begin_noise_tool_begin_same_chunk(kimi_k2_tool_parser):

570 - """

571 - Test that begin→noise→tool_begin within the SAME chunk suppresses

572 - the noise text correctly (not just across chunks).

573 - """

574 - kimi_k2_tool_parser.reset_streaming_state()

575 -

576 - section_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_begin|>")

577 - tool_call_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_call_begin|>")

578 -

579 - # Single delta containing: section_begin + spurious text + tool_call_begin

580 - combined_text = "<|tool_calls_section_begin|> noise text <|tool_call_begin|>"

581 -

582 - result = kimi_k2_tool_parser.extract_tool_calls_streaming(

583 - previous_text="Reasoning ",

584 - current_text="Reasoning " + combined_text,

585 - delta_text=combined_text,

586 - previous_token_ids=[1, 2],

587 - current_token_ids=[1, 2, section_begin_id, 3, 4, tool_call_begin_id],

588 - delta_token_ids=[section_begin_id, 3, 4, tool_call_begin_id],

589 - request=None,

590 - )

591 -

592 - # The noise text should NOT leak into content

593 - # Result should either be None/empty or start tool call parsing

594 - if result is not None and result.content is not None:

595 - # If content is returned, it should not contain the noise

596 - assert "noise text" not in result.content

597 - assert result.content == "" or result.content.strip() == ""

598 -

599 -

600 -def test_stream_ends_without_section_end_marker(kimi_k2_tool_parser):

601 - """

602 - Test that if the stream ends (EOF) without a proper section end marker,

603 - the parser doesn't leak text, doesn't crash, and resets state cleanly.

604 - """

605 - kimi_k2_tool_parser.reset_streaming_state()

606 -

607 - section_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_begin|>")

608 -

609 - # Enter tool section

610 - _result1 = kimi_k2_tool_parser.extract_tool_calls_streaming(

611 - previous_text="",

612 - current_text="<|tool_calls_section_begin|>",

613 - delta_text="<|tool_calls_section_begin|>",

614 - previous_token_ids=[],

615 - current_token_ids=[section_begin_id],

616 - delta_token_ids=[section_begin_id],

617 - request=None,

618 - )

619 - assert kimi_k2_tool_parser.in_tool_section is True

620 -

621 - # Some content in tool section

622 - result2 = kimi_k2_tool_parser.extract_tool_calls_streaming(

623 - previous_text="<|tool_calls_section_begin|>",

624 - current_text="<|tool_calls_section_begin|> partial content",

625 - delta_text=" partial content",

626 - previous_token_ids=[section_begin_id],

627 - current_token_ids=[section_begin_id, 10, 11],

628 - delta_token_ids=[10, 11],

629 - request=None,

630 - )

631 - # Content should be suppressed

632 - assert result2.content == "" or result2.content is None

633 -

634 - # Stream ends (EOF) - no more deltas, no section_end marker

635 - # Simulate this by manually checking state and resetting

636 - # (In real usage, the request handler would call reset_streaming_state)

637 - assert kimi_k2_tool_parser.in_tool_section is True # Still in section

638 -

639 - # Reset state (as would happen between requests)

640 - kimi_k2_tool_parser.reset_streaming_state()

641 -

642 - # Verify clean slate

643 - assert kimi_k2_tool_parser.in_tool_section is False

644 - assert kimi_k2_tool_parser.token_buffer == ""

645 -

646 - # Next request should work normally

647 - result3 = kimi_k2_tool_parser.extract_tool_calls_streaming(

648 - previous_text="",

649 - current_text="New reasoning",

650 - delta_text="New reasoning",

651 - previous_token_ids=[],

652 - current_token_ids=[20, 21],

653 - delta_token_ids=[20, 21],

654 - request=None,

655 - )

656 - assert result3 is not None

657 - assert result3.content == "New reasoning"

658 -

659 -

660 -def test_same_chunk_begin_and_end_markers(kimi_k2_tool_parser):

661 - """

662 - CRITICAL TEST: Verify that when both section_begin and section_end

663 - markers appear in the SAME chunk, the parser correctly:

664 - 1. Enters the tool section

665 - 2. Immediately exits the tool section

666 - 3. Does NOT get stuck in in_tool_section=True state

667 -

668 - This tests the bug fix where elif was changed to if to handle

669 - both state transitions in a single delta.

670 - """

671 - kimi_k2_tool_parser.reset_streaming_state()

672 -

673 - section_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_begin|>")

674 - section_end_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_end|>")

675 -

676 - # Single chunk with both markers (e.g., empty tool section)

677 - combined_delta = "<|tool_calls_section_begin|><|tool_calls_section_end|>"

678 -

679 - result = kimi_k2_tool_parser.extract_tool_calls_streaming(

680 - previous_text="Some reasoning ",

681 - current_text="Some reasoning " + combined_delta,

682 - delta_text=combined_delta,

683 - previous_token_ids=[1, 2],

684 - current_token_ids=[1, 2, section_begin_id, section_end_id],

685 - delta_token_ids=[section_begin_id, section_end_id],

686 - request=None,

687 - )

688 -

689 - # CRITICAL: Parser should NOT be stuck in tool section

690 - assert kimi_k2_tool_parser.in_tool_section is False, (

691 - "Parser stuck in tool section after processing both begin/end in same chunk. "

692 - "This indicates the elif bug was not fixed."

693 - )

694 -

695 - # Result should be empty or contain only stripped content

696 - assert result is not None

697 - assert result.content == "" or result.content is None

698 -

699 - # Verify subsequent content streams correctly (not suppressed)

700 - result2 = kimi_k2_tool_parser.extract_tool_calls_streaming(

701 - previous_text="Some reasoning " + combined_delta,

702 - current_text="Some reasoning " + combined_delta + " More reasoning",

703 - delta_text=" More reasoning",

704 - previous_token_ids=[1, 2, section_begin_id, section_end_id],

705 - current_token_ids=[1, 2, section_begin_id, section_end_id, 10, 11],

706 - delta_token_ids=[10, 11],

707 - request=None,

708 - )

709 -

710 - # This content should NOT be suppressed (we're out of tool section)

711 - assert result2 is not None

712 - assert result2.content == " More reasoning"

713 -

714 -

715 -def test_same_chunk_begin_content_end_markers(kimi_k2_tool_parser):

716 - """

717 - Test the same-chunk scenario with actual content between markers.

718 - Example: <|tool_calls_section_begin|> text <|tool_calls_section_end|>

719 - all arriving in one delta. The key is that the state machine correctly

720 - transitions in and out within the same chunk.

721 - """

722 - kimi_k2_tool_parser.reset_streaming_state()

723 -

724 - section_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_begin|>")

725 - section_end_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_end|>")

726 -

727 - # Chunk with begin, some whitespace/noise, and end all together

728 - # This simulates a tool section that opens and closes in the same chunk

729 - combined_delta = "<|tool_calls_section_begin|> <|tool_calls_section_end|>"

730 -

731 - _result = kimi_k2_tool_parser.extract_tool_calls_streaming(

732 - previous_text="Reasoning ",

733 - current_text="Reasoning " + combined_delta,

734 - delta_text=combined_delta,

735 - previous_token_ids=[1],

736 - current_token_ids=[1, section_begin_id, 100, section_end_id],

737 - delta_token_ids=[section_begin_id, 100, section_end_id],

738 - request=None,

739 - )

740 -

741 - # Parser should exit cleanly (not stuck in tool section)

742 - assert kimi_k2_tool_parser.in_tool_section is False

743 -

744 - # Verify the fix: next content should stream normally, not be suppressed

745 - result2 = kimi_k2_tool_parser.extract_tool_calls_streaming(

746 - previous_text="Reasoning " + combined_delta,

747 - current_text="Reasoning " + combined_delta + " Done",

748 - delta_text=" Done",

749 - previous_token_ids=[1, section_begin_id, 100, section_end_id],

750 - current_token_ids=[1, section_begin_id, 100, section_end_id, 200],

751 - delta_token_ids=[200],

752 - request=None,

753 - )

754 -

755 - # Content after section should be returned (not suppressed)

756 - assert result2 is not None

757 - assert result2.content == " Done"

758 -

759 -

760 -def test_tool_call_end_and_section_end_same_chunk(kimi_k2_tool_parser):

761 - """

762 - CRITICAL TEST (P1): Verify that when both <|tool_call_end|> and

763 - <|tool_calls_section_end|> appear in the SAME chunk, the parser:

764 - 1. Processes the tool_call_end first (emits final arguments)

765 - 2. THEN exits the section

766 - 3. Does NOT drop the final tool call update

767 - 4. Does NOT leak special tokens into reasoning

768 -

769 - This tests the deferred section exit fix.

770 - """

771 - kimi_k2_tool_parser.reset_streaming_state()

772 -

773 - section_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_begin|>")

774 - section_end_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_end|>")

775 - tool_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_call_begin|>")

776 - tool_end_id = kimi_k2_tool_parser.vocab.get("<|tool_call_end|>")

777 -

778 - # Simulate a streaming sequence for a SHORT tool call (all in one chunk):

779 - combined = (

780 - '<|tool_call_begin|>get_weather:0 <|tool_call_argument_begin|> {"city": "Paris"} '

781 - "<|tool_call_end|><|tool_calls_section_end|>"

782 - )

783 -

784 - deltas = [

785 - ("Let me help. ", [1, 2]),

786 - ("<|tool_calls_section_begin|>", [section_begin_id]),

787 - (combined, [tool_begin_id, 10, 11, 12, tool_end_id, section_end_id]),

788 - (" Done", [20]),

789 - ]

790 -

791 - results = run_streaming_sequence(kimi_k2_tool_parser, deltas)

792 -

793 - # CRITICAL: Parser should have exited section AFTER processing tool

794 - assert kimi_k2_tool_parser.in_tool_section is False

795 -

796 - # Tool call should have been emitted (not dropped)

797 - if results[2] is not None and results[2].content is not None:

798 - # Verify no special tokens leaked into content

799 - assert "<|tool_call_end|>" not in results[2].content

800 - assert "<|tool_calls_section_end|>" not in results[2].content

160 + content, tool_calls = run_tool_extraction(parser, model_output, streaming=False)

161 + assert len(tool_calls) == 2

162 + assert tool_calls[0].function.name == "bad"

163 + assert tool_calls[1].function.name == "good"

164 +

165 + def test_invalid_funcall_id_skipped(self, parser):

166 + """Tool calls with malformed id (no colon+digit) are skipped."""

167 + model_output = (

168 + "Help. "

169 + + SECTION_BEGIN

170 + + _tool("functions.invalid.0", '{"city": "Beijing"}')

171 + + _tool("functions.valid:1", '{"city": "Shanghai"}')

172 + + SECTION_END

173 + )

174 + content, tool_calls = run_tool_extraction(parser, model_output, streaming=False)

175 + assert len(tool_calls) == 1

176 + assert tool_calls[0].function.name == "valid"

177 +

178 + def test_native_id_extracted(self, parser):

179 + """Regression: parser extracts native ID onto ToolCall (PR #32768)."""

180 + model_output = "Checking weather. " + _wrap(

181 + _tool("functions.get_weather:0", '{"city": "Tokyo"}')

182 + )

183 + content, tool_calls = run_tool_extraction(parser, model_output, streaming=False)

184 + assert len(tool_calls) == 1

185 + assert tool_calls[0].id == "functions.get_weather:0"

186 + assert tool_calls[0].function.name == "get_weather"

187 + assert json.loads(tool_calls[0].function.arguments) == {"city": "Tokyo"}

188 +

189 + def test_multi_turn_native_id_continuity(self, kimi_k2_tokenizer):

190 + """Regression: native IDs from turn 1 preserved across turns (PR #32768)."""

191 + turn1_parser = KimiK2ToolParser(kimi_k2_tokenizer)

192 + turn1_output = "Let me check. " + _wrap(

193 + _tool("functions.get_weather:0", '{"city": "Beijing"}')

194 + )

195 + _, turn1_tools = run_tool_extraction(

196 + turn1_parser, turn1_output, streaming=False

197 + )

198 + assert len(turn1_tools) == 1

199 + assert turn1_tools[0].id == "functions.get_weather:0"

200

802 - # Content after tool section should stream normally

803 - assert results[3] is not None

804 - assert results[3].content == " Done"

201 + # Fresh parser for turn 2

202 + turn2_parser = KimiK2ToolParser(kimi_k2_tokenizer)

203 + turn2_output = "Now let me get news. " + _wrap(

204 + _tool("functions.get_news:0", '{"topic": "weather in Beijing"}')

205 + )

206 + _, turn2_tools = run_tool_extraction(

207 + turn2_parser, turn2_output, streaming=False

208 + )

209 + assert len(turn2_tools) == 1

210 + assert turn2_tools[0].id == "functions.get_news:0"

211

212

807 -def test_streaming_tool_call_markers_not_leaked(kimi_k2_tool_parser):

808 - """

809 - CRITICAL TEST: Verify that tool call markers (<|tool_call_begin|>,

810 - <|tool_call_end|>, <|tool_call_argument_begin|>) are NOT leaked

811 - into the content field during streaming.

213 +def _split_tool_output_to_deltas(

214 + content: str, tool_strs: list[tuple[str, str]]

215 +) -> list[str]:

216 + """Build a list of string deltas with special tokens as separate chunks.

217

813 - This reproduces the AWS Bedrock bug where tool call markers appeared

814 - in the 'text' field of responses.

218 + Args:

219 + content: text before tool section

220 + tool_strs: list of (tool_id, args_json)

221 """

816 - kimi_k2_tool_parser.reset_streaming_state()

817 -

818 - section_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_begin|>")

819 - section_end_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_end|>")

820 - tool_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_call_begin|>")

821 - tool_end_id = kimi_k2_tool_parser.vocab.get("<|tool_call_end|>")

822 -

823 - # List of markers that should NEVER appear in content

824 - forbidden_markers = [

825 - "<|tool_call_begin|>",

826 - "<|tool_call_end|>",

827 - "<|tool_call_argument_begin|>",

828 - "<|tool_calls_section_begin|>",

829 - "<|tool_calls_section_end|>",

830 - ]

831 -

832 - all_content = []

833 -

834 - # Steps: reasoning, section begin, tool call, section end, more reasoning

835 - tool_chunk = (

836 - "<|tool_call_begin|> functions.get_weather:0 "

837 - '<|tool_call_argument_begin|> {"city": "Tokyo"} <|tool_call_end|>'

838 - )

839 - deltas = [

840 - ("I'll check the weather. ", [1, 2, 3]),

841 - ("<|tool_calls_section_begin|>", [section_begin_id]),

842 - (tool_chunk, [tool_begin_id, 10, 11, tool_end_id]),

843 - ("<|tool_calls_section_end|>", [section_end_id]),

844 - (" Here's the result.", [20, 21]),

845 - ]

846 -

847 - results = run_streaming_sequence(kimi_k2_tool_parser, deltas)

848 -

849 - for res in results:

850 - if res and res.content:

851 - all_content.append(res.content)

852 -

853 - # CRITICAL ASSERTIONS: No forbidden markers in any content

854 - full_content = "".join(all_content)

855 - for marker in forbidden_markers:

856 - assert marker not in full_content, (

857 - f"MARKER LEAK DETECTED: '{marker}' found in content. "

858 - f"Full content: {repr(full_content)}"

222 + deltas = [content, SECTION_BEGIN]

223 + for tool_id, args_json in tool_strs:

224 + deltas.extend(

225 + [

226 + TOOL_BEGIN,

227 + f"{tool_id} ",

228 + ARG_BEGIN,

229 + f"{args_json} ",

230 + TOOL_END,

231 + ]

232 )

860 -

861 - # Also check that tool call content (function name, arguments) is not leaked

862 - assert "get_weather" not in full_content, (

863 - f"TOOL CALL CONTENT LEAKED: 'get_weather' found in content. "

864 - f"Full content: {repr(full_content)}"

865 - )

866 - assert "Tokyo" not in full_content, (

867 - f"TOOL CALL CONTENT LEAKED: 'Tokyo' found in content. "

868 - f"Full content: {repr(full_content)}"

869 - )

870 -

871 - # Verify that legitimate content was preserved

872 - assert "I'll check the weather." in full_content or len(all_content) > 0

233 + deltas.append(SECTION_END)

234 + return deltas

235

236

875 -def test_native_id_extracted_and_placed_on_tool_call(kimi_k2_tool_parser):

876 - """Regression: parser extracts native ID onto ToolCall (PR #32768)."""

877 - model_output = (

878 - "Checking weather. "

879 - "<|tool_calls_section_begin|>"

880 - "<|tool_call_begin|>functions.get_weather:0"

881 - '<|tool_call_argument_begin|>{"city": "Tokyo"}'

882 - "<|tool_call_end|>"

883 - "<|tool_calls_section_end|>"

884 - )

885 -

886 - result = kimi_k2_tool_parser.extract_tool_calls(model_output, request=None)

887 - assert result.tools_called

888 - assert len(result.tool_calls) == 1

889 -

890 - tc = result.tool_calls[0]

891 - # Native ID from model output must be used as the tool call ID

892 - assert tc.id == "functions.get_weather:0"

893 - assert tc.function.name == "get_weather"

894 - assert json.loads(tc.function.arguments) == {"city": "Tokyo"}

895 -

896 -

897 -def test_multi_turn_native_id_continuity(kimi_k2_tool_parser, kimi_k2_tokenizer):

898 - """Regression: native IDs from turn 1 preserved across turns (PR #32768)."""

899 - turn1_output = (

900 - "Let me check. "

901 - "<|tool_calls_section_begin|>"

902 - "<|tool_call_begin|>functions.get_weather:0"

903 - '<|tool_call_argument_begin|>{"city": "Beijing"}'

904 - "<|tool_call_end|>"

905 - "<|tool_calls_section_end|>"

237 +class TestStreamingHappyPath:

238 + def test_single_tool_call(self, parser):

239 + """Verify DeltaToolCall output: name, id, arguments for one tool."""

240 + deltas = _split_tool_output_to_deltas(

241 + "I'll help. ",

242 + [("functions.get_weather:0", '{"city": "Beijing"}')],

243 + )

244 + rec = run_tool_extraction_streaming(parser, deltas)

245 +

246 + assert len(rec.tool_calls) == 1

247 + tc = rec.tool_calls[0]

248 + assert tc.function.name == "get_weather"

249 + assert tc.id == "functions.get_weather:0"

250 + assert json.loads(tc.function.arguments) == {"city": "Beijing"}

251 +

252 + def test_multiple_tool_calls(self, parser):

253 + """Two tool calls emitted with correct indices, names, arguments."""

254 + deltas = _split_tool_output_to_deltas(

255 + "Compare weather. ",

256 + [

257 + ("functions.get_weather:0", '{"city": "Tokyo"}'),

258 + ("functions.get_weather:1", '{"city": "NYC"}'),

259 + ],

260 + )

261 + rec = run_tool_extraction_streaming(parser, deltas)

262 +

263 + assert len(rec.tool_calls) == 2

264 + assert rec.tool_calls[0].function.name == "get_weather"

265 + assert rec.tool_calls[0].id == "functions.get_weather:0"

266 + assert json.loads(rec.tool_calls[0].function.arguments) == {"city": "Tokyo"}

267 +

268 + assert rec.tool_calls[1].function.name == "get_weather"

269 + assert rec.tool_calls[1].id == "functions.get_weather:1"

270 + assert json.loads(rec.tool_calls[1].function.arguments) == {"city": "NYC"}

271 +

272 + def test_content_before_tools(self, parser):

273 + """Content before section is streamed; markers/args don't leak."""

274 + deltas = _split_tool_output_to_deltas(

275 + "I'll check the weather. ",

276 + [("functions.get_weather:0", '{"city": "Tokyo"}')],

277 + )

278 + rec = run_tool_extraction_streaming(parser, deltas)

279 +

280 + assert "check the weather" in rec.other_content

281 + # No markers or tool content leaked

282 + for marker in [SECTION_BEGIN, SECTION_END, TOOL_BEGIN, TOOL_END, ARG_BEGIN]:

283 + assert marker not in rec.other_content

284 + assert "get_weather" not in rec.other_content

285 + assert "Tokyo" not in rec.other_content

286 +

287 + def test_no_tool_calls(self, parser):

288 + """Plain text streaming returns content only."""

289 + deltas = ["This is just ", "regular text ", "without tools."]

290 + rec = run_tool_extraction_streaming(parser, deltas)

291 +

292 + assert rec.other_content == "This is just regular text without tools."

293 + assert rec.tool_calls == []

294 +

295 + def test_incremental_arguments(self, parser):

296 + """Arguments split across small chunks accumulate correctly."""

297 + deltas = [

298 + "Help. ",

299 + SECTION_BEGIN,

300 + TOOL_BEGIN,

301 + "functions.get_weather:0 ",

302 + ARG_BEGIN,

303 + '{"ci',

304 + 'ty": "Be',

305 + 'ijing"}',

306 + " ",

307 + TOOL_END,

308 + SECTION_END,

309 + ]

310 + rec = run_tool_extraction_streaming(parser, deltas)

311 +

312 + assert len(rec.tool_calls) == 1

313 + assert rec.tool_calls[0].function.name == "get_weather"

314 + assert json.loads(rec.tool_calls[0].function.arguments) == {"city": "Beijing"}

315 +

316 + @pytest.mark.parametrize(

317 + "model_output",

318 + [

319 + pytest.param(

320 + "Single. "

321 + + _wrap(_tool("functions.get_weather:0", '{"city": "Beijing"}')),

322 + id="single_tool",

323 + ),

324 + pytest.param(

325 + "Multi. "

326 + + _wrap(

327 + _tool("functions.get_weather:0", '{"city": "Tokyo"}'),

328 + _tool("functions.get_news:1", '{"topic": "tech"}'),

329 + ),

330 + id="parallel_tools",

331 + ),

332 + pytest.param(

333 + "No prefix id. " + _wrap(_tool("get_weather:0", '{"city": "NYC"}')),

334 + id="no_functions_prefix",

335 + ),

336 + ],

337 )

338 + def test_streaming_matches_nonstreaming(self, parser, model_output):

339 + """Streaming reconstruction matches non-streaming extraction."""

340 + content_non, tools_non = run_tool_extraction(

341 + parser, model_output, streaming=False

342 + )

343 + content_stream, tools_stream = run_tool_extraction(

344 + parser, model_output, streaming=True

345 + )

346

908 - turn1_result = kimi_k2_tool_parser.extract_tool_calls(turn1_output, request=None)

909 - assert turn1_result.tools_called

910 - assert turn1_result.tool_calls[0].id == "functions.get_weather:0"

911 -

912 - # Fresh parser for turn 2

913 - turn2_parser = KimiK2ToolParser(kimi_k2_tokenizer)

914 - turn2_output = (

915 - "Now let me get news. "

916 - "<|tool_calls_section_begin|>"

917 - "<|tool_call_begin|>functions.get_news:0"

918 - '<|tool_call_argument_begin|>{"topic": "weather in Beijing"}'

919 - "<|tool_call_end|>"

920 - "<|tool_calls_section_end|>"

921 - )

347 + assert len(tools_non) == len(tools_stream)

348 + for tc_non, tc_stream in zip(tools_non, tools_stream):

349 + assert tc_non.function.name == tc_stream.function.name

350 + assert json.loads(tc_non.function.arguments) == json.loads(

351 + tc_stream.function.arguments

352 + )

353

923 - turn2_result = turn2_parser.extract_tool_calls(turn2_output, request=None)

924 - assert turn2_result.tools_called

925 - assert turn2_result.tool_calls[0].id == "functions.get_news:0"

354

355 +class TestStreamingEdgeCases:

356 + def test_marker_suppression(self, parser):

357 + """No special-token markers appear in reconstructed content."""

358 + deltas = _split_tool_output_to_deltas(

359 + "I'll check. ",

360 + [("functions.get_weather:0", '{"city": "Tokyo"}')],

361 + )

362 + rec = run_tool_extraction_streaming(parser, deltas)

363 +

364 + forbidden = [SECTION_BEGIN, SECTION_END, TOOL_BEGIN, TOOL_END, ARG_BEGIN]

365 + for marker in forbidden:

366 + assert marker not in rec.other_content, (

367 + f"Marker leaked: {marker!r} in {rec.other_content!r}"

368 + )

369 +

370 + def test_noise_between_markers_suppressed(self, parser):

371 + """Text between section_begin and tool_call_begin doesn't leak."""

372 + deltas = [

373 + "Reasoning. ",

374 + SECTION_BEGIN,

375 + " spurious noise ",

376 + TOOL_BEGIN,

377 + "functions.test:0 ",

378 + ARG_BEGIN,

379 + '{"k": "v"} ',

380 + TOOL_END,

381 + SECTION_END,

382 + ]

383 + rec = run_tool_extraction_streaming(parser, deltas)

384 +

385 + assert "spurious" not in rec.other_content

386 + assert "noise" not in rec.other_content

387 +

388 + def test_empty_tool_section(self, parser):

389 + """Empty section (begin immediately followed by end) doesn't crash."""

390 + deltas = ["Reasoning. ", SECTION_BEGIN, SECTION_END]

391 + rec = run_tool_extraction_streaming(parser, deltas)

392 +

393 + assert rec.tool_calls == []

394 +

395 + def test_three_different_tools(self, parser):

396 + """Three tool calls with different functions stream correctly."""

397 + deltas = _split_tool_output_to_deltas(

398 + "Multiple tasks. ",

399 + [

400 + ("functions.get_weather:0", '{"city": "NYC"}'),

401 + ("functions.get_news:1", '{"topic": "tech"}'),

402 + ("functions.send_email:2", '{"to": "a@b.com"}'),

403 + ],

404 + )

405 + rec = run_tool_extraction_streaming(parser, deltas)

406 +

407 + assert len(rec.tool_calls) == 3

408 + names = [tc.function.name for tc in rec.tool_calls]

409 + assert names == ["get_weather", "get_news", "send_email"]

410 + ids = [tc.id for tc in rec.tool_calls]

411 + assert len(set(ids)) == 3 # unique ids

412 +

413 + def test_truncated_tool_call_no_end_marker(self, parser):

414 + """Stream ending mid-tool-call (max_tokens) doesn't crash."""

415 + deltas = [

416 + "I'll check. ",

417 + SECTION_BEGIN,

418 + TOOL_BEGIN,

419 + "functions.get_weather:0 ",

420 + ARG_BEGIN,

421 + '{"city": "Bei',

422 + # Stream ends here — no TOOL_END, no SECTION_END

423 + ]

424 + rec = run_tool_extraction_streaming(parser, deltas)

425 +

426 + # Should not crash; tool name and partial args extracted

427 + assert len(rec.tool_calls) == 1

428 + assert rec.tool_calls[0].function.name == "get_weather"

429 + assert rec.tool_calls[0].id == "functions.get_weather:0"

430 + assert rec.tool_calls[0].function.arguments == '{"city": "Bei'

431 + # No markers leaked into content

432 + for marker in [SECTION_BEGIN, SECTION_END, TOOL_BEGIN, TOOL_END, ARG_BEGIN]:

433 + assert marker not in rec.other_content

434 +

435 + def test_content_after_tool_section(self, parser):

436 + """Trailing text after section_end doesn't crash or leak markers."""

437 + deltas = [

438 + "Before. ",

439 + SECTION_BEGIN,

440 + TOOL_BEGIN,

441 + "functions.get_weather:0 ",

442 + ARG_BEGIN,

443 + '{"city": "Tokyo"} ',

444 + TOOL_END,

445 + SECTION_END,

446 + " After tools.",

447 + ]

448 + rec = run_tool_extraction_streaming(parser, deltas)

449 +

450 + # Tool call extracted correctly

451 + assert len(rec.tool_calls) == 1

452 + assert rec.tool_calls[0].function.name == "get_weather"

453 + assert json.loads(rec.tool_calls[0].function.arguments) == {"city": "Tokyo"}

454 + # Trailing content after tool section is dropped

455 + assert "After tools." not in rec.other_content

456 + # No markers leaked into content

457 + for marker in [SECTION_BEGIN, SECTION_END, TOOL_BEGIN, TOOL_END, ARG_BEGIN]:

458 + assert marker not in rec.other_content

459 +

460 +

461 +class TestAdjustRequest:

462 + def test_sets_skip_special_tokens_false(self, parser):

463 + request = MagicMock(spec=ChatCompletionRequest)

464 + request.tools = [{"type": "function", "function": {"name": "test"}}]

465 + request.tool_choice = "auto"

466 + request.skip_special_tokens = True

467 +

468 + result = parser.adjust_request(request)

469 + assert result.skip_special_tokens is False

470 +

471 + def test_no_change_when_tool_choice_none(self, parser):

472 + request = MagicMock(spec=ChatCompletionRequest)

473 + request.tools = [{"type": "function", "function": {"name": "test"}}]

474 + request.tool_choice = "none"

475 + request.skip_special_tokens = True

476 +

477 + result = parser.adjust_request(request)

478 + assert result.skip_special_tokens is True

479 +

480 + def test_no_change_when_no_tools(self, parser):

481 + request = MagicMock(spec=ChatCompletionRequest)

482 + request.tools = None

483 + request.tool_choice = "auto"

484 + request.skip_special_tokens = True

485 +

486 + result = parser.adjust_request(request)

487 + assert result.skip_special_tokens is True

488 +

489 +

490 +def _chunk_tokenized_deltas(tokenizer, text: str, stream_interval: int) -> list[str]:

491 + """Encode text, group tokens into chunks of stream_interval, decode each."""

492 + token_ids = tokenizer.encode(text, add_special_tokens=False)

493 + deltas = []

494 + prev = ""

495 + for i in range(0, len(token_ids), stream_interval):

496 + decoded = tokenizer.decode(

497 + token_ids[: i + stream_interval], skip_special_tokens=False

498 + )

499 + deltas.append(decoded[len(prev) :])

500 + prev = decoded

501 + return deltas

502

928 -def test_streaming_multiple_tool_calls_not_leaked(kimi_k2_tool_parser):

929 - """

930 - Test that MULTIPLE tool calls in streaming mode do not leak into content.

931 - This reproduces the AWS Bedrock scenario: "Compare weather in Tokyo and NYC".

932 - """

933 - kimi_k2_tool_parser.reset_streaming_state()

503

935 - section_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_begin|>")

936 - section_end_id = kimi_k2_tool_parser.vocab.get("<|tool_calls_section_end|>")

937 - tool_begin_id = kimi_k2_tool_parser.vocab.get("<|tool_call_begin|>")

938 - tool_end_id = kimi_k2_tool_parser.vocab.get("<|tool_call_end|>")

504 +class TestStreamingIntervals:

505 + """Test streaming at various token-chunk sizes to catch boundary bugs."""

506

940 - all_content = []

507 + @pytest.mark.parametrize("stream_interval", [1, 2, 3, 5, 8])

508 + def test_single_tool_call_at_interval(self, kimi_k2_tokenizer, stream_interval):

509 + text = "Help. " + _wrap(_tool("functions.get_weather:0", '{"city": "Beijing"}'))

510 + deltas = _chunk_tokenized_deltas(kimi_k2_tokenizer, text, stream_interval)

511 + parser = KimiK2ToolParser(kimi_k2_tokenizer)

512 + rec = run_tool_extraction_streaming(

513 + parser, deltas, assert_one_tool_per_delta=False

514 + )

515

942 - tool1 = '<|tool_call_begin|> get_weather:0 <|tool_call_argument_begin|> {"city": "Tokyo"} <|tool_call_end|>'

943 - tool2 = ' <|tool_call_begin|> get_weather:1 <|tool_call_argument_begin|> {"city": "New York"} <|tool_call_end|>'

516 + assert len(rec.tool_calls) == 1

517 + assert rec.tool_calls[0].function.name == "get_weather"

518 + assert json.loads(rec.tool_calls[0].function.arguments) == {"city": "Beijing"}

519

945 - deltas = [

946 - ("I'll compare the weather. ", [1, 2, 3]),

947 - ("<|tool_calls_section_begin|>", [section_begin_id]),

948 - (tool1, [tool_begin_id, 10, tool_end_id]),

949 - (tool2, [tool_begin_id, 20, tool_end_id]),

950 - ("<|tool_calls_section_end|>", [section_end_id]),

951 - (" Here's the comparison.", [30]),

952 - ]

520 + @pytest.mark.parametrize("stream_interval", [1, 2, 3, 5, 8])

521 + def test_content_then_tool_call_at_interval(

522 + self, kimi_k2_tokenizer, stream_interval

523 + ):

524 + text = "Sure, let me check. " + _wrap(

525 + _tool("functions.get_weather:0", '{"city": "Tokyo"}')

526 + )

527 + deltas = _chunk_tokenized_deltas(kimi_k2_tokenizer, text, stream_interval)

528 + parser = KimiK2ToolParser(kimi_k2_tokenizer)

529 + rec = run_tool_extraction_streaming(

530 + parser, deltas, assert_one_tool_per_delta=False

531 + )

532

954 - results = run_streaming_sequence(kimi_k2_tool_parser, deltas)

533 + assert "let me check" in rec.other_content

534 + assert "get_weather" not in rec.other_content

535 + assert len(rec.tool_calls) == 1

536 + assert rec.tool_calls[0].function.name == "get_weather"

537 + assert json.loads(rec.tool_calls[0].function.arguments) == {"city": "Tokyo"}

538 +

539 + @pytest.mark.parametrize("stream_interval", [1, 2, 3, 5, 8])

540 + def test_multiple_tool_calls_at_interval(self, kimi_k2_tokenizer, stream_interval):

541 + text = "Compare. " + _wrap(

542 + _tool("functions.search:0", '{"q": "cats"}'),

543 + _tool("functions.search:1", '{"q": "dogs"}'),

544 + )

545 + deltas = _chunk_tokenized_deltas(kimi_k2_tokenizer, text, stream_interval)

546 + parser = KimiK2ToolParser(kimi_k2_tokenizer)

547 + rec = run_tool_extraction_streaming(

548 + parser, deltas, assert_one_tool_per_delta=False

549 + )

550

956 - for res in results:

957 - if res and res.content:

958 - all_content.append(res.content)

551 + assert len(rec.tool_calls) == 2

552 + assert rec.tool_calls[0].function.name == "search"

553 + assert json.loads(rec.tool_calls[0].function.arguments) == {"q": "cats"}

554 + assert rec.tool_calls[1].function.name == "search"

555 + assert json.loads(rec.tool_calls[1].function.arguments) == {"q": "dogs"}

556 +

557 + @pytest.mark.parametrize("stream_interval", [1, 2, 3, 5, 8])

558 + def test_plain_text_at_interval(self, kimi_k2_tokenizer, stream_interval):

559 + text = "This is plain text with no tool calling involved."

560 + deltas = _chunk_tokenized_deltas(kimi_k2_tokenizer, text, stream_interval)

561 + parser = KimiK2ToolParser(kimi_k2_tokenizer)

562 + rec = run_tool_extraction_streaming(

563 + parser, deltas, assert_one_tool_per_delta=False

564 + )

565

960 - # Assertions

961 - full_content = "".join(all_content)

566 + assert rec.other_content == text

567 + assert rec.tool_calls == []

568

963 - # Check no markers leaked

964 - forbidden = ["<|tool_call", "<|tool_calls_section"]

965 - for marker in forbidden:

966 - assert marker not in full_content, (

967 - f"MARKER LEAKED: {marker} in {repr(full_content)}"

569 + def test_content_and_tool_call_in_single_chunk(self, kimi_k2_tokenizer):

570 + """Content + complete tool call in one chunk must both be emitted."""

571 + text = "Hi! " + _wrap(_tool("functions.get_weather:0", '{"city": "Beijing"}'))

572 + deltas = _chunk_tokenized_deltas(kimi_k2_tokenizer, text, stream_interval=9999)

573 + parser = KimiK2ToolParser(kimi_k2_tokenizer)

574 + rec = run_tool_extraction_streaming(

575 + parser, deltas, assert_one_tool_per_delta=False

576 )

577

970 - # Check no tool call content leaked (both tools)

971 - assert "get_weather" not in full_content, f"TOOL NAME LEAKED: {repr(full_content)}"

972 - assert "Tokyo" not in full_content, f"TOOL ARG LEAKED (Tokyo): {repr(full_content)}"

973 - assert "New York" not in full_content, (

974 - f"TOOL ARG LEAKED (NYC): {repr(full_content)}"

975 - )

976 -

977 - # Legitimate content preserved

978 - assert "compare" in full_content.lower() or len(all_content) > 0

578 + assert "Hi!" in rec.other_content

579 + assert "get_weather" not in rec.other_content

580 + assert len(rec.tool_calls) == 1

581 + assert rec.tool_calls[0].function.name == "get_weather"

582 + assert json.loads(rec.tool_calls[0].function.arguments) == {"city": "Beijing"}

vllm/tool_parsers/kimi_k2_tool_parser.py

+159 -484

@@ -1,6 +1,5 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -# code modified from deepseekv3_tool_parser.py

4 from collections.abc import Sequence

@@ -17,12 +16,14 @@ from vllm.entrypoints.openai.engine.protocol import (

16 FunctionCall,

17 ToolCall,

18 )

19 +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest

20 from vllm.logger import init_logger

21 from vllm.tokenizers import TokenizerLike

22 from vllm.tool_parsers.abstract_tool_parser import (

23 Tool,

24 ToolParser,

25 )

26 +from vllm.tool_parsers.utils import partial_tag_overlap

28 logger = init_logger(__name__)

@@ -30,124 +31,44 @@ logger = init_logger(__name__)

31 class KimiK2ToolParser(ToolParser):

32 def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None):

33 super().__init__(tokenizer, tools)

33 - self.current_tool_name_sent: bool = False

34 +

35 + # Streaming state

36 + self._sent_content_idx: int = 0

37 self.prev_tool_call_arr: list[dict] = []

35 - self.current_tool_id: int = -1

36 - self.streamed_args_for_tool: list[

37 - str

38 - ] = [] # map what has been streamed for each tool so far to a list

39 -

40 - # Section-level state management to prevent token leakage

41 - self.in_tool_section: bool = False

42 - self.token_buffer: str = ""

43 - # Buffer size: empirical worst-case for longest marker (~30 chars) * 2

44 - # + safety margin for unicode + partial overlap. Prevents unbounded growth.

45 - self.buffer_max_size: int = 1024

46 - self.section_char_count: int = 0 # Track characters processed in tool section

47 - self.max_section_chars: int = 8192 # Force exit if section exceeds this

48 - self._buffer_overflow_logged: bool = False # Log overflow once per session

49 -

50 - # Support both singular and plural variants

38 + self.streamed_args_for_tool: list[str] = []

39 +

40 + # Section marker

41 self.tool_calls_start_token: str = "<|tool_calls_section_begin|>"

52 - self.tool_calls_end_token: str = "<|tool_calls_section_end|>"

53 - self.tool_calls_start_token_variants: list[str] = [

54 - "<|tool_calls_section_begin|>",

55 - "<|tool_call_section_begin|>", # singular variant

56 - ]

57 - self.tool_calls_end_token_variants: list[str] = [

58 - "<|tool_calls_section_end|>",

59 - "<|tool_call_section_end|>", # singular variant

60 - ]

43 + # Individual tool call markers

44 self.tool_call_start_token: str = "<|tool_call_begin|>"

45 self.tool_call_end_token: str = "<|tool_call_end|>"

46 + self.tool_call_arg_token: str = "<|tool_call_argument_begin|>"

48 + # Regex for non-streaming extraction

49 self.tool_call_regex = re.compile(

50 + r"<\|tool_call_begin\|>\s*(?P<tool_call_id>[^<]+:\d+)\s*"

51 + r"<\|tool_call_argument_begin\|>\s*"

52 + r"(?P<function_arguments>(?:(?!<\|tool_call_begin\|>).)*?)\s*"

53 + r"<\|tool_call_end\|>",

54 re.DOTALL,

55 )

70 - self.stream_tool_call_portion_regex = re.compile(

71 - r"(?P<tool_call_id>.+:\d+)\s*<\|tool_call_argument_begin\|>\s*(?P<function_arguments>.*)"

72 - )

73 -

74 - self.stream_tool_call_name_regex = re.compile(r"(?P<tool_call_id>.+:\d+)\s*")

75 -

57 if not self.model_tokenizer:

58 raise ValueError(

59 "The model tokenizer must be passed to the ToolParser "

60 "constructor during construction."

61 )

81 - self.tool_calls_start_token_id = self.vocab.get(self.tool_calls_start_token)

82 - self.tool_calls_end_token_id = self.vocab.get(self.tool_calls_end_token)

83 -

84 - # Get token IDs for all variants

85 - self.tool_calls_start_token_ids: list[int] = [

86 - tid

87 - for variant in self.tool_calls_start_token_variants

88 - if (tid := self.vocab.get(variant)) is not None

89 - ]

90 - self.tool_calls_end_token_ids: list[int] = [

91 - tid

92 - for variant in self.tool_calls_end_token_variants

93 - if (tid := self.vocab.get(variant)) is not None

94 - ]

95 -

96 - self.tool_call_start_token_id = self.vocab.get(self.tool_call_start_token)

97 - self.tool_call_end_token_id = self.vocab.get(self.tool_call_end_token)

98 -

99 - if (

100 - self.tool_calls_start_token_id is None

101 - or self.tool_calls_end_token_id is None

102 - ):

103 - raise RuntimeError(

104 - "Kimi-K2 Tool parser could not locate tool call start/end "

105 - "tokens in the tokenizer!"

106 - )

108 - def _check_and_strip_markers(self, text: str) -> tuple[str, bool, bool]:

109 - """

110 - Check for section begin/end markers in text and strip them.

111 - Returns: (cleaned_text, found_section_begin, found_section_end)

112 - """

113 - found_begin = False

114 - found_end = False

115 - cleaned = text

116 -

117 - # Check for section begin markers (any variant)

118 - for variant in self.tool_calls_start_token_variants:

119 - if variant in cleaned:

120 - cleaned = cleaned.replace(variant, "")

121 - found_begin = True

122 -

123 - # Check for section end markers (any variant)

124 - for variant in self.tool_calls_end_token_variants:

125 - if variant in cleaned:

126 - cleaned = cleaned.replace(variant, "")

127 - found_end = True

128 - return cleaned, found_begin, found_end

129 -

130 - def _reset_section_state(self) -> None:

131 - """Reset state when exiting tool section."""

132 - self.in_tool_section = False

133 - self.token_buffer = ""

134 - self.section_char_count = 0

135 -

136 - def reset_streaming_state(self) -> None:

137 - """

138 - Reset all streaming state. Call this between requests to prevent

139 - state leakage when parser instance is reused.

140 - """

141 - # Reset section state

142 - self._reset_section_state()

143 -

144 - # Reset parent class state

145 - self.current_tool_name_sent = False

146 - self.prev_tool_call_arr = []

147 - self.current_tool_id = -1

148 - self.streamed_args_for_tool = []

149 -

150 - logger.debug("Streaming state reset")

63 + def adjust_request(

64 + self, request: ChatCompletionRequest | ResponsesRequest

65 + ) -> ChatCompletionRequest | ResponsesRequest:

66 + request = super().adjust_request(request)

67 + if request.tools and request.tool_choice != "none":

68 + # Ensure special-token markers appear as literal text in

69 + # current_text so we can do pure text-based parsing.

70 + request.skip_special_tokens = False

71 + return request

73 def extract_tool_calls(

74 self,

@@ -198,6 +119,95 @@ class KimiK2ToolParser(ToolParser):

119 tools_called=False, tool_calls=[], content=model_output

120 )

121

122 + def _extract_content(self, current_text: str) -> str | None:

123 + """Return unsent content before the tool-calls section, or None.

124 +

125 + Holds back any trailing suffix that partially matches

126 + ``<|tool_calls_section_begin|>`` to avoid leaking marker bytes.

127 + """

128 + if self.tool_calls_start_token not in current_text:

129 + overlap = partial_tag_overlap(current_text, self.tool_calls_start_token)

130 + sendable_idx = len(current_text) - overlap

131 + else:

132 + sendable_idx = current_text.index(self.tool_calls_start_token)

133 +

134 + if sendable_idx > self._sent_content_idx:

135 + content = current_text[self._sent_content_idx : sendable_idx]

136 + self._sent_content_idx = sendable_idx

137 + return content

138 + return None

139 +

140 + def _extract_tool_calls(self, current_text: str) -> list[str]:

141 + """Extract raw bodies from ``<|tool_call_begin|>…<|tool_call_end|>`` blocks."""

142 + if self.tool_calls_start_token not in current_text:

143 + return []

144 +

145 + results: list[str] = []

146 + pos = current_text.index(self.tool_calls_start_token)

147 + while True:

148 + start = current_text.find(self.tool_call_start_token, pos)

149 + if start == -1:

150 + break

151 + tc_start = start + len(self.tool_call_start_token)

152 + end = current_text.find(self.tool_call_end_token, tc_start)

153 +

154 + if end != -1:

155 + tool_call = current_text[tc_start:end]

156 + pos = end + len(self.tool_call_end_token)

157 + else:

158 + tool_call = current_text[tc_start:]

159 + overlap = partial_tag_overlap(tool_call, self.tool_call_end_token)

160 + if overlap:

161 + tool_call = tool_call[:-overlap]

162 +

163 + results.append(tool_call)

164 +

165 + if end == -1:

166 + break

167 + return results

168 +

169 + @staticmethod

170 + def _extract_tool_id_and_name(

171 + header: str | None,

172 + ) -> tuple[str | None, str | None]:

173 + """Parse ``(tool_id, tool_name)`` from a header

174 + like ``"functions.get_weather:0"``."""

175 + if header is None:

176 + return None, None

177 + match = re.match(r"(.+:\d+)", header)

178 + if not match:

179 + return None, None

180 +

181 + tool_id = match.group(1).strip()

182 + tool_name = tool_id.split(":")[0].split(".")[-1]

183 + return tool_id, tool_name

184 +

185 + def _split_tool_call(self, tool_call: str) -> tuple[str | None, str | None]:

186 + """Split a tool-call body into ``(header, arguments)`` at the argument marker.

187 +

188 + Example::

189 + 'get_weather:0 <|tool_call_argument_begin|>{"c'

190 + -> ("get_weather:0", '{"c')

191 + """

192 + arg_pos = tool_call.find(self.tool_call_arg_token)

193 + if arg_pos == -1:

194 + return None, None

195 + header = tool_call[:arg_pos].strip()

196 + tool_args = tool_call[arg_pos + len(self.tool_call_arg_token) :]

197 + return header, tool_args

198 +

199 + def _compute_args_diff(self, index: int, tool_args: str | None) -> str | None:

200 + """Return new argument text not yet sent for tool `index`, or None."""

201 + if tool_args is None:

202 + return None

203 + prev = self.streamed_args_for_tool[index]

204 + if len(tool_args) <= len(prev):

205 + return None

206 + diff = tool_args[len(prev) :]

207 + self.streamed_args_for_tool[index] = tool_args

208 + self.prev_tool_call_arr[index]["arguments"] = tool_args

209 + return diff

210 +

211 def extract_tool_calls_streaming(

212 self,

213 previous_text: str,

@@ -208,394 +218,59 @@ class KimiK2ToolParser(ToolParser):

218 delta_token_ids: Sequence[int],

219 request: ChatCompletionRequest,

220 ) -> DeltaMessage | None:

211 - logger.debug("delta_text: %s", delta_text)

212 - logger.debug("delta_token_ids: %s", delta_token_ids)

213 -

214 - # Flag to defer section exit until after tool parsing completes

215 - deferred_section_exit = False

216 -

217 - # Add delta to buffer for split marker detection

218 - self.token_buffer += delta_text

219 -

220 - # Enforce buffer size limit to prevent memory issues

221 - if len(self.token_buffer) > self.buffer_max_size:

222 - if not self._buffer_overflow_logged:

223 - logger.warning(

224 - "Token buffer exceeded max size (%d bytes), flushing excess. "

225 - "This may indicate very long markers or unusual tokenization.",

226 - self.buffer_max_size,

227 - )

228 - self._buffer_overflow_logged = True

229 - # Keep only the most recent content that might contain partial markers

230 - self.token_buffer = self.token_buffer[-self.buffer_max_size // 2 :]

231 -

232 - # Check buffer for section markers (handles split tokens)

233 - buffered_text, found_section_begin, found_section_end = (

234 - self._check_and_strip_markers(self.token_buffer)

235 - )

236 -

237 - # Track section state transitions

238 - if found_section_begin and not self.in_tool_section:

239 - logger.debug("Entering tool section")

240 - self.in_tool_section = True

241 - self.token_buffer = buffered_text # Use cleaned buffer

242 - self.section_char_count = 0 # Reset counter for new section

243 -

244 - if found_section_end and self.in_tool_section:

245 - logger.debug("Detected section end marker")

246 - # CRITICAL: Don't exit early if tool_call_end is in this chunk.

247 - # Tool parser must emit final arguments/close first to avoid dropping

248 - # the final tool update and leaking tokens into reasoning channel.

249 - has_tool_end = self.tool_call_end_token_id in delta_token_ids

250 - if has_tool_end:

251 - # Defer exit until after tool parsing completes

252 - deferred_section_exit = True

253 - logger.debug("Deferring section exit: tool_call_end in same chunk")

254 - self.token_buffer = buffered_text

255 - else:

256 - # No tool call ending, safe to exit immediately

257 - logger.debug("Exiting tool section")

258 - self._reset_section_state()

259 - # Extract any content AFTER the section end marker in delta_text

260 - # (don't use buffered_text as it contains tool call data)

261 - post_section_content = ""

262 - for variant in self.tool_calls_end_token_variants:

263 - if variant in delta_text:

264 - parts = delta_text.split(variant, 1)

265 - if len(parts) > 1:

266 - post_section_content = parts[1]

267 - break

268 - if post_section_content.strip():

269 - return DeltaMessage(content=post_section_content)

270 - return DeltaMessage(content="")

271 - else:

272 - self.token_buffer = buffered_text

273 -

274 - # Check if any variant of section start token is in current_token_ids

275 - has_section_token = any(

276 - tid in current_token_ids for tid in self.tool_calls_start_token_ids

277 - )

278 -

279 - # Early return: if no section token detected yet, return as reasoning content

280 - if not has_section_token and not self.in_tool_section:

281 - logger.debug("No tool call tokens found!")

282 - # Don't clear buffer - it needs to accumulate partial markers across deltas

283 - # Buffer overflow is already protected by lines 215-224

284 - return DeltaMessage(content=delta_text)

285 -

286 - # Strip section markers from delta_text for subsequent processing

287 - # NOTE: This preprocessing happens BEFORE the regex-based tool call

288 - # parsing (from PR #24847) to ensure markers are removed cleanly

289 - # before pattern matching. No double-stripping occurs because

290 - # section markers and tool call markers are distinct.

291 - delta_text, _, _ = self._check_and_strip_markers(delta_text)

292 -

293 - # Error recovery: If in tool section for too long, force exit

294 - if self.in_tool_section:

295 - self.section_char_count += len(delta_text)

296 - if self.section_char_count > self.max_section_chars:

297 - logger.warning(

298 - "Tool section exceeded max length (%d chars), forcing exit. "

299 - "This may indicate malformed model output.",

300 - self.max_section_chars,

301 - )

302 - self._reset_section_state()

303 - # Deferred exit already handled by forced exit above

304 - # Return remaining content as reasoning (or empty delta if no content)

305 - return DeltaMessage(content=delta_text if delta_text.strip() else "")

306 -

221 try:

308 - # figure out where we are in the parsing by counting tool call

309 - # start & end tags

310 - prev_tool_start_count = previous_token_ids.count(

311 - self.tool_call_start_token_id

312 - )

313 - prev_tool_end_count = previous_token_ids.count(self.tool_call_end_token_id)

314 - cur_tool_start_count = current_token_ids.count(

315 - self.tool_call_start_token_id

316 - )

317 - cur_tool_end_count = current_token_ids.count(self.tool_call_end_token_id)

318 - tool_call_portion = None

319 - text_portion = None

320 -

321 - # case: if we're generating text, OR rounding out a tool call

322 - if (

323 - cur_tool_start_count == cur_tool_end_count

324 - and prev_tool_end_count == cur_tool_end_count

325 - and self.tool_call_end_token not in delta_text

326 - ):

327 - # Suppress content between section begin and first tool begin

328 - # (header noise). Don't suppress content between tools to avoid

329 - # breaking potential delimiter characters.

330 - if self.in_tool_section and cur_tool_start_count == 0:

331 - logger.debug(

332 - "In tool section before first tool, suppressing: %s",

333 - delta_text,

334 - )

335 - # Return empty delta to maintain iterator contract

336 - return DeltaMessage(content="")

337 - logger.debug("Generating text content! skipping tool parsing.")

338 - return DeltaMessage(content=delta_text)

339 -

340 - if self.tool_call_end_token in delta_text:

341 - logger.debug("tool_call_end_token in delta_text")

342 - full_text = current_text + delta_text

343 - tool_call_portion = (

344 - full_text.split(self.tool_call_start_token)[-1]

345 - .split(self.tool_call_end_token)[0]

346 - .rstrip()

347 - )

348 - delta_text = delta_text.split(self.tool_call_end_token)[0].rstrip()

349 - text_portion = delta_text.split(self.tool_call_end_token)[-1].lstrip()

350 -

351 - # case -- we're starting a new tool call

352 - if (

353 - cur_tool_start_count > cur_tool_end_count

354 - and cur_tool_start_count > prev_tool_start_count

355 - ):

356 - if len(delta_token_ids) > 1:

357 - tool_call_portion = current_text.split(self.tool_call_start_token)[

358 - -1

359 - ]

360 - else:

361 - tool_call_portion = None

362 - delta = None

363 -

364 - text_portion = None

365 -

366 - # set cursors and state appropriately

367 - self.current_tool_id += 1

368 - self.current_tool_name_sent = False

369 - self.streamed_args_for_tool.append("")

370 - logger.debug("Starting on a new tool %s", self.current_tool_id)

371 -

372 - # case -- we're updating an existing tool call

373 - elif (

374 - cur_tool_start_count > cur_tool_end_count

375 - and cur_tool_start_count == prev_tool_start_count

376 - ):

377 - # get the portion of the text that's the tool call

378 - tool_call_portion = current_text.split(self.tool_call_start_token)[-1]

379 - text_portion = None

380 -

381 - # case -- the current tool call is being closed.

382 - elif (

383 - cur_tool_start_count == cur_tool_end_count

384 - and cur_tool_end_count >= prev_tool_end_count

385 - ):

386 - if self.prev_tool_call_arr is None or len(self.prev_tool_call_arr) == 0:

387 - logger.debug("attempting to close tool call, but no tool call")

388 - # Handle deferred section exit before returning

389 - if deferred_section_exit and self.in_tool_section:

390 - self._reset_section_state()

391 - return None

392 - diff = self.prev_tool_call_arr[self.current_tool_id].get("arguments")

393 - if diff:

394 - diff = (

395 - diff.encode("utf-8").decode("unicode_escape")

396 - if diff is str

397 - else diff

398 - )

399 - if '"}' not in delta_text:

400 - # Handle deferred section exit before returning

401 - if deferred_section_exit and self.in_tool_section:

402 - self._reset_section_state()

403 - return None

404 - end_loc = delta_text.rindex('"}')

405 - diff = delta_text[:end_loc] + '"}'

406 - logger.debug(

407 - "Finishing tool and found diff that had not "

408 - "been streamed yet: %s",

409 - diff,

410 - )

411 - self.streamed_args_for_tool[self.current_tool_id] += diff

412 - # Handle deferred section exit before returning

413 - if deferred_section_exit and self.in_tool_section:

414 - logger.debug("Completing deferred section exit")

415 - self._reset_section_state()

416 - return DeltaMessage(

417 - tool_calls=[

418 - DeltaToolCall(

419 - index=self.current_tool_id,

420 - function=DeltaFunctionCall(arguments=diff).model_dump(

421 - exclude_none=True

422 - ),

423 - )

424 - ]

425 - )

426 -

427 - # case -- otherwise we're just generating text

428 - else:

429 - # Check if we're in tool section - if so, suppress

430 - if self.in_tool_section:

431 - logger.debug("In tool section, suppressing text generation")

432 - # Handle deferred section exit before returning

433 - if deferred_section_exit:

434 - self._reset_section_state()

435 - return DeltaMessage(content="")

436 - text = delta_text.replace(self.tool_call_start_token, "")

437 - text = text.replace(self.tool_call_end_token, "")

438 - delta = DeltaMessage(tool_calls=[], content=text)

439 - # Handle deferred section exit before returning

440 - if deferred_section_exit and self.in_tool_section:

441 - self._reset_section_state()

442 - return delta

443 -

444 - current_tool_call = dict()

445 - if tool_call_portion:

446 - current_tool_call_matches = self.stream_tool_call_portion_regex.match(

447 - tool_call_portion

448 - )

449 - if current_tool_call_matches:

450 - tool_id, tool_args = current_tool_call_matches.groups()

451 - tool_name = tool_id.split(":")[0].split(".")[-1]

452 - current_tool_call["id"] = tool_id.strip()

453 - current_tool_call["name"] = tool_name

454 - current_tool_call["arguments"] = tool_args

455 - else:

456 - current_tool_call_name_matches = (

457 - self.stream_tool_call_name_regex.match(tool_call_portion)

458 - )

459 - if current_tool_call_name_matches:

460 - (tool_id_str,) = current_tool_call_name_matches.groups()

461 - tool_name = tool_id_str.split(":")[0].split(".")[-1]

462 - current_tool_call["id"] = tool_id_str.strip()

463 - current_tool_call["name"] = tool_name

464 - current_tool_call["arguments"] = ""

465 - else:

466 - logger.debug("Not enough token")

467 - return None

468 -

469 - # case - we haven't sent the tool name yet. If it's available, send

470 - # it. otherwise, wait until it's available.

471 - if not self.current_tool_name_sent:

472 - if current_tool_call is None:

473 - return None

474 - function_name: str | None = current_tool_call.get("name")

475 - tool_id = current_tool_call.get("id")

476 - if function_name:

477 - self.current_tool_name_sent = True

478 - return DeltaMessage(

479 - tool_calls=[

480 - DeltaToolCall(

481 - index=self.current_tool_id,

482 - type="function",

483 - id=tool_id,

484 - function=DeltaFunctionCall(

485 - name=function_name

486 - ).model_dump(exclude_none=True),

487 - )

488 - ]

222 + # Extract any content before tool calls.

223 + content = self._extract_content(current_text)

224 + tool_calls = self._extract_tool_calls(current_text)

225 + tool_call_deltas: list[DeltaToolCall] = []

226 +

227 + for i, tool_call in enumerate(tool_calls):

228 + # First time seeing tool call at index i.

229 + if i >= len(self.prev_tool_call_arr):

230 + # Initialize streaming state.

231 + self.prev_tool_call_arr.append({})

232 + self.streamed_args_for_tool.append("")

233 +

234 + header, tool_args = self._split_tool_call(tool_call)

235 +

236 + # Stream back tool name.

237 + if "name" not in self.prev_tool_call_arr[i]:

238 + tool_id, tool_name = self._extract_tool_id_and_name(header)

239 + if not tool_name:

240 + # Can't skip to tool i+1 if i isn't ready

241 + break

242 + self.prev_tool_call_arr[i]["name"] = tool_name

243 + self.prev_tool_call_arr[i]["id"] = tool_id

244 + tool_call_deltas.append(

245 + DeltaToolCall(

246 + index=i,

247 + type="function",

248 + id=tool_id,

249 + function=DeltaFunctionCall(name=tool_name).model_dump(

250 + exclude_none=True

251 + ),

252 + )

253 )

490 - else:

491 - return None

492 -

493 - # case -- otherwise, send the tool call delta

494 -

495 - # if the tool call portion is None, send the delta as text

496 - if tool_call_portion is None:

497 - # if there's text but not tool calls, send that -

498 - # otherwise None to skip chunk

499 - # CRITICAL: Never return content if we're in a tool section

500 - if self.in_tool_section:

501 - return None

502 - delta = (

503 - DeltaMessage(content=delta_text)

504 - if text_portion is not None

505 - else None

506 - )

507 - return delta

508 -

509 - # now, the nitty-gritty of tool calls

510 - # now we have the portion to parse as tool call.

511 -

512 - logger.debug(

513 - "Trying to parse current tool call with ID %s", self.current_tool_id

514 - )

515 -

516 - # if we're starting a new tool call, push an empty object in as

517 - # a placeholder for the arguments

518 - if len(self.prev_tool_call_arr) <= self.current_tool_id:

519 - self.prev_tool_call_arr.append({})

520 -

521 - # main logic for tool parsing here - compare prev. partially-parsed

522 - # JSON to the current partially-parsed JSON

523 - prev_arguments = self.prev_tool_call_arr[self.current_tool_id].get(

524 - "arguments"

525 - )

526 - cur_arguments = current_tool_call.get("arguments")

527 -

528 - logger.debug("diffing old arguments: %s", prev_arguments)

529 - logger.debug("against new ones: %s", cur_arguments)

530 -

531 - # case -- no arguments have been created yet. skip sending a delta.

532 - if not cur_arguments and not prev_arguments:

533 - logger.debug("Skipping text %s - no arguments", delta_text)

534 - delta = None

535 -

536 - # case -- prev arguments are defined, but non are now.

537 - # probably impossible, but not a fatal error - just keep going

538 - elif not cur_arguments and prev_arguments:

539 - logger.error(

540 - "should be impossible to have arguments reset "

541 - "mid-call. skipping streaming anything."

542 - )

543 - delta = None

254

545 - # case -- we now have the first info about arguments available from

546 - # autocompleting the JSON

547 - elif cur_arguments and not prev_arguments:

548 - delta = DeltaMessage(

549 - tool_calls=[

255 + # Stream back new tool args by diffing against what was sent.

256 + args_diff = self._compute_args_diff(i, tool_args)

257 + if args_diff:

258 + tool_call_deltas.append(

259 DeltaToolCall(

551 - index=self.current_tool_id,

552 - function=DeltaFunctionCall(

553 - arguments=cur_arguments

554 - ).model_dump(exclude_none=True),

260 + index=i,

261 + function=DeltaFunctionCall(arguments=args_diff).model_dump(

262 + exclude_none=True

263 + ),

264 )

556 - ]

557 - )

558 - self.streamed_args_for_tool[self.current_tool_id] = cur_arguments

559 -

560 - # last case -- we have an update to existing arguments.

561 - elif cur_arguments and prev_arguments:

562 - if (

563 - isinstance(delta_text, str)

564 - and cur_arguments != prev_arguments

565 - and len(cur_arguments) > len(prev_arguments)

566 - and cur_arguments.startswith(prev_arguments)

567 - ):

568 - delta_arguments = cur_arguments[len(prev_arguments) :]

569 - logger.debug("got diff %s", delta_text)

570 -

571 - delta = DeltaMessage(

572 - tool_calls=[

573 - DeltaToolCall(

574 - index=self.current_tool_id,

575 - function=DeltaFunctionCall(

576 - arguments=delta_arguments

577 - ).model_dump(exclude_none=True),

578 - )

579 - ]

265 )

581 - self.streamed_args_for_tool[self.current_tool_id] = cur_arguments

582 - else:

583 - delta = None

584 -

585 - # handle saving the state for the current tool into

586 - # the "prev" list for use in diffing for the next iteration

587 - if self.current_tool_id == len(self.prev_tool_call_arr) - 1:

588 - self.prev_tool_call_arr[self.current_tool_id] = current_tool_call

589 - else:

590 - self.prev_tool_call_arr.append(current_tool_call)

591 -

592 - # Handle deferred section exit after tool parsing completes

593 - if deferred_section_exit and self.in_tool_section:

594 - logger.debug("Completing deferred section exit")

595 - self._reset_section_state()

266

597 - return delta

267 + if content or tool_call_deltas:

268 + return DeltaMessage(

269 + content=content,

270 + tool_calls=tool_call_deltas,

271 + )

272 + return None

273

274 except Exception:

275 logger.exception("Error trying to handle streaming tool call.")

601 - return None # do not stream a delta. skip this token ID.

276 + return None