diff --git a/tensorrt_llm/serve/postprocess_handlers.py b/tensorrt_llm/serve/postprocess_handlers.py index 364f1cbf4e8f..63cd0297bad2 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( + 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 = "" 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..c3998762b6b8 100644 --- a/tensorrt_llm/serve/tool_parser/kimi_k3_tool_parser.py +++ b/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[^<]*?)<\|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, ) @@ -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,16 +142,19 @@ 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, @@ -146,6 +162,12 @@ 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( + 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) diff --git a/tests/unittest/llmapi/apps/test_tool_parsers.py b/tests/unittest/llmapi/apps/test_tool_parsers.py index 4109ff26c354..710f6dd7b2b8 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,398 @@ 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 == [] + + @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.""" + 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."""