[Kimi K3] Add reasoning, tool-call, and OpenAI serving support (#33025)
Co-authored-by: hnyls2002 <lsyincs@gmail.com> Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com> Co-authored-by: A-transformer <cl5743590921@gmail.com>
This commit is contained in:
co-authored by
hnyls2002
Liangsheng Yin
A-transformer
parent
f1b41a5b3d
commit
e2cf21b9e5
@@ -0,0 +1,221 @@
|
||||
import json
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
from sglang.srt.entrypoints.openai.protocol import Function, Tool
|
||||
from sglang.srt.function_call.core_types import ToolCallItem
|
||||
from sglang.srt.function_call.function_call_parser import FunctionCallParser
|
||||
from sglang.srt.function_call.kimik3_detector import KimiK3Detector
|
||||
from sglang.srt.function_call.kimik3_format import (
|
||||
MESSAGE_CLOSE,
|
||||
RESPONSE_CLOSE,
|
||||
RESPONSE_OPEN,
|
||||
TOOLS_CLOSE,
|
||||
TOOLS_OPEN,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
def _make_tool(name: str) -> Tool:
|
||||
return Tool(
|
||||
type="function",
|
||||
function=Function(
|
||||
name=name,
|
||||
description=f"{name} tool",
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {"code": {"type": "string"}},
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _call_block(tool: str, index: int, args: dict[str, tuple[str, str]]) -> str:
|
||||
parts = [f'<|open|>call tool="{tool}" index="{index}"<|sep|>']
|
||||
for key, (arg_type, value) in args.items():
|
||||
parts.append(
|
||||
f'<|open|>argument key="{key}" type="{arg_type}"<|sep|>'
|
||||
f"{value}<|close|>argument<|sep|>"
|
||||
)
|
||||
parts.append("<|close|>call<|sep|>")
|
||||
return "".join(parts)
|
||||
|
||||
|
||||
def _chunks(text: str, size: int) -> list[str]:
|
||||
return [text[index : index + size] for index in range(0, len(text), size)]
|
||||
|
||||
|
||||
def _stream(
|
||||
detector: KimiK3Detector, chunks: list[str], tools: list[Tool]
|
||||
) -> tuple[str, list[ToolCallItem]]:
|
||||
text = ""
|
||||
calls = []
|
||||
for chunk in chunks:
|
||||
result = detector.parse_streaming_increment(chunk, tools)
|
||||
text += result.normal_text
|
||||
calls.extend(result.calls)
|
||||
return text, calls
|
||||
|
||||
|
||||
def test_detect_and_parse_single_call() -> None:
|
||||
detector = KimiK3Detector()
|
||||
tools = [_make_tool("python")]
|
||||
text = (
|
||||
f"{RESPONSE_OPEN}Let me run it.{RESPONSE_CLOSE}{TOOLS_OPEN}"
|
||||
+ _call_block(
|
||||
"python",
|
||||
1,
|
||||
{"code": ("string", "print(1)"), "opts": ("object", '{"a": 1}')},
|
||||
)
|
||||
+ TOOLS_CLOSE
|
||||
)
|
||||
result = detector.detect_and_parse(text, tools)
|
||||
assert result.normal_text == "Let me run it."
|
||||
assert len(result.calls) == 1
|
||||
assert result.calls[0].name == "python"
|
||||
assert json.loads(result.calls[0].parameters) == {
|
||||
"code": "print(1)",
|
||||
"opts": {"a": 1},
|
||||
}
|
||||
|
||||
|
||||
def test_detect_and_parse_no_tools_channel() -> None:
|
||||
detector = KimiK3Detector()
|
||||
result = detector.detect_and_parse(
|
||||
f"{RESPONSE_OPEN}hi there{RESPONSE_CLOSE}{MESSAGE_CLOSE}",
|
||||
[_make_tool("python")],
|
||||
)
|
||||
assert result.normal_text == "hi there"
|
||||
assert result.calls == []
|
||||
|
||||
|
||||
def test_detect_and_parse_multiple_calls() -> None:
|
||||
detector = KimiK3Detector()
|
||||
text = (
|
||||
TOOLS_OPEN
|
||||
+ _call_block("python", 1, {"code": ("string", "a")})
|
||||
+ _call_block("python", 2, {"code": ("string", "b")})
|
||||
+ TOOLS_CLOSE
|
||||
)
|
||||
result = detector.detect_and_parse(text, [_make_tool("python")])
|
||||
assert [call.tool_index for call in result.calls] == [0, 1]
|
||||
assert json.loads(result.calls[1].parameters) == {"code": "b"}
|
||||
|
||||
|
||||
def test_detect_and_parse_unclosed_tools_section() -> None:
|
||||
detector = KimiK3Detector()
|
||||
text = TOOLS_OPEN + _call_block("python", 1, {"code": ("string", "x")})
|
||||
result = detector.detect_and_parse(text, [_make_tool("python")])
|
||||
assert len(result.calls) == 1
|
||||
assert json.loads(result.calls[0].parameters) == {"code": "x"}
|
||||
|
||||
|
||||
def test_attr_unescaping_and_raw_string_args() -> None:
|
||||
detector = KimiK3Detector()
|
||||
text = (
|
||||
f"{TOOLS_OPEN}"
|
||||
'<|open|>call tool="a&b" index="1"<|sep|>'
|
||||
'<|open|>argument key="q" type="string"<|sep|>'
|
||||
"say "hi"<|close|>argument<|sep|>"
|
||||
"<|close|>call<|sep|>"
|
||||
f"{TOOLS_CLOSE}"
|
||||
)
|
||||
result = detector.detect_and_parse(text, [_make_tool("python")])
|
||||
assert result.calls[0].name == "a&b"
|
||||
assert json.loads(result.calls[0].parameters) == {"q": "say "hi""}
|
||||
|
||||
|
||||
def test_non_string_arg_json_decoding() -> None:
|
||||
detector = KimiK3Detector()
|
||||
text = (
|
||||
TOOLS_OPEN
|
||||
+ _call_block(
|
||||
"python",
|
||||
1,
|
||||
{
|
||||
"n": ("number", "42"),
|
||||
"flag": ("boolean", "true"),
|
||||
"bad": ("object", "{not json"),
|
||||
},
|
||||
)
|
||||
+ TOOLS_CLOSE
|
||||
)
|
||||
result = detector.detect_and_parse(text, [_make_tool("python")])
|
||||
assert json.loads(result.calls[0].parameters) == {
|
||||
"n": 42,
|
||||
"flag": True,
|
||||
"bad": "{not json",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("chunk_size", [1, 7, 23])
|
||||
def test_streaming_split_markers(chunk_size: int) -> None:
|
||||
detector = KimiK3Detector()
|
||||
tools = [_make_tool("python")]
|
||||
text = (
|
||||
f"{RESPONSE_OPEN}Hello!{RESPONSE_CLOSE}{TOOLS_OPEN}"
|
||||
+ _call_block("python", 1, {"code": ("string", "print(2)")})
|
||||
+ TOOLS_CLOSE
|
||||
)
|
||||
normal_text, calls = _stream(detector, _chunks(text, chunk_size), tools)
|
||||
assert normal_text == "Hello!"
|
||||
assert len(calls) == 1
|
||||
assert calls[0].name == "python"
|
||||
assert json.loads(calls[0].parameters) == {"code": "print(2)"}
|
||||
|
||||
|
||||
def test_streaming_two_calls() -> None:
|
||||
detector = KimiK3Detector()
|
||||
tools = [_make_tool("python")]
|
||||
text = (
|
||||
TOOLS_OPEN
|
||||
+ _call_block("python", 1, {"code": ("string", "a")})
|
||||
+ _call_block("python", 2, {"code": ("string", "b")})
|
||||
+ TOOLS_CLOSE
|
||||
)
|
||||
_, calls = _stream(detector, _chunks(text, 7), tools)
|
||||
assert [call.tool_index for call in calls] == [0, 1]
|
||||
assert [json.loads(call.parameters) for call in calls] == [
|
||||
{"code": "a"},
|
||||
{"code": "b"},
|
||||
]
|
||||
|
||||
|
||||
def test_streaming_plain_text_only() -> None:
|
||||
detector = KimiK3Detector()
|
||||
text, calls = _stream(
|
||||
detector, ["just a ", "plain ", "reply"], [_make_tool("python")]
|
||||
)
|
||||
assert text == "just a plain reply"
|
||||
assert calls == []
|
||||
|
||||
|
||||
def test_streaming_bookkeeping_for_serving_layer() -> None:
|
||||
detector = KimiK3Detector()
|
||||
tools = [_make_tool("python")]
|
||||
text = (
|
||||
TOOLS_OPEN + _call_block("python", 1, {"code": ("string", "a")}) + TOOLS_CLOSE
|
||||
)
|
||||
_stream(detector, _chunks(text, 9), tools)
|
||||
assert detector.current_tool_id == 0
|
||||
assert detector.prev_tool_call_arr[0] == {
|
||||
"name": "python",
|
||||
"arguments": {"code": "a"},
|
||||
}
|
||||
assert json.loads(detector.streamed_args_for_tool[0]) == {"code": "a"}
|
||||
|
||||
|
||||
def test_detector_capabilities_and_registration() -> None:
|
||||
detector = KimiK3Detector()
|
||||
assert detector.supports_structural_tag()
|
||||
assert not detector.parses_required_natively()
|
||||
parser = FunctionCallParser([_make_tool("python")], "kimi_k3")
|
||||
assert isinstance(parser.detector, KimiK3Detector)
|
||||
assert parser.get_structure_constraint("required") is not None
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
@@ -359,30 +359,28 @@ class TestCreateGrammarBackend(unittest.TestCase):
|
||||
# encode must return a single-token list for think_start/end tokens
|
||||
tokenizer.encode.return_value = [42]
|
||||
|
||||
result = create_grammar_backend(args, tokenizer, 32000, think_end_id=42)
|
||||
result = create_grammar_backend(args, tokenizer, 32000, think_end_ids=[42])
|
||||
self.assertIsInstance(result, ReasonerGrammarBackend)
|
||||
self.assertIs(result.grammar_backend, mock_backend)
|
||||
|
||||
@patch("sglang.srt.constrained.outlines_backend.OutlinesGrammarBackend")
|
||||
def test_no_reasoner_wrapping_without_think_end_id(self, mock_outlines_cls):
|
||||
"""Without think_end_id passed in, no reasoner wrapping."""
|
||||
def test_no_reasoner_wrapping_without_think_end_ids(self, mock_outlines_cls):
|
||||
mock_backend = MagicMock(spec=BaseGrammarBackend)
|
||||
mock_outlines_cls.return_value = mock_backend
|
||||
args = self._make_server_args("outlines", reasoning_parser="deepseek-r1")
|
||||
tokenizer = MagicMock(spec=[]) # No think_end_id attribute
|
||||
tokenizer = MagicMock(spec=[])
|
||||
|
||||
result = create_grammar_backend(args, tokenizer, 32000, think_end_id=None)
|
||||
result = create_grammar_backend(args, tokenizer, 32000, think_end_ids=None)
|
||||
self.assertIs(result, mock_backend)
|
||||
|
||||
@patch("sglang.srt.constrained.outlines_backend.OutlinesGrammarBackend")
|
||||
def test_no_reasoner_wrapping_without_reasoning_parser(self, mock_outlines_cls):
|
||||
"""Without reasoning_parser, no reasoner wrapping even with think_end_id."""
|
||||
mock_backend = MagicMock(spec=BaseGrammarBackend)
|
||||
mock_outlines_cls.return_value = mock_backend
|
||||
args = self._make_server_args("outlines", reasoning_parser=None)
|
||||
tokenizer = MagicMock()
|
||||
|
||||
result = create_grammar_backend(args, tokenizer, 32000, think_end_id=42)
|
||||
result = create_grammar_backend(args, tokenizer, 32000, think_end_ids=[42])
|
||||
self.assertIs(result, mock_backend)
|
||||
|
||||
@patch("sglang.srt.constrained.xgrammar_backend.XGrammarGrammarBackend")
|
||||
|
||||
@@ -273,7 +273,7 @@ class TestProcessReqWithGrammar(unittest.TestCase):
|
||||
def test_cache_hit_applies_request_thinking_budget(self):
|
||||
mgr = self._make_mgr()
|
||||
grammar_obj = ReasonerGrammarObject(
|
||||
grammar=None, think_end_id=0, max_think_tokens=99
|
||||
grammar=None, think_end_ids=[0], max_think_tokens=99
|
||||
)
|
||||
mgr.grammar_backend.get_cached_or_future_value.return_value = (
|
||||
grammar_obj,
|
||||
@@ -292,7 +292,7 @@ class TestProcessReqWithGrammar(unittest.TestCase):
|
||||
mgr = self._make_mgr()
|
||||
mgr._enable_strict_thinking = True
|
||||
grammar_obj = ReasonerGrammarObject(
|
||||
grammar=None, think_end_id=0, max_think_tokens=99
|
||||
grammar=None, think_end_ids=[0], max_think_tokens=99
|
||||
)
|
||||
mgr.grammar_backend.init_strict_reasoning_grammar.return_value = grammar_obj
|
||||
|
||||
@@ -545,7 +545,7 @@ class TestGetReadyGrammarRequests(unittest.TestCase):
|
||||
mgr = self._make_mgr()
|
||||
|
||||
grammar_obj = ReasonerGrammarObject(
|
||||
grammar=None, think_end_id=0, max_think_tokens=99
|
||||
grammar=None, think_end_ids=[0], max_think_tokens=99
|
||||
)
|
||||
future = Future()
|
||||
future.set_result(grammar_obj)
|
||||
|
||||
@@ -13,6 +13,8 @@ from sglang.srt.constrained.reasoner_grammar_backend import (
|
||||
from sglang.srt.constrained.torch_ops.token_filter_torch_ops import (
|
||||
set_token_filter_torch,
|
||||
)
|
||||
from sglang.srt.function_call.kimik3_format import THINK_CLOSE
|
||||
from sglang.srt.parser.reasoning_parser import KimiK3Detector as KimiK3ReasoningDetector
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(2.0, "base-a-test-cpu")
|
||||
@@ -75,7 +77,7 @@ class TestReasonerGrammarObject(unittest.TestCase):
|
||||
def _make_strict_object(self):
|
||||
return ReasonerGrammarObject(
|
||||
grammar=None,
|
||||
think_end_id=7,
|
||||
think_end_ids=[7],
|
||||
think_excluded_token_ids=[3, 5],
|
||||
max_think_tokens=2,
|
||||
enable_token_filter=True,
|
||||
@@ -117,6 +119,26 @@ class TestReasonerGrammarObject(unittest.TestCase):
|
||||
self.assertIs(obj.move_vocab_mask(mask, "cpu"), mask)
|
||||
self.assertIsNotNone(obj.apply_vocab_mask)
|
||||
|
||||
def test_budget_exhaustion_walks_multi_token_end(self):
|
||||
obj = ReasonerGrammarObject(
|
||||
grammar=None,
|
||||
think_end_ids=[7, 8],
|
||||
max_think_tokens=1,
|
||||
enable_token_filter=True,
|
||||
token_filter_fn=set_token_filter_torch,
|
||||
)
|
||||
obj.maybe_init_reasoning(True)
|
||||
obj.accept_token(10)
|
||||
|
||||
first_mask = torch.zeros((1, 2), dtype=torch.int32)
|
||||
obj.fill_vocab_mask(first_mask, 0)
|
||||
self.assertEqual(_allowed_token_ids(first_mask, [7, 8, 10]), [7])
|
||||
|
||||
obj.accept_token(7)
|
||||
second_mask = torch.zeros((1, 2), dtype=torch.int32)
|
||||
obj.fill_vocab_mask(second_mask, 0)
|
||||
self.assertEqual(_allowed_token_ids(second_mask, [7, 8, 10]), [8])
|
||||
|
||||
|
||||
class TestReasonerGrammarBackend(unittest.TestCase):
|
||||
def setUp(self):
|
||||
@@ -163,6 +185,42 @@ class TestReasonerGrammarBackend(unittest.TestCase):
|
||||
self.assertEqual(obj.max_think_tokens, 2)
|
||||
self.assertEqual(obj.think_excluded_token_ids, [3, 4])
|
||||
|
||||
def test_kimi_k3_excluded_tokens_spare_the_xtml_control_tokens(self):
|
||||
"""Kimi K3 bans bare channel names, never the marker-composing tokens.
|
||||
|
||||
The excluded list is flattened into single token ids, so listing a whole
|
||||
marker such as "<|open|>response<|sep|>" would ban <|open|> and <|sep|>
|
||||
individually -- which also blocks the think-end sequence and the jump
|
||||
into the tools channel, leaving the model unable to stop thinking.
|
||||
"""
|
||||
control_ids = {"<|open|>": [1], "<|close|>": [2], "<|sep|>": [3]}
|
||||
think_end_ids = [2, 4, 3]
|
||||
tokenizer = _DummyTokenizer(
|
||||
{
|
||||
THINK_CLOSE: think_end_ids,
|
||||
"response": [10],
|
||||
"message": [11],
|
||||
"<|end_of_msg|>": [12],
|
||||
"[EOS]": [13],
|
||||
"[EOT]": [14],
|
||||
**control_ids,
|
||||
}
|
||||
)
|
||||
reasoner = ReasonerGrammarBackend(
|
||||
_DummyGrammarBackend(support_token_filter=True),
|
||||
SimpleNamespace(detector=KimiK3ReasoningDetector()),
|
||||
tokenizer,
|
||||
enable_strict_thinking=True,
|
||||
)
|
||||
|
||||
excluded = reasoner.think_excluded_token_ids
|
||||
|
||||
self.assertEqual(excluded, [10, 11, 12, 13, 14])
|
||||
for token, ids in control_ids.items():
|
||||
for token_id in ids:
|
||||
self.assertNotIn(token_id, excluded, f"{token} must stay generatable")
|
||||
self.assertEqual(set(think_end_ids) & set(excluded), set())
|
||||
|
||||
def test_init_strict_reasoning_grammar_none_when_strict_disabled(self):
|
||||
backend = _DummyGrammarBackend(support_token_filter=True)
|
||||
reasoner = ReasonerGrammarBackend(
|
||||
@@ -205,16 +263,15 @@ class TestReasonerGrammarBackend(unittest.TestCase):
|
||||
)
|
||||
self.assertIsNotNone(reasoner)
|
||||
|
||||
def test_rejects_multi_token_think_end_marker(self):
|
||||
def test_accepts_multi_token_think_end_marker(self):
|
||||
backend = _DummyGrammarBackend(support_token_filter=True)
|
||||
|
||||
with self.assertRaisesRegex(ValueError, "must encode to exactly one token"):
|
||||
ReasonerGrammarBackend(
|
||||
backend,
|
||||
self._make_parser(),
|
||||
self._make_tokenizer(end_ids=[2, 3]),
|
||||
enable_strict_thinking=True,
|
||||
)
|
||||
reasoner = ReasonerGrammarBackend(
|
||||
backend,
|
||||
self._make_parser(),
|
||||
self._make_tokenizer(end_ids=[2, 3]),
|
||||
enable_strict_thinking=True,
|
||||
)
|
||||
self.assertEqual(reasoner.think_end_ids, [2, 3])
|
||||
|
||||
def test_rejects_unencodable_excluded_token(self):
|
||||
backend = _DummyGrammarBackend(support_token_filter=True)
|
||||
@@ -255,7 +312,7 @@ class TestReasonerGrammarObjectRollback(unittest.TestCase):
|
||||
inner_grammar.is_terminated.return_value = False
|
||||
obj = ReasonerGrammarObject(
|
||||
grammar=inner_grammar,
|
||||
think_end_id=7,
|
||||
think_end_ids=[7],
|
||||
think_excluded_token_ids=[3, 5],
|
||||
max_think_tokens=-1,
|
||||
enable_token_filter=True,
|
||||
@@ -272,11 +329,10 @@ class TestReasonerGrammarObjectRollback(unittest.TestCase):
|
||||
obj, inner_grammar = self._make_object_with_mock_grammar()
|
||||
obj.maybe_init_reasoning(True)
|
||||
|
||||
# Accept 3 thinking tokens then think_end_id
|
||||
obj.accept_token(10)
|
||||
obj.accept_token(11)
|
||||
obj.accept_token(12)
|
||||
obj.accept_token(7) # think_end_id → tokens_after_end = 0
|
||||
obj.accept_token(7)
|
||||
|
||||
self.assertTrue(obj._is_generation())
|
||||
self.assertEqual(obj.tokens_after_end, 0)
|
||||
@@ -296,7 +352,7 @@ class TestReasonerGrammarObjectRollback(unittest.TestCase):
|
||||
# 2 thinking tokens + think_end + 3 generation tokens
|
||||
obj.accept_token(10) # think
|
||||
obj.accept_token(11) # think
|
||||
obj.accept_token(7) # think_end_id
|
||||
obj.accept_token(7)
|
||||
obj.accept_token(20) # gen 1
|
||||
obj.accept_token(21) # gen 2
|
||||
obj.accept_token(22) # gen 3
|
||||
@@ -315,7 +371,7 @@ class TestReasonerGrammarObjectRollback(unittest.TestCase):
|
||||
obj.maybe_init_reasoning(True)
|
||||
|
||||
obj.accept_token(10) # think
|
||||
obj.accept_token(7) # think_end_id
|
||||
obj.accept_token(7)
|
||||
obj.accept_token(20) # gen 1
|
||||
obj.accept_token(21) # gen 2
|
||||
|
||||
@@ -344,7 +400,7 @@ class TestReasonerGrammarObjectRollback(unittest.TestCase):
|
||||
obj.maybe_init_reasoning(True)
|
||||
|
||||
obj.accept_token(10)
|
||||
obj.accept_token(7) # think_end_id → GENERATION
|
||||
obj.accept_token(7)
|
||||
obj.accept_token(20)
|
||||
|
||||
self.assertEqual(obj.tokens_in_think, 1)
|
||||
@@ -370,6 +426,26 @@ class TestReasonerGrammarObjectRollback(unittest.TestCase):
|
||||
self.assertEqual(copy.tokens_after_end, -1)
|
||||
self.assertTrue(copy._is_thinking())
|
||||
|
||||
def test_multi_token_marker_survives_rollback(self):
|
||||
obj = ReasonerGrammarObject(grammar=None, think_end_ids=[2, 3])
|
||||
obj.maybe_init_reasoning(True)
|
||||
obj.accept_token(2)
|
||||
obj.accept_token(9)
|
||||
obj.rollback(1)
|
||||
obj.accept_token(3)
|
||||
self.assertTrue(obj._is_generation())
|
||||
|
||||
obj.rollback(1)
|
||||
self.assertTrue(obj._is_thinking())
|
||||
self.assertEqual(obj._matched_think_end_tokens, 1)
|
||||
|
||||
def test_self_overlapping_marker_is_matched(self):
|
||||
obj = ReasonerGrammarObject(grammar=None, think_end_ids=[2, 2, 3])
|
||||
obj.maybe_init_reasoning(True)
|
||||
for token in (2, 2, 2, 3):
|
||||
obj.accept_token(token)
|
||||
self.assertTrue(obj._is_generation())
|
||||
|
||||
|
||||
class TestReasonerGrammarObjectFillVocabMask(unittest.TestCase):
|
||||
"""Tests for fill_vocab_mask behavior in different states."""
|
||||
@@ -383,7 +459,7 @@ class TestReasonerGrammarObjectFillVocabMask(unittest.TestCase):
|
||||
)
|
||||
obj = ReasonerGrammarObject(
|
||||
grammar=inner_grammar,
|
||||
think_end_id=7,
|
||||
think_end_ids=[7],
|
||||
think_excluded_token_ids=[3, 5],
|
||||
max_think_tokens=-1,
|
||||
enable_token_filter=True,
|
||||
@@ -411,7 +487,7 @@ class TestReasonerGrammarObjectFillVocabMask(unittest.TestCase):
|
||||
)
|
||||
obj = ReasonerGrammarObject(
|
||||
grammar=inner_grammar,
|
||||
think_end_id=7,
|
||||
think_end_ids=[7],
|
||||
think_excluded_token_ids=[3, 5],
|
||||
max_think_tokens=-1,
|
||||
enable_token_filter=True,
|
||||
@@ -424,7 +500,7 @@ class TestReasonerGrammarObjectFillVocabMask(unittest.TestCase):
|
||||
)
|
||||
obj.maybe_init_reasoning(True)
|
||||
obj.accept_token(10)
|
||||
obj.accept_token(7) # think_end_id → GENERATION
|
||||
obj.accept_token(7)
|
||||
|
||||
mask = obj.allocate_vocab_mask(64, 1, "cpu")
|
||||
obj.fill_vocab_mask(mask, 0)
|
||||
@@ -435,7 +511,7 @@ class TestReasonerGrammarObjectFillVocabMask(unittest.TestCase):
|
||||
inner_grammar = MagicMock()
|
||||
obj = ReasonerGrammarObject(
|
||||
grammar=inner_grammar,
|
||||
think_end_id=7,
|
||||
think_end_ids=[7],
|
||||
think_excluded_token_ids=None,
|
||||
max_think_tokens=-1,
|
||||
enable_token_filter=False,
|
||||
@@ -463,7 +539,7 @@ class TestReasonerGrammarObjectCurrentToken(unittest.TestCase):
|
||||
inner_grammar.is_terminated.return_value = False
|
||||
obj = ReasonerGrammarObject(
|
||||
grammar=inner_grammar,
|
||||
think_end_id=7,
|
||||
think_end_ids=[7],
|
||||
think_excluded_token_ids=None,
|
||||
max_think_tokens=-1,
|
||||
enable_token_filter=False,
|
||||
@@ -480,7 +556,7 @@ class TestReasonerGrammarObjectCurrentToken(unittest.TestCase):
|
||||
obj, inner_grammar = self._make_object_with_mock_grammar()
|
||||
obj.maybe_init_reasoning(True)
|
||||
obj.accept_token(10) # thinking token
|
||||
obj.accept_token(7) # think_end_id -> GENERATION
|
||||
obj.accept_token(7)
|
||||
obj.accept_token(58) # generation token "["
|
||||
self.assertEqual(obj.current_token, 58)
|
||||
|
||||
@@ -497,7 +573,7 @@ class TestReasonerGrammarObjectCurrentToken(unittest.TestCase):
|
||||
must not be re-accepted; with current_token tracked, the guard skips."""
|
||||
obj, inner_grammar = self._make_object_with_mock_grammar()
|
||||
obj.maybe_init_reasoning(True)
|
||||
obj.accept_token(7) # think_end_id -> GENERATION
|
||||
obj.accept_token(7)
|
||||
obj.accept_token(58) # "[" accepted into inner grammar
|
||||
obj.accept_token(4913) # '{"' accepted into inner grammar
|
||||
inner_grammar.accept_token.reset_mock()
|
||||
|
||||
@@ -438,6 +438,68 @@ class TestChatCompletionRequest(unittest.TestCase):
|
||||
self.assertEqual(name, "VoiceNote")
|
||||
self.assertEqual(strict, True)
|
||||
|
||||
def test_schema_derived_strict_false_constraint_gated_on_renderer(self):
|
||||
"""A `strict` field on the user's model doubles as the protocol switch.
|
||||
|
||||
set_json_schema pops `strict` out of the schema's properties and feeds
|
||||
its default into response_format. strict=False drops the sampling
|
||||
constraint only when the renderer forwards response_format to the
|
||||
model; otherwise the schema would be silently ignored, so the
|
||||
constraint stays installed.
|
||||
"""
|
||||
|
||||
class Note(BaseModel):
|
||||
title: str
|
||||
strict: bool = False
|
||||
|
||||
request = ChatCompletionRequest(
|
||||
model="test-model",
|
||||
messages=[{"role": "user", "content": "Return JSON"}],
|
||||
response_format={
|
||||
"type": "json_schema",
|
||||
"schema": Note.model_json_schema(),
|
||||
},
|
||||
)
|
||||
|
||||
self.assertIs(request.response_format.json_schema.strict, False)
|
||||
self.assertNotIn(
|
||||
"strict", request.response_format.json_schema.schema_["properties"]
|
||||
)
|
||||
sampling_params = request.to_sampling_params(
|
||||
stop=[], model_generation_config={}
|
||||
)
|
||||
self.assertIn("json_schema", sampling_params)
|
||||
sampling_params = request.to_sampling_params(
|
||||
stop=[],
|
||||
model_generation_config={},
|
||||
renderer_handles_response_format=True,
|
||||
)
|
||||
self.assertNotIn("json_schema", sampling_params)
|
||||
|
||||
def test_non_strict_response_format_constraint_gated_on_renderer(self):
|
||||
request = ChatCompletionRequest(
|
||||
model="test-model",
|
||||
messages=[{"role": "user", "content": "Return JSON"}],
|
||||
response_format={
|
||||
"type": "json_schema",
|
||||
"json_schema": {
|
||||
"name": "answer",
|
||||
"schema": {"type": "object"},
|
||||
"strict": False,
|
||||
},
|
||||
},
|
||||
)
|
||||
sampling_params = request.to_sampling_params(
|
||||
stop=[], model_generation_config={}
|
||||
)
|
||||
self.assertIn("json_schema", sampling_params)
|
||||
sampling_params = request.to_sampling_params(
|
||||
stop=[],
|
||||
model_generation_config={},
|
||||
renderer_handles_response_format=True,
|
||||
)
|
||||
self.assertNotIn("json_schema", sampling_params)
|
||||
|
||||
|
||||
class TestModelSerialization(unittest.TestCase):
|
||||
"""Test model serialization with hidden states"""
|
||||
|
||||
@@ -27,6 +27,7 @@ from sglang.srt.entrypoints.openai.serving_chat import (
|
||||
OpenAIServingChat,
|
||||
normalize_tool_content,
|
||||
)
|
||||
from sglang.srt.function_call.kimik3_format import TOOLS_CLOSE
|
||||
from sglang.srt.managers.io_struct import GenerateReqInput
|
||||
from sglang.srt.parser.template_detection import ReasoningToggleConfig
|
||||
from sglang.srt.utils import get_or_create_event_loop
|
||||
@@ -301,12 +302,40 @@ class ServingChatTestCase(unittest.TestCase):
|
||||
[],
|
||||
[],
|
||||
None,
|
||||
require_reasoning=True,
|
||||
)
|
||||
|
||||
adapted, _ = self.chat._convert_to_internal_request(req)
|
||||
|
||||
self.assertTrue(adapted.require_reasoning)
|
||||
|
||||
def test_process_messages_records_template_reasoning_state(self):
|
||||
self.chat.default_chat_template_kwargs = {"thinking": True}
|
||||
self.template_manager.reasoning_config = ReasoningToggleConfig(
|
||||
toggle_param="thinking", default_enabled=False
|
||||
)
|
||||
self.chat.reasoning_parser = "deepseek-v3"
|
||||
request = ChatCompletionRequest(
|
||||
model="x",
|
||||
messages=[{"role": "user", "content": "What is 2+2?"}],
|
||||
)
|
||||
rendered = MessageProcessingResult(
|
||||
prompt="prompt",
|
||||
prompt_ids=[1, 2, 3],
|
||||
image_data=None,
|
||||
audio_data=None,
|
||||
video_data=None,
|
||||
modalities=[],
|
||||
stop=[],
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
self.chat, "_apply_conversation_template", return_value=rendered
|
||||
):
|
||||
processed = self.chat._process_messages(request, is_multimodal=False)
|
||||
|
||||
self.assertTrue(processed.require_reasoning)
|
||||
|
||||
def test_kimi_tool_call_respects_explicit_reasoning_disable(self):
|
||||
self.template_manager.reasoning_config = ReasoningToggleConfig(
|
||||
toggle_param="thinking", default_enabled=True
|
||||
@@ -632,6 +661,260 @@ class ServingChatTestCase(unittest.TestCase):
|
||||
parser.get_structure_constraint.call_args.kwargs["thinking_mode"]
|
||||
)
|
||||
|
||||
def test_kimi_k3_constraint_failure_keeps_native_stop_format(self):
|
||||
self.template_manager.chat_template_name = None
|
||||
self.template_manager.jinja_template_content_format = "string"
|
||||
self.chat.chat_encoding_spec = "kimi_k3"
|
||||
self.chat.tool_call_parser = "kimi_k3"
|
||||
self.tm.tokenizer.apply_chat_template.return_value = [1, 2, 3]
|
||||
tool = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"location": {"type": "string"}},
|
||||
"required": ["location"],
|
||||
},
|
||||
"strict": True,
|
||||
},
|
||||
}
|
||||
|
||||
cases = (
|
||||
(None, [TOOLS_CLOSE]),
|
||||
("USER_STOP", ["USER_STOP", TOOLS_CLOSE]),
|
||||
(["USER_STOP"], ["USER_STOP", TOOLS_CLOSE]),
|
||||
([TOOLS_CLOSE], [TOOLS_CLOSE]),
|
||||
)
|
||||
for request_stop, expected in cases:
|
||||
with (
|
||||
self.subTest(request_stop=request_stop),
|
||||
patch(
|
||||
"sglang.srt.entrypoints.openai.serving_chat.FunctionCallParser"
|
||||
) as parser_cls,
|
||||
):
|
||||
parser = parser_cls.return_value
|
||||
parser.detector.eot_token = TOOLS_CLOSE
|
||||
parser.detector.parses_required_natively.return_value = False
|
||||
parser.get_structure_constraint.return_value = None
|
||||
request = ChatCompletionRequest(
|
||||
model="x",
|
||||
messages=[{"role": "user", "content": "Weather in Paris?"}],
|
||||
tools=[tool],
|
||||
tool_choice="required",
|
||||
stop=request_stop,
|
||||
)
|
||||
original_stop = (
|
||||
list(request.stop)
|
||||
if isinstance(request.stop, list)
|
||||
else request.stop
|
||||
)
|
||||
|
||||
result = self.chat._process_messages(request, is_multimodal=False)
|
||||
|
||||
self.assertEqual(result.stop, expected)
|
||||
self.assertEqual(request.stop, original_stop)
|
||||
self.assertIsNone(result.tool_call_constraint)
|
||||
|
||||
def test_kimi_k3_tool_call_stop_is_scoped_to_active_tools(self):
|
||||
self.template_manager.chat_template_name = None
|
||||
self.template_manager.jinja_template_content_format = "string"
|
||||
self.chat.chat_encoding_spec = "kimi_k3"
|
||||
self.chat.tool_call_parser = "kimi_k3"
|
||||
self.tm.tokenizer.apply_chat_template.return_value = [1, 2, 3]
|
||||
request = ChatCompletionRequest(
|
||||
model="x",
|
||||
messages=[{"role": "user", "content": "Weather in Paris?"}],
|
||||
tools=[
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"parameters": {"type": "object"},
|
||||
},
|
||||
}
|
||||
],
|
||||
tool_choice="none",
|
||||
)
|
||||
|
||||
result = self.chat._process_messages(request, is_multimodal=False)
|
||||
|
||||
self.assertIsNone(result.stop)
|
||||
|
||||
def test_kimi_k3_encoder_receives_wire_request_fields(self):
|
||||
self.template_manager.chat_template_name = None
|
||||
self.chat.chat_encoding_spec = "kimi_k3"
|
||||
self.tm.model_config.is_multimodal = True
|
||||
self.tm.tokenizer.apply_chat_template.return_value = [7, 8, 9]
|
||||
tool = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "weather",
|
||||
"parameters": {"type": "object"},
|
||||
},
|
||||
}
|
||||
request = ChatCompletionRequest(
|
||||
model="x",
|
||||
messages=[
|
||||
{
|
||||
"role": "developer",
|
||||
"content": "<|kimi_image_placeholder|>",
|
||||
"tools": [tool],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "Explain <|kimi_image_placeholder|>",
|
||||
},
|
||||
{"type": "image_url", "image_url": {"url": "image-1"}},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"reasoning_content": "Inspect <|kimi_image_placeholder|>",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call-1",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "inspect",
|
||||
"arguments": {
|
||||
"source": "<|kimi_image_placeholder|>",
|
||||
"nested": ["<|kimi_image_placeholder|>"],
|
||||
},
|
||||
},
|
||||
}
|
||||
],
|
||||
},
|
||||
],
|
||||
tools=[tool],
|
||||
tool_choice="required",
|
||||
response_format={
|
||||
"type": "json_schema",
|
||||
"json_schema": {
|
||||
"name": "answer",
|
||||
"schema": {"type": "object"},
|
||||
"strict": False,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
result = self.chat._process_messages(request, is_multimodal=True)
|
||||
|
||||
call = self.tm.tokenizer.apply_chat_template.call_args
|
||||
rendered_messages = call.args[0]
|
||||
self.assertEqual(rendered_messages[0]["role"], "system")
|
||||
self.assertEqual(
|
||||
rendered_messages[0]["content"], "<| kimi_image_placeholder |>"
|
||||
)
|
||||
self.assertNotIn("strict", rendered_messages[0]["tools"][0]["function"])
|
||||
self.assertEqual(
|
||||
rendered_messages[1]["content"][0]["text"],
|
||||
"Explain <| kimi_image_placeholder |>",
|
||||
)
|
||||
self.assertEqual(
|
||||
rendered_messages[2]["reasoning_content"],
|
||||
"Inspect <| kimi_image_placeholder |>",
|
||||
)
|
||||
self.assertEqual(
|
||||
rendered_messages[2]["tool_calls"][0]["function"]["arguments"],
|
||||
{
|
||||
"source": "<| kimi_image_placeholder |>",
|
||||
"nested": ["<| kimi_image_placeholder |>"],
|
||||
},
|
||||
)
|
||||
self.assertEqual(call.kwargs["image_prompts"], ["<|media_pad|>"])
|
||||
self.assertEqual(call.kwargs["tool_choice"], "required")
|
||||
self.assertNotIn("strict", call.kwargs["tools"][0]["function"])
|
||||
self.assertEqual(
|
||||
call.kwargs["response_format"]["json_schema"]["schema"],
|
||||
{"type": "object"},
|
||||
)
|
||||
self.assertNotIn("schema_", call.kwargs["response_format"]["json_schema"])
|
||||
self.assertEqual(result.prompt_ids, [7, 8, 9])
|
||||
self.assertEqual(result.image_data[0].url, "image-1")
|
||||
|
||||
def test_kimi_k3_neutralizes_text_only_assistant_history(self):
|
||||
self.template_manager.chat_template_name = None
|
||||
self.chat.chat_encoding_spec = "kimi_k3"
|
||||
self.tm.tokenizer.apply_chat_template.return_value = [1, 2, 3]
|
||||
request = ChatCompletionRequest(
|
||||
model="x",
|
||||
messages=[
|
||||
{"role": "user", "content": "Run it"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"reasoning_content": "Read <|kimi_image_placeholder|>",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call-1",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "shell",
|
||||
"arguments": "not-json <|kimi_image_placeholder|>",
|
||||
},
|
||||
}
|
||||
],
|
||||
},
|
||||
],
|
||||
)
|
||||
|
||||
self.chat._process_messages(request, is_multimodal=False)
|
||||
|
||||
messages = self.tm.tokenizer.apply_chat_template.call_args.args[0]
|
||||
kwargs = self.tm.tokenizer.apply_chat_template.call_args.kwargs
|
||||
self.assertEqual(messages[-1]["role"], "assistant")
|
||||
self.assertEqual(
|
||||
messages[-1]["reasoning_content"],
|
||||
"Read <| kimi_image_placeholder |>",
|
||||
)
|
||||
self.assertEqual(
|
||||
messages[-1]["tool_calls"][0]["function"]["arguments"],
|
||||
"not-json <| kimi_image_placeholder |>",
|
||||
)
|
||||
self.assertNotIn("image_prompts", kwargs)
|
||||
|
||||
def test_message_tools_participate_in_validation_across_encodings(self):
|
||||
tool = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "weather",
|
||||
"parameters": {"type": "object"},
|
||||
},
|
||||
}
|
||||
messages = [
|
||||
{"role": "system", "content": "", "tools": [tool]},
|
||||
{"role": "user", "content": "Weather?"},
|
||||
]
|
||||
for chat_encoding_spec in (None, "dsv4", "dsv32", "kimi_k3"):
|
||||
with self.subTest(chat_encoding_spec=chat_encoding_spec):
|
||||
self.chat.chat_encoding_spec = chat_encoding_spec
|
||||
automatic = ChatCompletionRequest(
|
||||
model="x", messages=messages, tool_choice=None
|
||||
)
|
||||
self.assertEqual(automatic.tool_choice, "auto")
|
||||
self.assertIsNone(self.chat._validate_request(automatic))
|
||||
|
||||
required = ChatCompletionRequest(
|
||||
model="x", messages=messages, tool_choice="required"
|
||||
)
|
||||
self.assertIsNone(self.chat._validate_request(required))
|
||||
|
||||
duplicate = ChatCompletionRequest(
|
||||
model="x",
|
||||
messages=messages,
|
||||
tools=[tool],
|
||||
tool_choice="required",
|
||||
)
|
||||
self.assertEqual(
|
||||
self.chat._validate_request(duplicate),
|
||||
"Tool names must be unique across request and message tools.",
|
||||
)
|
||||
|
||||
def test_jinja_rejects_non_object_tool_call_arguments(self):
|
||||
"""History tool call arguments must parse to a JSON object."""
|
||||
self.template_manager.chat_template_name = None
|
||||
@@ -1256,6 +1539,19 @@ class ServingChatTestCase(unittest.TestCase):
|
||||
serving_chat = OpenAIServingChat(tm, TemplateManager())
|
||||
self.assertEqual(serving_chat.chat_encoding_spec, "dsv4")
|
||||
|
||||
def test_kimi_k3_encoding_detection(self):
|
||||
from sglang.srt.parser.template_manager import TemplateManager
|
||||
|
||||
tm = _MockTokenizerManager()
|
||||
tm.model_config.hf_config.architectures = ["KimiK3ForConditionalGeneration"]
|
||||
serving_chat = OpenAIServingChat(tm, TemplateManager())
|
||||
self.assertEqual(serving_chat.chat_encoding_spec, "kimi_k3")
|
||||
|
||||
tm.model_config.hf_config.architectures = ["LlamaForCausalLM"]
|
||||
tm.server_args.tool_call_parser = "kimi_k3"
|
||||
serving_chat = OpenAIServingChat(tm, TemplateManager())
|
||||
self.assertEqual(serving_chat.chat_encoding_spec, "kimi_k3")
|
||||
|
||||
# ------------- dsv4 task + latest_reminder -------------
|
||||
def test_dsv4_task_field_schema(self):
|
||||
"""Top-level `task` accepts the 6 DS task tokens and rejects others."""
|
||||
|
||||
@@ -18,6 +18,7 @@ from sglang.srt.entrypoints.openai.protocol import (
|
||||
)
|
||||
from sglang.srt.entrypoints.openai.serving_responses import OpenAIServingResponses
|
||||
from sglang.srt.function_call.core_types import ToolCallItem
|
||||
from sglang.srt.parser.template_detection import ReasoningToggleConfig
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=8, suite="base-a-test-cpu")
|
||||
@@ -205,6 +206,105 @@ class ChatToolForwardingTestCase(unittest.TestCase):
|
||||
result = asyncio.run(serving.create_responses(request, raw_request=None))
|
||||
self.assertEqual(getattr(result, "status_code", None), 400)
|
||||
|
||||
def test_kimi_k3_request_uses_chat_encoder_fields(self):
|
||||
serving = make_serving()
|
||||
serving.chat_encoding_spec = "kimi_k3"
|
||||
serving.default_chat_template_kwargs = {}
|
||||
serving.template_manager.chat_template_name = None
|
||||
serving.tokenizer_manager.tokenizer.apply_chat_template.return_value = [4, 5, 6]
|
||||
request = ResponsesRequest(
|
||||
model="x",
|
||||
input="Explain <|kimi_image_placeholder|>",
|
||||
tools=[
|
||||
{
|
||||
"type": "function",
|
||||
"name": "lookup",
|
||||
"parameters": {"type": "object"},
|
||||
}
|
||||
],
|
||||
tool_choice="required",
|
||||
reasoning={"effort": "high"},
|
||||
store=False,
|
||||
)
|
||||
|
||||
_, request_prompts, engine_prompts, _ = asyncio.run(
|
||||
serving._make_request(request, None, serving.tokenizer_manager.tokenizer)
|
||||
)
|
||||
|
||||
call = serving.tokenizer_manager.tokenizer.apply_chat_template.call_args
|
||||
self.assertEqual(
|
||||
call.args[0][0]["content"], "Explain <| kimi_image_placeholder |>"
|
||||
)
|
||||
self.assertEqual(call.kwargs["thinking_effort"], "high")
|
||||
self.assertEqual(call.kwargs["tool_choice"], "required")
|
||||
self.assertEqual(call.kwargs["tools"][0]["function"]["name"], "lookup")
|
||||
self.assertEqual(request_prompts, [[4, 5, 6]])
|
||||
self.assertEqual(engine_prompts, [[4, 5, 6]])
|
||||
|
||||
|
||||
class ReasoningRequestForwardingTestCase(unittest.TestCase):
|
||||
def test_create_responses_uses_processed_reasoning_state(self):
|
||||
serving = make_serving()
|
||||
serving.reasoning_parser = "deepseek-r1"
|
||||
serving.default_chat_template_kwargs = {"thinking": False}
|
||||
serving.template_manager.reasoning_config = ReasoningToggleConfig(
|
||||
toggle_param="thinking", default_enabled=True
|
||||
)
|
||||
rendered = MessageProcessingResult(
|
||||
prompt="prompt",
|
||||
prompt_ids=[1, 2, 3],
|
||||
image_data=None,
|
||||
audio_data=None,
|
||||
video_data=None,
|
||||
modalities=[],
|
||||
stop=[],
|
||||
)
|
||||
captured = {}
|
||||
|
||||
async def fake_generate(
|
||||
request_id,
|
||||
request_prompt,
|
||||
adapted_request,
|
||||
sampling_params,
|
||||
context,
|
||||
**kwargs,
|
||||
):
|
||||
captured["adapted_request"] = adapted_request
|
||||
context.append_output(
|
||||
{
|
||||
"text": "done",
|
||||
"meta_info": {
|
||||
"prompt_tokens": 3,
|
||||
"completion_tokens": 1,
|
||||
"cached_tokens": 0,
|
||||
},
|
||||
}
|
||||
)
|
||||
yield context
|
||||
|
||||
serving._generate_with_builtin_tools = fake_generate
|
||||
request = ResponsesRequest(
|
||||
model="x",
|
||||
input="answer",
|
||||
request_id="resp_reasoning",
|
||||
store=False,
|
||||
)
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
serving, "_apply_conversation_template", return_value=rendered
|
||||
),
|
||||
patch(
|
||||
"sglang.srt.entrypoints.openai.serving_responses.ReasoningParser"
|
||||
) as parser_cls,
|
||||
):
|
||||
parser_cls.return_value.parse_non_stream.return_value = (None, "done")
|
||||
response = asyncio.run(serving.create_responses(request))
|
||||
|
||||
self.assertEqual(response.status, "completed")
|
||||
self.assertFalse(captured["adapted_request"].require_reasoning)
|
||||
self.assertFalse(parser_cls.call_args.kwargs["force_reasoning"])
|
||||
|
||||
|
||||
class InputItemNormalizationTestCase(unittest.TestCase):
|
||||
def test_function_call_becomes_assistant_tool_call(self):
|
||||
@@ -293,6 +393,7 @@ class FullResponseUsageTestCase(unittest.TestCase):
|
||||
tokenizer=serving.tokenizer_manager.tokenizer,
|
||||
request_metadata=metadata,
|
||||
created_time=123,
|
||||
require_reasoning=False,
|
||||
)
|
||||
)
|
||||
|
||||
@@ -432,6 +533,7 @@ class OutputItemsTestCase(unittest.TestCase):
|
||||
self._function_tool_request(),
|
||||
"raw model output with <tool_call>",
|
||||
tokenizer=Mock(),
|
||||
require_reasoning=False,
|
||||
)
|
||||
|
||||
tool_calls = [
|
||||
@@ -464,7 +566,10 @@ class OutputItemsTestCase(unittest.TestCase):
|
||||
[fake_call],
|
||||
)
|
||||
output_items = serving._make_response_output_items(
|
||||
self._function_tool_request(), "raw model output", tokenizer=Mock()
|
||||
self._function_tool_request(),
|
||||
"raw model output",
|
||||
tokenizer=Mock(),
|
||||
require_reasoning=False,
|
||||
)
|
||||
|
||||
types = [type(item).__name__ for item in output_items]
|
||||
@@ -489,7 +594,7 @@ class OutputItemsTestCase(unittest.TestCase):
|
||||
raw = '[{"name": "get_weather", "parameters": {"city": "Beijing"}}]'
|
||||
|
||||
output_items = serving._make_response_output_items(
|
||||
request, raw, tokenizer=Mock()
|
||||
request, raw, tokenizer=Mock(), require_reasoning=False
|
||||
)
|
||||
|
||||
tool_calls = [
|
||||
@@ -524,7 +629,10 @@ class OutputItemsTestCase(unittest.TestCase):
|
||||
"sglang.srt.entrypoints.openai.serving_responses.FunctionCallParser"
|
||||
) as parser_cls:
|
||||
output_items = serving._make_response_output_items(
|
||||
request, "just a plain answer", tokenizer=Mock()
|
||||
request,
|
||||
"just a plain answer",
|
||||
tokenizer=Mock(),
|
||||
require_reasoning=False,
|
||||
)
|
||||
parser_cls.assert_not_called()
|
||||
|
||||
|
||||
@@ -20,9 +20,10 @@ register_cpu_ci(est_time=4, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
class _StreamFixture:
|
||||
def __init__(self, serving, request):
|
||||
def __init__(self, serving, request, *, require_reasoning=False):
|
||||
self.serving = serving
|
||||
self.request = request
|
||||
self.require_reasoning = require_reasoning
|
||||
self.request_metadata = RequestResponseMetadata(request_id=request.request_id)
|
||||
|
||||
def run(self, chunks):
|
||||
@@ -39,6 +40,7 @@ class _StreamFixture:
|
||||
model_name="x",
|
||||
tokenizer=Mock(),
|
||||
request_metadata=self.request_metadata,
|
||||
require_reasoning=self.require_reasoning,
|
||||
)
|
||||
)
|
||||
|
||||
@@ -60,6 +62,20 @@ def _engine_chunk(text, completion_tokens, *, finish=False):
|
||||
|
||||
|
||||
class NonHarmonyStreamTestCase(unittest.TestCase):
|
||||
def test_reasoning_parser_uses_processed_reasoning_state(self):
|
||||
serving = make_serving()
|
||||
serving.reasoning_parser = "deepseek-r1"
|
||||
request = ResponsesRequest(model="x", input="hi", stream=True, store=False)
|
||||
|
||||
with patch(
|
||||
"sglang.srt.entrypoints.openai.serving_responses.ReasoningParser"
|
||||
) as parser_cls:
|
||||
parser_cls.return_value.parse_stream_chunk.return_value = (None, "done")
|
||||
fixture = _StreamFixture(serving, request, require_reasoning=True)
|
||||
fixture.run([_engine_chunk("done", 1, finish=True)])
|
||||
|
||||
self.assertTrue(parser_cls.call_args.kwargs["force_reasoning"])
|
||||
|
||||
def test_emits_typed_sse_events_in_order(self):
|
||||
serving = make_serving()
|
||||
serving.reasoning_parser = None
|
||||
|
||||
@@ -0,0 +1,810 @@
|
||||
import json
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import xgrammar as xgr
|
||||
from xgrammar.testing import _is_grammar_accept_string
|
||||
|
||||
from sglang.srt.entrypoints.openai.protocol import (
|
||||
ChatCompletionRequest,
|
||||
Function,
|
||||
Tool,
|
||||
ToolChoice,
|
||||
ToolChoiceFuncName,
|
||||
)
|
||||
from sglang.srt.environ import ToolStrictLevel, envs
|
||||
from sglang.srt.function_call.function_call_parser import FunctionCallParser
|
||||
from sglang.srt.function_call.kimik3_detector import KimiK3Detector
|
||||
from sglang.srt.function_call.kimik3_format import (
|
||||
ARGUMENT_CLOSE,
|
||||
CALL_CLOSE,
|
||||
THINK_CLOSE,
|
||||
TOOLS_CLOSE,
|
||||
TOOLS_OPEN,
|
||||
)
|
||||
from sglang.srt.function_call.kimik3_structural_tag import (
|
||||
get_kimik3_auto_tool_call_structural_tag,
|
||||
get_kimik3_structural_tag,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=8, suite="base-a-test-cpu")
|
||||
|
||||
_CLOSE_TOKEN = "<|close|>"
|
||||
_CLOSE_TOKEN_ID = 256
|
||||
_TOKENIZER_INFO = xgr.TokenizerInfo(
|
||||
[bytes([token_id]) for token_id in range(256)] + [_CLOSE_TOKEN.encode()]
|
||||
)
|
||||
_TOKEN_COMPILER = xgr.GrammarCompiler(_TOKENIZER_INFO, cache_enabled=True)
|
||||
|
||||
|
||||
def _tool(name="weather", strict=True):
|
||||
return Tool(
|
||||
type="function",
|
||||
function=Function(
|
||||
name=name,
|
||||
strict=strict,
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"city": {
|
||||
"type": "string",
|
||||
"pattern": "[A-Z][A-Za-z ]+",
|
||||
},
|
||||
"days": {
|
||||
"type": "integer",
|
||||
"minimum": 1,
|
||||
"maximum": 10,
|
||||
},
|
||||
"unit": {
|
||||
"type": ["string", "null"],
|
||||
"enum": ["celsius", "fahrenheit", None],
|
||||
},
|
||||
"metadata": {
|
||||
"type": "object",
|
||||
"properties": {"source": {"type": "string"}},
|
||||
"required": ["source"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
"tags": {
|
||||
"type": "array",
|
||||
"items": {"type": "string"},
|
||||
},
|
||||
},
|
||||
"required": ["city", "days"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _argument(key, argument_type, value):
|
||||
return (
|
||||
f'<|open|>argument key="{key}" type="{argument_type}"<|sep|>'
|
||||
f"{value}{ARGUMENT_CLOSE}"
|
||||
)
|
||||
|
||||
|
||||
def _call(name, index, *arguments):
|
||||
return (
|
||||
f'<|open|>call tool="{name}" index="{index}"<|sep|>'
|
||||
+ "".join(arguments)
|
||||
+ CALL_CLOSE
|
||||
)
|
||||
|
||||
|
||||
def _tools_section(*calls):
|
||||
return TOOLS_OPEN + "".join(calls) + TOOLS_CLOSE
|
||||
|
||||
|
||||
def _grammar(tools, tool_choice="auto", thinking_mode=False, parallel_tool_calls=True):
|
||||
structural_tag = get_kimik3_structural_tag(
|
||||
tools,
|
||||
tool_choice=tool_choice,
|
||||
thinking_mode=thinking_mode,
|
||||
parallel_tool_calls=parallel_tool_calls,
|
||||
)
|
||||
return xgr.Grammar.from_structural_tag(structural_tag)
|
||||
|
||||
|
||||
def _accepts(grammar, value):
|
||||
return _is_grammar_accept_string(grammar, value)
|
||||
|
||||
|
||||
def _encode_with_close_token(value):
|
||||
token_ids = []
|
||||
start = 0
|
||||
while (index := value.find(_CLOSE_TOKEN, start)) != -1:
|
||||
token_ids.extend(value[start:index].encode())
|
||||
token_ids.append(_CLOSE_TOKEN_ID)
|
||||
start = index + len(_CLOSE_TOKEN)
|
||||
token_ids.extend(value[start:].encode())
|
||||
return token_ids
|
||||
|
||||
|
||||
def _token_accepts(structural_tag, value):
|
||||
compiled = _TOKEN_COMPILER.compile_structural_tag(structural_tag)
|
||||
matcher = xgr.GrammarMatcher(compiled)
|
||||
for token_id in _encode_with_close_token(value):
|
||||
if not matcher.accept_token(token_id):
|
||||
return False
|
||||
return matcher.is_completed()
|
||||
|
||||
|
||||
def _valid_weather_call(index=1):
|
||||
return _call(
|
||||
"weather",
|
||||
index,
|
||||
_argument("city", "string", "San Francisco"),
|
||||
_argument("days", "number", "3"),
|
||||
_argument("unit", "string", "celsius"),
|
||||
_argument("metadata", "object", '{"source":"forecast"}'),
|
||||
_argument("tags", "array", '["coastal","windy"]'),
|
||||
)
|
||||
|
||||
|
||||
def test_strict_schema_accepts_native_xtml_values():
|
||||
grammar = _grammar([_tool()], tool_choice="required")
|
||||
|
||||
assert _accepts(grammar, _tools_section(_valid_weather_call()))
|
||||
assert _accepts(
|
||||
grammar,
|
||||
_tools_section(
|
||||
_call(
|
||||
"weather",
|
||||
1,
|
||||
_argument("city", "string", "Paris"),
|
||||
_argument("days", "number", "1"),
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"arguments",
|
||||
[
|
||||
(_argument("city", "string", "paris"), _argument("days", "number", "3")),
|
||||
(
|
||||
_argument("city", "string", "Paris"),
|
||||
_argument("days", "number", "11"),
|
||||
),
|
||||
(
|
||||
_argument("city", "number", "3"),
|
||||
_argument("days", "number", "3"),
|
||||
),
|
||||
(
|
||||
_argument("city", "string", "Paris"),
|
||||
_argument("days", "number", "3"),
|
||||
_argument("unit", "string", "kelvin"),
|
||||
),
|
||||
(
|
||||
_argument("city", "string", "Paris"),
|
||||
_argument("days", "number", "3"),
|
||||
_argument("tags", "array", "[coastal]"),
|
||||
),
|
||||
(
|
||||
_argument("city", "string", "Paris"),
|
||||
_argument("days", "number", "3"),
|
||||
_argument("unknown", "string", "value"),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_strict_schema_rejects_invalid_parameters(arguments):
|
||||
grammar = _grammar([_tool()], tool_choice="required")
|
||||
|
||||
assert not _accepts(grammar, _tools_section(_call("weather", 1, *arguments)))
|
||||
|
||||
|
||||
def test_required_allows_response_prefix_but_requires_tools():
|
||||
grammar = _grammar([_tool()], tool_choice="required")
|
||||
response = "<|open|>response<|sep|>Checking." "<|close|>response<|sep|>"
|
||||
|
||||
assert _accepts(grammar, response + _tools_section(_valid_weather_call()))
|
||||
assert not _accepts(grammar, response)
|
||||
|
||||
|
||||
def test_auto_allows_plain_response_or_multiple_tool_calls():
|
||||
grammar = _grammar([_tool(), _tool("forecast")])
|
||||
plain = (
|
||||
"<|open|>response<|sep|>No tool needed."
|
||||
"<|close|>response<|sep|><|close|>message<|sep|>"
|
||||
)
|
||||
calls = _tools_section(
|
||||
_valid_weather_call(),
|
||||
_call(
|
||||
"forecast",
|
||||
2,
|
||||
_argument("city", "string", "London"),
|
||||
_argument("days", "number", "2"),
|
||||
),
|
||||
)
|
||||
|
||||
assert _accepts(grammar, plain)
|
||||
assert _accepts(grammar, calls)
|
||||
|
||||
|
||||
def test_named_tool_choice_forces_only_the_selected_tool():
|
||||
grammar = _grammar(
|
||||
[_tool(), _tool("forecast")],
|
||||
tool_choice=ToolChoice(function=ToolChoiceFuncName(name="forecast")),
|
||||
)
|
||||
forecast_call = _call(
|
||||
"forecast",
|
||||
1,
|
||||
_argument("city", "string", "London"),
|
||||
_argument("days", "number", "2"),
|
||||
)
|
||||
|
||||
assert _accepts(grammar, _tools_section(forecast_call))
|
||||
assert not _accepts(grammar, _tools_section(_valid_weather_call()))
|
||||
|
||||
|
||||
def test_function_call_parser_uses_native_tag_for_named_tool_choice():
|
||||
tool_choice = ToolChoice(function=ToolChoiceFuncName(name="forecast"))
|
||||
constraint = FunctionCallParser(
|
||||
[_tool(), _tool("forecast")], "kimi_k3"
|
||||
).get_structure_constraint(tool_choice)
|
||||
assert constraint is not None
|
||||
grammar = xgr.Grammar.from_structural_tag(constraint[1])
|
||||
forecast_call = _call(
|
||||
"forecast",
|
||||
1,
|
||||
_argument("city", "string", "London"),
|
||||
_argument("days", "number", "2"),
|
||||
)
|
||||
|
||||
assert _accepts(grammar, _tools_section(forecast_call))
|
||||
assert not _accepts(grammar, _tools_section(_valid_weather_call()))
|
||||
|
||||
|
||||
def test_non_strict_tool_keeps_xtml_structure_and_loose_parameters():
|
||||
grammar = _grammar([_tool(strict=False)], tool_choice="required")
|
||||
call = _call(
|
||||
"weather",
|
||||
1,
|
||||
_argument("custom", "array", '["x",1]'),
|
||||
_argument("other", "string", "raw text"),
|
||||
)
|
||||
|
||||
assert _accepts(grammar, _tools_section(call))
|
||||
assert not _accepts(grammar, _tools_section("unstructured"))
|
||||
|
||||
|
||||
def test_strict_schema_supports_refs_and_mixed_unions():
|
||||
tool = Tool(
|
||||
type="function",
|
||||
function=Function(
|
||||
name="convert",
|
||||
strict=True,
|
||||
parameters={
|
||||
"type": "object",
|
||||
"$defs": {
|
||||
"mode": {
|
||||
"type": "string",
|
||||
"enum": ["fast", "safe"],
|
||||
}
|
||||
},
|
||||
"properties": {
|
||||
"mode": {"$ref": "#/$defs/mode"},
|
||||
"value": {
|
||||
"anyOf": [
|
||||
{"type": "string", "enum": ["auto"]},
|
||||
{
|
||||
"type": "integer",
|
||||
"minimum": 2,
|
||||
"maximum": 3,
|
||||
},
|
||||
]
|
||||
},
|
||||
},
|
||||
"required": ["mode", "value"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
),
|
||||
)
|
||||
grammar = _grammar([tool], tool_choice="required")
|
||||
|
||||
assert _accepts(
|
||||
grammar,
|
||||
_tools_section(
|
||||
_call(
|
||||
"convert",
|
||||
1,
|
||||
_argument("mode", "string", "fast"),
|
||||
_argument("value", "number", "2"),
|
||||
)
|
||||
),
|
||||
)
|
||||
assert not _accepts(
|
||||
grammar,
|
||||
_tools_section(
|
||||
_call(
|
||||
"convert",
|
||||
1,
|
||||
_argument("mode", "string", "unsafe"),
|
||||
_argument("value", "number", "1"),
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_strict_schema_handles_number_enums_and_all_of_integer_constraints():
|
||||
tool = Tool(
|
||||
type="function",
|
||||
function=Function(
|
||||
name="score",
|
||||
strict=True,
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"value": {
|
||||
"type": "number",
|
||||
"enum": [1, 1.5],
|
||||
},
|
||||
"count": {
|
||||
"allOf": [
|
||||
{"type": "number"},
|
||||
{"type": "integer", "minimum": 1},
|
||||
]
|
||||
},
|
||||
},
|
||||
"required": ["value", "count"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
),
|
||||
)
|
||||
grammar = _grammar([tool], tool_choice="required")
|
||||
|
||||
assert _accepts(
|
||||
grammar,
|
||||
_tools_section(
|
||||
_call(
|
||||
"score",
|
||||
1,
|
||||
_argument("value", "number", "1"),
|
||||
_argument("count", "number", "2"),
|
||||
)
|
||||
),
|
||||
)
|
||||
assert not _accepts(
|
||||
grammar,
|
||||
_tools_section(
|
||||
_call(
|
||||
"score",
|
||||
1,
|
||||
_argument("value", "number", "2"),
|
||||
_argument("count", "number", "1.5"),
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_strict_schema_preserves_additional_properties_default():
|
||||
tool = Tool(
|
||||
type="function",
|
||||
function=Function(
|
||||
name="annotate",
|
||||
strict=True,
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"label": {"type": "string"},
|
||||
},
|
||||
"required": ["label"],
|
||||
},
|
||||
),
|
||||
)
|
||||
grammar = _grammar([tool], tool_choice="required")
|
||||
|
||||
assert _accepts(
|
||||
grammar,
|
||||
_tools_section(
|
||||
_call(
|
||||
"annotate",
|
||||
1,
|
||||
_argument("label", "string", "sample"),
|
||||
_argument("confidence", "number", "0.9"),
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_dynamic_argument_key_compiles_without_xgrammar_unicode_warning(capfd):
|
||||
tool = Tool(
|
||||
type="function",
|
||||
function=Function(
|
||||
name="annotate",
|
||||
strict=True,
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
},
|
||||
),
|
||||
)
|
||||
grammar = _grammar([tool], tool_choice="required")
|
||||
|
||||
assert _accepts(
|
||||
grammar,
|
||||
_tools_section(_call("annotate", 1, _argument("置信度", "number", "0.9"))),
|
||||
)
|
||||
assert "Negative Character class" not in capfd.readouterr().err
|
||||
|
||||
|
||||
def test_strict_empty_object_accepts_no_arguments_only():
|
||||
tool = Tool(
|
||||
type="function",
|
||||
function=Function(
|
||||
name="ping",
|
||||
strict=True,
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
"additionalProperties": False,
|
||||
},
|
||||
),
|
||||
)
|
||||
grammar = _grammar([tool], tool_choice="required")
|
||||
|
||||
assert _accepts(grammar, _tools_section(_call("ping", 1)))
|
||||
assert not _accepts(
|
||||
grammar,
|
||||
_tools_section(_call("ping", 1, _argument("unexpected", "string", "value"))),
|
||||
)
|
||||
|
||||
|
||||
def test_tool_strict_level_controls_native_tag_parameter_schema():
|
||||
invalid_call = _tools_section(
|
||||
_call(
|
||||
"weather",
|
||||
264,
|
||||
_argument("city", "string", "paris"),
|
||||
_argument("days", "number", "99"),
|
||||
)
|
||||
)
|
||||
empty_call = _tools_section(_call("weather", 264))
|
||||
|
||||
with envs.SGLANG_TOOL_STRICT_LEVEL.override(ToolStrictLevel.OFF.value):
|
||||
constraint = FunctionCallParser(
|
||||
[_tool(strict=False)], "kimi_k3"
|
||||
).get_structure_constraint("auto")
|
||||
assert constraint is not None
|
||||
assert _token_accepts(constraint[1], invalid_call)
|
||||
assert not _token_accepts(constraint[1], empty_call)
|
||||
|
||||
with envs.SGLANG_TOOL_STRICT_LEVEL.override(ToolStrictLevel.FUNCTION.value):
|
||||
constraint = FunctionCallParser(
|
||||
[_tool(strict=False)], "kimi_k3"
|
||||
).get_structure_constraint("auto")
|
||||
assert constraint is not None
|
||||
grammar = xgr.Grammar.from_structural_tag(constraint[1])
|
||||
assert _accepts(grammar, invalid_call)
|
||||
assert _accepts(grammar, empty_call)
|
||||
|
||||
with envs.SGLANG_TOOL_STRICT_LEVEL.override(ToolStrictLevel.PARAMETER.value):
|
||||
constraint = FunctionCallParser(
|
||||
[_tool(strict=False)], "kimi_k3"
|
||||
).get_structure_constraint("auto")
|
||||
assert constraint is not None
|
||||
assert not _accepts(
|
||||
xgr.Grammar.from_structural_tag(constraint[1]), invalid_call
|
||||
)
|
||||
|
||||
|
||||
def test_auto_hook_constrains_all_calls_and_requires_nonempty_values():
|
||||
structural_tag = get_kimik3_auto_tool_call_structural_tag([_tool(strict=False)])
|
||||
assert structural_tag is not None
|
||||
first = _call(
|
||||
"weather",
|
||||
3,
|
||||
_argument("city", "string", "Paris"),
|
||||
_argument("days", "number", "3"),
|
||||
)
|
||||
second = _call(
|
||||
"weather",
|
||||
264,
|
||||
_argument("city", "string", "London"),
|
||||
_argument("days", "number", "2"),
|
||||
)
|
||||
|
||||
assert _token_accepts(structural_tag, _tools_section(first, second))
|
||||
assert not _token_accepts(
|
||||
structural_tag,
|
||||
_tools_section(first, _call("weather", 264)),
|
||||
)
|
||||
assert not _token_accepts(
|
||||
structural_tag,
|
||||
_tools_section(
|
||||
_call(
|
||||
"weather",
|
||||
3,
|
||||
_argument("city", "string", ""),
|
||||
_argument("days", "number", "3"),
|
||||
)
|
||||
),
|
||||
)
|
||||
assert not _token_accepts(structural_tag, _tools_section(_call("weather", 3)))
|
||||
|
||||
|
||||
def test_auto_hook_rejects_unknown_or_unclosed_calls():
|
||||
structural_tag = get_kimik3_auto_tool_call_structural_tag([_tool(strict=False)])
|
||||
assert structural_tag is not None
|
||||
call = _call(
|
||||
"weather",
|
||||
49,
|
||||
_argument("city", "string", "Paris"),
|
||||
_argument("days", "number", "3"),
|
||||
)
|
||||
unknown = _call(
|
||||
"forecast",
|
||||
49,
|
||||
_argument("city", "string", "Paris"),
|
||||
_argument("days", "number", "3"),
|
||||
)
|
||||
|
||||
assert not _token_accepts(structural_tag, _tools_section(unknown))
|
||||
assert not _token_accepts(
|
||||
structural_tag, TOOLS_OPEN + call.removesuffix(CALL_CLOSE)
|
||||
)
|
||||
assert not _token_accepts(structural_tag, TOOLS_OPEN + call)
|
||||
|
||||
|
||||
def test_auto_hook_does_not_swallow_parser_visible_closes():
|
||||
structural_tag = get_kimik3_auto_tool_call_structural_tag([_tool(strict=False)])
|
||||
assert structural_tag is not None
|
||||
output = (
|
||||
TOOLS_OPEN
|
||||
+ '<|open|>call tool="weather" index="23"<|sep|>'
|
||||
+ _argument("city", "string", "Paris")
|
||||
+ '<|open|>argument key="days" type="number"<|sep|>'
|
||||
+ ARGUMENT_CLOSE
|
||||
+ CALL_CLOSE
|
||||
+ "3"
|
||||
+ ARGUMENT_CLOSE
|
||||
+ CALL_CLOSE
|
||||
+ TOOLS_CLOSE
|
||||
)
|
||||
|
||||
parsed = KimiK3Detector().detect_and_parse(output, [_tool(strict=False)])
|
||||
|
||||
assert json.loads(parsed.calls[0].parameters) == {"city": "Paris", "days": ""}
|
||||
assert not _token_accepts(structural_tag, output)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"tool_strict_level",
|
||||
[
|
||||
ToolStrictLevel.OFF,
|
||||
ToolStrictLevel.FUNCTION,
|
||||
ToolStrictLevel.PARAMETER,
|
||||
],
|
||||
)
|
||||
def test_parallel_tool_calls_false_rejects_second_call(tool_strict_level):
|
||||
with envs.SGLANG_TOOL_STRICT_LEVEL.override(tool_strict_level.value):
|
||||
constraint = FunctionCallParser(
|
||||
[_tool(strict=False)], "kimi_k3"
|
||||
).get_structure_constraint("auto", parallel_tool_calls=False)
|
||||
assert constraint is not None
|
||||
first = _call(
|
||||
"weather",
|
||||
165,
|
||||
_argument("city", "string", "Paris"),
|
||||
_argument("days", "number", "3"),
|
||||
)
|
||||
second = _call(
|
||||
"weather",
|
||||
166,
|
||||
_argument("city", "string", "London"),
|
||||
_argument("days", "number", "2"),
|
||||
)
|
||||
|
||||
assert _token_accepts(constraint[1], _tools_section(first))
|
||||
assert not _token_accepts(constraint[1], _tools_section(first, second))
|
||||
assert not _token_accepts(
|
||||
constraint[1], _tools_section(first) + _tools_section(second)
|
||||
)
|
||||
|
||||
|
||||
def test_parallel_tool_calls_true_constrains_every_call():
|
||||
grammar = _grammar(
|
||||
[_tool(strict=False)],
|
||||
parallel_tool_calls=True,
|
||||
)
|
||||
first = _call(
|
||||
"weather",
|
||||
7,
|
||||
_argument("city", "string", "Paris"),
|
||||
)
|
||||
second = _call(
|
||||
"weather",
|
||||
19,
|
||||
_argument("other", "array", '["loose"]'),
|
||||
)
|
||||
|
||||
assert _accepts(grammar, _tools_section(first, second))
|
||||
|
||||
|
||||
def test_parameter_level_constrains_every_parallel_call():
|
||||
with envs.SGLANG_TOOL_STRICT_LEVEL.override(ToolStrictLevel.PARAMETER.value):
|
||||
constraint = FunctionCallParser(
|
||||
[_tool(strict=False)], "kimi_k3"
|
||||
).get_structure_constraint("auto")
|
||||
assert constraint is not None
|
||||
grammar = xgr.Grammar.from_structural_tag(constraint[1])
|
||||
first = _valid_weather_call(index=3)
|
||||
valid_second = _call(
|
||||
"weather",
|
||||
49,
|
||||
_argument("city", "string", "London"),
|
||||
_argument("days", "number", "2"),
|
||||
)
|
||||
invalid_second = _call(
|
||||
"weather",
|
||||
49,
|
||||
_argument("city", "string", "london"),
|
||||
_argument("days", "number", "99"),
|
||||
)
|
||||
|
||||
assert _accepts(grammar, _tools_section(first, valid_second))
|
||||
assert not _accepts(grammar, _tools_section(first, invalid_second))
|
||||
|
||||
|
||||
def test_strict_tool_without_parameters_compiles_to_empty_arguments():
|
||||
"""SGLANG_TOOL_STRICT_LEVEL=2 marks every tool strict, including tools
|
||||
that declare no parameters; the grammar build must not fail for them."""
|
||||
tool = Tool(type="function", function=Function(name="ping", strict=True))
|
||||
grammar = _grammar([tool], tool_choice="required")
|
||||
|
||||
assert _accepts(grammar, _tools_section(_call("ping", 1)))
|
||||
assert not _accepts(
|
||||
grammar,
|
||||
_tools_section(_call("ping", 1, _argument("x", "string", "y"))),
|
||||
)
|
||||
|
||||
|
||||
def test_parameter_level_keeps_constraint_with_no_parameter_tool():
|
||||
"""A single no-parameter tool must not poison the whole request: a build
|
||||
error here was swallowed and required fell back to a JSON-only grammar
|
||||
the K3 parser cannot read."""
|
||||
tools = [
|
||||
_tool(strict=False),
|
||||
Tool(type="function", function=Function(name="ping")),
|
||||
]
|
||||
with envs.SGLANG_TOOL_STRICT_LEVEL.override(ToolStrictLevel.PARAMETER.value):
|
||||
constraint = FunctionCallParser(tools, "kimi_k3").get_structure_constraint(
|
||||
"required"
|
||||
)
|
||||
|
||||
assert constraint is not None
|
||||
assert constraint[0] == "structural_tag"
|
||||
|
||||
|
||||
def test_all_of_number_branches_do_not_narrow_to_integer():
|
||||
"""allOf with only number branches was intersected down to integer,
|
||||
silently dropping non-integer enum values."""
|
||||
tool = Tool(
|
||||
type="function",
|
||||
function=Function(
|
||||
name="scale",
|
||||
strict=True,
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"factor": {"allOf": [{"type": "number"}], "enum": [1.5, 2]},
|
||||
},
|
||||
"required": ["factor"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
),
|
||||
)
|
||||
grammar = _grammar([tool], tool_choice="required")
|
||||
|
||||
assert _accepts(
|
||||
grammar,
|
||||
_tools_section(_call("scale", 1, _argument("factor", "number", "1.5"))),
|
||||
)
|
||||
assert not _accepts(
|
||||
grammar,
|
||||
_tools_section(_call("scale", 1, _argument("factor", "number", "3"))),
|
||||
)
|
||||
|
||||
|
||||
def test_auto_hook_serializes_into_sampling_parameters():
|
||||
constraint = FunctionCallParser(
|
||||
[_tool(strict=False)], "kimi_k3"
|
||||
).get_structure_constraint("auto")
|
||||
assert constraint is not None
|
||||
request = ChatCompletionRequest(
|
||||
model="test",
|
||||
messages=[{"role": "user", "content": "Weather?"}],
|
||||
max_completion_tokens=16,
|
||||
)
|
||||
|
||||
sampling_params = request.to_sampling_params(
|
||||
stop=[],
|
||||
model_generation_config={},
|
||||
tool_call_constraint=constraint,
|
||||
)
|
||||
|
||||
serialized = json.loads(sampling_params["structural_tag"])
|
||||
assert serialized["type"] == "structural_tag"
|
||||
assert serialized["format"]["type"] == "triggered_tags"
|
||||
|
||||
|
||||
def test_auto_hook_forces_one_typed_property_when_none_are_required():
|
||||
tool = Tool(
|
||||
type="function",
|
||||
function=Function(
|
||||
name="search",
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"query": {"type": "string"},
|
||||
"limit": {"type": "integer"},
|
||||
},
|
||||
},
|
||||
),
|
||||
)
|
||||
structural_tag = get_kimik3_auto_tool_call_structural_tag([tool])
|
||||
assert structural_tag is not None
|
||||
|
||||
assert _token_accepts(
|
||||
structural_tag,
|
||||
_tools_section(_call("search", 1, _argument("limit", "number", "3"))),
|
||||
)
|
||||
assert not _token_accepts(structural_tag, _tools_section(_call("search", 1)))
|
||||
|
||||
|
||||
def test_auto_hook_keeps_structure_for_ambiguous_required_argument_type():
|
||||
tool = Tool(
|
||||
type="function",
|
||||
function=Function(
|
||||
name="lookup",
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"key": {"type": ["string", "integer"]},
|
||||
},
|
||||
"required": ["key"],
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
structural_tag = get_kimik3_auto_tool_call_structural_tag([tool])
|
||||
assert structural_tag is not None
|
||||
grammar = xgr.Grammar.from_structural_tag(structural_tag)
|
||||
|
||||
assert _accepts(
|
||||
grammar,
|
||||
_tools_section(_call("lookup", 9254, _argument("key", "number", "3"))),
|
||||
)
|
||||
assert not _accepts(
|
||||
grammar,
|
||||
TOOLS_OPEN + _call("lookup", 9254, _argument("key", "number", "3")),
|
||||
)
|
||||
|
||||
|
||||
def test_parameter_level_applies_to_other_model_native_tags():
|
||||
with envs.SGLANG_TOOL_STRICT_LEVEL.override(ToolStrictLevel.PARAMETER.value):
|
||||
constraint = FunctionCallParser(
|
||||
[_tool(strict=False)], "kimi_k2"
|
||||
).get_structure_constraint("auto")
|
||||
|
||||
assert constraint is not None
|
||||
serialized = constraint[1].model_dump_json()
|
||||
assert '"properties"' in serialized
|
||||
assert '"city"' in serialized
|
||||
|
||||
|
||||
def test_reasoning_prefix_is_owned_by_exactly_one_layer():
|
||||
tool = _tool()
|
||||
wrapped_by_xgrammar = _grammar([tool], tool_choice="required", thinking_mode=True)
|
||||
post_reasoning_only = _grammar([tool], tool_choice="required", thinking_mode=False)
|
||||
output = "reasoning" + THINK_CLOSE + _tools_section(_valid_weather_call())
|
||||
|
||||
assert _accepts(wrapped_by_xgrammar, output)
|
||||
assert not _accepts(post_reasoning_only, output)
|
||||
assert _accepts(post_reasoning_only, _tools_section(_valid_weather_call()))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
@@ -66,7 +66,7 @@ def _make_processor() -> SchedulerBatchResultProcessor:
|
||||
enable_overlap=False,
|
||||
enable_overlap_mlx=False,
|
||||
server_args=SimpleNamespace(enable_metrics=False),
|
||||
model_config=SimpleNamespace(think_end_id=None),
|
||||
model_config=SimpleNamespace(think_end_ids=None),
|
||||
token_to_kv_pool_allocator=None,
|
||||
tree_cache=None,
|
||||
hisparse_coordinator=None,
|
||||
@@ -134,5 +134,19 @@ class TestSpecV2GrammarTruncation(CustomTestCase):
|
||||
self.assertEqual(req.kv_committed_len, 3)
|
||||
|
||||
|
||||
class TestReasoningTokenAccounting(CustomTestCase):
|
||||
def test_multi_token_end_can_span_decode_steps(self):
|
||||
req = _make_req(terminate_after=99)
|
||||
req.require_reasoning = True
|
||||
processor = _make_processor()
|
||||
processor.model_config.think_end_ids = [7, 8]
|
||||
|
||||
processor._maybe_update_reasoning_tokens(req, [10, 7])
|
||||
processor._maybe_update_reasoning_tokens(req, [8, 11])
|
||||
|
||||
self.assertEqual(req.reasoning_tokens, 3)
|
||||
self.assertTrue(req._is_reasoning_over)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -0,0 +1,162 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
from sglang.srt.function_call.kimik3_format import (
|
||||
MESSAGE_CLOSE,
|
||||
RESPONSE_CLOSE,
|
||||
RESPONSE_OPEN,
|
||||
THINK_CLOSE,
|
||||
THINK_OPEN,
|
||||
TOOLS_CLOSE,
|
||||
TOOLS_OPEN,
|
||||
)
|
||||
from sglang.srt.parser.reasoning_parser import KimiK3Detector, ReasoningParser
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
def _stream(detector: KimiK3Detector, chunks: list[str]) -> tuple[str, str]:
|
||||
reasoning = ""
|
||||
content = ""
|
||||
for chunk in chunks:
|
||||
result = detector.parse_streaming_increment(chunk)
|
||||
reasoning += result.reasoning_text
|
||||
content += result.normal_text
|
||||
return reasoning, content
|
||||
|
||||
|
||||
def _chunks(text: str, size: int) -> list[str]:
|
||||
return [text[index : index + size] for index in range(0, len(text), size)]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "reasoning", "content"),
|
||||
[
|
||||
(
|
||||
f"{THINK_OPEN}deep thought{THINK_CLOSE}"
|
||||
f"{RESPONSE_OPEN}the answer{RESPONSE_CLOSE}{MESSAGE_CLOSE}",
|
||||
"deep thought",
|
||||
"the answer",
|
||||
),
|
||||
(
|
||||
f"thinking...{THINK_CLOSE}{RESPONSE_OPEN}done{RESPONSE_CLOSE}",
|
||||
"thinking...",
|
||||
"done",
|
||||
),
|
||||
(
|
||||
f"{RESPONSE_OPEN}plain reply{RESPONSE_CLOSE}{MESSAGE_CLOSE}",
|
||||
"",
|
||||
"plain reply",
|
||||
),
|
||||
("still going", "still going", ""),
|
||||
],
|
||||
)
|
||||
def test_non_stream_reasoning_channels(text: str, reasoning: str, content: str) -> None:
|
||||
detector = KimiK3Detector(force_reasoning=True)
|
||||
result = detector.detect_and_parse(text)
|
||||
assert result.reasoning_text == reasoning
|
||||
assert result.normal_text == content
|
||||
|
||||
|
||||
def test_non_stream_tools_channel_passthrough() -> None:
|
||||
tools_channel = (
|
||||
f'{TOOLS_OPEN}<|open|>call tool="python" index="1"<|sep|>'
|
||||
"<|close|>call<|sep|>"
|
||||
f"{TOOLS_CLOSE}"
|
||||
)
|
||||
detector = KimiK3Detector(force_reasoning=True)
|
||||
result = detector.detect_and_parse(
|
||||
f"thought{THINK_CLOSE}{RESPONSE_OPEN}reply{RESPONSE_CLOSE}{tools_channel}"
|
||||
)
|
||||
assert result.reasoning_text == "thought"
|
||||
assert result.normal_text == f"reply{tools_channel}"
|
||||
|
||||
|
||||
def test_non_stream_recovers_missing_think_separator() -> None:
|
||||
detector = KimiK3Detector(force_reasoning=True)
|
||||
result = detector.detect_and_parse(
|
||||
f"thought{THINK_CLOSE.removesuffix('<|sep|>')}{RESPONSE_OPEN}"
|
||||
f"reply{RESPONSE_CLOSE}"
|
||||
)
|
||||
assert result.reasoning_text == "thought"
|
||||
assert result.normal_text == "reply"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "reasoning", "content"),
|
||||
[
|
||||
("deep thought<|close|>", "deep thought", ""),
|
||||
("deep thought<|close|>think", "deep thought", ""),
|
||||
(f"{THINK_CLOSE}<|open|>", "", ""),
|
||||
(f"{THINK_CLOSE}<|open|>response", "", ""),
|
||||
(
|
||||
f"{THINK_CLOSE}{RESPONSE_OPEN}the answer<|close|>response",
|
||||
"",
|
||||
"the answer",
|
||||
),
|
||||
(
|
||||
f"{THINK_CLOSE}{RESPONSE_OPEN}the answer{RESPONSE_CLOSE}<|close|>message",
|
||||
"",
|
||||
"the answer",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_non_stream_strips_partial_marker_suffixes(
|
||||
text: str, reasoning: str, content: str
|
||||
) -> None:
|
||||
result = KimiK3Detector(force_reasoning=True).detect_and_parse(text)
|
||||
assert result.reasoning_text == reasoning
|
||||
assert result.normal_text == content
|
||||
|
||||
|
||||
def test_non_stream_preserves_non_marker_angle_bracket_suffix() -> None:
|
||||
result = KimiK3Detector(force_reasoning=True).detect_and_parse(
|
||||
f"{THINK_CLOSE}{RESPONSE_OPEN}answer <3"
|
||||
)
|
||||
assert result.normal_text == "answer <3"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("chunk_size", [1, 4, 13])
|
||||
def test_streaming_split_markers(chunk_size: int) -> None:
|
||||
detector = KimiK3Detector(force_reasoning=True)
|
||||
text = (
|
||||
f"{THINK_OPEN}deep thought{THINK_CLOSE}"
|
||||
f"{RESPONSE_OPEN}the answer{RESPONSE_CLOSE}{MESSAGE_CLOSE}"
|
||||
)
|
||||
reasoning, content = _stream(detector, _chunks(text, chunk_size))
|
||||
assert reasoning == "deep thought"
|
||||
assert content == "the answer"
|
||||
|
||||
|
||||
def test_streaming_tools_channel_passthrough() -> None:
|
||||
tools_channel = (
|
||||
f'{TOOLS_OPEN}<|open|>call tool="python" index="1"<|sep|>'
|
||||
"<|close|>call<|sep|>"
|
||||
f"{TOOLS_CLOSE}"
|
||||
)
|
||||
detector = KimiK3Detector(force_reasoning=True)
|
||||
text = f"thought{THINK_CLOSE}{RESPONSE_OPEN}reply{RESPONSE_CLOSE}{tools_channel}"
|
||||
reasoning, content = _stream(detector, _chunks(text, 5))
|
||||
assert reasoning == "thought"
|
||||
assert content == f"reply{tools_channel}"
|
||||
|
||||
|
||||
def test_streaming_recovers_missing_think_separator() -> None:
|
||||
detector = KimiK3Detector(force_reasoning=True)
|
||||
text = (
|
||||
f"thought{THINK_CLOSE.removesuffix('<|sep|>')}{RESPONSE_OPEN}"
|
||||
f"reply{RESPONSE_CLOSE}"
|
||||
)
|
||||
reasoning, content = _stream(detector, _chunks(text, 3))
|
||||
assert reasoning == "thought"
|
||||
assert content == "reply"
|
||||
|
||||
|
||||
def test_reasoning_parser_registration() -> None:
|
||||
assert isinstance(ReasoningParser("kimi_k3").detector, KimiK3Detector)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
@@ -15,7 +15,7 @@ from sglang.srt.parser.template_detection import (
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(2.0, "base-a-test-cpu")
|
||||
register_cpu_ci(est_time=2.0, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
class _DummyTokenizer:
|
||||
@@ -872,6 +872,34 @@ class TestResolveAutoParsers(unittest.TestCase):
|
||||
self.assertEqual(args.reasoning_parser, "deepseek-v4")
|
||||
self.assertEqual(args.tool_call_parser, "deepseekv4")
|
||||
|
||||
def test_kimi_k3_arch_without_chat_template_uses_custom_encoder(self):
|
||||
args = self._make_server_args(reasoning_parser="auto", tool_call_parser="auto")
|
||||
tokenizer = _DummyTokenizer([])
|
||||
config = SimpleNamespace(
|
||||
architectures=["KimiK3ForConditionalGeneration"], model_type="kimi_k3"
|
||||
)
|
||||
|
||||
with _patch_hf_transformers_utils(
|
||||
Mock(return_value=tokenizer), Mock(return_value=config)
|
||||
):
|
||||
resolve_auto_parsers(args)
|
||||
|
||||
self.assertEqual(args.reasoning_parser, "kimi_k3")
|
||||
self.assertEqual(args.tool_call_parser, "kimi_k3")
|
||||
|
||||
def test_kimi_k3_model_type_without_architecture_uses_custom_encoder(self):
|
||||
args = self._make_server_args(reasoning_parser="auto", tool_call_parser="auto")
|
||||
tokenizer = _DummyTokenizer([])
|
||||
config = SimpleNamespace(architectures=None, model_type="kimi_k3")
|
||||
|
||||
with _patch_hf_transformers_utils(
|
||||
Mock(return_value=tokenizer), Mock(return_value=config)
|
||||
):
|
||||
resolve_auto_parsers(args)
|
||||
|
||||
self.assertEqual(args.reasoning_parser, "kimi_k3")
|
||||
self.assertEqual(args.tool_call_parser, "kimi_k3")
|
||||
|
||||
def test_deepseek_arch_fallback_runs_when_tokenizer_load_fails(self):
|
||||
args = self._make_server_args(reasoning_parser="auto", tool_call_parser="auto")
|
||||
config = SimpleNamespace(architectures=["DeepseekV32ForCausalLM"])
|
||||
|
||||
Reference in New Issue
Block a user