[TRTLLM-15176][fix] Harden Kimi K3 tool-call parsing

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(
+            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(
+            f"Forced tool_choice '{forced_name}' but the model emitted a call "
+            f"to '{parsed_name}'; using the forced 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(
+                    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)
         else:
             if text is None:
                 text = ""

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:
         """

tensorrt_llm/serve/tool_parser/kimi_k3_tool_parser.py

@@ -50,29 +50,43 @@ 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__()
         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(
             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<attrs>[^<]*?)<\|sep\|>"
+            r"<\|open\|>call(?P<attrs>" + attrs_pattern + r")<\|sep\|>"
             r"(?P<body>.*?)<\|close\|>call<\|sep\|>",
             re.DOTALL,
         )
         self._argument_regex = re.compile(
-            r"<\|open\|>argument(?P<attrs>[^<]*?)<\|sep\|>"
+            r"<\|open\|>argument(?P<attrs>" + attrs_pattern + r")<\|sep\|>"
             r"(?P<value>.*?)<\|close\|>argument<\|sep\|>",
             re.DOTALL,
         )
         self._json_regex = re.compile(
-            r"<\|open\|>json(?P<attrs>[^<]*?)<\|sep\|>"
+            r"<\|open\|>json(?P<attrs>" + attrs_pattern + r")<\|sep\|>"
             r"(?P<value>.*?)<\|close\|>json<\|sep\|>",
             re.DOTALL,
         )
@@ -99,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
 
@@ -129,23 +142,32 @@ 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:
                 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,
                     name=name,
                     parameters=self._parse_call_arguments(match.group("body")),
                 )
             )
+        if matched_calls < opened_calls:
+            logger.warning(
+                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
 
     def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult:
@@ -174,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)
@@ -193,9 +222,34 @@ 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
         )
+
+    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:
+            # 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"
+        )
+        return self.detect_and_parse(buffer, tools)