| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | |
| 5 | # copy from https://huggingface.co/deepseek-ai/DeepSeek-V3.2/blob/main/encoding/encoding_dsv32.py |
| 6 | import copy |
| 7 | import json |
| 8 | from typing import Any |
| 9 | |
| 10 | # flake8: noqa: E501 |
| 11 | TOOLS_SYSTEM_TEMPLATE = """## Tools |
| 12 | You have access to a set of tools you can use to answer the user's question. |
| 13 | You can invoke functions by writing a "<{dsml_token}function_calls>" block like the following as part of your reply to the user: |
| 14 | <{dsml_token}function_calls> |
| 15 | <{dsml_token}invoke name="$FUNCTION_NAME"> |
| 16 | <{dsml_token}parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</{dsml_token}parameter> |
| 17 | ... |
| 18 | </{dsml_token}invoke> |
| 19 | <{dsml_token}invoke name="$FUNCTION_NAME2"> |
| 20 | ... |
| 21 | </{dsml_token}invoke> |
| 22 | </{dsml_token}function_calls> |
| 23 | String and scalar parameters should be specified as is without any escaping or quotes, while lists and objects should use JSON format. The "string" attribute should be set to "true" for string type parameters and "false" for other types (numbers, booleans, arrays, objects). |
| 24 | If the thinking_mode is enabled, then after function results you should strongly consider outputting a thinking block. Here is an example: |
| 25 | <{dsml_token}function_calls> |
| 26 | ... |
| 27 | </{dsml_token}function_calls> |
| 28 | <function_results> |
| 29 | ... |
| 30 | </function_results> |
| 31 | {thinking_start_token}...thinking about results{thinking_end_token} |
| 32 | Here are the functions available in JSONSchema format: |
| 33 | <functions> |
| 34 | {tool_schemas} |
| 35 | </functions> |
| 36 | """ |
| 37 | |
| 38 | bos_token: str = "<｜begin▁of▁sentence｜>" |
| 39 | eos_token: str = "<｜end▁of▁sentence｜>" |
| 40 | thinking_start_token: str = "<think>" |
| 41 | thinking_end_token: str = "</think>" |
| 42 | dsml_token: str = "｜DSML｜" |
| 43 | system_msg_template: str = "{content}" |
| 44 | user_msg_template: str = "<｜User｜>{content}<｜Assistant｜>" |
| 45 | assistant_msg_template: str = "{reasoning}{content}{tool_calls}<｜end▁of▁sentence｜>" |
| 46 | thinking_template = "{reasoning}" |
| 47 | |
| 48 | response_format_template: str = "## Response Format:\n\nYou MUST strictly adhere to the following schema to reply:\n{schema}" |
| 49 | tool_call_template: str = ( |
| 50 | '<{dsml_token}invoke name="{name}">\n{arguments}\n</{dsml_token}invoke>' |
| 51 | ) |
| 52 | tool_calls_template = ( |
| 53 | "<{dsml_token}function_calls>\n{tool_calls}\n</{dsml_token}function_calls>" |
| 54 | ) |
| 55 | |
| 56 | tool_output_template: str = "\n<result>{content}</result>" |
| 57 | |
| 58 | |
| 59 | def to_json(value: Any) -> str: |
| 60 | try: |
| 61 | return json.dumps(value, ensure_ascii=False) |
| 62 | except Exception: |
| 63 | return json.dumps(value, ensure_ascii=True) |
| 64 | |
| 65 | |
| 66 | def tools_from_openai_format(tools): |
| 67 | return [tool["function"] for tool in tools] |
| 68 | |
| 69 | |
| 70 | def tool_calls_from_openai_format(tool_calls): |
| 71 | return [ |
| 72 | { |
| 73 | "name": tool_call["function"]["name"], |
| 74 | "arguments": tool_call["function"]["arguments"], |
| 75 | } |
| 76 | for tool_call in tool_calls |
| 77 | ] |
| 78 | |
| 79 | |
| 80 | def encode_arguments_to_dsml(tool_call: dict[str, str]) -> str: |
| 81 | p_dsml_template = """<{dsml_token}parameter name="{key}" string="{is_str}">{value}</{dsml_token}parameter>""" |
| 82 | P_dsml_strs = [] |
| 83 | if isinstance(tool_call["arguments"], str): |
| 84 | arguments = json.loads(tool_call["arguments"]) |
| 85 | else: |
| 86 | arguments = tool_call["arguments"] |
| 87 | |
| 88 | for k, v in arguments.items(): |
| 89 | p_dsml_str = p_dsml_template.format( |
| 90 | dsml_token=dsml_token, |
| 91 | key=k, |
| 92 | is_str="true" if isinstance(v, str) else "false", |
| 93 | value=v if isinstance(v, str) else to_json(v), |
| 94 | ) |
| 95 | |
| 96 | P_dsml_strs.append(p_dsml_str) |
| 97 | |
| 98 | return "\n".join(P_dsml_strs) |
| 99 | |
| 100 | |
| 101 | def render_tools(tools: list[dict[str, str | dict[str, Any]]]) -> str: |
| 102 | tools_json = [to_json(t) for t in tools] |
| 103 | |
| 104 | return TOOLS_SYSTEM_TEMPLATE.format( |
| 105 | tool_schemas="\n".join(tools_json), |
| 106 | dsml_token=dsml_token, |
| 107 | thinking_start_token=thinking_start_token, |
| 108 | thinking_end_token=thinking_end_token, |
| 109 | ) |
| 110 | |
| 111 | |
| 112 | def find_last_user_index(messages: list[dict[str, Any]]) -> int: |
| 113 | last_user_index = -1 |
| 114 | for idx in range(len(messages) - 1, -1, -1): |
| 115 | if messages[idx].get("role") in ["user", "developer"]: |
| 116 | last_user_index = idx |
| 117 | break |
| 118 | return last_user_index |
| 119 | |
| 120 | |
| 121 | def render_message( |
| 122 | index: int, |
| 123 | messages: list[dict[str, Any]], |
| 124 | thinking_mode: str, |
| 125 | last_user_idx: int | None = None, |
| 126 | ) -> str: |
| 127 | if not (0 <= index < len(messages)): |
| 128 | raise ValueError( |
| 129 | f"Index {index} out of range for messages list of length {len(messages)}" |
| 130 | ) |
| 131 | if thinking_mode not in ["chat", "thinking"]: |
| 132 | raise ValueError(f"Invalid thinking_mode `{thinking_mode}`") |
| 133 | |
| 134 | prompt = "" |
| 135 | msg = messages[index] |
| 136 | last_user_idx = ( |
| 137 | find_last_user_index(messages) if last_user_idx is None else last_user_idx |
| 138 | ) |
| 139 | |
| 140 | role = msg.get("role") |
| 141 | content = msg.get("content") |
| 142 | tools = msg.get("tools") |
| 143 | response_format = msg.get("response_format") |
| 144 | tool_calls = msg.get("tool_calls") |
| 145 | reasoning = msg.get("reasoning") |
| 146 | is_prefix = msg.get("prefix", False) |
| 147 | |
| 148 | if tools: |
| 149 | tools = tools_from_openai_format(tools) |
| 150 | if tool_calls: |
| 151 | tool_calls = tool_calls_from_openai_format(tool_calls) |
| 152 | |
| 153 | if role == "system": |
| 154 | prompt += system_msg_template.format(content=content or "") |
| 155 | if tools: |
| 156 | prompt += "\n\n" + render_tools(tools) |
| 157 | |
| 158 | if response_format: |
| 159 | prompt += "\n\n" + response_format_template.format( |
| 160 | schema=to_json(response_format) |
| 161 | ) |
| 162 | |
| 163 | elif role == "developer": |
| 164 | if not content: |
| 165 | raise ValueError(f"Invalid message for role `{role}`: {msg}") |
| 166 | content_developer = "" |
| 167 | if tools: |
| 168 | content_developer += "\n\n" + render_tools(tools) |
| 169 | |
| 170 | if response_format: |
| 171 | content_developer += "\n\n" + response_format_template.format( |
| 172 | schema=to_json(response_format) |
| 173 | ) |
| 174 | |
| 175 | content_developer += "\n\n# The user's message is: {}".format(content) |
| 176 | |
| 177 | prompt += user_msg_template.format(content=content_developer) |
| 178 | if index == last_user_idx and thinking_mode == "thinking": |
| 179 | prompt += thinking_start_token |
| 180 | else: |
| 181 | prompt += thinking_end_token |
| 182 | |
| 183 | elif role == "user": |
| 184 | prompt += user_msg_template.format(content=content) |
| 185 | |
| 186 | if index == last_user_idx and thinking_mode == "thinking": |
| 187 | prompt += thinking_start_token |
| 188 | else: |
| 189 | prompt += thinking_end_token |
| 190 | |
| 191 | elif role == "tool": |
| 192 | prev_assistant_idx = index - 1 |
| 193 | assistant_msg = messages[prev_assistant_idx] |
| 194 | while prev_assistant_idx >= 0 and assistant_msg.get("role") == "tool": |
| 195 | prev_assistant_idx -= 1 |
| 196 | assistant_msg = messages[prev_assistant_idx] |
| 197 | |
| 198 | if not ( |
| 199 | index == 0 |
| 200 | or prev_assistant_idx >= 0 |
| 201 | and assistant_msg.get("role") == "assistant" |
| 202 | ): |
| 203 | raise ValueError(f"Invalid messages at {index}:\n{assistant_msg}") |
| 204 | |
| 205 | tool_call_order = index - prev_assistant_idx |
| 206 | assistant_tool_calls = assistant_msg.get("tool_calls") |
| 207 | if not (assistant_tool_calls and len(assistant_tool_calls) >= tool_call_order): |
| 208 | raise ValueError("No tool calls but found tool output") |
| 209 | |
| 210 | if tool_call_order == 1: |
| 211 | prompt += "\n\n<function_results>" |
| 212 | |
| 213 | prompt += tool_output_template.format(content=content) |
| 214 | |
| 215 | if tool_call_order == len(assistant_tool_calls): |
| 216 | prompt += "\n</function_results>" |
| 217 | |
| 218 | if index >= last_user_idx and thinking_mode == "thinking": |
| 219 | prompt += "\n\n" + thinking_start_token |
| 220 | else: |
| 221 | prompt += "\n\n" + thinking_end_token |
| 222 | |
| 223 | elif role == "assistant": |
| 224 | prev_assistant_idx = index |
| 225 | thinking_part = "" |
| 226 | |
| 227 | tool_calls_content = "" |
| 228 | if tool_calls: |
| 229 | tool_calls = [ |
| 230 | tool_call_template.format( |
| 231 | dsml_token=dsml_token, |
| 232 | name=tool_call.get("name"), |
| 233 | arguments=encode_arguments_to_dsml(tool_call), |
| 234 | ) |
| 235 | for tool_call in tool_calls |
| 236 | ] |
| 237 | tool_calls_content += "\n\n" + tool_calls_template.format( |
| 238 | dsml_token=dsml_token, tool_calls="\n".join(tool_calls) |
| 239 | ) |
| 240 | |
| 241 | summary_content = content or "" |
| 242 | |
| 243 | if thinking_mode == "thinking" and index > last_user_idx: |
| 244 | if not (reasoning or tool_calls): |
| 245 | raise ValueError( |
| 246 | f"ThinkingMode: {thinking_mode}, invalid message without reasoning/tool_calls `{msg}` after last user message" |
| 247 | ) |
| 248 | thinking_part = ( |
| 249 | thinking_template.format(reasoning=reasoning or "") + thinking_end_token |
| 250 | ) |
| 251 | |
| 252 | if not tool_calls and is_prefix: |
| 253 | prompt += summary_content |
| 254 | else: |
| 255 | prompt += assistant_msg_template.format( |
| 256 | reasoning=thinking_part, |
| 257 | content=summary_content, |
| 258 | tool_calls=tool_calls_content, |
| 259 | ) |
| 260 | else: |
| 261 | raise ValueError(f"Invalid role: {role}") |
| 262 | |
| 263 | return prompt |
| 264 | |
| 265 | |
| 266 | def drop_thinking_messages( |
| 267 | messages: list[dict[str, Any]], last_user_idx: int | None = None |
| 268 | ) -> list[dict[str, Any]]: |
| 269 | messages_wo_thinking: list[dict[str, Any]] = [] |
| 270 | last_user_idx = ( |
| 271 | find_last_user_index(messages) if last_user_idx is None else last_user_idx |
| 272 | ) |
| 273 | for idx, msg in enumerate(messages): |
| 274 | role = msg.get("role") |
| 275 | if role in ["user", "system", "tool"] or idx >= last_user_idx: |
| 276 | messages_wo_thinking.append(msg) |
| 277 | continue |
| 278 | |
| 279 | elif role == "assistant": |
| 280 | msg_wo_thinking = copy.copy(msg) |
| 281 | msg_wo_thinking.pop("reasoning", None) |
| 282 | messages_wo_thinking.append(msg_wo_thinking) |
| 283 | |
| 284 | return messages_wo_thinking |
| 285 | |
| 286 | |
| 287 | def encode_messages( |
| 288 | messages: list[dict[str, Any]], |
| 289 | thinking_mode: str, |
| 290 | context: list[dict[str, Any]] | None = None, |
| 291 | drop_thinking: bool = True, |
| 292 | add_default_bos_token: bool = True, |
| 293 | ) -> str: |
| 294 | context = context if context else [] |
| 295 | full_messages = context + messages |
| 296 | |
| 297 | prompt = bos_token if add_default_bos_token and len(context) == 0 else "" |
| 298 | |
| 299 | if thinking_mode == "thinking" and drop_thinking: |
| 300 | full_messages = drop_thinking_messages(full_messages) |
| 301 | |
| 302 | last_user_idx = find_last_user_index(full_messages) |
| 303 | |
| 304 | for idx in range(len(messages)): |
| 305 | prompt += render_message( |
| 306 | idx + len(context), |
| 307 | full_messages, |
| 308 | thinking_mode=thinking_mode, |
| 309 | last_user_idx=last_user_idx, |
| 310 | ) |
| 311 | |
| 312 | return prompt |