forked from Karylab-cklius/vllm
[Model] Add Kimi K3 support: Python frontend [2/2] (#50093)
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
This commit is contained in:
Bugen Zhao
@@ -0,0 +1,250 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.entrypoints.openai.chat_completion.protocol import (
|
||||
ChatCompletionRequest,
|
||||
)
|
||||
from vllm.entrypoints.openai.engine.protocol import DeltaMessage
|
||||
from vllm.parser.kimi_k3 import KimiK3Parser
|
||||
from vllm.parser.parser_manager import ParserManager
|
||||
from vllm.reasoning.kimi_k3_reasoning_parser import KimiK3ReasoningParser
|
||||
|
||||
pytestmark = pytest.mark.skip_global_cleanup
|
||||
|
||||
OPEN = "<|open|>"
|
||||
CLOSE = "<|close|>"
|
||||
SEP = "<|sep|>"
|
||||
THINK_OPEN = f"{OPEN}think{SEP}"
|
||||
THINK_CLOSE = f"{CLOSE}think{SEP}"
|
||||
RESPONSE_OPEN = f"{OPEN}response{SEP}"
|
||||
|
||||
|
||||
class DummyTokenizer:
|
||||
def get_vocab(self) -> dict[str, int]:
|
||||
return {}
|
||||
|
||||
def encode(self, text: str, add_special_tokens: bool = False) -> list[int]:
|
||||
if text == THINK_OPEN:
|
||||
return [1, 2, 3]
|
||||
if text == THINK_CLOSE:
|
||||
return [4, 2, 3]
|
||||
return [ord(ch) for ch in text]
|
||||
|
||||
|
||||
class ReasoningOnlyParser(KimiK3Parser):
|
||||
reasoning_parser_cls = KimiK3ReasoningParser
|
||||
|
||||
|
||||
def test_parser_manager_selects_kimi_k3_parser_for_reasoning_only():
|
||||
parser_cls = ParserManager.get_parser(reasoning_parser_name="kimi_k3")
|
||||
|
||||
assert parser_cls is not None
|
||||
assert issubclass(parser_cls, KimiK3Parser)
|
||||
assert parser_cls.reasoning_parser_cls is KimiK3ReasoningParser
|
||||
assert parser_cls.tool_parser_cls is None
|
||||
|
||||
|
||||
def test_parser_selection_thinking_disabled():
|
||||
parser = KimiK3ReasoningParser(
|
||||
DummyTokenizer(), chat_template_kwargs={"thinking": False}
|
||||
)
|
||||
|
||||
assert parser._thinking_enabled is False
|
||||
|
||||
|
||||
def test_extract_reasoning_with_xtml_tags():
|
||||
parser = KimiK3ReasoningParser(DummyTokenizer())
|
||||
request = ChatCompletionRequest(model="test-model", messages=[])
|
||||
|
||||
reasoning, content = parser.extract_reasoning_content(
|
||||
f"{THINK_OPEN}step{THINK_CLOSE}{RESPONSE_OPEN}answer",
|
||||
request,
|
||||
)
|
||||
|
||||
assert reasoning == "step"
|
||||
assert content == "answer"
|
||||
|
||||
|
||||
def test_extract_reasoning_with_generation_prefix_consumed():
|
||||
parser = KimiK3ReasoningParser(DummyTokenizer())
|
||||
request = ChatCompletionRequest(model="test-model", messages=[])
|
||||
|
||||
reasoning, content = parser.extract_reasoning_content(
|
||||
f"step{THINK_CLOSE}{RESPONSE_OPEN}answer",
|
||||
request,
|
||||
)
|
||||
|
||||
assert reasoning == "step"
|
||||
assert content == "answer"
|
||||
|
||||
|
||||
def test_delegating_parser_strips_response_wrapper_without_tool_parser():
|
||||
parser = ReasoningOnlyParser(DummyTokenizer())
|
||||
request = ChatCompletionRequest(model="test-model", messages=[])
|
||||
|
||||
reasoning, content, tool_calls = parser.parse(
|
||||
f"{THINK_OPEN}step{THINK_CLOSE}{RESPONSE_OPEN}answer",
|
||||
request,
|
||||
)
|
||||
|
||||
assert reasoning == "step"
|
||||
assert content == "answer"
|
||||
assert tool_calls == []
|
||||
|
||||
|
||||
def test_is_reasoning_end_uses_full_input_ids():
|
||||
parser = KimiK3ReasoningParser(DummyTokenizer())
|
||||
|
||||
assert not parser.is_reasoning_end([4, 2])
|
||||
assert parser.is_reasoning_end([4, 2, 3])
|
||||
|
||||
|
||||
def test_is_reasoning_end_ignores_stale_close_from_prior_turn():
|
||||
# DummyTokenizer: THINK_OPEN -> [1, 2, 3], THINK_CLOSE -> [4, 2, 3].
|
||||
# Multi-turn / agent continuation: a prior turn's think channel (its close
|
||||
# marker) is kept in the prompt, then the current turn opens a new think
|
||||
# block that has not closed yet. Reasoning must read as NOT ended, otherwise
|
||||
# the structured-output gate constrains the current turn's reasoning.
|
||||
parser = KimiK3ReasoningParser(DummyTokenizer())
|
||||
|
||||
stale_close = [4, 2, 3]
|
||||
new_open = [1, 2, 3]
|
||||
# prior close, then current-turn open still unclosed -> not ended
|
||||
assert not parser.is_reasoning_end([*stale_close, *new_open])
|
||||
# ...then the current turn emits its own close -> ended
|
||||
assert parser.is_reasoning_end([*stale_close, *new_open, *stale_close])
|
||||
# open with no close yet -> not ended
|
||||
assert not parser.is_reasoning_end([*new_open])
|
||||
|
||||
|
||||
def test_streaming_split_open_marker_is_held_back():
|
||||
parser = KimiK3ReasoningParser(DummyTokenizer())
|
||||
|
||||
first = parser.extract_reasoning_content_streaming(
|
||||
previous_text="",
|
||||
current_text=OPEN,
|
||||
delta_text=OPEN,
|
||||
previous_token_ids=[],
|
||||
current_token_ids=[1],
|
||||
delta_token_ids=[1],
|
||||
)
|
||||
second = parser.extract_reasoning_content_streaming(
|
||||
previous_text=OPEN,
|
||||
current_text=f"{OPEN}think",
|
||||
delta_text="think",
|
||||
previous_token_ids=[1],
|
||||
current_token_ids=[1, 2],
|
||||
delta_token_ids=[2],
|
||||
)
|
||||
third = parser.extract_reasoning_content_streaming(
|
||||
previous_text=f"{OPEN}think",
|
||||
current_text=THINK_OPEN + "step",
|
||||
delta_text=f"{SEP}step",
|
||||
previous_token_ids=[1, 2],
|
||||
current_token_ids=[1, 2, 3, 9],
|
||||
delta_token_ids=[3, 9],
|
||||
)
|
||||
|
||||
assert first is None
|
||||
assert second is None
|
||||
assert isinstance(third, DeltaMessage)
|
||||
assert third.reasoning == "step"
|
||||
|
||||
|
||||
def test_streaming_split_close_marker_hands_content_downstream():
|
||||
parser = KimiK3ReasoningParser(DummyTokenizer())
|
||||
|
||||
previous_text = f"{THINK_OPEN}step"
|
||||
partial_close = parser.extract_reasoning_content_streaming(
|
||||
previous_text=previous_text,
|
||||
current_text=previous_text + CLOSE,
|
||||
delta_text=CLOSE,
|
||||
previous_token_ids=[1, 2, 3, 9],
|
||||
current_token_ids=[1, 2, 3, 9, 4],
|
||||
delta_token_ids=[4],
|
||||
)
|
||||
closed = parser.extract_reasoning_content_streaming(
|
||||
previous_text=previous_text + CLOSE,
|
||||
current_text=previous_text + f"{THINK_CLOSE}{RESPONSE_OPEN}answer",
|
||||
delta_text=f"think{SEP}{RESPONSE_OPEN}answer",
|
||||
previous_token_ids=[1, 2, 3, 9, 4],
|
||||
current_token_ids=[1, 2, 3, 9, 4, 2, 3, 10],
|
||||
delta_token_ids=[2, 3, 10],
|
||||
)
|
||||
|
||||
assert partial_close is None
|
||||
assert isinstance(closed, DeltaMessage)
|
||||
assert closed.reasoning is None
|
||||
assert closed.content == f"{RESPONSE_OPEN}answer"
|
||||
assert parser.extract_content_ids([2, 3, 10]) == [10]
|
||||
|
||||
|
||||
def test_thinking_disabled_streams_content():
|
||||
parser = KimiK3ReasoningParser(
|
||||
DummyTokenizer(), chat_template_kwargs={"enable_thinking": False}
|
||||
)
|
||||
|
||||
delta = parser.extract_reasoning_content_streaming(
|
||||
previous_text="",
|
||||
current_text=f"{RESPONSE_OPEN}answer",
|
||||
delta_text=f"{RESPONSE_OPEN}answer",
|
||||
previous_token_ids=[],
|
||||
current_token_ids=[1],
|
||||
delta_token_ids=[1],
|
||||
)
|
||||
|
||||
assert isinstance(delta, DeltaMessage)
|
||||
assert delta.content == f"{RESPONSE_OPEN}answer"
|
||||
assert delta.reasoning is None
|
||||
|
||||
|
||||
def test_delegating_parser_thinking_false_streams_response_content():
|
||||
parser = ReasoningOnlyParser(
|
||||
DummyTokenizer(), chat_template_kwargs={"thinking": False}
|
||||
)
|
||||
request = ChatCompletionRequest(
|
||||
model="test-model",
|
||||
messages=[],
|
||||
chat_template_kwargs={"thinking": False},
|
||||
)
|
||||
|
||||
first = parser.parse_delta(
|
||||
delta_text="OK",
|
||||
delta_token_ids=[10],
|
||||
request=request,
|
||||
prompt_token_ids=[1],
|
||||
finished=False,
|
||||
)
|
||||
partial_close = parser.parse_delta(
|
||||
delta_text=CLOSE,
|
||||
delta_token_ids=[2],
|
||||
request=request,
|
||||
prompt_token_ids=[1],
|
||||
finished=False,
|
||||
)
|
||||
closed = parser.parse_delta(
|
||||
delta_text=f"response{SEP}",
|
||||
delta_token_ids=[3, 4],
|
||||
request=request,
|
||||
prompt_token_ids=[1],
|
||||
finished=False,
|
||||
)
|
||||
|
||||
assert first is not None
|
||||
assert first.content == "OK"
|
||||
assert first.reasoning is None
|
||||
assert partial_close is None
|
||||
assert closed is None
|
||||
|
||||
|
||||
def test_adjust_request_keeps_xtml_markers_contiguous():
|
||||
parser = KimiK3ReasoningParser(DummyTokenizer())
|
||||
request = ChatCompletionRequest(model="test-model", messages=[])
|
||||
|
||||
adjusted = parser.adjust_request(request)
|
||||
|
||||
assert adjusted.skip_special_tokens is False
|
||||
if hasattr(adjusted, "spaces_between_special_tokens"):
|
||||
assert adjusted.spaces_between_special_tokens is False
|
||||
@@ -0,0 +1,353 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.exceptions import VLLMValidationError
|
||||
from vllm.renderers import ChatParams
|
||||
from vllm.renderers.kimi_k3 import KimiK3Renderer, _merge_k3_media_io_kwargs
|
||||
from vllm.renderers.registry import RENDERER_REGISTRY
|
||||
from vllm.tokenizers.registry import TokenizerRegistry
|
||||
|
||||
|
||||
class StubTokenizer:
|
||||
"""Stands in for the model's TikTokenTokenizer.
|
||||
|
||||
Records the kwargs it is called with and returns fixed token ids, so tests
|
||||
can assert how the renderer drives ``apply_chat_template`` without the real
|
||||
(git-LFS) tiktoken vocabulary.
|
||||
"""
|
||||
|
||||
def __init__(self, token_ids: list[int]) -> None:
|
||||
self.token_ids = token_ids
|
||||
self.calls: list[dict[str, Any]] = []
|
||||
self.conversations: list[list[dict[str, Any]]] = []
|
||||
|
||||
def apply_chat_template(self, conversation, **kwargs) -> list[int]:
|
||||
self.conversations.append(conversation)
|
||||
self.calls.append(kwargs)
|
||||
return list(self.token_ids)
|
||||
|
||||
|
||||
@dataclass
|
||||
class MockHFConfig:
|
||||
model_type: str = "kimi_k3"
|
||||
|
||||
|
||||
@dataclass
|
||||
class MockModelConfig:
|
||||
runner_type: str = "generate"
|
||||
is_multimodal_model: bool = False
|
||||
multimodal_config: Any = None
|
||||
hf_config: MockHFConfig = field(default_factory=MockHFConfig)
|
||||
allowed_local_media_path: str = ""
|
||||
allowed_media_domains: Any = None
|
||||
enable_prompt_embeds: bool = False
|
||||
renderer_num_workers: int = 1
|
||||
|
||||
|
||||
@dataclass
|
||||
class MockParallelConfig:
|
||||
_api_process_rank: int = 0
|
||||
|
||||
|
||||
@dataclass
|
||||
class MockVllmConfig:
|
||||
model_config: MockModelConfig
|
||||
parallel_config: MockParallelConfig
|
||||
|
||||
|
||||
def _make_renderer(tokenizer: StubTokenizer) -> KimiK3Renderer:
|
||||
config = MockVllmConfig(MockModelConfig(), MockParallelConfig())
|
||||
return KimiK3Renderer(config, tokenizer)
|
||||
|
||||
|
||||
def test_kimi_k3_registered():
|
||||
assert RENDERER_REGISTRY.load_renderer_cls("kimi_k3").__name__ == "KimiK3Renderer"
|
||||
assert (
|
||||
TokenizerRegistry.load_tokenizer_cls("kimi_k3").__name__ == "CachedHfTokenizer"
|
||||
)
|
||||
|
||||
|
||||
def test_k3_media_io_defaults_preserve_original_mode():
|
||||
# Default: K3 keeps the original image mode (no background flattening).
|
||||
assert _merge_k3_media_io_kwargs(None) == {"image": {"image_mode": None}}
|
||||
|
||||
# Server-/request-level values take precedence over the K3 default.
|
||||
assert _merge_k3_media_io_kwargs({"image": {"image_mode": "RGB"}}) == {
|
||||
"image": {"image_mode": "RGB"}
|
||||
}
|
||||
|
||||
# Unrelated image kwargs are merged with the default.
|
||||
assert _merge_k3_media_io_kwargs(
|
||||
{"image": {"rgba_background_color": (0, 0, 0)}}
|
||||
) == {"image": {"image_mode": None, "rgba_background_color": (0, 0, 0)}}
|
||||
|
||||
|
||||
def test_apply_chat_template_forces_tokenize_and_pins_return_dict():
|
||||
tokenizer = StubTokenizer([7, 8, 9])
|
||||
renderer = _make_renderer(tokenizer)
|
||||
tools = [{"type": "function", "function": {"name": "search"}}]
|
||||
params = ChatParams(
|
||||
chat_template_kwargs={"tools": tools, "tokenize": False, "thinking": True}
|
||||
)
|
||||
|
||||
token_ids = renderer._apply_chat_template(
|
||||
[{"role": "user", "content": "hi"}], params
|
||||
)
|
||||
|
||||
assert token_ids == [7, 8, 9]
|
||||
kwargs = tokenizer.calls[-1]
|
||||
# tokenize is forced on even though the request asked for False, so K3 keeps
|
||||
# the special-vs-ordinary token distinction instead of re-tokenizing a string.
|
||||
assert kwargs["tokenize"] is True
|
||||
# return_dict is pinned False so we always get a flat list of ids.
|
||||
assert kwargs["return_dict"] is False
|
||||
assert kwargs["tools"] == tools
|
||||
assert kwargs["thinking"] is True
|
||||
|
||||
|
||||
def test_apply_chat_template_translates_standard_thinking_kwargs():
|
||||
# Standard enable_thinking/reasoning_effort kwargs must be translated
|
||||
# to K3's native thinking/thinking_effort.
|
||||
tokenizer = StubTokenizer([7, 8, 9])
|
||||
renderer = _make_renderer(tokenizer)
|
||||
params = ChatParams(
|
||||
chat_template_kwargs={"enable_thinking": False, "reasoning_effort": "none"}
|
||||
)
|
||||
|
||||
renderer._apply_chat_template([{"role": "user", "content": "hi"}], params)
|
||||
|
||||
kwargs = tokenizer.calls[-1]
|
||||
assert kwargs["thinking"] is False
|
||||
assert "thinking_effort" not in kwargs
|
||||
assert "enable_thinking" not in kwargs
|
||||
assert "reasoning_effort" not in kwargs
|
||||
|
||||
|
||||
@pytest.mark.parametrize("reasoning_effort", ["low", "high", "max"])
|
||||
def test_apply_chat_template_translates_supported_reasoning_effort(
|
||||
reasoning_effort: str,
|
||||
):
|
||||
tokenizer = StubTokenizer([7, 8, 9])
|
||||
renderer = _make_renderer(tokenizer)
|
||||
params = ChatParams(chat_template_kwargs={"reasoning_effort": reasoning_effort})
|
||||
|
||||
renderer._apply_chat_template([{"role": "user", "content": "hi"}], params)
|
||||
|
||||
kwargs = tokenizer.calls[-1]
|
||||
assert kwargs["thinking_effort"] == reasoning_effort
|
||||
assert "reasoning_effort" not in kwargs
|
||||
|
||||
|
||||
@pytest.mark.parametrize("reasoning_effort", ["minimal", "medium", "xhigh"])
|
||||
def test_apply_chat_template_rejects_unsupported_reasoning_effort(
|
||||
reasoning_effort: str,
|
||||
):
|
||||
tokenizer = StubTokenizer([7, 8, 9])
|
||||
renderer = _make_renderer(tokenizer)
|
||||
params = ChatParams(chat_template_kwargs={"reasoning_effort": reasoning_effort})
|
||||
|
||||
with pytest.raises(VLLMValidationError, match="thinking_effort") as exc_info:
|
||||
renderer._apply_chat_template([{"role": "user", "content": "hi"}], params)
|
||||
|
||||
assert exc_info.value.parameter == "thinking_effort"
|
||||
assert exc_info.value.value == reasoning_effort
|
||||
assert tokenizer.calls == []
|
||||
|
||||
|
||||
def test_apply_chat_template_validates_canonical_native_thinking_effort():
|
||||
tokenizer = StubTokenizer([7, 8, 9])
|
||||
renderer = _make_renderer(tokenizer)
|
||||
params = ChatParams(
|
||||
chat_template_kwargs={
|
||||
"thinking_effort": "low",
|
||||
"reasoning_effort": "medium",
|
||||
}
|
||||
)
|
||||
|
||||
renderer._apply_chat_template([{"role": "user", "content": "hi"}], params)
|
||||
|
||||
assert tokenizer.calls[-1]["thinking_effort"] == "low"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("thinking_effort", ["none", "minimal", "medium", "xhigh"])
|
||||
def test_apply_chat_template_rejects_unsupported_native_thinking_effort(
|
||||
thinking_effort: str,
|
||||
):
|
||||
tokenizer = StubTokenizer([7, 8, 9])
|
||||
renderer = _make_renderer(tokenizer)
|
||||
params = ChatParams(chat_template_kwargs={"thinking_effort": thinking_effort})
|
||||
|
||||
with pytest.raises(VLLMValidationError, match="thinking_effort") as exc_info:
|
||||
renderer._apply_chat_template([{"role": "user", "content": "hi"}], params)
|
||||
|
||||
assert exc_info.value.parameter == "thinking_effort"
|
||||
assert exc_info.value.value == thinking_effort
|
||||
assert tokenizer.calls == []
|
||||
|
||||
|
||||
def test_apply_chat_template_native_k3_kwargs_take_precedence():
|
||||
tokenizer = StubTokenizer([7, 8, 9])
|
||||
renderer = _make_renderer(tokenizer)
|
||||
params = ChatParams(
|
||||
chat_template_kwargs={
|
||||
"thinking": True,
|
||||
"enable_thinking": False,
|
||||
"thinking_effort": "low",
|
||||
"reasoning_effort": "high",
|
||||
}
|
||||
)
|
||||
|
||||
renderer._apply_chat_template([{"role": "user", "content": "hi"}], params)
|
||||
|
||||
kwargs = tokenizer.calls[-1]
|
||||
assert kwargs["thinking"] is True
|
||||
assert kwargs["thinking_effort"] == "low"
|
||||
|
||||
|
||||
def test_apply_chat_template_adds_k3_api_metadata():
|
||||
tokenizer = StubTokenizer([7, 8, 9])
|
||||
renderer = _make_renderer(tokenizer)
|
||||
response_format = {"type": "json_object"}
|
||||
params = ChatParams(
|
||||
tool_choice="required",
|
||||
response_format=response_format,
|
||||
)
|
||||
|
||||
renderer._apply_chat_template([{"role": "user", "content": "hi"}], params)
|
||||
|
||||
kwargs = tokenizer.calls[-1]
|
||||
assert kwargs["tool_choice"] == "required"
|
||||
assert kwargs["response_format"] == response_format
|
||||
|
||||
|
||||
def test_apply_chat_template_auto_tool_choice_keeps_template_kwarg():
|
||||
tokenizer = StubTokenizer([7, 8, 9])
|
||||
renderer = _make_renderer(tokenizer)
|
||||
params = ChatParams(
|
||||
chat_template_kwargs={"tool_choice": "required"},
|
||||
tool_choice="auto",
|
||||
)
|
||||
|
||||
renderer._apply_chat_template([{"role": "user", "content": "hi"}], params)
|
||||
|
||||
assert tokenizer.calls[-1]["tool_choice"] == "required"
|
||||
|
||||
|
||||
def test_apply_chat_template_omits_tool_choice_without_tools():
|
||||
tokenizer = StubTokenizer([7, 8, 9])
|
||||
renderer = _make_renderer(tokenizer)
|
||||
|
||||
renderer._apply_chat_template(
|
||||
[{"role": "user", "content": "hi"}], ChatParams(tool_choice=None)
|
||||
)
|
||||
|
||||
assert "tool_choice" not in tokenizer.calls[-1]
|
||||
|
||||
|
||||
def test_render_messages_returns_token_prompt():
|
||||
renderer = _make_renderer(StubTokenizer([1, 2, 3]))
|
||||
|
||||
conversation, prompt = renderer.render_messages(
|
||||
[{"role": "user", "content": "hi"}], ChatParams()
|
||||
)
|
||||
|
||||
assert prompt == {"prompt_token_ids": [1, 2, 3]}
|
||||
assert "multi_modal_data" not in prompt
|
||||
assert conversation[0]["role"] == "user"
|
||||
|
||||
|
||||
def test_render_messages_derives_private_xtml_tool_attrs():
|
||||
tokenizer = StubTokenizer([1, 2, 3])
|
||||
renderer = _make_renderer(tokenizer)
|
||||
|
||||
conversation, _ = renderer.render_messages(
|
||||
[
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "lookup:0",
|
||||
"type": "function",
|
||||
"function": {"name": "lookup", "arguments": "{}"},
|
||||
},
|
||||
{
|
||||
"id": "lookup:1",
|
||||
"type": "function",
|
||||
"function": {"name": "lookup", "arguments": "{}"},
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": "lookup:1",
|
||||
"tool": "client-supplied-name",
|
||||
"index": 99,
|
||||
"content": "second",
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": "lookup:0",
|
||||
"content": "first",
|
||||
},
|
||||
],
|
||||
ChatParams(),
|
||||
)
|
||||
|
||||
assert [message["content"] for message in conversation[1:]] == [
|
||||
"first",
|
||||
"second",
|
||||
]
|
||||
assert conversation[1]["tool"] == "lookup"
|
||||
assert conversation[1]["index"] == 1
|
||||
assert conversation[2]["tool"] == "lookup"
|
||||
assert conversation[2]["index"] == 2
|
||||
assert tokenizer.conversations[-1] == conversation
|
||||
|
||||
|
||||
def test_render_messages_ignores_client_supplied_xtml_tool_attrs():
|
||||
tokenizer = StubTokenizer([1, 2, 3])
|
||||
renderer = _make_renderer(tokenizer)
|
||||
|
||||
conversation, _ = renderer.render_messages(
|
||||
[
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "lookup:0",
|
||||
"type": "function",
|
||||
"function": {"name": "lookup", "arguments": "{}"},
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": "unknown",
|
||||
"tool": "lookup",
|
||||
"index": 1,
|
||||
"content": "result",
|
||||
},
|
||||
],
|
||||
ChatParams(),
|
||||
)
|
||||
|
||||
assert "tool" not in conversation[1]
|
||||
assert "index" not in conversation[1]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_render_messages_async_returns_token_prompt():
|
||||
renderer = _make_renderer(StubTokenizer([4, 5]))
|
||||
|
||||
conversation, prompt = await renderer.render_messages_async(
|
||||
[{"role": "user", "content": "hi"}], ChatParams()
|
||||
)
|
||||
|
||||
assert prompt == {"prompt_token_ids": [4, 5]}
|
||||
assert conversation[0]["role"] == "user"
|
||||
@@ -0,0 +1,62 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Named tool choice for Kimi K3: allowed when the XTML structural tag is
|
||||
attached (strict tool calling), rejected otherwise."""
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.entrypoints.openai.chat_completion.protocol import (
|
||||
ChatCompletionRequest,
|
||||
)
|
||||
from vllm.exceptions import VLLMValidationError
|
||||
from vllm.sampling_params import StructuredOutputsParams
|
||||
|
||||
|
||||
class _DummyTokenizer:
|
||||
def get_vocab(self):
|
||||
return {}
|
||||
|
||||
def encode(self, text, add_special_tokens=False):
|
||||
return [ord(c) for c in text]
|
||||
|
||||
|
||||
def _request(with_tag: bool) -> ChatCompletionRequest:
|
||||
req = ChatCompletionRequest(
|
||||
model="k3",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
tools=[
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"city": {"type": "string"}},
|
||||
"required": ["city"],
|
||||
},
|
||||
},
|
||||
}
|
||||
],
|
||||
tool_choice={"type": "function", "function": {"name": "get_weather"}},
|
||||
)
|
||||
if with_tag:
|
||||
req.structured_outputs = StructuredOutputsParams(
|
||||
structural_tag='{"type": "structural_tag", "format": {}}'
|
||||
)
|
||||
return req
|
||||
|
||||
|
||||
def _parser():
|
||||
from vllm.tool_parsers.kimi_k3_tool_parser import KimiK3ToolParser
|
||||
|
||||
return KimiK3ToolParser(_DummyTokenizer())
|
||||
|
||||
|
||||
def test_named_choice_allowed_with_structural_tag():
|
||||
req = _parser().adjust_request(_request(with_tag=True))
|
||||
assert req.skip_special_tokens is False
|
||||
|
||||
|
||||
def test_named_choice_rejected_without_structural_tag():
|
||||
with pytest.raises(VLLMValidationError):
|
||||
_parser().adjust_request(_request(with_tag=False))
|
||||
@@ -5,7 +5,8 @@ from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from xgrammar import StructuralTag
|
||||
from xgrammar import Grammar, StructuralTag
|
||||
from xgrammar.testing import _is_grammar_accept_string
|
||||
|
||||
from vllm.entrypoints.openai.chat_completion.protocol import (
|
||||
ChatCompletionNamedFunction,
|
||||
@@ -24,6 +25,7 @@ from vllm.tool_parsers.deepseekv32_engine_tool_parser import (
|
||||
from vllm.tool_parsers.glm47_moe_tool_parser import Glm47MoeModelToolParser
|
||||
from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser
|
||||
from vllm.tool_parsers.kimi_k2_tool_parser import KimiK2ToolParser
|
||||
from vllm.tool_parsers.kimi_k3_tool_parser import KimiK3ToolParser
|
||||
from vllm.tool_parsers.llama_tool_parser import Llama3JsonToolParser
|
||||
from vllm.tool_parsers.minimax_m2_tool_parser import MinimaxM2ToolParser
|
||||
from vllm.tool_parsers.qwen3_engine_tool_parser import Qwen3EngineToolParser
|
||||
@@ -165,6 +167,196 @@ def test_hermes_required_tool_calls_use_empty_separator():
|
||||
assert tag.format.separator == ""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Kimi K3 (XTML channel format) structural tag
|
||||
# ---------------------------------------------------------------------------
|
||||
_K3_RESPONSE_OPEN = "<|open|>response<|sep|>"
|
||||
_K3_RESPONSE_CLOSE = "<|close|>response<|sep|>"
|
||||
_K3_TOOLS_OPEN = "<|open|>tools<|sep|>"
|
||||
_K3_TOOLS_CLOSE = "<|close|>tools<|sep|>"
|
||||
_K3_CALL_CLOSE = "<|close|>call<|sep|>"
|
||||
_K3_ARG_CLOSE = "<|close|>argument<|sep|>"
|
||||
_K3_MESSAGE_CLOSE = "<|close|>message<|sep|>"
|
||||
|
||||
|
||||
def _k3_tools_by_name() -> list[ChatCompletionToolsParam]:
|
||||
return [
|
||||
ChatCompletionToolsParam(
|
||||
type="function",
|
||||
function={
|
||||
"name": "get_weather",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"city": {"type": "string"},
|
||||
"days": {"type": "integer"},
|
||||
},
|
||||
"required": ["city"],
|
||||
},
|
||||
},
|
||||
),
|
||||
ChatCompletionToolsParam(
|
||||
type="function",
|
||||
function={
|
||||
"name": "run_command",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"command": {"type": "string"}},
|
||||
"required": ["command"],
|
||||
},
|
||||
},
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def _k3_arg(key: str, typ: str, val: str) -> str:
|
||||
return f'<|open|>argument key="{key}" type="{typ}"<|sep|>{val}{_K3_ARG_CLOSE}'
|
||||
|
||||
|
||||
def _k3_call(name: str, args: str, idx: int = 1) -> str:
|
||||
return f'<|open|>call tool="{name}" index="{idx}"<|sep|>{args}{_K3_CALL_CLOSE}'
|
||||
|
||||
|
||||
def _k3_response(content: str = "") -> str:
|
||||
return f"{_K3_RESPONSE_OPEN}{content}{_K3_RESPONSE_CLOSE}"
|
||||
|
||||
|
||||
def _k3_tools(*calls: str) -> str:
|
||||
return f"{_K3_TOOLS_OPEN}{''.join(calls)}{_K3_TOOLS_CLOSE}"
|
||||
|
||||
|
||||
def _k3_grammar(tool_choice, tools=None):
|
||||
tag = get_model_structural_tag(
|
||||
model="kimi_k3",
|
||||
tools=tools if tools is not None else _k3_tools_by_name(),
|
||||
tool_choice=tool_choice,
|
||||
reasoning=False,
|
||||
)
|
||||
assert isinstance(tag, StructuralTag)
|
||||
return Grammar.from_structural_tag(tag)
|
||||
|
||||
|
||||
def test_kimi_k3_registered_as_vllm_builtin():
|
||||
assert "kimi_k3" in VLLM_BUILTIN_STRUCTURAL_TAG_MODELS
|
||||
assert KimiK3ToolParser.structural_tag_model == "kimi_k3"
|
||||
|
||||
|
||||
def test_kimi_k3_auto_without_strict_is_unconstrained():
|
||||
# auto + no strict tool => no structural tag (matches the strict gate).
|
||||
tag = get_model_structural_tag(
|
||||
model="kimi_k3",
|
||||
tools=_k3_tools_by_name(),
|
||||
tool_choice="auto",
|
||||
reasoning=False,
|
||||
)
|
||||
assert tag is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"body",
|
||||
[
|
||||
# single required arg
|
||||
_k3_response()
|
||||
+ _k3_tools(_k3_call("get_weather", _k3_arg("city", "string", "Paris"))),
|
||||
# response content + two args (string + number)
|
||||
_k3_response("Checking.")
|
||||
+ _k3_tools(
|
||||
_k3_call(
|
||||
"get_weather",
|
||||
_k3_arg("city", "string", "Paris") + _k3_arg("days", "number", "3"),
|
||||
)
|
||||
),
|
||||
# args in reverse order (parser is order-agnostic)
|
||||
_k3_response()
|
||||
+ _k3_tools(
|
||||
_k3_call(
|
||||
"get_weather",
|
||||
_k3_arg("days", "number", "3") + _k3_arg("city", "string", "Paris"),
|
||||
)
|
||||
),
|
||||
# two calls, second tool
|
||||
_k3_response()
|
||||
+ _k3_tools(
|
||||
_k3_call("get_weather", _k3_arg("city", "string", "Paris"), 1),
|
||||
_k3_call("run_command", _k3_arg("command", "string", "ls -la"), 2),
|
||||
),
|
||||
# string value with regex metacharacters / spaces
|
||||
_k3_response()
|
||||
+ _k3_tools(
|
||||
_k3_call(
|
||||
"run_command", _k3_arg("command", "string", "grep -E 'a|b{2,}' x.py")
|
||||
)
|
||||
),
|
||||
# trailing message-close marker (model's natural turn terminator)
|
||||
_k3_response()
|
||||
+ _k3_tools(_k3_call("get_weather", _k3_arg("city", "string", "Paris")))
|
||||
+ _K3_MESSAGE_CLOSE,
|
||||
# non-thinking mode: response-open is the prompt prefix, so it is absent
|
||||
_K3_RESPONSE_CLOSE
|
||||
+ _k3_tools(_k3_call("get_weather", _k3_arg("city", "string", "Paris"))),
|
||||
],
|
||||
)
|
||||
def test_kimi_k3_required_accepts_valid_tool_calls(body: str):
|
||||
assert _is_grammar_accept_string(_k3_grammar("required"), body)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"body",
|
||||
[
|
||||
# unknown tool name
|
||||
_k3_response()
|
||||
+ _k3_tools(_k3_call("get_temperature", _k3_arg("city", "string", "x"))),
|
||||
# number arg given a non-numeric JSON value
|
||||
_k3_response()
|
||||
+ _k3_tools(_k3_call("get_weather", _k3_arg("days", "number", "abc"))),
|
||||
# undeclared argument key
|
||||
_k3_response()
|
||||
+ _k3_tools(
|
||||
_k3_call(
|
||||
"get_weather",
|
||||
_k3_arg("city", "string", "Paris") + _k3_arg("zzz", "string", "x"),
|
||||
)
|
||||
),
|
||||
# required schema but no argument tags
|
||||
_k3_response() + _k3_tools(_k3_call("get_weather", "")),
|
||||
# missing tools close marker
|
||||
_k3_response()
|
||||
+ _K3_TOOLS_OPEN
|
||||
+ _k3_call("get_weather", _k3_arg("city", "string", "Paris")),
|
||||
# required but no tool call
|
||||
_k3_response("hello"),
|
||||
],
|
||||
)
|
||||
def test_kimi_k3_required_rejects_invalid(body: str):
|
||||
assert not _is_grammar_accept_string(_k3_grammar("required"), body)
|
||||
|
||||
|
||||
def test_kimi_k3_schema_without_required_accepts_empty_call():
|
||||
tools = [
|
||||
ChatCompletionToolsParam(
|
||||
type="function",
|
||||
function={
|
||||
"name": "get_weather",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"city": {"type": "string"}},
|
||||
},
|
||||
},
|
||||
)
|
||||
]
|
||||
grammar = _k3_grammar("required", tools=tools)
|
||||
body = _k3_response() + _k3_tools(_k3_call("get_weather", ""))
|
||||
|
||||
assert _is_grammar_accept_string(grammar, body)
|
||||
|
||||
|
||||
def test_kimi_k3_auto_strict_allows_response_only(sample_tools_strict):
|
||||
# With a strict tool the tag is built; the tools channel is optional so a
|
||||
# plain response (no tool call) is still valid.
|
||||
grammar = _k3_grammar("auto", tools=sample_tools_strict)
|
||||
assert _is_grammar_accept_string(grammar, _k3_response("Just answering."))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", sorted(XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS))
|
||||
def test_get_model_structural_tag_supports_named_tool_choice(
|
||||
model: str,
|
||||
@@ -339,3 +531,151 @@ def test_get_function_parameters_relaxes_function_strict_false():
|
||||
)
|
||||
|
||||
assert _get_function_parameters(function) is True
|
||||
|
||||
|
||||
def _k3_tools_with_root_defs() -> list[ChatCompletionToolsParam]:
|
||||
return [
|
||||
ChatCompletionToolsParam(
|
||||
type="function",
|
||||
function={
|
||||
"name": "make_config",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"config": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"build": {"$ref": "#/$defs/build"},
|
||||
"index": {"type": "string"},
|
||||
},
|
||||
"required": ["index"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
"required": ["config"],
|
||||
"$defs": {
|
||||
"build": {
|
||||
"type": "object",
|
||||
"properties": {"outDir": {"type": "string"}},
|
||||
"additionalProperties": False,
|
||||
}
|
||||
},
|
||||
},
|
||||
},
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def test_kimi_k3_property_ref_to_root_defs_compiles_and_accepts():
|
||||
# Root-level $defs referenced from inside a property schema (the walle
|
||||
# TestReferences shape). Slicing the property out of the parameters
|
||||
# document orphans "#/$defs/..." unless the builder re-attaches $defs;
|
||||
# before the fix Grammar.from_structural_tag raised on the dangling ref.
|
||||
grammar = _k3_grammar("required", tools=_k3_tools_with_root_defs())
|
||||
|
||||
body = _k3_response() + _k3_tools(
|
||||
_k3_call(
|
||||
"make_config",
|
||||
_k3_arg(
|
||||
"config",
|
||||
"object",
|
||||
'{"build": {"outDir": "dist"}, "index": "a.html"}',
|
||||
),
|
||||
)
|
||||
)
|
||||
assert _is_grammar_accept_string(grammar, body)
|
||||
|
||||
|
||||
def _k3_tools_with_string_enum() -> list[ChatCompletionToolsParam]:
|
||||
return [
|
||||
ChatCompletionToolsParam(
|
||||
type="function",
|
||||
function={
|
||||
"name": "set_unit",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"unit": {
|
||||
"type": "string",
|
||||
"enum": ["celsius", " fahrenheit", "\tkelvin"],
|
||||
},
|
||||
},
|
||||
"required": ["unit"],
|
||||
},
|
||||
},
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", ["celsius", " fahrenheit", "\tkelvin"])
|
||||
def test_kimi_k3_string_enum_accepts_exact_values(value: str):
|
||||
# Raw string channel with enum: constrained to the exact enum values,
|
||||
# including leading-whitespace variants the model otherwise flubs.
|
||||
grammar = _k3_grammar("required", tools=_k3_tools_with_string_enum())
|
||||
body = _k3_response() + _k3_tools(
|
||||
_k3_call("set_unit", _k3_arg("unit", "string", value))
|
||||
)
|
||||
assert _is_grammar_accept_string(grammar, body)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", ["kelvin", "Celsius", "celsius ", ""])
|
||||
def test_kimi_k3_string_enum_rejects_non_members(value: str):
|
||||
grammar = _k3_grammar("required", tools=_k3_tools_with_string_enum())
|
||||
body = _k3_response() + _k3_tools(
|
||||
_k3_call("set_unit", _k3_arg("unit", "string", value))
|
||||
)
|
||||
assert not _is_grammar_accept_string(grammar, body)
|
||||
|
||||
|
||||
def _k3_tools_with_maxlen() -> list[ChatCompletionToolsParam]:
|
||||
return [
|
||||
ChatCompletionToolsParam(
|
||||
type="function",
|
||||
function={
|
||||
"name": "set_note",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"note": {"type": "string", "maxLength": 8, "minLength": 2},
|
||||
},
|
||||
"required": ["note"],
|
||||
},
|
||||
},
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def test_kimi_k3_string_maxlength_bounds_raw_channel():
|
||||
# Raw string channel with maxLength/minLength: enforced via a bounded
|
||||
# regex that keeps the "<|" marker prefix unambiguous but still allows
|
||||
# a bare '<' inside values.
|
||||
grammar = _k3_grammar("required", tools=_k3_tools_with_maxlen())
|
||||
|
||||
def body(val: str) -> str:
|
||||
return _k3_response() + _k3_tools(
|
||||
_k3_call("set_note", _k3_arg("note", "string", val))
|
||||
)
|
||||
|
||||
assert _is_grammar_accept_string(grammar, body("ab"))
|
||||
assert _is_grammar_accept_string(grammar, body("a<b then"))
|
||||
assert not _is_grammar_accept_string(grammar, body("way too long note"))
|
||||
assert not _is_grammar_accept_string(grammar, body("a")) # under minLength
|
||||
|
||||
|
||||
def test_kimi_k3_forced_tool_choice_builds_single_mandatory_call():
|
||||
# Named tool choice normalizes to "forced": the tag must require exactly
|
||||
# the named tool's call (no response-only escape).
|
||||
grammar = _k3_grammar(
|
||||
ChatCompletionNamedToolChoiceParam(
|
||||
type="function",
|
||||
function=ChatCompletionNamedFunction(name="get_weather"),
|
||||
),
|
||||
tools=_k3_tools_by_name(),
|
||||
)
|
||||
|
||||
ok = _k3_response() + _k3_tools(
|
||||
_k3_call("get_weather", _k3_arg("city", "string", "Paris"))
|
||||
)
|
||||
response_only = _k3_response("no call here")
|
||||
assert _is_grammar_accept_string(grammar, ok)
|
||||
assert not _is_grammar_accept_string(grammar, response_only)
|
||||
|
||||
@@ -0,0 +1,600 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.entrypoints.openai.chat_completion.protocol import (
|
||||
ChatCompletionRequest,
|
||||
)
|
||||
from vllm.entrypoints.openai.engine.protocol import DeltaMessage
|
||||
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
|
||||
from vllm.entrypoints.openai.responses.utils import build_response_output_items
|
||||
from vllm.exceptions import VLLMValidationError
|
||||
from vllm.parser.kimi_k3 import KimiK3Parser
|
||||
from vllm.parser.parser_manager import ParserManager
|
||||
from vllm.reasoning.kimi_k3_reasoning_parser import KimiK3ReasoningParser
|
||||
from vllm.tool_parsers.kimi_k3_tool_parser import KimiK3ToolParser
|
||||
|
||||
OPEN = "<|open|>"
|
||||
CLOSE = "<|close|>"
|
||||
SEP = "<|sep|>"
|
||||
THINK_OPEN = f"{OPEN}think{SEP}"
|
||||
THINK_CLOSE = f"{CLOSE}think{SEP}"
|
||||
RESPONSE_CLOSE = f"{CLOSE}response{SEP}"
|
||||
|
||||
|
||||
class DummyTokenizer:
|
||||
def get_vocab(self) -> dict[str, int]:
|
||||
return {}
|
||||
|
||||
def encode(self, text: str, add_special_tokens: bool = False) -> list[int]:
|
||||
if text == THINK_OPEN:
|
||||
return [1, 2, 3]
|
||||
if text == THINK_CLOSE:
|
||||
return [4, 2, 3]
|
||||
return [ord(ch) for ch in text]
|
||||
|
||||
|
||||
class KimiK3DelegatingParser(KimiK3Parser):
|
||||
reasoning_parser_cls = KimiK3ReasoningParser
|
||||
tool_parser_cls = KimiK3ToolParser
|
||||
|
||||
|
||||
def test_parser_manager_selects_kimi_k3_parser():
|
||||
parser_cls = ParserManager.get_parser(
|
||||
tool_parser_name="kimi_k3",
|
||||
reasoning_parser_name="kimi_k3",
|
||||
enable_auto_tools=True,
|
||||
)
|
||||
|
||||
assert parser_cls is not None
|
||||
assert issubclass(parser_cls, KimiK3Parser)
|
||||
assert parser_cls.reasoning_parser_cls is KimiK3ReasoningParser
|
||||
assert parser_cls.tool_parser_cls is KimiK3ToolParser
|
||||
|
||||
|
||||
def _request() -> ChatCompletionRequest:
|
||||
return ChatCompletionRequest(
|
||||
model="test-model",
|
||||
messages=[],
|
||||
tools=[
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "calc",
|
||||
"parameters": {"type": "object", "properties": {}},
|
||||
},
|
||||
}
|
||||
],
|
||||
tool_choice="auto",
|
||||
)
|
||||
|
||||
|
||||
def _named_request() -> ChatCompletionRequest:
|
||||
return ChatCompletionRequest(
|
||||
model="test-model",
|
||||
messages=[],
|
||||
tools=[
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "calc",
|
||||
"parameters": {"type": "object", "properties": {}},
|
||||
},
|
||||
}
|
||||
],
|
||||
tool_choice={"type": "function", "function": {"name": "calc"}},
|
||||
)
|
||||
|
||||
|
||||
def _responses_request(*, tool_choice="auto") -> ResponsesRequest:
|
||||
return ResponsesRequest.model_validate(
|
||||
{
|
||||
"model": "test-model",
|
||||
"input": "Call the calc tool.",
|
||||
"tools": [
|
||||
{
|
||||
"type": "function",
|
||||
"name": "calc",
|
||||
"parameters": {"type": "object", "properties": {}},
|
||||
}
|
||||
],
|
||||
"tool_choice": tool_choice,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _arg(key: str, typ: str, value: str) -> str:
|
||||
return f'{OPEN}argument key="{key}" type="{typ}"{SEP}{value}{CLOSE}argument{SEP}'
|
||||
|
||||
|
||||
def _call(tool: str, index: int, *args: str) -> str:
|
||||
body = "".join(args)
|
||||
return f'{OPEN}call tool="{tool}" index="{index}"{SEP}{body}{CLOSE}call{SEP}'
|
||||
|
||||
|
||||
def _response(content: str) -> str:
|
||||
return f"{OPEN}response{SEP}{content}{RESPONSE_CLOSE}"
|
||||
|
||||
|
||||
def _tools(*calls: str) -> str:
|
||||
return f"{OPEN}tools{SEP}{''.join(calls)}{CLOSE}tools{SEP}"
|
||||
|
||||
|
||||
def test_extract_tool_calls_with_response_and_typed_arguments():
|
||||
parser = KimiK3ToolParser(DummyTokenizer())
|
||||
|
||||
output = _response("answer") + _tools(
|
||||
_call(
|
||||
"calc",
|
||||
1,
|
||||
_arg("x", "number", "1"),
|
||||
_arg("flag", "boolean", "true"),
|
||||
_arg("text", "string", "raw"),
|
||||
)
|
||||
)
|
||||
extracted = parser.extract_tool_calls(output, _request())
|
||||
|
||||
assert extracted.tools_called is True
|
||||
assert extracted.content == "answer"
|
||||
assert len(extracted.tool_calls) == 1
|
||||
tool_call = extracted.tool_calls[0]
|
||||
assert tool_call.id == "calc:0"
|
||||
assert tool_call.function.name == "calc"
|
||||
assert json.loads(tool_call.function.arguments) == {
|
||||
"x": 1,
|
||||
"flag": True,
|
||||
"text": "raw",
|
||||
}
|
||||
|
||||
|
||||
def test_delegating_parser_preserves_tool_calls_after_reasoning():
|
||||
parser = KimiK3DelegatingParser(DummyTokenizer())
|
||||
output = (
|
||||
f"{THINK_OPEN}step{THINK_CLOSE}"
|
||||
+ _response("answer")
|
||||
+ _tools(_call("calc", 1, _arg("x", "number", "1")))
|
||||
)
|
||||
|
||||
reasoning, content, tool_calls = parser.parse(
|
||||
output,
|
||||
_request(),
|
||||
enable_auto_tools=True,
|
||||
)
|
||||
|
||||
assert reasoning == "step"
|
||||
assert content == "answer"
|
||||
assert tool_calls is not None
|
||||
assert len(tool_calls) == 1
|
||||
assert tool_calls[0].id == "calc:0"
|
||||
assert tool_calls[0].name == "calc"
|
||||
assert json.loads(tool_calls[0].arguments) == {"x": 1}
|
||||
|
||||
|
||||
def test_delegating_parser_required_tool_choice_uses_xtml_parser():
|
||||
parser = KimiK3DelegatingParser(DummyTokenizer())
|
||||
request = _request().model_copy(update={"tool_choice": "required"})
|
||||
output = (
|
||||
f"{THINK_OPEN}step{THINK_CLOSE}"
|
||||
+ _response("")
|
||||
+ _tools(_call("calc", 1, _arg("x", "number", "1")))
|
||||
)
|
||||
|
||||
reasoning, content, tool_calls = parser.parse(
|
||||
output,
|
||||
request,
|
||||
enable_auto_tools=True,
|
||||
)
|
||||
|
||||
assert reasoning == "step"
|
||||
assert content is None
|
||||
assert tool_calls is not None
|
||||
assert len(tool_calls) == 1
|
||||
assert tool_calls[0].name == "calc"
|
||||
assert json.loads(tool_calls[0].arguments) == {"x": 1}
|
||||
|
||||
|
||||
def test_delegating_parser_named_tool_choice_uses_xtml_parser():
|
||||
parser = KimiK3DelegatingParser(DummyTokenizer())
|
||||
output = (
|
||||
f"{THINK_OPEN}step{THINK_CLOSE}"
|
||||
+ _response("")
|
||||
+ _tools(_call("calc", 1, _arg("x", "number", "1")))
|
||||
)
|
||||
|
||||
reasoning, content, tool_calls = parser.parse(
|
||||
output,
|
||||
_named_request(),
|
||||
enable_auto_tools=True,
|
||||
)
|
||||
|
||||
assert reasoning == "step"
|
||||
assert content is None
|
||||
assert tool_calls is not None
|
||||
assert len(tool_calls) == 1
|
||||
assert tool_calls[0].name == "calc"
|
||||
assert json.loads(tool_calls[0].arguments) == {"x": 1}
|
||||
|
||||
|
||||
def test_delegating_parser_auto_no_call_strips_consumed_response_prefix():
|
||||
parser = KimiK3DelegatingParser(
|
||||
DummyTokenizer(), chat_template_kwargs={"thinking": False}
|
||||
)
|
||||
request = _request().model_copy(
|
||||
update={"chat_template_kwargs": {"thinking": False}}
|
||||
)
|
||||
|
||||
reasoning, content, tool_calls = parser.parse(
|
||||
f"answer{RESPONSE_CLOSE}",
|
||||
request,
|
||||
enable_auto_tools=True,
|
||||
)
|
||||
|
||||
assert reasoning is None
|
||||
assert content == "answer"
|
||||
assert tool_calls is None
|
||||
|
||||
|
||||
def test_delegating_parser_required_call_strips_consumed_response_prefix():
|
||||
parser = KimiK3DelegatingParser(
|
||||
DummyTokenizer(), chat_template_kwargs={"thinking": False}
|
||||
)
|
||||
request = _request().model_copy(
|
||||
update={
|
||||
"tool_choice": "required",
|
||||
"chat_template_kwargs": {"thinking": False},
|
||||
}
|
||||
)
|
||||
output = RESPONSE_CLOSE + _tools(_call("calc", 1, _arg("x", "number", "1")))
|
||||
|
||||
reasoning, content, tool_calls = parser.parse(
|
||||
output,
|
||||
request,
|
||||
enable_auto_tools=True,
|
||||
)
|
||||
|
||||
assert reasoning is None
|
||||
assert content is None
|
||||
assert tool_calls is not None
|
||||
assert len(tool_calls) == 1
|
||||
assert tool_calls[0].name == "calc"
|
||||
assert json.loads(tool_calls[0].arguments) == {"x": 1}
|
||||
|
||||
|
||||
def test_delegating_parser_truncated_tools_do_not_leak_xtml():
|
||||
parser = KimiK3DelegatingParser(
|
||||
DummyTokenizer(), chat_template_kwargs={"thinking": False}
|
||||
)
|
||||
request = _request().model_copy(
|
||||
update={
|
||||
"tool_choice": "required",
|
||||
"chat_template_kwargs": {"thinking": False},
|
||||
}
|
||||
)
|
||||
|
||||
reasoning, content, tool_calls = parser.parse(
|
||||
(f'{RESPONSE_CLOSE}{OPEN}tools{SEP}{OPEN}call tool="calc" index="1"'),
|
||||
request,
|
||||
enable_auto_tools=True,
|
||||
)
|
||||
|
||||
assert reasoning is None
|
||||
assert content is None
|
||||
assert tool_calls is None
|
||||
|
||||
|
||||
def test_extract_tool_calls_unescapes_attributes():
|
||||
parser = KimiK3ToolParser(DummyTokenizer())
|
||||
|
||||
output = _tools(_call("a&b"c", 1, _arg("k&q", "string", "v")))
|
||||
extracted = parser.extract_tool_calls(output, _request())
|
||||
|
||||
assert extracted.tools_called is True
|
||||
assert extracted.tool_calls[0].function.name == 'a&b"c'
|
||||
assert json.loads(extracted.tool_calls[0].function.arguments) == {"k&q": "v"}
|
||||
|
||||
|
||||
def test_extract_tool_calls_allows_less_than_in_attributes():
|
||||
parser = KimiK3ToolParser(DummyTokenizer())
|
||||
|
||||
output = _tools(_call("calc<beta", 1, _arg("foo<bar", "string", "raw")))
|
||||
extracted = parser.extract_tool_calls(output, _request())
|
||||
|
||||
assert extracted.tools_called is True
|
||||
assert extracted.tool_calls[0].function.name == "calc<beta"
|
||||
assert json.loads(extracted.tool_calls[0].function.arguments) == {"foo<bar": "raw"}
|
||||
|
||||
|
||||
def test_extract_content_from_whitespace_degraded_markers():
|
||||
parser = KimiK3ToolParser(DummyTokenizer())
|
||||
|
||||
extracted = parser.extract_tool_calls(
|
||||
f"{OPEN} response {SEP}answer{CLOSE} response {SEP}",
|
||||
_request(),
|
||||
)
|
||||
|
||||
assert extracted.tools_called is False
|
||||
assert extracted.content == "answer"
|
||||
|
||||
|
||||
def test_streaming_split_markers_do_not_leak():
|
||||
parser = KimiK3ToolParser(DummyTokenizer())
|
||||
request = _request()
|
||||
previous_text = ""
|
||||
previous_ids: list[int] = []
|
||||
messages: list[DeltaMessage] = []
|
||||
chunks = [
|
||||
OPEN,
|
||||
"response",
|
||||
f"{SEP}Hi",
|
||||
OPEN,
|
||||
"tools",
|
||||
SEP,
|
||||
f'{OPEN}call tool="calc" index="1"{SEP}',
|
||||
_arg("x", "number", "1"),
|
||||
f"{CLOSE}call",
|
||||
SEP,
|
||||
]
|
||||
|
||||
for i, chunk in enumerate(chunks, start=1):
|
||||
current_text = previous_text + chunk
|
||||
current_ids = previous_ids + [i]
|
||||
delta = parser.extract_tool_calls_streaming(
|
||||
previous_text=previous_text,
|
||||
current_text=current_text,
|
||||
delta_text=chunk,
|
||||
previous_token_ids=previous_ids,
|
||||
current_token_ids=current_ids,
|
||||
delta_token_ids=[i],
|
||||
request=request,
|
||||
)
|
||||
if delta is not None:
|
||||
messages.append(delta)
|
||||
previous_text = current_text
|
||||
previous_ids = current_ids
|
||||
|
||||
content = "".join(message.content or "" for message in messages)
|
||||
tool_deltas = [
|
||||
tool_call for message in messages for tool_call in (message.tool_calls or [])
|
||||
]
|
||||
|
||||
assert content == "Hi"
|
||||
assert OPEN not in content
|
||||
assert SEP not in content
|
||||
assert len(tool_deltas) == 1
|
||||
assert tool_deltas[0].id == "calc:0"
|
||||
assert tool_deltas[0].function.name == "calc"
|
||||
assert json.loads(tool_deltas[0].function.arguments) == {"x": 1}
|
||||
|
||||
|
||||
def test_streaming_consumed_response_prefix_no_call_keeps_content():
|
||||
parser = KimiK3ToolParser(DummyTokenizer())
|
||||
request = _request()
|
||||
previous_text = ""
|
||||
previous_ids: list[int] = []
|
||||
messages: list[DeltaMessage] = []
|
||||
chunks = ["O", "K", CLOSE, f"response{SEP}"]
|
||||
|
||||
for i, chunk in enumerate(chunks, start=1):
|
||||
current_text = previous_text + chunk
|
||||
current_ids = previous_ids + [i]
|
||||
delta = parser.extract_tool_calls_streaming(
|
||||
previous_text=previous_text,
|
||||
current_text=current_text,
|
||||
delta_text=chunk,
|
||||
previous_token_ids=previous_ids,
|
||||
current_token_ids=current_ids,
|
||||
delta_token_ids=[i],
|
||||
request=request,
|
||||
)
|
||||
if delta is not None:
|
||||
messages.append(delta)
|
||||
previous_text = current_text
|
||||
previous_ids = current_ids
|
||||
|
||||
assert "".join(message.content or "" for message in messages) == "OK"
|
||||
assert all(CLOSE not in (message.content or "") for message in messages)
|
||||
|
||||
|
||||
def test_delegating_parser_tool_choice_none_strips_xtml_and_suppresses_calls():
|
||||
parser = KimiK3DelegatingParser(
|
||||
DummyTokenizer(), chat_template_kwargs={"thinking": False}
|
||||
)
|
||||
request = _request().model_copy(
|
||||
update={
|
||||
"tool_choice": "none",
|
||||
"chat_template_kwargs": {"thinking": False},
|
||||
}
|
||||
)
|
||||
messages: list[DeltaMessage] = []
|
||||
chunks = [
|
||||
OPEN,
|
||||
"response",
|
||||
f"{SEP}answer",
|
||||
RESPONSE_CLOSE,
|
||||
_tools(_call("calc", 1, _arg("x", "number", "1"))),
|
||||
]
|
||||
|
||||
for index, chunk in enumerate(chunks, start=1):
|
||||
delta = parser.parse_delta(
|
||||
delta_text=chunk,
|
||||
delta_token_ids=[index],
|
||||
request=request,
|
||||
prompt_token_ids=[1],
|
||||
finished=index == len(chunks),
|
||||
)
|
||||
if delta is not None:
|
||||
messages.append(delta)
|
||||
|
||||
content = "".join(message.content or "" for message in messages)
|
||||
assert content == "answer"
|
||||
assert OPEN not in content
|
||||
assert CLOSE not in content
|
||||
assert SEP not in content
|
||||
assert all(not message.tool_calls for message in messages)
|
||||
|
||||
|
||||
def test_adjust_request_keeps_xtml_markers_contiguous():
|
||||
parser = KimiK3ToolParser(DummyTokenizer())
|
||||
request = _request()
|
||||
|
||||
adjusted = parser.adjust_request(request)
|
||||
|
||||
assert adjusted.skip_special_tokens is False
|
||||
if hasattr(adjusted, "spaces_between_special_tokens"):
|
||||
assert adjusted.spaces_between_special_tokens is False
|
||||
assert KimiK3ToolParser.supports_required_and_named is False
|
||||
|
||||
|
||||
def test_adjust_request_required_uses_xtml_parser_not_json_guidance():
|
||||
parser = KimiK3ToolParser(DummyTokenizer())
|
||||
request = _request().model_copy(update={"tool_choice": "required"})
|
||||
|
||||
adjusted = parser.adjust_request(request)
|
||||
|
||||
assert adjusted.structured_outputs is None
|
||||
assert adjusted.skip_special_tokens is False
|
||||
if hasattr(adjusted, "spaces_between_special_tokens"):
|
||||
assert adjusted.spaces_between_special_tokens is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"tool_request",
|
||||
[
|
||||
_named_request(),
|
||||
_responses_request(
|
||||
tool_choice={"type": "function", "name": "calc"},
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_adjust_request_rejects_named_tool_choice(tool_request):
|
||||
parser = KimiK3ToolParser(DummyTokenizer())
|
||||
|
||||
with pytest.raises(VLLMValidationError) as exc_info:
|
||||
parser.adjust_request(tool_request)
|
||||
|
||||
assert exc_info.value.parameter == "tool_choice"
|
||||
assert "requires strict tool calling" in str(exc_info.value)
|
||||
|
||||
|
||||
def test_responses_chat_params_carries_tool_choice_metadata():
|
||||
request = _responses_request(tool_choice="required")
|
||||
|
||||
chat_params = request.build_chat_params(
|
||||
default_template=None,
|
||||
default_template_content_format="auto",
|
||||
)
|
||||
|
||||
assert chat_params.tool_choice == "required"
|
||||
|
||||
|
||||
def test_responses_chat_params_keeps_template_tool_choice_when_api_auto():
|
||||
request = _responses_request().model_copy(
|
||||
update={"chat_template_kwargs": {"tool_choice": "required"}}
|
||||
)
|
||||
|
||||
chat_params = request.build_chat_params(
|
||||
default_template=None,
|
||||
default_template_content_format="auto",
|
||||
)
|
||||
|
||||
assert chat_params.chat_template_kwargs["tool_choice"] == "required"
|
||||
assert chat_params.tool_choice == "auto"
|
||||
|
||||
|
||||
def test_responses_required_tool_choice_uses_xtml_parser():
|
||||
parser = KimiK3DelegatingParser(
|
||||
DummyTokenizer(), chat_template_kwargs={"thinking": False}
|
||||
)
|
||||
request = _responses_request(tool_choice="required").model_copy(
|
||||
update={"chat_template_kwargs": {"thinking": False}}
|
||||
)
|
||||
output = RESPONSE_CLOSE + _tools(_call("calc", 1, _arg("x", "number", "1")))
|
||||
|
||||
reasoning, content, tool_calls = parser.parse(
|
||||
output, request, enable_auto_tools=True, model_output_token_ids=[]
|
||||
)
|
||||
response_outputs = build_response_output_items(
|
||||
reasoning=reasoning,
|
||||
content=content,
|
||||
tool_calls=tool_calls,
|
||||
tools=request.tools,
|
||||
)
|
||||
|
||||
assert len(response_outputs) == 1
|
||||
tool_call = response_outputs[0]
|
||||
assert tool_call.type == "function_call"
|
||||
assert tool_call.name == "calc"
|
||||
assert json.loads(tool_call.arguments) == {"x": 1}
|
||||
|
||||
|
||||
def test_responses_named_tool_choice_uses_xtml_parser():
|
||||
parser = KimiK3DelegatingParser(
|
||||
DummyTokenizer(), chat_template_kwargs={"thinking": False}
|
||||
)
|
||||
request = _responses_request(
|
||||
tool_choice={"type": "function", "name": "calc"}
|
||||
).model_copy(update={"chat_template_kwargs": {"thinking": False}})
|
||||
output = RESPONSE_CLOSE + _tools(_call("calc", 1, _arg("x", "number", "1")))
|
||||
|
||||
reasoning, content, tool_calls = parser.parse(
|
||||
output, request, enable_auto_tools=True, model_output_token_ids=[]
|
||||
)
|
||||
response_outputs = build_response_output_items(
|
||||
reasoning=reasoning,
|
||||
content=content,
|
||||
tool_calls=tool_calls,
|
||||
tools=request.tools,
|
||||
)
|
||||
|
||||
assert len(response_outputs) == 1
|
||||
tool_call = response_outputs[0]
|
||||
assert tool_call.type == "function_call"
|
||||
assert tool_call.name == "calc"
|
||||
assert json.loads(tool_call.arguments) == {"x": 1}
|
||||
|
||||
|
||||
def test_chat_params_carries_tool_choice_metadata():
|
||||
request = _request().model_copy(update={"tool_choice": "required"})
|
||||
|
||||
chat_params = request.build_chat_params(
|
||||
default_template=None,
|
||||
default_template_content_format="auto",
|
||||
)
|
||||
|
||||
assert chat_params.tool_choice == "required"
|
||||
|
||||
|
||||
def test_chat_params_carries_response_format_metadata():
|
||||
request = ChatCompletionRequest(
|
||||
model="test-model",
|
||||
messages=[],
|
||||
response_format={"type": "json_object"},
|
||||
)
|
||||
|
||||
chat_params = request.build_chat_params(
|
||||
default_template=None,
|
||||
default_template_content_format="auto",
|
||||
)
|
||||
|
||||
assert chat_params.response_format is request.response_format
|
||||
assert chat_params.tool_choice is None
|
||||
|
||||
|
||||
def test_chat_params_keeps_template_tool_choice_when_api_auto():
|
||||
request = _request().model_copy(
|
||||
update={
|
||||
"tool_choice": "auto",
|
||||
"chat_template_kwargs": {"tool_choice": "required"},
|
||||
}
|
||||
)
|
||||
|
||||
chat_params = request.build_chat_params(
|
||||
default_template=None,
|
||||
default_template_content_format="auto",
|
||||
)
|
||||
|
||||
assert chat_params.chat_template_kwargs["tool_choice"] == "required"
|
||||
assert chat_params.tool_choice == "auto"
|
||||
+12
-1
@@ -84,7 +84,14 @@ RunnerOption = Literal["auto", RunnerType]
|
||||
ConvertType = Literal["none", "embed", "classify"]
|
||||
ConvertOption = Literal["auto", ConvertType]
|
||||
TokenizerMode = Literal[
|
||||
"auto", "hf", "slow", "mistral", "deepseek_v32", "deepseek_v4", "inkling"
|
||||
"auto",
|
||||
"hf",
|
||||
"slow",
|
||||
"mistral",
|
||||
"deepseek_v32",
|
||||
"deepseek_v4",
|
||||
"inkling",
|
||||
"kimi_k3",
|
||||
]
|
||||
ModelDType = Literal["auto", "half", "float16", "bfloat16", "float", "float32"]
|
||||
LogprobsMode = Literal[
|
||||
@@ -141,6 +148,8 @@ class ModelConfig:
|
||||
- "mistral" will always use the tokenizer from `mistral_common`.
|
||||
- "deepseek_v32" will always use the tokenizer from `deepseek_v32`.
|
||||
- "deepseek_v4" will always use the tokenizer from `deepseek_v4`.
|
||||
- "kimi_k3" will always use the "hf" tokenizer but render chat prompts
|
||||
with Kimi K3's Python XTML encoding instead of a Jinja template.
|
||||
- Other custom values can be supported via plugins.
|
||||
|
||||
To swap the Rust BPE backend that powers HF fast tokenizers for the
|
||||
@@ -634,6 +643,8 @@ class ModelConfig:
|
||||
self.tokenizer_mode = "terratorch"
|
||||
elif arch == "MoonshotKimiaForCausalLM":
|
||||
self.tokenizer_mode = "kimi_audio"
|
||||
elif arch == "KimiK3ForConditionalGeneration":
|
||||
self.tokenizer_mode = "kimi_k3"
|
||||
elif arch == "DeepseekV32ForCausalLM":
|
||||
self.tokenizer_mode = "deepseek_v32"
|
||||
elif arch == "DeepseekV4ForCausalLM":
|
||||
|
||||
@@ -2024,13 +2024,19 @@ def get_history_tool_calls_cnt(conversation: list[ConversationMessage]):
|
||||
return idx
|
||||
|
||||
|
||||
_KIMI_MODEL_TYPES = ("kimi_k2", "kimi_k25")
|
||||
_KIMI_MODEL_TYPES = ("kimi_k2", "kimi_k25", "kimi_k3")
|
||||
|
||||
|
||||
def get_tool_call_id_type(model_config: ModelConfig) -> str:
|
||||
"""Return the tool-call ID type for a given model configuration."""
|
||||
hf_overrides = getattr(model_config, "hf_overrides", None)
|
||||
if model_config.hf_text_config.model_type in _KIMI_MODEL_TYPES or (
|
||||
hf_config = getattr(model_config, "hf_config", None)
|
||||
hf_text_config = getattr(model_config, "hf_text_config", None)
|
||||
model_types = (
|
||||
getattr(hf_config, "model_type", None),
|
||||
getattr(hf_text_config, "model_type", None),
|
||||
)
|
||||
if any(model_type in _KIMI_MODEL_TYPES for model_type in model_types) or (
|
||||
isinstance(hf_overrides, dict)
|
||||
and hf_overrides.get("model_type") in _KIMI_MODEL_TYPES
|
||||
):
|
||||
|
||||
@@ -565,6 +565,12 @@ class ChatCompletionRequest(OpenAIBaseModel):
|
||||
),
|
||||
media_io_kwargs=self.media_io_kwargs,
|
||||
return_assistant_tokens_mask=bool(self.return_assistant_tokens_mask),
|
||||
# No-tools requests default to tool_choice="none" at the API
|
||||
# layer. Collapse that default before rendering, so K3 emits a
|
||||
# model-visible tool-choice instruction only for requests with a
|
||||
# tools block.
|
||||
tool_choice=self.tool_choice if self.tools else None,
|
||||
response_format=self.response_format,
|
||||
)
|
||||
|
||||
def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
|
||||
|
||||
@@ -337,6 +337,7 @@ class ResponsesRequest(OpenAIBaseModel):
|
||||
extra_kwargs,
|
||||
),
|
||||
media_io_kwargs=self.media_io_kwargs,
|
||||
tool_choice=self.tool_choice if self.tools else None,
|
||||
)
|
||||
|
||||
def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
|
||||
|
||||
@@ -0,0 +1,135 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from vllm.entrypoints.openai.engine.protocol import (
|
||||
DeltaMessage,
|
||||
FunctionCall,
|
||||
)
|
||||
from vllm.parser.abstract_parser import DelegatingParser
|
||||
from vllm.reasoning.kimi_k3_reasoning_parser import KimiK3ReasoningParser
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm.entrypoints.openai.chat_completion.protocol import (
|
||||
ChatCompletionRequest,
|
||||
)
|
||||
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
|
||||
|
||||
|
||||
class KimiK3Parser(DelegatingParser):
|
||||
"""Compose the Kimi K3 reasoning and tool parsers for XTML output."""
|
||||
|
||||
# TODO: Switch Kimi K3 to the parser engine once its XTML reasoning/tool
|
||||
# path is covered there.
|
||||
def _extract_tool_calls(
|
||||
self,
|
||||
content: str | None,
|
||||
request: ChatCompletionRequest | ResponsesRequest,
|
||||
enable_auto_tools: bool = False,
|
||||
) -> tuple[list[FunctionCall] | None, str | None]:
|
||||
if self._tool_parser is None or not enable_auto_tools:
|
||||
return super()._extract_tool_calls(content, request, enable_auto_tools)
|
||||
|
||||
tool_call_info = self.extract_tool_calls(content or "", request=request)
|
||||
if request.tool_choice == "none":
|
||||
return [], tool_call_info.content
|
||||
if not tool_call_info.tools_called:
|
||||
return None, tool_call_info.content
|
||||
|
||||
tool_calls = [
|
||||
FunctionCall(
|
||||
id=tool_call.id,
|
||||
name=tool_call.function.name,
|
||||
arguments=tool_call.function.arguments,
|
||||
)
|
||||
for tool_call in tool_call_info.tool_calls
|
||||
]
|
||||
parsed_content = tool_call_info.content
|
||||
if parsed_content and parsed_content.strip() == "":
|
||||
parsed_content = None
|
||||
return tool_calls, parsed_content
|
||||
|
||||
def _extract_tool_calls_streaming(
|
||||
self,
|
||||
previous_text: str,
|
||||
current_text: str,
|
||||
delta_text: str,
|
||||
previous_token_ids: Sequence[int],
|
||||
current_token_ids: Sequence[int],
|
||||
delta_token_ids: Sequence[int],
|
||||
request: ChatCompletionRequest | ResponsesRequest,
|
||||
tool_call_idx: int | None = None,
|
||||
tool_call_id_type: str = "random",
|
||||
function_name_returned: bool = False,
|
||||
) -> tuple[DeltaMessage | None, bool]:
|
||||
if request.tool_choice != "none":
|
||||
return super()._extract_tool_calls_streaming(
|
||||
previous_text,
|
||||
current_text,
|
||||
delta_text,
|
||||
previous_token_ids,
|
||||
current_token_ids,
|
||||
delta_token_ids,
|
||||
request,
|
||||
tool_call_idx=tool_call_idx,
|
||||
tool_call_id_type=tool_call_id_type,
|
||||
function_name_returned=function_name_returned,
|
||||
)
|
||||
|
||||
delta_message = self.extract_tool_calls_streaming(
|
||||
previous_text,
|
||||
current_text,
|
||||
delta_text,
|
||||
previous_token_ids,
|
||||
current_token_ids,
|
||||
delta_token_ids,
|
||||
request,
|
||||
)
|
||||
if delta_message is not None:
|
||||
delta_message.tool_calls = []
|
||||
return delta_message, False
|
||||
|
||||
def parse_delta(
|
||||
self,
|
||||
delta_text: str,
|
||||
delta_token_ids: list[int],
|
||||
request: ChatCompletionRequest | ResponsesRequest,
|
||||
prompt_token_ids: list[int] | None = None,
|
||||
*,
|
||||
finished: bool,
|
||||
) -> DeltaMessage | None:
|
||||
state = self._stream_state
|
||||
previous_content = state.previous_text if state.reasoning_ended else ""
|
||||
delta_message = super().parse_delta(
|
||||
delta_text,
|
||||
delta_token_ids,
|
||||
request,
|
||||
prompt_token_ids,
|
||||
finished=finished,
|
||||
)
|
||||
|
||||
if (
|
||||
self._tool_parser is not None
|
||||
or not isinstance(self._reasoning_parser, KimiK3ReasoningParser)
|
||||
or not state.reasoning_ended
|
||||
or delta_message is None
|
||||
):
|
||||
return delta_message
|
||||
|
||||
stripped = self._reasoning_parser.strip_content_streaming(
|
||||
previous_text=previous_content,
|
||||
current_text=state.previous_text,
|
||||
)
|
||||
delta_message.content = stripped.content if stripped is not None else None
|
||||
if (
|
||||
delta_message.role is None
|
||||
and delta_message.content is None
|
||||
and delta_message.reasoning is None
|
||||
and not delta_message.tool_calls
|
||||
):
|
||||
return None
|
||||
return delta_message
|
||||
@@ -125,6 +125,18 @@ class ParserManager:
|
||||
MistralParser.tool_parser_cls = tool_parser_cls
|
||||
return MistralParser
|
||||
|
||||
if reasoning_parser_name == "kimi_k3" or tool_parser_name == "kimi_k3":
|
||||
from vllm.parser.kimi_k3 import KimiK3Parser
|
||||
|
||||
r_cls = reasoning_parser_cls
|
||||
t_cls = tool_parser_cls
|
||||
|
||||
class _KimiK3Parser(KimiK3Parser):
|
||||
reasoning_parser_cls = r_cls
|
||||
tool_parser_cls = t_cls
|
||||
|
||||
return _KimiK3Parser
|
||||
|
||||
from vllm.parser.abstract_parser import DelegatingParser
|
||||
|
||||
r_cls = reasoning_parser_cls
|
||||
|
||||
@@ -84,6 +84,10 @@ _REASONING_PARSERS_TO_REGISTER = {
|
||||
"kimi_k2_reasoning_parser",
|
||||
"KimiK2ReasoningParser",
|
||||
),
|
||||
"kimi_k3": (
|
||||
"kimi_k3_reasoning_parser",
|
||||
"KimiK3ReasoningParser",
|
||||
),
|
||||
"mimo": (
|
||||
"qwen3_engine_reasoning_parser",
|
||||
"Qwen3ParserReasoningAdapter",
|
||||
|
||||
@@ -0,0 +1,371 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Reasoning parser for the Kimi K3 (XTML) chat format.
|
||||
|
||||
This strips the ``think`` channel out of generated text and hands the remainder
|
||||
(``response`` + ``tools`` channels) downstream. Kimi K3 wraps the thinking
|
||||
channel as an XTML element built from special tokens::
|
||||
|
||||
<|open|>think<|sep|> <reasoning> <|close|>think<|sep|>
|
||||
|
||||
Two subtleties drive the implementation:
|
||||
* Unlike Kimi-K2 (a single ``<think>`` token), each K3 marker is a 3-token
|
||||
sequence, so the token-id helpers search for the marker *subsequence*
|
||||
rather than a single id.
|
||||
* In thinking mode the serving layer may feed ``<|open|>think<|sep|>`` as the
|
||||
generation prefix, so the model's output can begin *inside* the think
|
||||
channel with no open marker. The text paths therefore treat a missing open
|
||||
marker as "reasoning starts at offset 0".
|
||||
|
||||
When thinking is disabled (``chat_template_kwargs={"thinking": False}`` or
|
||||
``{"enable_thinking": False}``, i.e. instruct mode) the parser returns every
|
||||
delta as normal content; there is simply no think channel to extract.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import regex as re
|
||||
from transformers import PreTrainedTokenizerBase
|
||||
|
||||
from vllm.entrypoints.openai.engine.protocol import DeltaMessage
|
||||
from vllm.reasoning import ReasoningParser
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
|
||||
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
|
||||
|
||||
|
||||
def _subseq_index(haystack: Sequence[int], needle: Sequence[int]) -> int:
|
||||
"""Return start index of the last occurrence of needle in haystack, or -1."""
|
||||
n = len(needle)
|
||||
if n == 0:
|
||||
return -1
|
||||
for i in range(len(haystack) - n, -1, -1):
|
||||
if list(haystack[i : i + n]) == list(needle):
|
||||
return i
|
||||
return -1
|
||||
|
||||
|
||||
class KimiK3ReasoningParser(ReasoningParser):
|
||||
"""Reasoning parser for the Kimi K3 (XTML) think channel."""
|
||||
|
||||
def __init__(self, tokenizer: PreTrainedTokenizerBase, *args, **kwargs):
|
||||
super().__init__(tokenizer)
|
||||
|
||||
if not self.model_tokenizer:
|
||||
raise ValueError(
|
||||
"The model tokenizer must be passed to the ReasoningParser "
|
||||
"constructor during construction."
|
||||
)
|
||||
|
||||
# thinking can be disabled via chat_template_kwargs -> identity fallthrough
|
||||
chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {}
|
||||
thinking = chat_kwargs.get("thinking", None)
|
||||
if thinking is None:
|
||||
thinking = chat_kwargs.get("enable_thinking", True)
|
||||
self._thinking_enabled = bool(thinking)
|
||||
|
||||
# XTML markers as literal strings (skip_special_tokens=False at serve time)
|
||||
self._think_open = "<|open|>think<|sep|>"
|
||||
self._think_close = "<|close|>think<|sep|>"
|
||||
|
||||
# Content-channel markers that must be stripped from the final output.
|
||||
# The tool parser handles these when active, but when tool_choice is
|
||||
# "none" (the default for requests without tools) the tool parser is
|
||||
# bypassed and the reasoning parser must strip them itself.
|
||||
self._response_open = "<|open|>response<|sep|>"
|
||||
self._response_close = "<|close|>response<|sep|>"
|
||||
self._message_close = "<|close|>message<|sep|>"
|
||||
|
||||
# Tolerant matchers: defense-in-depth against vLLM's added-token spacing
|
||||
# ("<|open|> think <|sep|>"). The `\s*` is a no-op on clean input, so the
|
||||
# normal path stays byte-exact; it only helps if some serving path leaves
|
||||
# spaces_between_special_tokens on. See adjust_request for the real fix.
|
||||
open_marker = r"<\|open\|>"
|
||||
close_marker = r"<\|close\|>"
|
||||
sep_marker = r"<\|sep\|>"
|
||||
self._think_open_re = re.compile(open_marker + r"\s*think\s*" + sep_marker)
|
||||
self._think_close_re = re.compile(close_marker + r"\s*think\s*" + sep_marker)
|
||||
self._response_open_re = re.compile(
|
||||
open_marker + r"\s*response\s*" + sep_marker
|
||||
)
|
||||
self._response_close_re = re.compile(
|
||||
close_marker + r"\s*response\s*" + sep_marker
|
||||
)
|
||||
self._message_close_re = re.compile(
|
||||
close_marker + r"\s*message\s*" + sep_marker
|
||||
)
|
||||
|
||||
# marker token-id subsequences (each marker is 3 tokens)
|
||||
self._think_open_ids = tokenizer.encode(
|
||||
self._think_open, add_special_tokens=False
|
||||
)
|
||||
self._think_close_ids = tokenizer.encode(
|
||||
self._think_close, add_special_tokens=False
|
||||
)
|
||||
self._last_streaming_delta_token_ids: tuple[int, ...] | None = None
|
||||
self._last_streaming_content_token_ids: list[int] | None = None
|
||||
|
||||
@property
|
||||
def reasoning_start_str(self) -> str | None:
|
||||
return self._think_open
|
||||
|
||||
@property
|
||||
def reasoning_end_str(self) -> str | None:
|
||||
return self._think_close
|
||||
|
||||
def adjust_request(
|
||||
self,
|
||||
request: "ChatCompletionRequest | ResponsesRequest",
|
||||
) -> "ChatCompletionRequest | ResponsesRequest":
|
||||
request.skip_special_tokens = False
|
||||
if hasattr(request, "spaces_between_special_tokens"):
|
||||
request.spaces_between_special_tokens = False
|
||||
return request
|
||||
|
||||
def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:
|
||||
if not self._thinking_enabled:
|
||||
return True
|
||||
# Reasoning has ended only if the *most recent* think block is closed:
|
||||
# the last close marker must come after the last open marker. A plain
|
||||
# "a close marker exists anywhere" check false-positives in multi-turn /
|
||||
# agent continuations, where the chat template keeps a prior turn's
|
||||
# think channel (with its <|close|>think<|sep|>) in the prompt while the
|
||||
# current turn is still reasoning (its <|open|>think<|sep|> is the newest
|
||||
# marker). See _subseq_index (returns the last occurrence).
|
||||
last_close = _subseq_index(input_ids, self._think_close_ids)
|
||||
last_open = _subseq_index(input_ids, self._think_open_ids)
|
||||
if last_open == -1:
|
||||
# No open marker in scope (e.g. the open was consumed as the
|
||||
# generation prefix): a close marker means reasoning has ended.
|
||||
return last_close != -1
|
||||
return last_close > last_open
|
||||
|
||||
def _extract_content_ids(self, input_ids: list[int]) -> list[int]:
|
||||
if not self._thinking_enabled:
|
||||
return input_ids
|
||||
idx = _subseq_index(input_ids, self._think_close_ids)
|
||||
if idx == -1:
|
||||
return [] # still reasoning
|
||||
return input_ids[idx + len(self._think_close_ids) :]
|
||||
|
||||
def extract_content_ids(self, input_ids: list[int]) -> list[int]:
|
||||
cached_delta_ids = self._last_streaming_delta_token_ids
|
||||
cached_content_ids = self._last_streaming_content_token_ids
|
||||
self._last_streaming_delta_token_ids = None
|
||||
self._last_streaming_content_token_ids = None
|
||||
if cached_delta_ids == tuple(input_ids) and cached_content_ids is not None:
|
||||
return cached_content_ids
|
||||
return self._extract_content_ids(input_ids)
|
||||
|
||||
def _strip_content_wrapper(self, text: str) -> str:
|
||||
"""Strip ``<|open|>response<|sep|>…<|close|>response<|sep|>`` wrapper and
|
||||
``<|close|>message<|sep|>`` from *text*.
|
||||
|
||||
When Kimi K3 tool parsing is active it needs the raw XTML
|
||||
``response`` + ``tools`` channels. Otherwise the reasoning parser cleans
|
||||
up the response wrapper itself so API users do not see XTML markers.
|
||||
"""
|
||||
# Try structured unwrap first: extract body of response channel
|
||||
m_ro = self._response_open_re.search(text)
|
||||
m_rc = self._response_close_re.search(text, m_ro.end() if m_ro else 0)
|
||||
if m_ro is not None and m_rc is not None:
|
||||
text = text[m_ro.end() : m_rc.start()]
|
||||
elif m_ro is not None:
|
||||
# response opened but not closed (truncated)
|
||||
text = text[m_ro.end() :]
|
||||
else:
|
||||
# No response wrapper — strip stray markers if present
|
||||
text = self._response_open_re.sub("", text)
|
||||
text = self._response_close_re.sub("", text)
|
||||
text = self._message_close_re.sub("", text)
|
||||
return text
|
||||
|
||||
@staticmethod
|
||||
def _should_preserve_tool_channels(
|
||||
request: "ChatCompletionRequest | ResponsesRequest",
|
||||
) -> bool:
|
||||
return bool(getattr(request, "tools", None)) and (
|
||||
getattr(request, "tool_choice", None) != "none"
|
||||
)
|
||||
|
||||
def _content_after_reasoning(
|
||||
self,
|
||||
text: str,
|
||||
request: "ChatCompletionRequest | ResponsesRequest",
|
||||
) -> str | None:
|
||||
if self._should_preserve_tool_channels(request):
|
||||
return text or None
|
||||
return self._strip_content_wrapper(text) or None
|
||||
|
||||
def extract_reasoning(
|
||||
self, model_output: str, request: "ChatCompletionRequest | ResponsesRequest"
|
||||
) -> tuple[str | None, str | None]:
|
||||
"""Split full text into ``(reasoning, rest)`` for the non-streaming path.
|
||||
|
||||
Handles three shapes:
|
||||
* no think channel at all -> ``(None, model_output)`` (all content)
|
||||
* open marker present -> reasoning starts after ``<|open|>think<|sep|>``
|
||||
* open marker absent but a close marker exists (gen-prefix consumed
|
||||
the open) -> reasoning starts at offset 0
|
||||
``rest`` is whatever follows the close marker, fed on to the tool parser.
|
||||
"""
|
||||
if not self._thinking_enabled:
|
||||
return None, self._content_after_reasoning(model_output, request)
|
||||
|
||||
m_open = self._think_open_re.search(model_output)
|
||||
# reasoning content begins right after think-open (or at start if the
|
||||
# open marker was already consumed as a generation prefix)
|
||||
content_start = m_open.end() if m_open is not None else 0
|
||||
# if there is no think channel at all, everything is content
|
||||
if m_open is None and self._think_close_re.search(model_output) is None:
|
||||
return None, self._content_after_reasoning(model_output, request)
|
||||
|
||||
m_close = self._think_close_re.search(model_output, content_start)
|
||||
if m_close is not None:
|
||||
reasoning = model_output[content_start : m_close.start()]
|
||||
rest = model_output[m_close.end() :]
|
||||
return (reasoning or None, self._content_after_reasoning(rest, request))
|
||||
# think not closed -> still reasoning, no content yet
|
||||
return (model_output[content_start:] or None, None)
|
||||
|
||||
def _reasoning_text_ready_to_emit(self, text: str) -> str:
|
||||
"""Return the reasoning prefix that is safe to stream now.
|
||||
|
||||
Work from accumulated text, not from the current delta alone. That turns
|
||||
split-open and split-close handling into the same prefix-diff problem.
|
||||
|
||||
Example 1: split think-open marker.
|
||||
chunks: ``<|open|>`` / ``think`` / ``<|sep|>reasoning``
|
||||
current text after chunk 1: ``<|open|>`` -> emit ``""``
|
||||
current text after chunk 2: ``<|open|>think`` -> emit ``""``
|
||||
current text after chunk 3: ``<|open|>think<|sep|>reasoning``
|
||||
-> emit ``reasoning``
|
||||
|
||||
Example 2: split think-close marker.
|
||||
chunks: ``reasoning`` / ``<|close|>`` / ``think<|sep|>...``
|
||||
after ``<|close|>``, the suffix is only a partial close marker, so the
|
||||
sendable reasoning is still ``reasoning`` and the delta is empty. Once
|
||||
``think<|sep|>`` arrives, the close branch hands the following response
|
||||
or tools text to the downstream parser.
|
||||
"""
|
||||
m_open = self._think_open_re.search(text)
|
||||
if m_open is not None:
|
||||
text = text[m_open.end() :]
|
||||
overlap = 0
|
||||
for marker in (self._think_open, self._think_close):
|
||||
max_check = min(len(marker) - 1, len(text))
|
||||
for n in range(max_check, 0, -1):
|
||||
if text.endswith(marker[:n]):
|
||||
overlap = max(overlap, n)
|
||||
break
|
||||
return text[:-overlap] if overlap else text
|
||||
|
||||
def _content_ready_to_emit(self, text: str) -> str:
|
||||
"""Return the content prefix that is safe to stream now.
|
||||
|
||||
Mirrors ``_reasoning_text_ready_to_emit`` but for the post-reasoning
|
||||
content phase. Strips the ``<|open|>response<|sep|>`` prefix, holds back
|
||||
any partial marker suffix, and removes complete
|
||||
``<|close|>response<|sep|>`` / ``<|close|>message<|sep|>`` markers.
|
||||
"""
|
||||
# Strip response-open prefix
|
||||
m_open = self._response_open_re.search(text)
|
||||
if m_open is not None:
|
||||
text = text[m_open.end() :]
|
||||
|
||||
# Remove complete close/message markers
|
||||
text = self._response_close_re.sub("", text)
|
||||
text = self._message_close_re.sub("", text)
|
||||
|
||||
# Hold back partial markers at the end
|
||||
overlap = 0
|
||||
for marker in (
|
||||
self._response_open,
|
||||
self._response_close,
|
||||
self._message_close,
|
||||
):
|
||||
max_check = min(len(marker) - 1, len(text))
|
||||
for n in range(max_check, 0, -1):
|
||||
if text.endswith(marker[:n]):
|
||||
overlap = max(overlap, n)
|
||||
break
|
||||
return text[:-overlap] if overlap else text
|
||||
|
||||
def strip_content_streaming(
|
||||
self,
|
||||
previous_text: str,
|
||||
current_text: str,
|
||||
) -> DeltaMessage | None:
|
||||
"""Strip XTML content wrappers from streaming deltas after reasoning.
|
||||
|
||||
Called by KimiK3Parser when no tool parser is configured, so the
|
||||
reasoning parser handles ``<|open|>response<|sep|>`` /
|
||||
``<|close|>response<|sep|>`` / ``<|close|>message<|sep|>`` stripping itself.
|
||||
|
||||
Works from accumulated text (``previous_text`` / ``current_text``
|
||||
already contain only post-reasoning content).
|
||||
"""
|
||||
current_safe = self._content_ready_to_emit(current_text)
|
||||
previous_safe = self._content_ready_to_emit(previous_text)
|
||||
if current_safe.startswith(previous_safe):
|
||||
delta = current_safe[len(previous_safe) :]
|
||||
else:
|
||||
delta = current_safe
|
||||
return DeltaMessage(content=delta) if delta else None
|
||||
|
||||
def extract_reasoning_streaming(
|
||||
self,
|
||||
previous_text: str,
|
||||
current_text: str,
|
||||
delta_text: str,
|
||||
previous_token_ids: Sequence[int],
|
||||
current_token_ids: Sequence[int],
|
||||
delta_token_ids: Sequence[int],
|
||||
) -> DeltaMessage | None:
|
||||
self._last_streaming_delta_token_ids = None
|
||||
self._last_streaming_content_token_ids = None
|
||||
if not self._thinking_enabled:
|
||||
return DeltaMessage(content=delta_text)
|
||||
|
||||
# reasoning already ended -> downstream content
|
||||
if self._think_close_re.search(previous_text):
|
||||
return DeltaMessage(content=delta_text)
|
||||
|
||||
# the close marker completes within this delta's accumulated text:
|
||||
# split the buffer at the close marker into reasoning vs trailing content.
|
||||
m_close = self._think_close_re.search(current_text)
|
||||
if m_close is not None:
|
||||
self._last_streaming_delta_token_ids = tuple(delta_token_ids)
|
||||
self._last_streaming_content_token_ids = self._extract_content_ids(
|
||||
list(current_token_ids)
|
||||
)
|
||||
m_open = self._think_open_re.search(current_text)
|
||||
r_start = m_open.end() if m_open is not None else 0
|
||||
reasoning = current_text[r_start : m_close.start()]
|
||||
already_sent = self._reasoning_text_ready_to_emit(previous_text)
|
||||
if reasoning.startswith(already_sent):
|
||||
reasoning_delta = reasoning[len(already_sent) :]
|
||||
else:
|
||||
reasoning_delta = reasoning
|
||||
content = current_text[m_close.end() :]
|
||||
return DeltaMessage(
|
||||
reasoning=reasoning_delta or None,
|
||||
content=content or None,
|
||||
)
|
||||
|
||||
current_reasoning = self._reasoning_text_ready_to_emit(current_text)
|
||||
previous_reasoning = self._reasoning_text_ready_to_emit(previous_text)
|
||||
if current_reasoning.startswith(previous_reasoning):
|
||||
reasoning_delta = current_reasoning[len(previous_reasoning) :]
|
||||
else:
|
||||
reasoning_delta = current_reasoning
|
||||
if not reasoning_delta:
|
||||
return None
|
||||
return DeltaMessage(reasoning=reasoning_delta)
|
||||
|
||||
# Backward-compatible aliases for existing unit tests and downstream users
|
||||
# that still call the pre-split method names.
|
||||
extract_reasoning_content = extract_reasoning
|
||||
extract_reasoning_content_streaming = extract_reasoning_streaming
|
||||
@@ -0,0 +1,220 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from typing import Any, cast
|
||||
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.entrypoints.chat_utils import (
|
||||
ChatCompletionMessageParam,
|
||||
ConversationMessage,
|
||||
parse_chat_messages,
|
||||
parse_chat_messages_async,
|
||||
)
|
||||
from vllm.exceptions import VLLMValidationError
|
||||
from vllm.multimodal.media.connector import merge_media_io_kwargs
|
||||
from vllm.tokenizers.hf import HfTokenizer
|
||||
from vllm.utils.async_utils import make_async
|
||||
|
||||
from .base import BaseRenderer
|
||||
from .inputs import DictPrompt
|
||||
from .inputs.preprocess import parse_dec_only_prompt
|
||||
from .params import ChatParams
|
||||
|
||||
# Keep the original image mode (including the alpha channel) for K3 instead of
|
||||
# flattening images onto a background color. Server-level (--media-io-kwargs)
|
||||
# and request-level media_io_kwargs still take precedence over this default.
|
||||
_K3_MEDIA_IO_DEFAULTS: dict[str, dict[str, Any]] = {"image": {"image_mode": None}}
|
||||
_K3_THINKING_EFFORTS = ("low", "high", "max")
|
||||
|
||||
|
||||
def _merge_k3_media_io_kwargs(
|
||||
media_io_kwargs: dict[str, dict[str, Any]] | None,
|
||||
) -> dict[str, dict[str, Any]] | None:
|
||||
return merge_media_io_kwargs(_K3_MEDIA_IO_DEFAULTS, media_io_kwargs)
|
||||
|
||||
|
||||
def _dump_k3_template_value(value: Any) -> Any:
|
||||
model_dump = getattr(value, "model_dump", None)
|
||||
if callable(model_dump):
|
||||
return model_dump(mode="json", exclude_none=True)
|
||||
|
||||
to_dict = getattr(value, "to_dict", None)
|
||||
if callable(to_dict):
|
||||
return to_dict()
|
||||
|
||||
return value
|
||||
|
||||
|
||||
def _apply_k3_thinking_kwargs(kwargs: dict[str, Any]) -> None:
|
||||
if (enable_thinking := kwargs.pop("enable_thinking", None)) is not None:
|
||||
kwargs.setdefault("thinking", enable_thinking)
|
||||
|
||||
reasoning_effort = kwargs.pop("reasoning_effort", None)
|
||||
if reasoning_effort == "none":
|
||||
kwargs.setdefault("thinking", False)
|
||||
elif reasoning_effort is not None:
|
||||
kwargs.setdefault("thinking_effort", reasoning_effort)
|
||||
|
||||
thinking_effort = kwargs.get("thinking_effort")
|
||||
if thinking_effort is not None and thinking_effort not in _K3_THINKING_EFFORTS:
|
||||
supported = ", ".join(_K3_THINKING_EFFORTS)
|
||||
raise VLLMValidationError(
|
||||
f"Kimi K3 supports thinking_effort values: {supported}",
|
||||
parameter="thinking_effort",
|
||||
value=thinking_effort,
|
||||
)
|
||||
|
||||
|
||||
def _normalize_k3_tool_messages(
|
||||
conversation: list[ConversationMessage],
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Reorder tool-result messages to match assistant tool_call order.
|
||||
|
||||
Supports matching by ``tool_call_id`` or by the synthetic
|
||||
``"{tool}:{zero_based_index}"`` alias. When any tool message in a
|
||||
block cannot be resolved, the whole block is left in its original
|
||||
order (graceful fallback).
|
||||
|
||||
Returns a new list; caller-owned message dicts are not mutated.
|
||||
"""
|
||||
normalized: list[dict[str, Any]] = []
|
||||
i = 0
|
||||
|
||||
while i < len(conversation):
|
||||
message = conversation[i]
|
||||
normalized.append(dict(message))
|
||||
i += 1
|
||||
|
||||
if message.get("role") != "assistant":
|
||||
continue
|
||||
|
||||
tool_calls = message.get("tool_calls")
|
||||
if not tool_calls:
|
||||
continue
|
||||
|
||||
# Build the lookup table from the assistant's tool_calls.
|
||||
targets_by_id: dict[str, tuple[int, str]] = {}
|
||||
for position, tool_call in enumerate(tool_calls):
|
||||
function = tool_call.get("function")
|
||||
if not isinstance(function, dict):
|
||||
continue
|
||||
name = function.get("name")
|
||||
if not isinstance(name, str) or not name:
|
||||
continue
|
||||
|
||||
call_target = (position, name)
|
||||
aliases = [f"{name}:{position}"]
|
||||
if tool_call_id := tool_call.get("id"):
|
||||
aliases.insert(0, str(tool_call_id))
|
||||
for alias in aliases:
|
||||
targets_by_id.setdefault(alias, call_target)
|
||||
|
||||
# Collect consecutive tool messages.
|
||||
block_start = i
|
||||
while i < len(conversation) and conversation[i].get("role") == "tool":
|
||||
i += 1
|
||||
if i == block_start:
|
||||
continue
|
||||
|
||||
# Try to resolve every tool message in the block.
|
||||
resolved: list[tuple[int, int, dict[str, Any]]] = []
|
||||
for order, tool_message in enumerate(conversation[block_start:i]):
|
||||
tool_call_id = tool_message.get("tool_call_id")
|
||||
resolved_target = (
|
||||
targets_by_id.get(str(tool_call_id))
|
||||
if tool_call_id is not None
|
||||
else None
|
||||
)
|
||||
if resolved_target is None:
|
||||
normalized.extend(dict(item) for item in conversation[block_start:i])
|
||||
break
|
||||
|
||||
position, name = resolved_target
|
||||
enriched = dict(tool_message)
|
||||
enriched["tool"] = name
|
||||
enriched["index"] = position + 1
|
||||
resolved.append((position, order, enriched))
|
||||
else:
|
||||
resolved.sort(key=lambda item: (item[0], item[1]))
|
||||
normalized.extend(item for _, _, item in resolved)
|
||||
|
||||
return normalized
|
||||
|
||||
|
||||
class KimiK3Renderer(BaseRenderer[HfTokenizer]):
|
||||
"""Render chat prompts with Kimi K3's Python XTML encoding.
|
||||
|
||||
K3 ships no Jinja chat template; its tokenizer renders messages through
|
||||
``encoding_k3`` instead. We tokenize eagerly so the structural markers keep
|
||||
their special-token ids while user- and tool-supplied text stays ordinary.
|
||||
"""
|
||||
|
||||
def __init__(self, config: VllmConfig, tokenizer: HfTokenizer | None) -> None:
|
||||
super().__init__(config, tokenizer)
|
||||
|
||||
self._apply_chat_template_async = make_async(
|
||||
self._apply_chat_template, executor=self._executor
|
||||
)
|
||||
|
||||
def _apply_chat_template(
|
||||
self,
|
||||
conversation: list[dict[str, Any]],
|
||||
params: ChatParams,
|
||||
) -> list[int]:
|
||||
# Tokenize eagerly: K3 encodes structural markers as special tokens and
|
||||
# user/tool text as ordinary tokens, so we cannot defer to a plain
|
||||
# re-tokenization of the rendered string downstream.
|
||||
kwargs = params.get_apply_chat_template_kwargs()
|
||||
_apply_k3_thinking_kwargs(kwargs)
|
||||
if params.tool_choice not in (None, "auto"):
|
||||
kwargs["tool_choice"] = _dump_k3_template_value(params.tool_choice)
|
||||
if params.response_format is not None:
|
||||
kwargs["response_format"] = _dump_k3_template_value(params.response_format)
|
||||
kwargs["tokenize"] = True
|
||||
return self.get_tokenizer().apply_chat_template(conversation, **kwargs)
|
||||
|
||||
def render_messages(
|
||||
self,
|
||||
messages: list[ChatCompletionMessageParam],
|
||||
params: ChatParams,
|
||||
) -> tuple[list[ConversationMessage], DictPrompt]:
|
||||
conversation, mm_data, mm_uuids = parse_chat_messages(
|
||||
messages,
|
||||
self.model_config,
|
||||
content_format="string",
|
||||
media_io_kwargs=_merge_k3_media_io_kwargs(params.media_io_kwargs),
|
||||
mm_processor_kwargs=params.mm_processor_kwargs,
|
||||
)
|
||||
|
||||
rendered_conversation = _normalize_k3_tool_messages(conversation)
|
||||
prompt = parse_dec_only_prompt(
|
||||
self._apply_chat_template(rendered_conversation, params)
|
||||
)
|
||||
if mm_data is not None:
|
||||
prompt["multi_modal_data"] = mm_data
|
||||
if mm_uuids is not None:
|
||||
prompt["multi_modal_uuids"] = mm_uuids
|
||||
|
||||
return cast(list[ConversationMessage], rendered_conversation), prompt
|
||||
|
||||
async def render_messages_async(
|
||||
self,
|
||||
messages: list[ChatCompletionMessageParam],
|
||||
params: ChatParams,
|
||||
) -> tuple[list[ConversationMessage], DictPrompt]:
|
||||
conversation, mm_data, mm_uuids = await parse_chat_messages_async(
|
||||
messages,
|
||||
self.model_config,
|
||||
content_format="string",
|
||||
media_io_kwargs=_merge_k3_media_io_kwargs(params.media_io_kwargs),
|
||||
mm_processor_kwargs=params.mm_processor_kwargs,
|
||||
)
|
||||
|
||||
rendered_conversation = _normalize_k3_tool_messages(conversation)
|
||||
token_ids = await self._apply_chat_template_async(rendered_conversation, params)
|
||||
prompt = parse_dec_only_prompt(token_ids)
|
||||
if mm_data is not None:
|
||||
prompt["multi_modal_data"] = mm_data
|
||||
if mm_uuids is not None:
|
||||
prompt["multi_modal_uuids"] = mm_uuids
|
||||
|
||||
return cast(list[ConversationMessage], rendered_conversation), prompt
|
||||
@@ -90,6 +90,12 @@ class ChatParams:
|
||||
return_assistant_tokens_mask: bool = False
|
||||
"""Request a per-token assistant mask from apply_chat_template."""
|
||||
|
||||
tool_choice: Any | None = None
|
||||
"""Request-level tool choice for renderers that need API metadata."""
|
||||
|
||||
response_format: Any | None = None
|
||||
"""Request-level response format for renderers that need API metadata."""
|
||||
|
||||
def with_defaults(
|
||||
self,
|
||||
default_chat_template_kwargs: dict[str, Any] | None = None,
|
||||
@@ -119,6 +125,8 @@ class ChatParams:
|
||||
self.mm_processor_kwargs,
|
||||
),
|
||||
return_assistant_tokens_mask=self.return_assistant_tokens_mask,
|
||||
tool_choice=self.tool_choice,
|
||||
response_format=self.response_format,
|
||||
)
|
||||
|
||||
def get_apply_chat_template_kwargs(self) -> dict[str, Any]:
|
||||
|
||||
@@ -24,6 +24,7 @@ _VLLM_RENDERERS = {
|
||||
"deepseek_v4": ("deepseek_v4", "DeepseekV4Renderer"),
|
||||
"hf": ("hf", "HfRenderer"),
|
||||
"kimi_audio": ("hf", "HfRenderer"),
|
||||
"kimi_k3": ("kimi_k3", "KimiK3Renderer"),
|
||||
"mistral": ("mistral", "MistralRenderer"),
|
||||
"terratorch": ("terratorch", "TerratorchRenderer"),
|
||||
"inkling": ("inkling", "InklingRenderer"),
|
||||
|
||||
@@ -44,6 +44,7 @@ _VLLM_TOKENIZERS = {
|
||||
"deepseek_v4": ("deepseek_v4", "DeepseekV4Tokenizer"),
|
||||
"hf": ("hf", "CachedHfTokenizer"),
|
||||
"kimi_audio": ("kimi_audio", "KimiAudioTokenizer"),
|
||||
"kimi_k3": ("hf", "CachedHfTokenizer"),
|
||||
"mistral": ("mistral", "MistralTokenizer"),
|
||||
# Inkling uses the plain HF tokenizer for token operations; the "inkling"
|
||||
# mode exists to select the InklingRenderer, which renders chat to
|
||||
|
||||
@@ -102,6 +102,10 @@ _TOOL_PARSERS_TO_REGISTER = {
|
||||
"kimi_k2_tool_parser",
|
||||
"KimiK2ToolParser",
|
||||
),
|
||||
"kimi_k3": (
|
||||
"kimi_k3_tool_parser",
|
||||
"KimiK3ToolParser",
|
||||
),
|
||||
"llama3_json": (
|
||||
"llama_tool_parser",
|
||||
"Llama3JsonToolParser",
|
||||
|
||||
@@ -0,0 +1,405 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tool-call parser for the Kimi K3 (XTML) chat format.
|
||||
|
||||
This turns the generated XTML ``response`` and ``tools`` channels back into
|
||||
OpenAI-compatible ``content`` and ``tool_calls``.
|
||||
|
||||
K3 assistant tool calls live in a nested ``tools`` channel::
|
||||
|
||||
<|open|>tools<|sep|>
|
||||
<|open|>call tool="python" index="1"<|sep|>
|
||||
<|open|>argument key="code" type="string"<|sep|>print(1)<|close|>argument<|sep|>
|
||||
<|open|>argument key="opts" type="object"<|sep|>{"a":1}<|close|>argument<|sep|>
|
||||
<|close|>call<|sep|>
|
||||
<|close|>tools<|sep|>
|
||||
|
||||
where ``<|open|>``, ``<|close|>``, ``<|sep|>`` are dedicated special tokens. The plain
|
||||
reply lives in a sibling ``response`` channel which is unwrapped into content.
|
||||
``response`` always precedes ``tools`` (see the protocol channel grammar).
|
||||
|
||||
Argument decoding mirrors the template's type tagging (inverse encoding):
|
||||
* type="string" -> value is the RAW text, no unescaping (template emits it raw)
|
||||
* other types -> value is JSON-decoded (number/boolean/null/object/array)
|
||||
|
||||
Attribute values (``tool=``, ``index=``, ``key=``, ``type=``) ARE escaped on the
|
||||
encode side (``&`` -> ``&``, ``"`` -> ``"``), so ``_attrs`` reverses
|
||||
them on decode (``"`` before ``&`` -- reverse of encode order).
|
||||
|
||||
Known limitation: because string argument and response bodies are emitted raw,
|
||||
a value that literally contains ``<|close|>argument<|sep|>`` or
|
||||
``<|close|>response<|sep|>`` is indistinguishable from a real closing marker.
|
||||
"""
|
||||
|
||||
import json
|
||||
from collections.abc import Sequence
|
||||
|
||||
import regex as re
|
||||
from openai.types.responses import ToolChoiceFunction
|
||||
|
||||
from vllm.entrypoints.openai.chat_completion.protocol import (
|
||||
ChatCompletionNamedToolChoiceParam,
|
||||
ChatCompletionRequest,
|
||||
)
|
||||
from vllm.entrypoints.openai.engine.protocol import (
|
||||
DeltaFunctionCall,
|
||||
DeltaMessage,
|
||||
DeltaToolCall,
|
||||
ExtractedToolCallInformation,
|
||||
FunctionCall,
|
||||
ToolCall,
|
||||
)
|
||||
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
|
||||
from vllm.exceptions import VLLMValidationError
|
||||
from vllm.logger import init_logger
|
||||
from vllm.tokenizers import TokenizerLike
|
||||
from vllm.tool_parsers.abstract_tool_parser import Tool, ToolParser
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
_O, _C, _S = r"<\|open\|>", r"<\|close\|>", r"<\|sep\|>"
|
||||
_TEXT_UNTIL_SEP = r"(?:(?!" + _S + r").)*?"
|
||||
|
||||
|
||||
def _partial_tag_overlap(text: str, tag: str) -> int:
|
||||
max_len = min(len(text), len(tag) - 1)
|
||||
for n in range(max_len, 0, -1):
|
||||
if text.endswith(tag[:n]):
|
||||
return n
|
||||
return 0
|
||||
|
||||
|
||||
class KimiK3ToolParser(ToolParser):
|
||||
supports_required_and_named = False
|
||||
# Enables the vLLM-side XTML structural tag builder
|
||||
# (``get_kimi_k3_structural_tag`` in ``structural_tag_registry``). With
|
||||
# ``VLLM_ENFORCE_STRICT_TOOL_CALLING`` on (default), ``_apply_structural_tag``
|
||||
# constrains generation to K3's ``<|open|>tools<|sep|>`` channel for
|
||||
# ``required`` (and ``auto`` when a tool sets ``strict``) instead of the
|
||||
# generic JSON guided decoding, which conflicts with the XTML format.
|
||||
structural_tag_model = "kimi_k3"
|
||||
|
||||
def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None):
|
||||
super().__init__(tokenizer, tools)
|
||||
|
||||
self.tools_open = "<|open|>tools<|sep|>"
|
||||
self.tools_close = "<|close|>tools<|sep|>"
|
||||
self.response_open = "<|open|>response<|sep|>"
|
||||
self.response_close = "<|close|>response<|sep|>"
|
||||
|
||||
# Regexes operate on detokenized text. The XTML markers reach us as the
|
||||
# literal strings <|open|>/<|close|>/<|sep|>. adjust_request keeps them
|
||||
# from being stripped and suppresses spacing between adjacent token
|
||||
# pieces. As DEFENSE-IN-DEPTH we still tolerate optional whitespace
|
||||
# WITHIN a marker (e.g. "<|open|> tools <|sep|>") in case some serving
|
||||
# path leaves vLLM's added-token spacing on -- this `\s*` is a no-op on
|
||||
# clean input, so the normal byte-exact path is unaffected. We do NOT
|
||||
# strip body whitespace, so clean content stays byte-exact.
|
||||
# All bodies use non-greedy (.*?) so each block stops at its own first
|
||||
# closing marker -- see the module-level KNOWN LIMITATION about
|
||||
# literal-marker values.
|
||||
self._tools_open_re = re.compile(_O + r"\s*tools\s*" + _S)
|
||||
self._tools_close_re = re.compile(_C + r"\s*tools\s*" + _S)
|
||||
self._response_open_re = re.compile(_O + r"\s*response\s*" + _S)
|
||||
self._response_close_re = re.compile(_C + r"\s*response\s*" + _S)
|
||||
self._message_close_re = re.compile(_C + r"\s*message\s*" + _S)
|
||||
self._call_re = re.compile(
|
||||
_O
|
||||
+ r"\s*call\s+(?P<attrs>"
|
||||
+ _TEXT_UNTIL_SEP
|
||||
+ r")"
|
||||
+ _S
|
||||
+ r"(?P<body>.*?)"
|
||||
+ _C
|
||||
+ r"\s*call\s*"
|
||||
+ _S,
|
||||
re.DOTALL,
|
||||
)
|
||||
self._arg_re = re.compile(
|
||||
_O
|
||||
+ r"\s*argument\s+(?P<attrs>"
|
||||
+ _TEXT_UNTIL_SEP
|
||||
+ r")"
|
||||
+ _S
|
||||
+ r"(?P<val>.*?)"
|
||||
+ _C
|
||||
+ r"\s*argument\s*"
|
||||
+ _S,
|
||||
re.DOTALL,
|
||||
)
|
||||
# attr segment: key="value" (value already escaped on the encode side)
|
||||
self._attr_re = re.compile(r'(?P<k>\w+)="(?P<v>[^"]*)"')
|
||||
self._response_re = re.compile(
|
||||
_O + r"\s*response\s*" + _S + r"(?P<c>.*?)" + _C + r"\s*response\s*" + _S,
|
||||
re.DOTALL,
|
||||
)
|
||||
|
||||
# streaming state
|
||||
self._sent_content_idx = 0
|
||||
self._sent_tool_call_count = 0
|
||||
|
||||
if not self.model_tokenizer:
|
||||
raise ValueError(
|
||||
"The model tokenizer must be passed to the ToolParser "
|
||||
"constructor during construction."
|
||||
)
|
||||
|
||||
def adjust_request(
|
||||
self, request: ChatCompletionRequest | ResponsesRequest
|
||||
) -> ChatCompletionRequest | ResponsesRequest:
|
||||
named = isinstance(
|
||||
request.tool_choice,
|
||||
(ChatCompletionNamedToolChoiceParam, ToolChoiceFunction),
|
||||
)
|
||||
structured_outputs = getattr(request, "structured_outputs", None)
|
||||
has_structural_tag = (
|
||||
structured_outputs is not None
|
||||
and structured_outputs.structural_tag is not None
|
||||
)
|
||||
if named and not has_structural_tag:
|
||||
# Without the XTML structural tag there is no way to force the
|
||||
# named call (the generic JSON guided-decoding path conflicts
|
||||
# with the XTML channel format).
|
||||
raise VLLMValidationError(
|
||||
"Named tool choice for Kimi K3 requires strict tool calling "
|
||||
"(VLLM_ENFORCE_STRICT_TOOL_CALLING) so the XTML structural "
|
||||
"tag can force the call. Otherwise use `tool_choice` set to "
|
||||
'"auto", "required", or "none".',
|
||||
parameter="tool_choice",
|
||||
value=request.tool_choice,
|
||||
)
|
||||
|
||||
if request.tools and (request.tool_choice == "required" or named):
|
||||
# K3 emits tool calls in XTML. When strict tool calling is enabled,
|
||||
# DelegatingParser._apply_structural_tag has already attached the K3
|
||||
# XTML structural tag (structural_tag_model = "kimi_k3"). We return
|
||||
# early to skip the generic parent path, which would otherwise attach
|
||||
# JSON guided decoding when strict calling is off -- that JSON
|
||||
# constraint conflicts with the <|open|>tools<|sep|> channel.
|
||||
request.skip_special_tokens = False
|
||||
if hasattr(request, "spaces_between_special_tokens"):
|
||||
request.spaces_between_special_tokens = False
|
||||
return request
|
||||
|
||||
request = super().adjust_request(request)
|
||||
# The XTML markers (<|open|>/<|close|>/<|sep|>,
|
||||
# <|open|>response<|sep|> ...) must
|
||||
# reach this parser as CONTIGUOUS literal text. Two request flags govern
|
||||
# that, and vLLM's detokenizer couples them:
|
||||
# effective_spaces_between =
|
||||
# skip_special_tokens OR spaces_between_special_tokens
|
||||
# (vllm/v1/engine/detokenizer.py).
|
||||
# The detokenizer treats these control tokens as separate sub-texts and,
|
||||
# when the effective flag is True, joins them with spaces ->
|
||||
# "<|open|> response <|sep|>", which the regexes below would NOT match.
|
||||
# Forcing BOTH flags off is the only way to suppress that spacing. We
|
||||
# set both unconditionally: the response channel is unwrapped from these
|
||||
# markers even with no tools.
|
||||
request.skip_special_tokens = False
|
||||
if hasattr(request, "spaces_between_special_tokens"):
|
||||
request.spaces_between_special_tokens = False
|
||||
return request
|
||||
|
||||
def _attrs(self, s: str) -> dict[str, str]:
|
||||
return {
|
||||
m["k"]: m["v"].replace(""", '"').replace("&", "&")
|
||||
for m in self._attr_re.finditer(s)
|
||||
}
|
||||
|
||||
def _decode_call(self, attrs: str, body: str) -> ToolCall | None:
|
||||
"""Decode one ``call`` block into a :class:`ToolCall`.
|
||||
|
||||
``attrs`` is the text between ``<|open|>call `` and ``<|sep|>`` (carries
|
||||
``tool=`` / ``index=``); ``body`` is the sequence of ``argument`` blocks.
|
||||
Each argument is re-typed per its ``type=`` tag: strings pass through
|
||||
raw, everything else is JSON-decoded (falling back to raw text if the
|
||||
JSON is malformed, so a partial stream never raises). Returns ``None``
|
||||
when no tool name is present (an empty/garbage block is dropped).
|
||||
"""
|
||||
call_attrs = self._attrs(attrs)
|
||||
tool_name = call_attrs.get("tool", "")
|
||||
tool_index = call_attrs.get("index", "")
|
||||
arguments: dict = {}
|
||||
for arg_match in self._arg_re.finditer(body):
|
||||
arg_attrs = self._attrs(arg_match["attrs"])
|
||||
key = arg_attrs.get("key", "")
|
||||
arg_type = arg_attrs.get("type", "string")
|
||||
raw_value = arg_match["val"]
|
||||
if arg_type == "string":
|
||||
arguments[key] = raw_value
|
||||
else:
|
||||
try:
|
||||
arguments[key] = json.loads(raw_value)
|
||||
except json.JSONDecodeError:
|
||||
arguments[key] = raw_value
|
||||
if not tool_name:
|
||||
return None
|
||||
tool_call_id = tool_name
|
||||
if tool_index:
|
||||
try:
|
||||
tool_call_id = f"{tool_name}:{int(tool_index) - 1}"
|
||||
except ValueError:
|
||||
tool_call_id = f"{tool_name}:{tool_index}"
|
||||
# id uses the API-side zero-based call ordinal; XTML's message index
|
||||
# stays one-based when rendering tool result messages.
|
||||
return ToolCall(
|
||||
id=tool_call_id,
|
||||
type="function",
|
||||
function=FunctionCall(
|
||||
name=tool_name,
|
||||
arguments=json.dumps(arguments, ensure_ascii=False),
|
||||
),
|
||||
)
|
||||
|
||||
def _strip_response_content(self, text: str) -> str | None:
|
||||
"""Strip XTML response/message markers from generated response text.
|
||||
|
||||
In chat serving, ``<|open|>response<|sep|>`` is often part of the prompt
|
||||
generation prefix, so the model output may only contain the body plus
|
||||
``<|close|>response<|sep|>``. Handle both that consumed-prefix shape and a
|
||||
complete ``<|open|>response<|sep|>... <|close|>response<|sep|>`` wrapper.
|
||||
"""
|
||||
m_open = self._response_open_re.search(text)
|
||||
if m_open is not None:
|
||||
m_close = self._response_close_re.search(text, m_open.end())
|
||||
if m_close is not None:
|
||||
text = text[m_open.end() : m_close.start()]
|
||||
else:
|
||||
text = text[m_open.end() :]
|
||||
else:
|
||||
text = self._response_close_re.sub("", text)
|
||||
text = self._message_close_re.sub("", text)
|
||||
return text or None
|
||||
|
||||
def _content(self, model_output: str, before: str) -> str | None:
|
||||
# prefer the unwrapped response channel; else the text before the tools
|
||||
m = self._response_re.search(model_output)
|
||||
if m is not None:
|
||||
return m["c"] or None
|
||||
return self._strip_response_content(before)
|
||||
|
||||
def _extract_response_content(self, current_text: str) -> str | None:
|
||||
# Streaming response text is computed from the accumulated text. This is
|
||||
# what keeps split markers from leaking:
|
||||
# <|open|> / response / <|sep|>Hi -> emit only "Hi" after open closes
|
||||
# Hi<|open|> / tools / <|sep|>... -> emit "Hi", hold the tools marker
|
||||
# Hi<|close|> / response / <|sep|>... -> emit "Hi", hold the close marker
|
||||
# Tool calls are even simpler: they are not emitted until a full
|
||||
# <|close|>call<|sep|> is present in current_text.
|
||||
m_open = self._response_open_re.search(current_text)
|
||||
# In the normal chat path, the response-open marker may be consumed as
|
||||
# the generation prefix. Then generated text starts directly with the
|
||||
# response body or with <|close|>response<|sep|> before a tools channel.
|
||||
body_start = m_open.end() if m_open is not None else 0
|
||||
m_tools = self._tools_open_re.search(current_text, body_start)
|
||||
m_rclose = self._response_close_re.search(current_text, body_start)
|
||||
tools_start = m_tools.start() if m_tools else -1
|
||||
response_end = m_rclose.start() if m_rclose else -1
|
||||
|
||||
candidates = [i for i in (tools_start, response_end) if i != -1]
|
||||
if candidates:
|
||||
sendable_idx = min(candidates)
|
||||
else:
|
||||
overlap = max(
|
||||
_partial_tag_overlap(current_text, self.response_open),
|
||||
_partial_tag_overlap(current_text, self.response_close),
|
||||
_partial_tag_overlap(current_text, self.tools_open),
|
||||
)
|
||||
sendable_idx = len(current_text) - overlap
|
||||
|
||||
if sendable_idx <= body_start:
|
||||
return None
|
||||
if self._sent_content_idx < body_start:
|
||||
self._sent_content_idx = body_start
|
||||
if sendable_idx <= self._sent_content_idx:
|
||||
return None
|
||||
|
||||
content = current_text[self._sent_content_idx : sendable_idx]
|
||||
self._sent_content_idx = sendable_idx
|
||||
return content or None
|
||||
|
||||
def extract_tool_calls(
|
||||
self, model_output: str, request: ChatCompletionRequest
|
||||
) -> ExtractedToolCallInformation:
|
||||
m_open = self._tools_open_re.search(model_output)
|
||||
if m_open is None:
|
||||
# no tools channel -> content is the response channel (unwrapped)
|
||||
return ExtractedToolCallInformation(
|
||||
tools_called=False,
|
||||
tool_calls=[],
|
||||
content=self._content(model_output, model_output),
|
||||
)
|
||||
try:
|
||||
before = model_output[: m_open.start()]
|
||||
start = m_open.end()
|
||||
m_close = self._tools_close_re.search(model_output, start)
|
||||
section = (
|
||||
model_output[start:]
|
||||
if m_close is None
|
||||
else model_output[start : m_close.start()]
|
||||
)
|
||||
|
||||
tool_calls = [
|
||||
tc
|
||||
for m in self._call_re.finditer(section)
|
||||
if (tc := self._decode_call(m["attrs"], m["body"])) is not None
|
||||
]
|
||||
if not tool_calls:
|
||||
return ExtractedToolCallInformation(
|
||||
tools_called=False,
|
||||
tool_calls=[],
|
||||
content=self._content(model_output, before),
|
||||
)
|
||||
return ExtractedToolCallInformation(
|
||||
tools_called=True,
|
||||
tool_calls=tool_calls,
|
||||
content=self._content(model_output, before),
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Error extracting K3 tool calls.")
|
||||
return ExtractedToolCallInformation(
|
||||
tools_called=False, tool_calls=[], content=model_output
|
||||
)
|
||||
|
||||
def extract_tool_calls_streaming(
|
||||
self,
|
||||
previous_text: str,
|
||||
current_text: str,
|
||||
delta_text: str,
|
||||
previous_token_ids: Sequence[int],
|
||||
current_token_ids: Sequence[int],
|
||||
delta_token_ids: Sequence[int],
|
||||
request: ChatCompletionRequest,
|
||||
) -> DeltaMessage | None:
|
||||
# Conservative streaming: stream unwrapped response-channel text, then
|
||||
# buffer tool calls and emit each call once its block closes.
|
||||
content = self._extract_response_content(current_text)
|
||||
|
||||
# tools channel is open: parse fully-closed calls we have not emitted yet
|
||||
m_tools = self._tools_open_re.search(current_text)
|
||||
if m_tools is None:
|
||||
return DeltaMessage(content=content) if content else None
|
||||
|
||||
section = current_text[m_tools.end() :]
|
||||
calls = [
|
||||
tc
|
||||
for m in self._call_re.finditer(section)
|
||||
if (tc := self._decode_call(m["attrs"], m["body"])) is not None
|
||||
]
|
||||
if len(calls) <= self._sent_tool_call_count:
|
||||
return DeltaMessage(content=content) if content else None
|
||||
new = calls[self._sent_tool_call_count :]
|
||||
|
||||
deltas = [
|
||||
DeltaToolCall(
|
||||
index=self._sent_tool_call_count + i,
|
||||
id=tc.id,
|
||||
type="function",
|
||||
function=DeltaFunctionCall(
|
||||
name=tc.function.name, arguments=tc.function.arguments
|
||||
).model_dump(exclude_none=True),
|
||||
)
|
||||
for i, tc in enumerate(new)
|
||||
]
|
||||
self._sent_tool_call_count = len(calls)
|
||||
return DeltaMessage(content=content, tool_calls=deltas)
|
||||
@@ -19,7 +19,12 @@ from xgrammar.structural_tag import (
|
||||
AnyTextFormat,
|
||||
ConstStringFormat,
|
||||
JSONSchemaFormat,
|
||||
OptionalFormat,
|
||||
OrFormat,
|
||||
PlusFormat,
|
||||
RegexFormat,
|
||||
SequenceFormat,
|
||||
StarFormat,
|
||||
TagFormat,
|
||||
TagsWithSeparatorFormat,
|
||||
TriggeredTagsFormat,
|
||||
@@ -66,7 +71,7 @@ XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS = frozenset(
|
||||
"deepseek_v4",
|
||||
}
|
||||
)
|
||||
VLLM_BUILTIN_STRUCTURAL_TAG_MODELS = frozenset({"hermes"})
|
||||
VLLM_BUILTIN_STRUCTURAL_TAG_MODELS = frozenset({"hermes", "kimi_k3"})
|
||||
SUPPORTED_STRUCTURAL_TAG_MODELS = (
|
||||
XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS | VLLM_BUILTIN_STRUCTURAL_TAG_MODELS
|
||||
)
|
||||
@@ -345,3 +350,259 @@ def get_minimax_structural_tag(
|
||||
)
|
||||
|
||||
return StructuralTag(format=suffix_tag)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Kimi K3 (XTML channel format)
|
||||
# ---------------------------------------------------------------------------
|
||||
# K3 assistant output after the reasoning gate (``<|close|>think<|sep|>``):
|
||||
# <|open|>response<|sep|> <content> <|close|>response<|sep|>
|
||||
# [ <|open|>tools<|sep|>
|
||||
# <|open|>call tool="NAME" index="1"<|sep|>
|
||||
# <|open|>argument key="K" type="TYPE"<|sep|>VALUE<|close|>argument<|sep|>
|
||||
# <|close|>call<|sep|> ...
|
||||
# <|close|>tools<|sep|> ]
|
||||
# See ``encoding_k3.py`` (_render_assistant_segments / _open_tag / _attr) and
|
||||
# ``kimi_k3_tool_parser.py`` for the exact byte-level encoding this mirrors.
|
||||
_K3_OPEN = "<|open|>"
|
||||
_K3_CLOSE = "<|close|>"
|
||||
_K3_SEP = "<|sep|>"
|
||||
_K3_RESPONSE_OPEN = f"{_K3_OPEN}response{_K3_SEP}"
|
||||
_K3_RESPONSE_CLOSE = f"{_K3_CLOSE}response{_K3_SEP}"
|
||||
_K3_TOOLS_OPEN = f"{_K3_OPEN}tools{_K3_SEP}"
|
||||
_K3_TOOLS_CLOSE = f"{_K3_CLOSE}tools{_K3_SEP}"
|
||||
_K3_CALL_CLOSE = f"{_K3_CLOSE}call{_K3_SEP}"
|
||||
_K3_ARG_CLOSE = f"{_K3_CLOSE}argument{_K3_SEP}"
|
||||
# The model closes the assistant turn with <|close|>message<|sep|> right before
|
||||
# the end-of-message token. It is generated (not part of the prompt prefix), so
|
||||
# the tag must permit it or the FSM would mask the model's natural terminator.
|
||||
_K3_MESSAGE_CLOSE = f"{_K3_CLOSE}message{_K3_SEP}"
|
||||
|
||||
# JSON-schema type -> K3 XTML ``type=`` attribute value. Mirrors
|
||||
# ``encoding_k3._xtml_type`` (integer collapses onto number).
|
||||
_K3_JSON_TO_XTML_TYPE = {
|
||||
"string": "string",
|
||||
"integer": "number",
|
||||
"number": "number",
|
||||
"boolean": "boolean",
|
||||
"null": "null",
|
||||
"object": "object",
|
||||
"array": "array",
|
||||
}
|
||||
|
||||
|
||||
def _k3_escape_attr(value: str) -> str:
|
||||
"""Mirror ``encoding_k3._escape_attr_value`` (``&`` then ``"``)."""
|
||||
return str(value).replace("&", "&").replace('"', """)
|
||||
|
||||
|
||||
_K3_STRING_ATOM = r"(?:[^<]|<[^|])"
|
||||
"""One raw-string character: anything but the ambiguous "<|" marker prefix.
|
||||
Allows '<' inside values (e.g. HTML snippets); a value *ending* in '<' or
|
||||
containing a literal "<|" is not expressible and falls back to AnyText via
|
||||
the pattern checks below never matching those cases at build time (schemas
|
||||
cannot know values, so the only build-time effect is the length bound)."""
|
||||
|
||||
|
||||
def _k3_bounded_string_regex(prop: dict[str, Any]) -> str | None:
|
||||
"""Length/pattern constraint for the raw string channel, if expressible.
|
||||
|
||||
The XTML string channel emits values raw (not JSON-quoted), so
|
||||
JSONSchemaFormat cannot enforce string constraints there; unconstrained
|
||||
AnyText lets maxLength/pattern violations through (observed on the walle
|
||||
verifier: over-long junk strings pass the grammar and fail validation).
|
||||
xgrammar's regex engine has no lookahead, so the close marker is kept
|
||||
unambiguous by excluding the "<|" prefix from value characters.
|
||||
|
||||
Returns a regex for the value, or None to keep permissive AnyText.
|
||||
"""
|
||||
max_len = prop.get("maxLength")
|
||||
min_len = prop.get("minLength", 0)
|
||||
if not isinstance(max_len, int) or max_len < 0 or max_len > 4096:
|
||||
return None
|
||||
if not isinstance(min_len, int) or min_len < 0 or min_len > max_len:
|
||||
min_len = 0
|
||||
return _K3_STRING_ATOM + f"{{{min_len},{max_len}}}"
|
||||
|
||||
|
||||
def _k3_argument_tag(
|
||||
key: str,
|
||||
schema: dict[str, Any],
|
||||
root_defs: dict[str, Any] | None = None,
|
||||
) -> TagFormat | None:
|
||||
"""Build one ``argument`` XTML tag for property ``key``.
|
||||
|
||||
``string`` values are emitted raw (bounded by the close marker); every other
|
||||
JSON type is emitted as JSON and validated against the property schema. A
|
||||
property whose type is a union / missing is left permissive (any XTML type,
|
||||
raw value) so a valid call is never rejected.
|
||||
|
||||
``root_defs`` carries the tool parameters' root-level ``$defs`` /
|
||||
``definitions``: slicing a property out of the parameters document orphans
|
||||
its ``#/$defs/...`` references, so those tables must be re-attached to keep
|
||||
the embedded schema self-contained.
|
||||
"""
|
||||
prop = schema if isinstance(schema, dict) else {}
|
||||
json_type = prop.get("type")
|
||||
xtml_type = (
|
||||
_K3_JSON_TO_XTML_TYPE.get(json_type) if isinstance(json_type, str) else None
|
||||
)
|
||||
if xtml_type is None:
|
||||
# Unknown / union type: constrain the key but keep the value permissive.
|
||||
return None
|
||||
begin = (
|
||||
f'{_K3_OPEN}argument key="{_k3_escape_attr(key)}" type="{xtml_type}"{_K3_SEP}'
|
||||
)
|
||||
if xtml_type == "string":
|
||||
# Raw string channel: JSONSchemaFormat can't apply (values are not
|
||||
# JSON-quoted), but an enum/const of strings is a finite set that can
|
||||
# be enforced exactly with const-string alternation. Enum semantics
|
||||
# are exclusive, so this never over-rejects. Fall back to permissive
|
||||
# AnyText for open-ended strings or non-representable enums.
|
||||
enum_values = prop.get("enum")
|
||||
if enum_values is None and isinstance(prop.get("const"), str):
|
||||
enum_values = [prop["const"]]
|
||||
if (
|
||||
isinstance(enum_values, list)
|
||||
and enum_values
|
||||
and len(enum_values) <= 256
|
||||
and all(isinstance(v, str) for v in enum_values)
|
||||
and not any("<|" in v for v in enum_values)
|
||||
):
|
||||
branches = [ConstStringFormat(value=v) for v in enum_values]
|
||||
content: Any = (
|
||||
branches[0] if len(branches) == 1 else OrFormat(elements=branches)
|
||||
)
|
||||
elif (bounded := _k3_bounded_string_regex(prop)) is not None:
|
||||
content = RegexFormat(pattern=bounded)
|
||||
else:
|
||||
content = AnyTextFormat(excludes=[_K3_CLOSE])
|
||||
else:
|
||||
embedded = prop
|
||||
if root_defs:
|
||||
embedded = dict(prop)
|
||||
for defs_key, defs_value in root_defs.items():
|
||||
embedded.setdefault(defs_key, defs_value)
|
||||
content = JSONSchemaFormat(json_schema=embedded)
|
||||
return TagFormat(begin=begin, content=content, end=_K3_ARG_CLOSE)
|
||||
|
||||
|
||||
def _k3_permissive_argument_tag() -> TagFormat:
|
||||
"""A key/type-agnostic ``argument`` tag: any attributes, raw value.
|
||||
|
||||
Used as a fallback so tools with union/loose schemas still get the XTML
|
||||
skeleton constrained without over-rejecting the value.
|
||||
"""
|
||||
return TagFormat(
|
||||
begin=_K3_OPEN + "argument ",
|
||||
content=SequenceFormat(
|
||||
elements=[
|
||||
RegexFormat(pattern=r"[^<]*" + _K3_SEP.replace("|", r"\|")),
|
||||
AnyTextFormat(excludes=[_K3_CLOSE]),
|
||||
]
|
||||
),
|
||||
end=_K3_ARG_CLOSE,
|
||||
)
|
||||
|
||||
|
||||
def _k3_arguments_block(parameters: dict[str, Any] | bool) -> Any:
|
||||
"""Build ``argument`` tags for a tool's parameter schema.
|
||||
|
||||
Require at least one tag when the root schema declares required properties.
|
||||
Otherwise, keep accepting zero-or-more tags. Arguments remain order-agnostic
|
||||
and non-unique.
|
||||
"""
|
||||
if not isinstance(parameters, dict):
|
||||
return StarFormat(content=_k3_permissive_argument_tag())
|
||||
props = parameters.get("properties")
|
||||
if not isinstance(props, dict) or not props:
|
||||
# No declared properties: allow any argument blocks (or none).
|
||||
return StarFormat(content=_k3_permissive_argument_tag())
|
||||
root_defs = {
|
||||
defs_key: parameters[defs_key]
|
||||
for defs_key in ("$defs", "definitions")
|
||||
if isinstance(parameters.get(defs_key), dict)
|
||||
}
|
||||
tags: list[TagFormat] = []
|
||||
for key, prop in props.items():
|
||||
tag = _k3_argument_tag(key, prop, root_defs)
|
||||
tags.append(tag if tag is not None else _k3_permissive_argument_tag())
|
||||
inner = tags[0] if len(tags) == 1 else OrFormat(elements=list(tags))
|
||||
required = parameters.get("required")
|
||||
if isinstance(required, list) and required:
|
||||
return PlusFormat(content=inner)
|
||||
return StarFormat(content=inner)
|
||||
|
||||
|
||||
def _k3_call_tag(tool: FunctionToolParam) -> TagFormat:
|
||||
"""One ``call`` tag: ``<|open|>call tool="N" index="<digits>"<|sep|> args``."""
|
||||
function = tool.function
|
||||
parameters = _get_function_parameters(function)
|
||||
begin = f'{_K3_OPEN}call tool="{_k3_escape_attr(function.name)}" index="'
|
||||
return TagFormat(
|
||||
begin=begin,
|
||||
content=SequenceFormat(
|
||||
elements=[
|
||||
RegexFormat(pattern=r"[0-9]+"),
|
||||
ConstStringFormat(value=f'"{_K3_SEP}'),
|
||||
_k3_arguments_block(parameters),
|
||||
]
|
||||
),
|
||||
end=_K3_CALL_CLOSE,
|
||||
)
|
||||
|
||||
|
||||
def _k3_response_prefix() -> list[Any]:
|
||||
"""The response channel that always precedes the tools channel.
|
||||
|
||||
``response`` is generated in thinking mode (prefix ends at
|
||||
``<|open|>think<|sep|>``) but is part of the generation prefix in
|
||||
non-thinking mode, so its open marker is optional. The body is bounded by
|
||||
the response close marker.
|
||||
"""
|
||||
return [
|
||||
OptionalFormat(content=ConstStringFormat(value=_K3_RESPONSE_OPEN)),
|
||||
TagFormat(begin="", content=AnyTextFormat(), end=_K3_RESPONSE_CLOSE),
|
||||
]
|
||||
|
||||
|
||||
def _k3_tools_channel(tools: list[FunctionToolParam]) -> TagFormat:
|
||||
return TagFormat(
|
||||
begin=_K3_TOOLS_OPEN,
|
||||
content=TagsWithSeparatorFormat(
|
||||
tags=[_k3_call_tag(tool) for tool in tools],
|
||||
separator="",
|
||||
at_least_one=True,
|
||||
),
|
||||
end=_K3_TOOLS_CLOSE,
|
||||
)
|
||||
|
||||
|
||||
@register_vllm_structural_tag("kimi_k3")
|
||||
def get_kimi_k3_structural_tag(
|
||||
tools: list[FunctionToolParam],
|
||||
builtin_tools: list[BuiltinToolParam],
|
||||
tool_choice: SimplifiedToolChoice,
|
||||
reasoning: bool,
|
||||
) -> StructuralTag:
|
||||
del builtin_tools, reasoning
|
||||
|
||||
trailer = OptionalFormat(content=ConstStringFormat(value=_K3_MESSAGE_CLOSE))
|
||||
|
||||
if not tools:
|
||||
return StructuralTag(
|
||||
format=SequenceFormat(elements=[*_k3_response_prefix(), trailer])
|
||||
)
|
||||
|
||||
if tool_choice == "auto":
|
||||
tools_part: Any = OptionalFormat(content=_k3_tools_channel(tools))
|
||||
elif tool_choice == "forced":
|
||||
# K3 rejects named tool choice upstream; treat defensively as a single
|
||||
# mandatory call of the first tool.
|
||||
tools_part = _k3_tools_channel(tools[:1])
|
||||
else: # required
|
||||
tools_part = _k3_tools_channel(tools)
|
||||
|
||||
return StructuralTag(
|
||||
format=SequenceFormat(elements=[*_k3_response_prefix(), tools_part, trailer])
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user