@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / c91fe7b1b9

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 09:45 UTC

[Frontend][Bug Fix] Update llama4 pythonic jinja template and llama4_pythonic parser (#17917)

Signed-off-by: Kai Wu <kaiwu@meta.com>

Kai Wu committed May 22, 2025 at 16:44 UTC c91fe7b1b9c4398c6d4c980fc480ada0da8a0b23

6 files changed +541 -72

docs/source/features/tool_calling.md

+5 -6

@@ -158,13 +158,13 @@ All Llama 3.1, 3.2 and 4 models should be supported.

158 * `meta-llama/Llama-3.2-*`

159 * `meta-llama/Llama-4-*`

160

161 -The tool calling that is supported is the [JSON based tool calling](https://llama.meta.com/docs/model-cards-and-prompt-formats/llama3_1/#json-based-tool-calling). For [pythonic tool calling](https://github.com/meta-llama/llama-models/blob/main/models/llama3_2/text_prompt_format.md#zero-shot-function-calling) introduced by the Llama-3.2 models, see the `pythonic` tool parser below.

161 +The tool calling that is supported is the [JSON based tool calling](https://llama.meta.com/docs/model-cards-and-prompt-formats/llama3_1/#json-based-tool-calling). For [pythonic tool calling](https://github.com/meta-llama/llama-models/blob/main/models/llama3_2/text_prompt_format.md#zero-shot-function-calling) introduced by the Llama-3.2 models, see the `pythonic` tool parser below. As for llama 4 models, it is recommended to use the `llama4_pythonic` tool parser.

162

163 Other tool calling formats like the built in python tool calling or custom tool calling are not supported.

164

165 Known issues:

166

167 -1. Parallel tool calls are not supported.

167 +1. Parallel tool calls are not supported for llama 3, but it is supported in llama 4 models.

168 2. The model can generate parameters with a wrong format, such as generating

169 an array serialized as string instead of an array.

170

@@ -177,11 +177,10 @@ images.

177

178 Recommended flags: `--tool-call-parser llama3_json --chat-template {see_above}`

179

180 -VLLM also provides a JSON based chat template for Llama 4:

181 -* <gh-file:examples/tool_chat_template_llama4_json.jinja> - this is based on the "official" chat template for the Llama 4

182 -models, but tweaked so that it works better with vLLM.

180 +VLLM also provides a pythonic and JSON based chat template for Llama 4, but pythonic tool calling is recommended:

181 +* <gh-file:examples/tool_chat_template_llama4_pythonic.jinja> - this is based on the [official chat template](https://www.llama.com/docs/model-cards-and-prompt-formats/llama4/) for the Llama 4 models.

182

184 -For Llama 4 use `--tool-call-parser llama4_json examples/tool_chat_template_llama4_json.jinja`.

183 +For Llama 4 model, use `--tool-call-parser llama4_pythonic --chat-template examples/tool_chat_template_llama4_pythonic.jinja`.

184

185 #### IBM Granite

186

examples/tool_chat_template_llama4_pythonic.jinja

+36 -64

@@ -1,16 +1,17 @@

1 {{- bos_token }}

2 -{%- if custom_tools is defined %}

2 +{%- if custom_tools is defined and custom_tools%}

3 {%- set tools = custom_tools %}

4 {%- endif %}

5 -{%- if not tools_in_user_message is defined %}

6 - {%- set tools_in_user_message = false %}

7 -{%- endif %}

8 -{%- if not tools is defined %}

5 +{%- if tools is defined and tools %}

6 + {%- set tool_definition = tool_definition ~ (tools | tojson(indent=4)) %}

7 +{%- else %}

8 {%- set tools = none %}

9 {%- endif %}

11 +

12 {#- This block extracts the system message, so we can slot it into the right place. #}

13 {%- if messages[0]['role'] == 'system' %}

14 + {%- set user_provided_system_message = true %}

15 {%- if messages[0]['content'] is string %}

16 {%- set system_message = messages[0]['content']|trim %}

17 {%- else %}

@@ -18,68 +19,33 @@

19 {%- endif %}

20 {%- set messages = messages[1:] %}

21 {%- else %}

21 - {%- if tools is not none %}

22 - {#- Add default tool system message when tools are provided #}

23 - {%- set system_message = "You are a helpful assistant with tool calling "

24 - "capabilities. Only reply with a tool call if the function exists in the "

25 - "library provided by the user. If it doesn't exist, just reply directly in "

26 - "natural language. When you receive a tool call response, use the output to "

27 - "format an answer to the original user question." %}

22 + {%- if tools is not none %}

23 + {#- Since not system_message was provided by user, if tool is provided, system_message is now default tool system message #}

24 + {#- This system message is from llama website:https://www.llama.com/docs/model-cards-and-prompt-formats/llama4/ #}

25 + {%- set system_message = "You are a helpful assistant and an expert in function composition. You can answer general questions using your internal knowledge OR invoke functions when necessary. Follow these strict guidelines:\n\n1. FUNCTION CALLS:\n- ONLY use functions that are EXPLICITLY listed in the function list below\n- If NO functions are listed (empty function list []), respond ONLY with internal knowledge or \"I don't have access to [Unavailable service] information\"\n- If a function is not in the list, respond ONLY with internal knowledge or \"I don't have access to [Unavailable service] information\"\n- If ALL required parameters are present AND the query EXACTLY matches a listed function's purpose: output ONLY the function call(s)\n- Use exact format: [func_name1(param1=value1, param2=value2), func_name2(...)]\nExamples:\nCORRECT: [get_weather(location=\"Vancouver\"), calculate_route(start=\"Boston\", end=\"New York\")] <- Only if get_weather and calculate_route are in function list\nINCORRECT: get_weather(location=\"New York\")\nINCORRECT: Let me check the weather: [get_weather(location=\"New York\")]\nINCORRECT: [get_events(location=\"Singapore\")] <- If function not in list\n\n2. RESPONSE RULES:\n- For pure function requests matching a listed function: ONLY output the function call(s)\n- For knowledge questions: ONLY output text\n- For missing parameters: ONLY request the specific missing parameters\n- For unavailable services (not in function list): output ONLY with internal knowledge or \"I don't have access to [Unavailable service] information\". Do NOT execute a function call.\n- If the query asks for information beyond what a listed function provides: output ONLY with internal knowledge about your limitations\n- NEVER combine text and function calls in the same response\n- NEVER suggest alternative functions when the requested service is unavailable\n- NEVER create or invent new functions not listed below\n\n3. STRICT BOUNDARIES:\n- ONLY use functions from the list below - no exceptions\n- NEVER use a function as an alternative to unavailable information\n- NEVER call functions not present in the function list\n- NEVER add explanatory text to function calls\n- NEVER respond with empty brackets\n- Use proper Python/JSON syntax for function calls\n- Check the function list carefully before responding\n\n4. TOOL RESPONSE HANDLING:\n- When receiving tool responses: provide concise, natural language responses\n- Don't repeat tool response verbatim\n- Don't add supplementary information\n\nHere is a list of functions in JSON format that you can invoke:\n" %}

26 {%- else %}

27 {%- set system_message = "" %}

28 {%- endif %}

29 {%- endif %}

32 -

33 -{#- System message if the user supplied one, or if tools are used (default tool system message) #}

30 +{#- Now writing the system message: use the user provided system message if user_provided_system_message, else default tool system message if tools presented #}

31 {%- if system_message %}

32 {#- always use user provided system message to override default tool system message #}

33 {{- "<|header_start|>system<|header_end|>\n\n" }}

34 {{- system_message }}

38 - {%- if tools is not none and not tools_in_user_message %}

39 - {{- "Tools: You have access to the following tools. You might need to use one "

40 - "or more function/tool calls to fulfill the task. \n"

41 - "If none are needed, then proceed to the response.\n\n"

42 - "Tool Call Syntax: You can call tools using the following syntax:\n"

43 - "[func_name1(params_name1=params_value1, params_name2=params_value2, ...), ...]\n"

44 - "Do not include anything else when calling the tools with the syntax above.\n\n"

45 - "Here is a list of functions in JSON format that you can invoke.\n " }}

46 - {%- for t in tools %}

47 - {{- t | tojson(indent=4) }}

48 - {{- "\n\n" }}

49 - {%- endfor %}

35 + {%- if user_provided_system_message and tools %}

36 + {{- "\nHere is a list of functions in JSON format that you can invoke. Use exact format: [func_name1(param1=value1, param2=value2), func_name2(...)]\n" }}

37 + {{- tool_definition -}}

38 + {%- elif tool_definition %}

39 + {{- tool_definition -}}

40 {%- endif %}

41 {{- "<|eot|>" }}

42 {%- endif %}

54 -{#- Custom tools are passed in a user message with some extra guidance #}

55 -{%- if tools_in_user_message and tools is not none %}

56 - {#- Extract the first user message so we can plug it in here #}

57 - {%- if messages | length != 0 %}

58 - {%- if messages[0]['content'] is string %}

59 - {%- set first_user_message = messages[0]['content']|trim %}

60 - {%- else %}

61 - {%- set first_user_message = messages[0]['content'] | selectattr('type', 'equalto', 'text') | map(attribute='text') | map('trim') | join('\n') %}

62 - {%- endif %}

63 - {%- set messages = messages[1:] %}

64 - {%- else %}

65 - {{- raise_exception("Cannot put tools in the first user message when there's no first user message!") }}

66 - {%- endif %}

67 - {{- '<|header_start|>user<|header_end|>\n\n' -}}

68 - {{- first_user_message}}

69 - {{- "\nHere is a list of functions in JSON format that you can invoke:"}}

70 - {%- for t in tools %}

71 - {{- t | tojson(indent=4) }}

72 - {{- "\n\n" }}

73 - {%- endfor %}

74 - {{- "Should you decide to return the function call(s), put them in the format "

75 - "of [func_name1(params_name1=params_value1, params_name2=params_value2, "

76 - "...), ...]\nDo not include anything else when calling the tools with the "

77 - "syntax above." }}

78 -{%- endif %}

79 -

44 +{#- Now deal with all other messages #}

45 {%- for message in messages %}

81 - {%- if not (message.role == 'ipython' or message.role == 'tool' or 'tool_calls' in message) %}

82 - {{- '<|header_start|>' + message['role'] + '<|header_end|>\n\n' }}

46 + {#- Base case: messages that are not from tool role and has empty tool_call list #}

47 + {%- if not (message.role == 'ipython' or message.role == 'tool' or ('tool_calls' in message and message.tool_calls|length != 0 )) %}

48 + {{- '<|header_start|>' + message['role'] + '<|header_end|>\n\n' }}

49 {%- if message['content'] is string %}

50 {{- message['content'] }}

51 {%- else %}

@@ -91,10 +57,12 @@

57 {%- endif %}

58 {%- endfor %}

59 {%- endif %}

94 - {{- "<|eot|>" }}

95 - {%- elif 'tool_calls' in message and message.tool_calls|length > 0 %}

96 - {%- set tool_call = message.tool_calls[0].function %}

97 - {{- '<|header_start|>assistant<|header_end|>\n\n' -}}

60 + {{- "<|eot|>" }}

61 + {#- Tool case: messages has non-empty tool_call list, must from assistant #}

62 + {%- elif 'tool_calls' in message %}

63 + {#- assume tool_calls are always coming from assistant #}

64 + {%- if message.role == 'assistant' %}

65 + {{- '<|header_start|>assistant<|header_end|>\n\n' -}}

66 {%- if message['content'] is string %}

67 {{- message['content'] }}

68 {%- else %}

@@ -106,32 +74,36 @@

74 {%- endif %}

75 {%- endfor %}

76 {%- endif %}

77 + {{- "[" }}

78 {%- for tool_call in message.tool_calls %}

79 {%- if tool_call.function is defined %}

80 {%- set tool_call = tool_call.function %}

81 {%- endif %}

113 - {{- tool_call.name + '(' -}}

82 + {{- tool_call.name + '(' -}}

83 {%- for param in tool_call.arguments %}

115 - {{- param + '=' -}}

84 + {{- param + '="' -}}

85 {{- "%s" | format(tool_call.arguments[param]) -}}

86 + {{- '"' -}}

87 {% if not loop.last %}, {% endif %}

88 {%- endfor %}

89 {{- ')' -}}

90 {% if not loop.last %}, {% endif %}

91 {%- endfor %}

122 - {{- "<|eom|>" }}

92 + {{- "]<|eot|>" }}

93 +{%- endif %}

94 +{#- Tool_response case: messages are from tool_response #}

95 {%- elif message.role == "tool" or message.role == "ipython" %}

96 {{- "<|header_start|>ipython<|header_end|>\n\n" }}

97 {%- if message.content is string %}

126 - {{- message.content | tojson }}

98 + {{- message.content | tojson }}

99 {%- else %}

100 {%- for content in message['content'] %}

101 {%- if content['type'] == 'text' %}

130 - {{- content['text'] | tojson }}

102 + {{- content['text'] | tojson }}

103 {%- endif %}

104 {%- endfor %}

105 {%- endif %}

134 - {{- "<|eom|>" }}

106 + {{- "<|eot|>" }}

107 {%- endif %}

108 {%- endfor %}

109 {%- if add_generation_prompt %}

tests/entrypoints/openai/tool_parsers/test_llama4_pythonic_tool_parser.py new

+193

@@ -0,0 +1,193 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +

3 +from unittest.mock import MagicMock

4 +

5 +import pytest

6 +

7 +from tests.entrypoints.openai.tool_parsers.utils import (

8 + run_tool_extraction, run_tool_extraction_streaming)

9 +from vllm.entrypoints.openai.protocol import FunctionCall

10 +from vllm.entrypoints.openai.tool_parsers import ToolParser, ToolParserManager

11 +

12 +# Test cases similar to pythonic parser but with Llama4 specific format

13 +SIMPLE_FUNCTION_OUTPUT = "[get_weather(city='LA', metric='C')]"

14 +SIMPLE_FUNCTION_CALL = FunctionCall(

15 + name="get_weather",

16 + arguments='{"city": "LA", "metric": "C"}',

17 +)

18 +MORE_TYPES_FUNCTION_OUTPUT = ("[register_user(name='Doe', "

19 + "age=9, "

20 + "address={'city': 'LA', 'state': 'CA'}, "

21 + "role=None, "

22 + "passed_test=True, "

23 + "aliases=['John', 'Johnny'])]")

24 +MORE_TYPES_FUNCTION_CALL = FunctionCall(

25 + name="register_user",

26 + arguments='{"name": "Doe", '

27 + '"age": 9, '

28 + '"address": {"city": "LA", "state": "CA"}, '

29 + '"role": null, '

30 + '"passed_test": true, '

31 + '"aliases": ["John", "Johnny"]}',

32 +)

33 +PARAMETERLESS_FUNCTION_OUTPUT = "[get_weather()]"

34 +PARAMETERLESS_FUNCTION_CALL = FunctionCall(

35 + name="get_weather",

36 + arguments='{}',

37 +)

38 +EMPTY_DICT_FUNCTION_OUTPUT = "[do_something_cool(additional_data={})]"

39 +EMPTY_DICT_FUNCTION_CALL = FunctionCall(

40 + name="do_something_cool",

41 + arguments='{"additional_data": {}}',

42 +)

43 +EMPTY_LIST_FUNCTION_OUTPUT = "[do_something_cool(steps=[])]"

44 +EMPTY_LIST_FUNCTION_CALL = FunctionCall(

45 + name="do_something_cool",

46 + arguments='{"steps": []}',

47 +)

48 +ESCAPED_STRING_FUNCTION_OUTPUT = (

49 + r"[get_weather(city='Martha\'s Vineyard', metric='\"cool units\"')]")

50 +ESCAPED_STRING_FUNCTION_CALL = FunctionCall(

51 + name="get_weather",

52 + arguments='{"city": "Martha\'s Vineyard", "metric": "\\"cool units\\""}',

53 +)

54 +PYTHON_TAG_FUNCTION_OUTPUT = (

55 + "<|python_start|>[get_weather(city='LA', metric='C')]<|python_end|>")

56 +

57 +

58 +@pytest.mark.parametrize("streaming", [True, False])

59 +def test_no_tool_call(streaming: bool):

60 + mock_tokenizer = MagicMock()

61 + tool_parser: ToolParser = ToolParserManager.get_tool_parser(

62 + "llama4_pythonic")(mock_tokenizer)

63 + model_output = "How can I help you today?"

64 +

65 + content, tool_calls = run_tool_extraction(tool_parser,

66 + model_output,

67 + streaming=streaming)

68 +

69 + assert content == model_output

70 + assert len(tool_calls) == 0

71 +

72 +

73 +test_str = "<|python_start|>"

74 +test_str += "[get_weather(city='LA', metric='C'),"

75 +test_str += "register_user(name='Doe', age=9)]"

76 +TEST_CASES = [

77 + pytest.param(True,

78 + ESCAPED_STRING_FUNCTION_OUTPUT,

79 + [ESCAPED_STRING_FUNCTION_CALL],

80 + id="simple_streaming"),

81 + pytest.param(False,

82 + SIMPLE_FUNCTION_OUTPUT, [SIMPLE_FUNCTION_CALL],

83 + id="simple_nonstreaming"),

84 + pytest.param(True,

85 + MORE_TYPES_FUNCTION_OUTPUT, [MORE_TYPES_FUNCTION_CALL],

86 + id="more_types_streaming"),

87 + pytest.param(False,

88 + MORE_TYPES_FUNCTION_OUTPUT, [MORE_TYPES_FUNCTION_CALL],

89 + id="more_types_nonstreaming"),

90 + pytest.param(True,

91 + PARAMETERLESS_FUNCTION_OUTPUT, [PARAMETERLESS_FUNCTION_CALL],

92 + id="parameterless_streaming"),

93 + pytest.param(False,

94 + PARAMETERLESS_FUNCTION_OUTPUT, [PARAMETERLESS_FUNCTION_CALL],

95 + id="parameterless_nonstreaming"),

96 + pytest.param(True,

97 + EMPTY_DICT_FUNCTION_OUTPUT, [EMPTY_DICT_FUNCTION_CALL],

98 + id="empty_dict_streaming"),

99 + pytest.param(False,

100 + EMPTY_DICT_FUNCTION_OUTPUT, [EMPTY_DICT_FUNCTION_CALL],

101 + id="empty_dict_nonstreaming"),

102 + pytest.param(True,

103 + EMPTY_LIST_FUNCTION_OUTPUT, [EMPTY_LIST_FUNCTION_CALL],

104 + id="empty_list_streaming"),

105 + pytest.param(False,

106 + EMPTY_LIST_FUNCTION_OUTPUT, [EMPTY_LIST_FUNCTION_CALL],

107 + id="empty_list_nonstreaming"),

108 + pytest.param(True,

109 + ESCAPED_STRING_FUNCTION_OUTPUT,

110 + [ESCAPED_STRING_FUNCTION_CALL],

111 + id="escaped_string_streaming"),

112 + pytest.param(False,

113 + ESCAPED_STRING_FUNCTION_OUTPUT,

114 + [ESCAPED_STRING_FUNCTION_CALL],

115 + id="escaped_string_nonstreaming"),

116 + pytest.param(

117 + True,

118 + "[get_weather(city='LA',metric='C'),register_user(name='Doe',age=9)]",

119 + [

120 + SIMPLE_FUNCTION_CALL,

121 + FunctionCall(name="register_user",

122 + arguments='{"name": "Doe", "age": 9}')

123 + ],

124 + id="parallel_calls_streaming"),

125 + pytest.param(

126 + False,

127 + "[get_weather(city='LA',metric='C'),register_user(name='Doe',age=9)]",

128 + [

129 + SIMPLE_FUNCTION_CALL,

130 + FunctionCall(name="register_user",

131 + arguments='{"name": "Doe", "age": 9}')

132 + ],

133 + id="parallel_calls_nonstreaming"),

134 + pytest.param(True,

135 + PYTHON_TAG_FUNCTION_OUTPUT, [SIMPLE_FUNCTION_CALL],

136 + id="python_tag_streaming"),

137 + pytest.param(False,

138 + PYTHON_TAG_FUNCTION_OUTPUT, [SIMPLE_FUNCTION_CALL],

139 + id="python_tag_nonstreaming"),

140 + pytest.param(True,

141 + test_str, [

142 + SIMPLE_FUNCTION_CALL,

143 + FunctionCall(name="register_user",

144 + arguments='{"name": "Doe", "age": 9}')

145 + ],

146 + id="parallel_calls_streaming"),

147 + pytest.param(False,

148 + "<|python_start|>[get_weather(city='LA', metric='C'), " +

149 + "register_user(name='Doe', age=9)]", [

150 + SIMPLE_FUNCTION_CALL,

151 + FunctionCall(name="register_user",

152 + arguments='{"name": "Doe", "age": 9}')

153 + ],

154 + id="parallel_calls_nonstreaming"),

155 +]

156 +

157 +

158 +@pytest.mark.parametrize("streaming, model_output, expected_tool_calls",

159 + TEST_CASES)

160 +def test_tool_call(streaming: bool, model_output: str,

161 + expected_tool_calls: list[FunctionCall]):

162 + mock_tokenizer = MagicMock()

163 + tool_parser: ToolParser = ToolParserManager.get_tool_parser(

164 + "llama4_pythonic")(mock_tokenizer)

165 +

166 + content, tool_calls = run_tool_extraction(tool_parser,

167 + model_output,

168 + streaming=streaming)

169 +

170 + assert len(tool_calls) == len(expected_tool_calls)

171 + for actual, expected in zip(tool_calls, expected_tool_calls):

172 + assert actual.type == "function"

173 + assert actual.function == expected

174 +

175 +

176 +def test_streaming_tool_call_with_large_steps():

177 + mock_tokenizer = MagicMock()

178 + tool_parser: ToolParser = ToolParserManager.get_tool_parser(

179 + "llama4_pythonic")(mock_tokenizer)

180 + model_output_deltas = [

181 + "<|python_start|>[get_weather(city='LA', metric='C'), "

182 + "get_weather(), "

183 + "do_something_cool(steps=[])]<|python_end|>",

184 + ]

185 +

186 + reconstructor = run_tool_extraction_streaming(

187 + tool_parser, model_output_deltas, assert_one_tool_per_delta=False)

188 +

189 + assert reconstructor.other_content == ""

190 + assert len(reconstructor.tool_calls) == 3

191 + assert reconstructor.tool_calls[0].function == SIMPLE_FUNCTION_CALL

192 + assert reconstructor.tool_calls[1].function == PARAMETERLESS_FUNCTION_CALL

193 + assert reconstructor.tool_calls[2].function == EMPTY_LIST_FUNCTION_CALL

tests/tool_use/utils.py

+1 -1

@@ -88,7 +88,7 @@ CONFIGS: dict[str, ServerConfig] = {

88 "meta-llama/Llama-4-Scout-17B-16E-Instruct",

89 "arguments": [

90 "--enforce-eager", "--no-enable-prefix-caching",

91 - "--tool-call-parser", "pythonic", "--chat-template",

91 + "--tool-call-parser", "llama4_pythonic", "--chat-template",

92 str(VLLM_PATH /

93 "examples/tool_chat_template_llama4_pythonic.jinja"), "-tp",

94 "4"

vllm/entrypoints/openai/tool_parsers/__init__.py

+3 -1

@@ -7,6 +7,7 @@ from .granite_tool_parser import GraniteToolParser

7 from .hermes_tool_parser import Hermes2ProToolParser

8 from .internlm2_tool_parser import Internlm2ToolParser

9 from .jamba_tool_parser import JambaToolParser

10 +from .llama4_pythonic_tool_parser import Llama4PythonicToolParser

11 from .llama_tool_parser import Llama3JsonToolParser

12 from .mistral_tool_parser import MistralToolParser

13 from .phi4mini_tool_parser import Phi4MiniJsonToolParser

@@ -16,5 +17,6 @@ __all__ = [

17 "ToolParser", "ToolParserManager", "Granite20bFCToolParser",

18 "GraniteToolParser", "Hermes2ProToolParser", "MistralToolParser",

19 "Internlm2ToolParser", "Llama3JsonToolParser", "JambaToolParser",

19 - "PythonicToolParser", "Phi4MiniJsonToolParser", "DeepSeekV3ToolParser"

20 + "Llama4PythonicToolParser", "PythonicToolParser", "Phi4MiniJsonToolParser",

21 + "DeepSeekV3ToolParser"

22 ]

vllm/entrypoints/openai/tool_parsers/llama4_pythonic_tool_parser.py new

+303

@@ -0,0 +1,303 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +

3 +import ast

4 +import json

5 +import re

6 +from collections.abc import Sequence

7 +from typing import Any, Union

8 +

9 +from transformers import PreTrainedTokenizerBase

10 +

11 +from vllm.entrypoints.openai.protocol import (ChatCompletionRequest,

12 + DeltaFunctionCall, DeltaMessage,

13 + DeltaToolCall,

14 + ExtractedToolCallInformation,

15 + FunctionCall, ToolCall)

16 +from vllm.entrypoints.openai.tool_parsers.abstract_tool_parser import (

17 + ToolParser, ToolParserManager)

18 +from vllm.logger import init_logger

19 +

20 +logger = init_logger(__name__)

21 +

22 +

23 +class _UnexpectedAstError(Exception):

24 + pass

25 +

26 +

27 +@ToolParserManager.register_module("llama4_pythonic")

28 +class Llama4PythonicToolParser(ToolParser):

29 + """

30 + Toolcall parser for Llama4 that produce tool calls in a pythonic style

31 + Use --enable-auto-tool-choice --tool-call-parser llama4_pythonic

32 + """

33 + # TODO(mdepinet): Possible future improvements:

34 + # 1. Support text + tools separated by either <|python_tag|> or \n\n

35 + # 2. Support tools outside of a list (or separated by a semicolon).

36 + # This depends on item 1 for consistent streaming.

37 + # Neither of these are necessary for e.g. ToolACE, but both would help make

38 + # Llama3.2 models more reliable.

39 +

40 + TOOL_CALL_REGEX = re.compile(

41 + r"\[([a-zA-Z]+\w*\(([a-zA-Z]+\w*=.*,\s*)*([a-zA-Z]+\w*=.*\s)?\),\s*)*([a-zA-Z]+\w*\(([a-zA-Z]+\w*=.*,\s*)*([a-zA-Z]+\w*=.*\s*)?\)\s*)+\]",

42 + re.DOTALL)

43 +

44 + def __init__(self, tokenizer: PreTrainedTokenizerBase):

45 + super().__init__(tokenizer)

46 +

47 + # Rename for readability. This is NOT a tool id.

48 + @property

49 + def current_tool_index(self) -> int:

50 + return self.current_tool_id

51 +

52 + @current_tool_index.setter

53 + def current_tool_index(self, value: int) -> None:

54 + self.current_tool_id = value

55 +

56 + def extract_tool_calls(

57 + self, model_output: str,

58 + request: ChatCompletionRequest) -> ExtractedToolCallInformation:

59 + """

60 + Extract the tool calls from a complete model response.

61 + """

62 +

63 + # remove <|python_start|> and <|python_end|>

64 + # as Llama 4 model sometime will output those tokens

65 + if model_output.startswith("<|python_start|>"):

66 + model_output = model_output[len("<|python_start|>"):]

67 + model_output = model_output.replace("<|python_end|>", "")

68 + if not (self.TOOL_CALL_REGEX.match(model_output)):

69 + return ExtractedToolCallInformation(tools_called=False,

70 + tool_calls=[],

71 + content=model_output)

72 +

73 + try:

74 + module = ast.parse(model_output)

75 + parsed = getattr(module.body[0], "value", None)

76 + if isinstance(parsed, ast.List) and all(

77 + isinstance(e, ast.Call) for e in parsed.elts):

78 + return ExtractedToolCallInformation(

79 + tools_called=True,

80 + tool_calls=[

81 + _handle_single_tool(e) # type: ignore

82 + for e in parsed.elts

83 + ],

84 + content=None)

85 + else:

86 + raise _UnexpectedAstError(

87 + "Tool output must be a list of function calls")

88 + except Exception:

89 + logger.exception("Error in extracting tool call from response.")

90 + # Treat as regular text

91 + return ExtractedToolCallInformation(tools_called=False,

92 + tool_calls=[],

93 + content=model_output)

94 +

95 + def extract_tool_calls_streaming(

96 + self,

97 + previous_text: str,

98 + current_text: str,

99 + delta_text: str,

100 + previous_token_ids: Sequence[int],

101 + current_token_ids: Sequence[int],

102 + delta_token_ids: Sequence[int],

103 + request: ChatCompletionRequest,

104 + ) -> Union[DeltaMessage, None]:

105 +

106 + if not current_text.startswith("[") and not current_text.startswith(

107 + "<|python_start|>"):

108 + return DeltaMessage(content=delta_text)

109 +

110 + try:

111 + # remove <|python_start|> and <|python_end|>

112 + if current_text.startswith("<|python_start|>"):

113 + current_text = current_text[len("<|python_start|>"):]

114 + if current_text.endswith("<|python_end|>"):

115 + current_text = current_text[:current_text.

116 + rfind("<|python_end|>")]

117 + valid_and_added_text = _make_valid_python(current_text)

118 + if valid_and_added_text is None:

119 + return None

120 + valid_text, added_text = valid_and_added_text

121 +

122 + module = ast.parse(valid_text)

123 + parsed = getattr(module.body[0], "value", None)

124 + if not isinstance(parsed, ast.List) or not all(

125 + isinstance(e, ast.Call) for e in parsed.elts):

126 + raise _UnexpectedAstError(

127 + "Tool output must be a list of function calls")

128 + tool_calls = [

129 + _handle_single_tool(e) # type: ignore

130 + for e in parsed.elts

131 + ]

132 +

133 + tool_deltas = []

134 + for index, new_call in enumerate(tool_calls):

135 + if index < self.current_tool_index:

136 + continue

137 +

138 + self.current_tool_index = index

139 + if len(self.streamed_args_for_tool) == index:

140 + self.streamed_args_for_tool.append("")

141 +

142 + new_call_complete = index < len(

143 + tool_calls) - 1 or ")]" not in added_text

144 + if new_call_complete:

145 + self.current_tool_index += 1

146 +

147 + withheld_suffix = (added_text[:-2]

148 + if not new_call_complete else "")

149 + if not new_call_complete and added_text[-2] == ")":

150 + # Function call is incomplete. Withhold the closing bracket.

151 + withheld_suffix = withheld_suffix + "}"

152 + # Strings get single quotes in the model-produced string.

153 + # JSON requires double quotes.

154 + withheld_suffix = withheld_suffix.replace("'", '"')

155 + delta = _compute_tool_delta(self.streamed_args_for_tool[index],

156 + new_call, index, withheld_suffix)

157 +

158 + if delta is not None:

159 + tool_deltas.append(delta)

160 + if (delta.function is not None

161 + and delta.function.arguments is not None):

162 + self.streamed_args_for_tool[

163 + index] += delta.function.arguments

164 +

165 + # HACK: serving_chat.py inspects the internal state of tool parsers

166 + # when determining it's final streaming delta, automatically

167 + # adding autocompleted JSON.

168 + # These two lines avoid that nonsense while ensuring finish_reason

169 + # is set to tool_calls when at least one tool is called.

170 + if tool_deltas and not self.prev_tool_call_arr:

171 + self.prev_tool_call_arr = [{"arguments": {}}]

172 +

173 + if tool_deltas:

174 + return DeltaMessage(tool_calls=tool_deltas)

175 + elif not added_text and self.current_tool_id > 0:

176 + # Return an empty DeltaMessage once the tool calls are all done

177 + # so that finish_reason gets set.

178 + return DeltaMessage(content='')

179 + else:

180 + return None

181 + except Exception:

182 + logger.exception("Error trying to handle streaming tool call.")

183 + logger.debug(

184 + "Skipping chunk as a result of tool streaming extraction "

185 + "error")

186 + return None

187 +

188 +

189 +def _get_parameter_value(val: ast.expr) -> Any:

190 + if isinstance(val, ast.Constant):

191 + return val.value

192 + elif isinstance(val, ast.Dict):

193 + if not all(isinstance(k, ast.Constant) for k in val.keys):

194 + raise _UnexpectedAstError(

195 + "Dict tool call arguments must have literal keys")

196 + return {

197 + k.value: _get_parameter_value(v) # type: ignore

198 + for k, v in zip(val.keys, val.values)

199 + }

200 + elif isinstance(val, ast.List):

201 + return [_get_parameter_value(v) for v in val.elts]

202 + else:

203 + raise _UnexpectedAstError("Tool call arguments must be literals")

204 +

205 +

206 +def _handle_single_tool(call: ast.Call) -> ToolCall:

207 + if not isinstance(call.func, ast.Name):

208 + raise _UnexpectedAstError("Invalid tool call name")

209 + function_name = call.func.id

210 + arguments = {}

211 + for keyword in call.keywords:

212 + arguments[keyword.arg] = _get_parameter_value(keyword.value)

213 + return ToolCall(type="function",

214 + function=FunctionCall(name=function_name,

215 + arguments=json.dumps(arguments)))

216 +

217 +

218 +def _make_valid_python(text: str) -> Union[tuple[str, str], None]:

219 + bracket_stack = []

220 + for index, char in enumerate(text):

221 + if char in {"[", "(", "{"}:

222 + bracket_stack.append(char)

223 + elif char == "]":

224 + if not bracket_stack or bracket_stack.pop() != "[":

225 + raise _UnexpectedAstError("Mismatched square brackets")

226 + elif char == ")":

227 + if not bracket_stack or bracket_stack.pop() != "(":

228 + raise _UnexpectedAstError("Mismatched parentheses")

229 + elif char == "}":

230 + if not bracket_stack or bracket_stack.pop() != "{":

231 + raise _UnexpectedAstError("Mismatched curly braces")

232 + elif char in {"'", '"'}:

233 + if bracket_stack and bracket_stack[-1] == char:

234 + if index > 0 and text[index - 1] == "\\":

235 + # Treat an escaped quote as a regular character

236 + pass

237 + else:

238 + bracket_stack.pop()

239 + elif bracket_stack and bracket_stack[-1] in {"'", '"'}:

240 + # Double quote within a single quote string or vice versa.

241 + pass

242 + else:

243 + bracket_stack.append(char)

244 +

245 + text = text.rstrip()

246 + if text.endswith("=") or text.endswith(":"):

247 + # Since we have no type information for this property/parameter value,

248 + # we can't fill in a valid value.

249 + return None

250 + if bracket_stack and bracket_stack[-1] == "{":

251 + trailing_dict_text = text[:text.rfind("{")]

252 + num_keys = trailing_dict_text.count(":")

253 + num_values = trailing_dict_text.count(",")

254 + if num_keys <= num_values:

255 + return None # Incomplete property name within parameter value

256 + if bracket_stack and bracket_stack[-1] == "(":

257 + trailing_params_text = text[:text.rfind("(")]

258 + num_full_param_names = trailing_params_text.count("=")

259 + num_full_param_values = trailing_params_text.count(",")

260 + if num_full_param_names <= num_full_param_values:

261 + return None # Incomplete parameter name

262 + if text.endswith(","):

263 + text = text[:-1]

264 + if bracket_stack and bracket_stack[-1] == "[" and not text.endswith(

265 + "[") and not text.endswith(")"):

266 + return None # Incomplete function name

267 +

268 + added_text = ""

269 + for char in reversed(bracket_stack):

270 + if char == "[":

271 + added_text += "]"

272 + elif char == "(":

273 + added_text += ")"

274 + elif char == "{":

275 + added_text += "}"

276 + elif char == "'":

277 + added_text += "'"

278 + elif char == '"':

279 + added_text += '"'

280 +

281 + return text + added_text, added_text

282 +

283 +

284 +def _compute_tool_delta(previously_sent_args: str, new_call: ToolCall,

285 + index: int,

286 + withheld_suffix: str) -> Union[DeltaToolCall, None]:

287 + new_call_args = new_call.function.arguments

288 + if withheld_suffix:

289 + assert new_call_args.endswith(withheld_suffix)

290 + new_call_args = new_call_args[:-len(withheld_suffix)]

291 + if not previously_sent_args:

292 + return DeltaToolCall(id=new_call.id,

293 + type="function",

294 + index=index,

295 + function=DeltaFunctionCall(

296 + name=new_call.function.name,

297 + arguments=new_call_args,

298 + ))

299 +

300 + arg_diff = new_call_args[len(previously_sent_args):]

301 + return DeltaToolCall(

302 + id=None, index=index, function=DeltaFunctionCall(

303 + arguments=arg_diff)) if arg_diff else None