main
py 61 lines 1.86 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import pytest
5
6 from vllm.tokenizers import get_tokenizer
7 from vllm.tool_parsers.deepseekv31_tool_parser import (
8 DeepSeekV31ToolParser,
9 )
10
11 MODEL = "deepseek-ai/DeepSeek-V3.1"
12
13
14 @pytest.fixture(scope="module")
15 def deepseekv31_tokenizer():
16 return get_tokenizer(tokenizer_name=MODEL)
17
18
19 @pytest.fixture
20 def parser(deepseekv31_tokenizer):
21 return DeepSeekV31ToolParser(deepseekv31_tokenizer)
22
23
24 def test_extract_tool_calls_with_tool(parser):
25 model_output = (
26 "normal text"
27 "<|tool▁calls▁begin|>"
28 '<|tool▁call▁begin|>foo<|tool▁sep|>{"x":1}<|tool▁call▁end|>'
29 "<|tool▁calls▁end|>"
30 )
31 result = parser.extract_tool_calls(model_output, None)
32 assert result.tools_called
33 assert len(result.tool_calls) == 1
34 assert result.tool_calls[0].function.name == "foo"
35 assert result.tool_calls[0].function.arguments == '{"x":1}'
36 assert result.content == "normal text"
37
38
39 def test_extract_tool_calls_with_multiple_tools(parser):
40 model_output = (
41 "some prefix text"
42 "<|tool▁calls▁begin|>"
43 '<|tool▁call▁begin|>foo<|tool▁sep|>{"x":1}<|tool▁call▁end|>'
44 '<|tool▁call▁begin|>bar<|tool▁sep|>{"y":2}<|tool▁call▁end|>'
45 "<|tool▁calls▁end|>"
46 " some suffix text"
47 )
48
49 result = parser.extract_tool_calls(model_output, None)
50
51 assert result.tools_called
52 assert len(result.tool_calls) == 2
53
54 assert result.tool_calls[0].function.name == "foo"
55 assert result.tool_calls[0].function.arguments == '{"x":1}'
56
57 assert result.tool_calls[1].function.name == "bar"
58 assert result.tool_calls[1].function.arguments == '{"y":2}'
59
60 # prefix is content
61 assert result.content == "some prefix text"