From b5b35195f94674be2c7b24a30bdb811f90764fee Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Wed, 19 Aug 2026 18:17:45 +0000 Subject: [PATCH 1/4] [TRTLLM-15176][fix] Harden Kimi K3 tool-call parsing Forced/named tool_choice previously returned the raw model output in tool_calls[0].function.arguments and finish_reason=stop. K3 has no structural-tag grammar for its XTML format, so nothing constrains the model and the raw text carries a free-text preamble plus markup. - Add BaseToolParser.extracts_forced_tool_calls (default False) so parsers whose forced output still carries native markup can opt in to serve-level extraction; KimiK3ToolParser opts in. For opted-in parsers, the named-choice paths (streaming and non-streaming) now run the tool parser: extracted JSON becomes arguments, the preamble becomes content, and the name always comes from the request (a model that calls a different tool is logged). If no markup arrived, the text is returned as content instead of garbage arguments. - Named tool_choice now reports finish_reason=tool_calls for the legacy raw-passthrough path too. - Add BaseToolParser.finish() (default no-op), called at end of stream, so a stream that ends before <|close|>tools<|sep|> emits the buffered K3 section instead of silently dropping it; complete call blocks are salvaged from a truncated section. - Tag-header regexes now stop at special tokens instead of any '<', so a literal '<' in an attribute value no longer silently drops the call; unparsable call blocks are counted and logged. - Extend the unit suite: forced-choice-with-preamble (streaming and non-streaming, driven through the real postprocessors), name mismatch, no-markup fallback, raw passthrough, early-end flush, multi-call streaming, and the '<'-in-attribute case. Signed-off-by: Brian Nguyen --- tensorrt_llm/serve/postprocess_handlers.py | 121 +++++- .../serve/tool_parser/base_tool_parser.py | 16 + .../serve/tool_parser/kimi_k3_tool_parser.py | 48 ++- .../unittest/llmapi/apps/test_tool_parsers.py | 371 ++++++++++++++++++ 4 files changed, 541 insertions(+), 15 deletions(-) diff --git a/tensorrt_llm/serve/postprocess_handlers.py b/tensorrt_llm/serve/postprocess_handlers.py index 364f1cbf4e8f..9ba09664df57 100644 --- a/tensorrt_llm/serve/postprocess_handlers.py +++ b/tensorrt_llm/serve/postprocess_handlers.py @@ -2,6 +2,7 @@ from dataclasses import dataclass, field from typing import Any, List, Literal, Optional, Tuple, Union +from tensorrt_llm.logger import logger from tensorrt_llm.serve.responses_utils import ResponsesStreamingProcessor from tensorrt_llm.serve.responses_utils import \ create_response_non_store as responses_api_create_response_non_store @@ -40,7 +41,7 @@ ResponsesResponse, StreamOptions, ToolCall, UsageInfo, to_disaggregated_params) from .tool_parser.base_tool_parser import BaseToolParser -from .tool_parser.core_types import ToolCallItem +from .tool_parser.core_types import StreamingParseResult, ToolCallItem from .tool_parser.tool_parser_factory import ToolParserFactory # yapf: enable @@ -187,8 +188,11 @@ def apply_reasoning_parser(args: ChatPostprocArgs, return content, reasoning_content -def apply_tool_parser(args: ChatPostprocArgs, output_index: int, text: str, - streaming: bool) -> Tuple[str, List[ToolCallItem]]: +def apply_tool_parser(args: ChatPostprocArgs, + output_index: int, + text: str, + streaming: bool, + finished: bool = False) -> Tuple[str, List[ToolCallItem]]: tool_parser = None tools = args.tools if args.tool_parser is not None and tools is not None: @@ -203,6 +207,11 @@ def apply_tool_parser(args: ChatPostprocArgs, output_index: int, text: str, result = tool_parser.detect_and_parse(text, tools) else: result = tool_parser.parse_streaming_increment(text, tools) + if finished: + finish_result = tool_parser.finish(tools) + result = StreamingParseResult( + normal_text=result.normal_text + finish_result.normal_text, + calls=result.calls + finish_result.calls) normal_text, calls = result.normal_text, result.calls if result.calls: args.has_tool_call[output_index] = True @@ -212,6 +221,47 @@ def apply_tool_parser(args: ChatPostprocArgs, output_index: int, text: str, return normal_text, calls +def _forced_tool_choice( + args: ChatPostprocArgs) -> Optional[ChatCompletionNamedToolChoiceParam]: + """Return the named tool_choice param if the request forces a tool.""" + if isinstance(args.tool_choice, ChatCompletionNamedToolChoiceParam): + return args.tool_choice + return None + + +def _forced_choice_uses_tool_parser(args: ChatPostprocArgs) -> bool: + """Whether the configured tool parser extracts forced/named tool calls. + + Parsers whose forced output is grammar-constrained bare JSON keep the + raw-passthrough behavior; parsers that opt in (see + ``BaseToolParser.extracts_forced_tool_calls``) still emit their native + markup on forced calls and need the extraction path. + """ + if args.tool_parser is None or args.tools is None: + return False + parser_cls = ToolParserFactory.parsers.get(args.tool_parser.lower()) + return bool(parser_cls + and getattr(parser_cls, "extracts_forced_tool_calls", False)) + + +def _forced_call_name(calls: List[ToolCallItem], forced_name: str) -> str: + """Validate parser-extracted calls against the request's forced tool. + + The name always comes from the request (the caller chose it); the parser + output is only checked so a disagreeing model is visible in the logs. + """ + if len(calls) > 1: + logger.warning( + "Forced tool_choice '%s' produced %d tool calls; keeping the first.", + forced_name, len(calls)) + parsed_name = calls[0].name + if parsed_name and parsed_name != forced_name: + logger.warning( + "Forced tool_choice '%s' but the model emitted a call to '%s'; " + "using the forced name.", forced_name, parsed_name) + return forced_name + + @nvtx_range_debug("chat_stream_post_processor") def chat_stream_post_processor(rsp: GenerationResultBase, args: ChatPostprocArgs) -> List[str]: @@ -279,18 +329,32 @@ def yield_first_chat(num_tokens: int, True, finished=(output.finish_reason is not None)) - if args.tool_choice and type( - args.tool_choice) is ChatCompletionNamedToolChoiceParam: + forced_tool = _forced_tool_choice(args) + if forced_tool and not _forced_choice_uses_tool_parser(args): + # Forced calls constrained to bare JSON arguments: raw deltas are + # the arguments stream. The response carries a tool call, so the + # final chunk must report finish_reason="tool_calls". + args.has_tool_call[i] = True delta_message = DeltaMessage(tool_calls=[ DeltaToolCall( - function=DeltaFunctionCall( - name=args.tool_choice.function.name, - arguments=delta_text), + function=DeltaFunctionCall(name=forced_tool.function.name, + arguments=delta_text), index=i, ), ], ) else: - delta_text, calls = apply_tool_parser(args, i, delta_text, True) + delta_text, calls = apply_tool_parser(args, + i, + delta_text, + True, + finished=(output.finish_reason + is not None)) + if forced_tool and calls: + forced_name = _forced_call_name(calls, + forced_tool.function.name) + calls = calls[:1] + if calls[0].name: + calls[0].name = forced_name tool_calls = [] for call_item in calls: # Tool call ID should be generated only once per tool call @@ -392,15 +456,48 @@ def chat_response_post_processor( text, reasoning_text = apply_reasoning_parser(args, output.index, output.text, False) - if args.tool_choice and isinstance(args.tool_choice, - ChatCompletionNamedToolChoiceParam): + forced_tool = _forced_tool_choice(args) + if forced_tool and not _forced_choice_uses_tool_parser(args): + # Forced calls constrained to bare JSON arguments: the whole text + # is the arguments payload. The response carries a tool call, so + # finish_reason must be "tool_calls". + args.has_tool_call[output.index] = True message = ChatMessage( role=role, content="", tool_calls=[ ToolCall(function=FunctionCall( - name=args.tool_choice.function.name, arguments=text)) + name=forced_tool.function.name, arguments=text)) ]) + elif forced_tool: + # The parser extracts the forced call from the model's native + # markup; any free-text preamble becomes content per OpenAI + # semantics. + text, calls = apply_tool_parser(args, output.index, text or "", + False) + if calls: + forced_name = _forced_call_name(calls, + forced_tool.function.name) + message = ChatMessage( + role=role, + content=text, + reasoning_content=reasoning_text, + tool_calls=[ + ToolCall(function=FunctionCall( + name=forced_name, arguments=calls[0].parameters)) + ]) + else: + # No tool markup despite the forced choice (nothing + # constrains the model for these parsers). Returning the text + # as arguments would hand the caller garbage JSON, so return + # it as content and keep finish_reason honest. + logger.warning( + "Forced tool_choice '%s' but the model emitted no tool-" + "call markup; returning the text as content.", + forced_tool.function.name) + message = ChatMessage(role=role, + content=text, + reasoning_content=reasoning_text) else: if text is None: text = "" diff --git a/tensorrt_llm/serve/tool_parser/base_tool_parser.py b/tensorrt_llm/serve/tool_parser/base_tool_parser.py index ece736749d2f..2d70d2fa572f 100644 --- a/tensorrt_llm/serve/tool_parser/base_tool_parser.py +++ b/tensorrt_llm/serve/tool_parser/base_tool_parser.py @@ -17,6 +17,11 @@ class BaseToolParser(ABC): """Base class providing two sets of interfaces: one-time and streaming incremental.""" needs_raw_special_tokens: bool = False + # Parsers whose forced/named ``tool_choice`` output still carries the + # model's native tool-call markup (rather than grammar-constrained bare + # JSON arguments) set this to True so the serving layer runs extraction + # on the named-choice path instead of passing raw text through. + extracts_forced_tool_calls: bool = False def __init__(self): # Streaming state management @@ -110,6 +115,17 @@ def _ends_with_partial_token(self, buffer: str, bot_token: str) -> int: return i return 0 + def finish(self, tools: List[Tool]) -> StreamingParseResult: + """Finalize a stream that ended before the format's closing markers. + + Called once by the serving layer when generation finishes, after the + last ``parse_streaming_increment`` call, so parsers that buffer whole + sections can emit whatever the truncated stream still holds. The + default is a no-op to preserve the existing behavior of parsers that + manage ``self._buffer`` incrementally. + """ + return StreamingParseResult() + def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> StreamingParseResult: """ diff --git a/tensorrt_llm/serve/tool_parser/kimi_k3_tool_parser.py b/tensorrt_llm/serve/tool_parser/kimi_k3_tool_parser.py index b434e019be61..d699a65da323 100644 --- a/tensorrt_llm/serve/tool_parser/kimi_k3_tool_parser.py +++ b/tensorrt_llm/serve/tool_parser/kimi_k3_tool_parser.py @@ -50,6 +50,10 @@ class KimiK3ToolParser(BaseToolParser): """Detector for the Kimi K3 XTML function-call format.""" needs_raw_special_tokens = True + # Forced/named tool_choice has no grammar for XTML (no structural-tag + # support), so the model output still carries preamble + markup and the + # serving layer must extract instead of passing raw text through. + extracts_forced_tool_calls = True def __init__(self): super().__init__() @@ -61,18 +65,24 @@ def __init__(self): r"(?:<\|close\|>message<\|sep\|>|<\|end_of_msg\|>)+\s*$" ) + # Tag headers run to the next special token. The encoder escapes only + # ``&`` and ``"`` in attribute values, so a literal ``<`` (or ``>``) + # can appear inside one; only ``<|`` is impossible without ending the + # header, so headers match any text that doesn't contain ``<|``. + attrs_pattern = r"(?:(?!<\|).)*?" + self._call_open_regex = re.compile(r"<\|open\|>call(?![a-zA-Z])") self._call_regex = re.compile( - r"<\|open\|>call(?P[^<]*?)<\|sep\|>" + r"<\|open\|>call(?P" + attrs_pattern + r")<\|sep\|>" r"(?P.*?)<\|close\|>call<\|sep\|>", re.DOTALL, ) self._argument_regex = re.compile( - r"<\|open\|>argument(?P[^<]*?)<\|sep\|>" + r"<\|open\|>argument(?P" + attrs_pattern + r")<\|sep\|>" r"(?P.*?)<\|close\|>argument<\|sep\|>", re.DOTALL, ) self._json_regex = re.compile( - r"<\|open\|>json(?P[^<]*?)<\|sep\|>" + r"<\|open\|>json(?P" + attrs_pattern + r")<\|sep\|>" r"(?P.*?)<\|close\|>json<\|sep\|>", re.DOTALL, ) @@ -129,7 +139,10 @@ def _parse_call_arguments(self, body: str) -> str: def _parse_tools_section(self, section: str, tools: List[Tool]) -> List[ToolCallItem]: tool_indices = self._get_tool_indices(tools) calls: List[ToolCallItem] = [] + opened_calls = len(self._call_open_regex.findall(section)) + matched_calls = 0 for position, match in enumerate(self._call_regex.finditer(section)): + matched_calls += 1 attrs = _parse_attrs(match.group("attrs")) name = attrs.get("tool") if not name: @@ -146,6 +159,13 @@ def _parse_tools_section(self, section: str, tools: List[Tool]) -> List[ToolCall parameters=self._parse_call_arguments(match.group("body")), ) ) + if matched_calls < opened_calls: + logger.warning( + "kimi_k3 tool parser: %d of %d call blocks were malformed or " + "truncated and could not be parsed", + opened_calls - matched_calls, + opened_calls, + ) return calls def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult: @@ -199,3 +219,25 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami return StreamingParseResult( normal_text=normal_text + result.normal_text, calls=result.calls ) + + def finish(self, tools: List[Tool]) -> StreamingParseResult: + """Emit whatever the buffer holds when the stream ends early. + + ``parse_streaming_increment`` buffers the whole tools section until + ``<|close|>tools<|sep|>``; if generation stops first (length limit, + cancellation), the buffered content would otherwise be dropped. + Complete call blocks are salvaged; a call truncated mid-block is + reported by the malformed-call warning in ``_parse_tools_section``. + """ + buffer, self._buffer = self._buffer, "" + if not buffer: + return StreamingParseResult() + if self.bot_token not in buffer: + # Only a partial bot_token prefix could be held back here; the + # stream is over, so it is plain text after all. + return StreamingParseResult(normal_text=buffer) + logger.warning( + "kimi_k3 tool parser: stream ended before %s; parsing the partial tools section", + self.eot_token, + ) + return self.detect_and_parse(buffer, tools) diff --git a/tests/unittest/llmapi/apps/test_tool_parsers.py b/tests/unittest/llmapi/apps/test_tool_parsers.py index 4109ff26c354..2ce9ebe5869d 100644 --- a/tests/unittest/llmapi/apps/test_tool_parsers.py +++ b/tests/unittest/llmapi/apps/test_tool_parsers.py @@ -344,6 +344,25 @@ def test_structure_info(self): assert "test_function" in info.begin assert info.trigger == "[TOOL_CALLS]" + def test_finish_default_noop(self, sample_tools): + """The default finalization hook emits nothing and keeps state. + + Existing parsers manage ``self._buffer`` incrementally; the + end-of-stream hook must not change their behavior. + """ + parser = ConcreteToolParser() + parser._buffer = '[TOOL_CALLS] {"name":"get_weather"' + + result = parser.finish(sample_tools) + + assert result.normal_text == "" + assert result.calls == [] + assert parser._buffer == '[TOOL_CALLS] {"name":"get_weather"' + + def test_extracts_forced_tool_calls_default_false(self): + """Forced tool_choice extraction is opt-in per parser.""" + assert ConcreteToolParser.extracts_forced_tool_calls is False + # ============================================================================ # Qwen3ToolParser Tests @@ -4443,6 +4462,358 @@ def test_composes_with_kimi_k3_reasoning_parser(self, sample_tools, parser): assert stage2.calls[0].name == "get_weather" assert json.loads(stage2.calls[0].parameters) == {"location": "NYC"} + def test_extracts_forced_tool_calls(self): + """K3 opts in to serve-level extraction on forced tool_choice. + + XTML has no structural-tag grammar, so a forced call still arrives + as preamble + markup and must not be passed through raw. + """ + assert KimiK3ToolParser.extracts_forced_tool_calls is True + + def test_streaming_early_end_flush(self, sample_tools, parser): + """Emit a buffered complete call when the stream ends before EOT. + + Without the finalization hook the buffered section was silently + dropped. + """ + result = parser.parse_streaming_increment( + "Sure. " + self.BOT + self._call( + "get_weather", 1, self._argument("location", "string", "NYC")), + sample_tools) + assert result.normal_text == "Sure. " + assert result.calls == [] + + result = parser.finish(sample_tools) + + assert len(result.calls) == 1 + assert result.calls[0].name == "get_weather" + assert json.loads(result.calls[0].parameters) == {"location": "NYC"} + assert parser._buffer == "" + + def test_streaming_early_end_flush_salvages_complete_calls( + self, sample_tools, parser): + """A stream truncated mid-call still emits the calls that completed.""" + parser.parse_streaming_increment( + self.BOT + self._call("get_weather", 1, + self._argument("location", "string", "LA")) + + '<|open|>call tool="search_web" index="2"<|sep|>' + '<|open|>argument key="query" type="str', sample_tools) + + result = parser.finish(sample_tools) + + assert [call.name for call in result.calls] == ["get_weather"] + assert json.loads(result.calls[0].parameters) == {"location": "LA"} + + def test_finish_after_complete_section_is_empty(self, sample_tools, parser): + """A cleanly closed section leaves nothing for the flush to emit.""" + result = parser.parse_streaming_increment( + self._section( + self._call("get_weather", 1, + self._argument("location", "string", "NYC"))), + sample_tools) + assert len(result.calls) == 1 + + result = parser.finish(sample_tools) + + assert result.normal_text == "" + assert result.calls == [] + + def test_finish_flushes_held_partial_bot_token_as_text( + self, sample_tools, parser): + """A held-back bot_token prefix is plain text once the stream ends.""" + result = parser.parse_streaming_increment("A <|open|>too", sample_tools) + assert result.normal_text == "A " + + result = parser.finish(sample_tools) + + assert result.normal_text == "<|open|>too" + assert result.calls == [] + + def test_streaming_multiple_calls_split_across_chunks( + self, sample_tools, parser): + """Both calls of a two-call section arrive once EOT lands.""" + first = self._call("get_weather", 1, + self._argument("location", "string", "LA")) + second = self._call("search_web", 2, + self._argument("query", "string", "AI")) + section = self.BOT + first + second + self.EOT + # Split inside the second call header and inside the EOT token. + chunks = [ + section[:len(self.BOT) + len(first) + 14], + section[len(self.BOT) + len(first) + 14:-7], + section[-7:], + ] + + collected = [] + for chunk in chunks: + collected.extend( + parser.parse_streaming_increment(chunk, sample_tools).calls) + + assert [call.name + for call in collected] == ["get_weather", "search_web"] + assert json.loads(collected[0].parameters) == {"location": "LA"} + assert json.loads(collected[1].parameters) == {"query": "AI"} + + def test_literal_lt_in_attribute_value(self, sample_tools, parser): + """Attribute values may contain a literal ``<``. + + The K3 encoder escapes only ``&`` and ``"``, so ``<`` reaches the + parser raw; the old header pattern (``[^<]*?``) silently dropped the + whole call. + """ + text = self._section( + self._call("acall tool="get_weather" index="1"<|sep|>' + '<|open|>argument key="location" type="string"<|sep|>NYC' + '<|close|>argument<|sep|>' + '<|close|>call<|sep|>' + EOT) + + @staticmethod + def _make_args(sample_tools, + tool_parser=None, + forced_tool_name=None, + stream=False): + from tensorrt_llm.serve.openai_protocol import ChatCompletionRequest + from tensorrt_llm.serve.postprocess_handlers import ChatPostprocArgs + + request_kwargs = dict( + model="test-model", + messages=[{ + "role": "user", + "content": "What is the weather in NYC?" + }], + tools=sample_tools, + stream=stream, + ) + if forced_tool_name is not None: + request_kwargs["tool_choice"] = { + "type": "function", + "function": { + "name": forced_tool_name + }, + } + req = ChatCompletionRequest(**request_kwargs) + args = ChatPostprocArgs.from_request(req) + args.tool_parser = tool_parser + args.num_prompt_tokens = 5 + return args + + @staticmethod + def _fake_response(text, finish_reason="stop"): + from types import SimpleNamespace + output = SimpleNamespace(index=0, + text=text, + token_ids=[1, 2, 3], + finish_reason=finish_reason, + stop_reason=None, + disaggregated_params=None) + return SimpleNamespace(outputs=[output], cached_tokens=0) + + @staticmethod + def _stream_deltas(args, chunks, finish_reason="stop"): + """Feed text chunks through the streaming postprocessor. + + Returns the decoded SSE payloads (one dict per emitted chunk). + """ + from types import SimpleNamespace + + from tensorrt_llm.serve.postprocess_handlers import \ + chat_stream_post_processor + + payloads = [] + for chunk_index, chunk in enumerate(chunks): + last = chunk_index == len(chunks) - 1 + output = SimpleNamespace( + index=0, + text_diff=chunk, + token_ids_diff=[chunk_index], + logprobs_diff=[], + finish_reason=finish_reason if last else None, + stop_reason=None, + length=chunk_index + 1, + disaggregated_params=None, + ) + rsp = SimpleNamespace(outputs=[output], + cached_tokens=0, + id="test-request", + _done=last) + for line in chat_stream_post_processor(rsp, args): + payloads.append(json.loads(line[len("data: "):])) + return payloads + + def test_forced_choice_k3_extracts_arguments(self, sample_tools): + """The headline TRTLLM-15176 bug. + + A forced tool_choice must return the extracted JSON, not the raw + preamble + XTML markup. + """ + from tensorrt_llm.serve.postprocess_handlers import \ + chat_response_post_processor + + args = self._make_args(sample_tools, + tool_parser="kimi_k3", + forced_tool_name="get_weather") + rsp = self._fake_response("Let me check. " + self.WEATHER_CALL_SECTION) + + choice = chat_response_post_processor(rsp, args).choices[0] + + assert choice.finish_reason == "tool_calls" + assert len(choice.message.tool_calls) == 1 + function = choice.message.tool_calls[0].function + assert function.name == "get_weather" + assert json.loads(function.arguments) == {"location": "NYC"} + assert choice.message.content == "Let me check. " + + def test_forced_choice_k3_name_mismatch_uses_request_name( + self, sample_tools): + """The caller chose the tool; a disagreeing model only earns a log.""" + from tensorrt_llm.serve.postprocess_handlers import \ + chat_response_post_processor + + args = self._make_args(sample_tools, + tool_parser="kimi_k3", + forced_tool_name="search_web") + rsp = self._fake_response(self.WEATHER_CALL_SECTION) + + choice = chat_response_post_processor(rsp, args).choices[0] + + function = choice.message.tool_calls[0].function + assert function.name == "search_web" + assert json.loads(function.arguments) == {"location": "NYC"} + + def test_forced_choice_k3_no_markup_returns_content(self, sample_tools): + """Return content when a K3 forced call produced no markup. + + Nothing constrains K3 forced calls; if the model emits no markup, + return the text as content rather than garbage arguments. + """ + from tensorrt_llm.serve.postprocess_handlers import \ + chat_response_post_processor + + args = self._make_args(sample_tools, + tool_parser="kimi_k3", + forced_tool_name="get_weather") + rsp = self._fake_response("I cannot help with that.") + + choice = chat_response_post_processor(rsp, args).choices[0] + + assert choice.finish_reason == "stop" + assert not choice.message.tool_calls + assert choice.message.content == "I cannot help with that." + + def test_forced_choice_raw_passthrough_keeps_text_as_arguments( + self, sample_tools): + """Ungated parsers keep the legacy raw passthrough on forced calls. + + Parsers without ``extracts_forced_tool_calls`` pass text through as + arguments, but now report finish_reason="tool_calls". + """ + from tensorrt_llm.serve.postprocess_handlers import \ + chat_response_post_processor + + args = self._make_args(sample_tools, + tool_parser="qwen3", + forced_tool_name="get_weather") + rsp = self._fake_response('{"location": "NYC"}') + + choice = chat_response_post_processor(rsp, args).choices[0] + + assert choice.finish_reason == "tool_calls" + function = choice.message.tool_calls[0].function + assert function.name == "get_weather" + assert function.arguments == '{"location": "NYC"}' + + def test_forced_choice_k3_streaming_extracts(self, sample_tools): + """Extract the forced call from a streamed K3 response. + + The preamble streams as content deltas and the extracted call as a + tool_calls delta with the forced name. + """ + args = self._make_args(sample_tools, + tool_parser="kimi_k3", + forced_tool_name="get_weather", + stream=True) + section = self.WEATHER_CALL_SECTION + payloads = self._stream_deltas( + args, ["Let me check. ", section[:40], section[40:]]) + + deltas = [p["choices"][0]["delta"] for p in payloads if p["choices"]] + content = "".join(d.get("content") or "" for d in deltas) + assert content == "Let me check. " + tool_deltas = [d for d in deltas if d.get("tool_calls")] + assert len(tool_deltas) == 1 + function = tool_deltas[0]["tool_calls"][0]["function"] + assert function["name"] == "get_weather" + assert json.loads(function["arguments"]) == {"location": "NYC"} + finish_reasons = [ + p["choices"][0].get("finish_reason") for p in payloads + if p["choices"] + ] + assert finish_reasons[-1] == "tool_calls" + + def test_streaming_early_end_flushes_buffered_call(self, sample_tools): + """Flush the buffered call when an auto-choice stream ends early. + + A stream that ends before EOT still emits the buffered call via the + finalization hook instead of dropping it. + """ + args = self._make_args(sample_tools, tool_parser="kimi_k3", stream=True) + truncated = self.WEATHER_CALL_SECTION[:-len(self.EOT)] + payloads = self._stream_deltas( + args, ["Sure. ", truncated[:30], truncated[30:]]) + + deltas = [p["choices"][0]["delta"] for p in payloads if p["choices"]] + tool_deltas = [d for d in deltas if d.get("tool_calls")] + assert len(tool_deltas) == 1 + function = tool_deltas[0]["tool_calls"][0]["function"] + assert function["name"] == "get_weather" + assert json.loads(function["arguments"]) == {"location": "NYC"} + finish_reasons = [ + p["choices"][0].get("finish_reason") for p in payloads + if p["choices"] + ] + assert finish_reasons[-1] == "tool_calls" + + def test_forced_choice_k3_streaming_no_markup_is_content( + self, sample_tools): + """Streaming honest fallback: no markup means content, not a call.""" + args = self._make_args(sample_tools, + tool_parser="kimi_k3", + forced_tool_name="get_weather", + stream=True) + payloads = self._stream_deltas(args, ["I cannot ", "help with that."]) + + deltas = [p["choices"][0]["delta"] for p in payloads if p["choices"]] + assert not any(d.get("tool_calls") for d in deltas) + content = "".join(d.get("content") or "" for d in deltas) + assert content == "I cannot help with that." + finish_reasons = [ + p["choices"][0].get("finish_reason") for p in payloads + if p["choices"] + ] + assert finish_reasons[-1] == "stop" + class TestConfigureParserSpecialTokenDecoding: """Test parser-specific detokenization settings in the OpenAI server.""" From 6ac459ec1a5399a9a506f2e3681a84feb0f2fac2 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Wed, 19 Aug 2026 19:34:30 +0000 Subject: [PATCH 2/4] Address trivial review comments Signed-off-by: Brian Nguyen --- tensorrt_llm/serve/postprocess_handlers.py | 14 +++++++------- .../serve/tool_parser/kimi_k3_tool_parser.py | 11 +++++------ 2 files changed, 12 insertions(+), 13 deletions(-) diff --git a/tensorrt_llm/serve/postprocess_handlers.py b/tensorrt_llm/serve/postprocess_handlers.py index 9ba09664df57..63cd0297bad2 100644 --- a/tensorrt_llm/serve/postprocess_handlers.py +++ b/tensorrt_llm/serve/postprocess_handlers.py @@ -252,13 +252,13 @@ def _forced_call_name(calls: List[ToolCallItem], forced_name: str) -> str: """ if len(calls) > 1: logger.warning( - "Forced tool_choice '%s' produced %d tool calls; keeping the first.", - forced_name, len(calls)) + f"Forced tool_choice '{forced_name}' produced {len(calls)} tool " + "calls; keeping the first.") parsed_name = calls[0].name if parsed_name and parsed_name != forced_name: logger.warning( - "Forced tool_choice '%s' but the model emitted a call to '%s'; " - "using the forced name.", forced_name, parsed_name) + f"Forced tool_choice '{forced_name}' but the model emitted a call " + f"to '{parsed_name}'; using the forced name.") return forced_name @@ -492,9 +492,9 @@ def chat_response_post_processor( # as arguments would hand the caller garbage JSON, so return # it as content and keep finish_reason honest. logger.warning( - "Forced tool_choice '%s' but the model emitted no tool-" - "call markup; returning the text as content.", - forced_tool.function.name) + f"Forced tool_choice '{forced_tool.function.name}' but the " + "model emitted no tool-call markup; returning the text as " + "content.") message = ChatMessage(role=role, content=text, reasoning_content=reasoning_text) diff --git a/tensorrt_llm/serve/tool_parser/kimi_k3_tool_parser.py b/tensorrt_llm/serve/tool_parser/kimi_k3_tool_parser.py index d699a65da323..ca579cc3a480 100644 --- a/tensorrt_llm/serve/tool_parser/kimi_k3_tool_parser.py +++ b/tensorrt_llm/serve/tool_parser/kimi_k3_tool_parser.py @@ -161,10 +161,9 @@ def _parse_tools_section(self, section: str, tools: List[Tool]) -> List[ToolCall ) if matched_calls < opened_calls: logger.warning( - "kimi_k3 tool parser: %d of %d call blocks were malformed or " - "truncated and could not be parsed", - opened_calls - matched_calls, - opened_calls, + f"kimi_k3 tool parser: {opened_calls - matched_calls} of " + f"{opened_calls} call blocks were malformed or truncated and " + "could not be parsed" ) return calls @@ -237,7 +236,7 @@ def finish(self, tools: List[Tool]) -> StreamingParseResult: # stream is over, so it is plain text after all. return StreamingParseResult(normal_text=buffer) logger.warning( - "kimi_k3 tool parser: stream ended before %s; parsing the partial tools section", - self.eot_token, + f"kimi_k3 tool parser: stream ended before {self.eot_token}; " + "parsing the partial tools section" ) return self.detect_and_parse(buffer, tools) From 58ae929909ed1d5da154bfabbec55f3d2c6d8665 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Thu, 20 Aug 2026 07:37:50 +0000 Subject: [PATCH 3/4] Address trivial review comments Signed-off-by: Brian Nguyen --- tensorrt_llm/serve/tool_parser/kimi_k3_tool_parser.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/serve/tool_parser/kimi_k3_tool_parser.py b/tensorrt_llm/serve/tool_parser/kimi_k3_tool_parser.py index ca579cc3a480..60e9bd353910 100644 --- a/tensorrt_llm/serve/tool_parser/kimi_k3_tool_parser.py +++ b/tensorrt_llm/serve/tool_parser/kimi_k3_tool_parser.py @@ -232,9 +232,11 @@ def finish(self, tools: List[Tool]) -> StreamingParseResult: if not buffer: return StreamingParseResult() if self.bot_token not in buffer: - # Only a partial bot_token prefix could be held back here; the - # stream is over, so it is plain text after all. - return StreamingParseResult(normal_text=buffer) + # The buffer holds either a partial bot_token prefix or the + # structural residue left after a completed tools section; the + # stream is over, so it is plain text after stripping any + # trailing structural tokens (matching detect_and_parse). + return StreamingParseResult(normal_text=self._trailing_structural.sub("", buffer)) logger.warning( f"kimi_k3 tool parser: stream ended before {self.eot_token}; " "parsing the partial tools section" From 0a65fb63bf7fd35c1e22f726496e9f4b1541dc4e Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Thu, 20 Aug 2026 11:22:46 -0500 Subject: [PATCH 4/4] [TRTLLM-15176][fix] Stop K3 streaming leaking post-section structural tokens parse_streaming_increment left the structural framing that trails a completed tools section (<|close|>message<|sep|>, <|end_of_msg|>) in the buffer. When that residue arrived in a later increment, the no-bot_token path emitted it verbatim as content, so streaming disagreed with detect_and_parse (which strips it via _trailing_structural) and regressed against main on the ordinary success path. Track a _section_done flag: once a complete section is emitted the K3 message is over, so later text is buffered for finish() to strip instead of being streamed as content. Add a parametrized test that splits each residue variant one character per increment. Also switch the remaining %s-style logger.warning calls in the parser to f-strings; tensorrt_llm.logger joins its args rather than interpolating, so those diagnostics printed the raw format string. Signed-off-by: Brian Nguyen --- .../serve/tool_parser/kimi_k3_tool_parser.py | 25 ++++++++---- .../unittest/llmapi/apps/test_tool_parsers.py | 40 +++++++++++++++++++ 2 files changed, 58 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/serve/tool_parser/kimi_k3_tool_parser.py b/tensorrt_llm/serve/tool_parser/kimi_k3_tool_parser.py index 60e9bd353910..c3998762b6b8 100644 --- a/tensorrt_llm/serve/tool_parser/kimi_k3_tool_parser.py +++ b/tensorrt_llm/serve/tool_parser/kimi_k3_tool_parser.py @@ -59,6 +59,10 @@ def __init__(self): super().__init__() self.bot_token = "<|open|>tools<|sep|>" # nosec B105 self.eot_token = "<|close|>tools<|sep|>" # nosec B105 + # Set once a complete tools section has been emitted. A K3 tools + # section terminates the message, so anything streamed afterwards is + # structural framing, not content. + self._section_done = False # Structural leftovers that may trail the tools section when the # reasoning parser is not in front of this parser. self._trailing_structural = re.compile( @@ -109,9 +113,8 @@ def _coerce_value(value: str, value_type: str) -> Any: return json.loads(value) except json.JSONDecodeError: logger.warning( - "kimi_k3 tool parser: argument declared type=%s but body is " - "not valid JSON; keeping raw text", - value_type, + f"kimi_k3 tool parser: argument declared type={value_type} but " + "body is not valid JSON; keeping raw text" ) return value @@ -147,11 +150,11 @@ def _parse_tools_section(self, section: str, tools: List[Tool]) -> List[ToolCall name = attrs.get("tool") if not name: logger.warning( - "kimi_k3 tool parser: call without tool attribute: %s", match.group("attrs") + f"kimi_k3 tool parser: call without tool attribute: {match.group('attrs')}" ) continue if name not in tool_indices: - logger.warning("Model attempted to call undefined function: %s", name) + logger.warning(f"Model attempted to call undefined function: {name}") calls.append( ToolCallItem( tool_index=position, @@ -193,6 +196,13 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> StreamingParseResult: self._buffer += new_text + if self._section_done: + # The completed tools section terminated the K3 message; any later + # text is structural framing (``<|close|>message<|sep|>``, + # ``<|end_of_msg|>``), never user content. Buffer it so ``finish`` + # strips it — matching ``detect_and_parse`` — instead of emitting + # protocol tokens as content. + return StreamingParseResult() bot_idx = self._buffer.find(self.bot_token) if bot_idx == -1: hold = self._ends_with_partial_token(self._buffer, self.bot_token) @@ -212,9 +222,10 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami return StreamingParseResult(normal_text=normal_text) section_end = eot_idx + len(self.eot_token) result = self.detect_and_parse(self._buffer[:section_end], tools) - # Anything after the section (normally empty) is re-examined on the - # next increment rather than dropped. + # The section terminates the message; hold any trailing framing in the + # buffer for ``finish`` to strip rather than emitting it as content. self._buffer = self._buffer[section_end:] + self._section_done = True return StreamingParseResult( normal_text=normal_text + result.normal_text, calls=result.calls ) diff --git a/tests/unittest/llmapi/apps/test_tool_parsers.py b/tests/unittest/llmapi/apps/test_tool_parsers.py index 2ce9ebe5869d..710f6dd7b2b8 100644 --- a/tests/unittest/llmapi/apps/test_tool_parsers.py +++ b/tests/unittest/llmapi/apps/test_tool_parsers.py @@ -4518,6 +4518,46 @@ def test_finish_after_complete_section_is_empty(self, sample_tools, parser): assert result.normal_text == "" assert result.calls == [] + @pytest.mark.parametrize("residue", [ + "<|close|>message<|sep|>", "<|end_of_msg|>", + "<|close|>message<|sep|><|end_of_msg|>" + ]) + def test_streaming_post_section_residue_never_leaks(self, sample_tools, + parser, residue): + """Structural framing after the section is stripped, not streamed. + + ``parse_streaming_increment`` leaves post-section residue in the + buffer; if it arrives in a later increment the no-``bot_token`` path + must not emit it as content. Streaming must match non-streaming, which + strips the same residue via ``_trailing_structural``. Regression: + without the ``_section_done`` guard the residue leaked as ``content``. + """ + preamble = "Sure. " + section = self._section( + self._call("get_weather", 1, + self._argument("location", "string", "NYC"))) + # Section completes in the first increment; residue trickles in one + # character at a time in later increments (worst case for partial + # structural tokens). + chunks = [preamble + section] + list(residue) + + normal_text = "" + calls = [] + for chunk in chunks: + result = parser.parse_streaming_increment(chunk, sample_tools) + normal_text += result.normal_text + calls.extend(result.calls) + result = parser.finish(sample_tools) + normal_text += result.normal_text + calls.extend(result.calls) + + assert normal_text == preamble + assert [call.name for call in calls] == ["get_weather"] + # Streaming agrees with the non-streaming path on the same full text. + non_streaming = self.make_parser().detect_and_parse( + preamble + section + residue, sample_tools) + assert non_streaming.normal_text == preamble + def test_finish_flushes_held_partial_bot_token_as_text( self, sample_tools, parser): """A held-back bot_token prefix is plain text once the stream ends."""