deepseek_v4_encoding.py 28 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766
  1. # SPDX-License-Identifier: Apache-2.0
  2. # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
  3. # ruff: noqa
  4. # fmt: off
  5. """
  6. DeepSeek-V4 Encoding
  7. A self-contained implementation for encoding/decoding DeepSeek-V4 chat messages
  8. with tool calling, thinking mode, and quick instruction task support.
  9. """
  10. from typing import Any, Dict, List, Union, Optional, Tuple
  11. import copy
  12. import json
  13. import regex as re
  14. # ============================================================
  15. # Special Tokens
  16. # ============================================================
  17. bos_token: str = "<|begin▁of▁sentence|>"
  18. eos_token: str = "<|end▁of▁sentence|>"
  19. thinking_start_token: str = "<think>"
  20. thinking_end_token: str = "</think>"
  21. dsml_token: str = "|DSML|"
  22. USER_SP_TOKEN = "<|User|>"
  23. ASSISTANT_SP_TOKEN = "<|Assistant|>"
  24. LATEST_REMINDER_SP_TOKEN = "<|latest_reminder|>"
  25. # Task special tokens for internal classification tasks
  26. DS_TASK_SP_TOKENS = {
  27. "action": "<|action|>",
  28. "query": "<|query|>",
  29. "authority": "<|authority|>",
  30. "domain": "<|domain|>",
  31. "title": "<|title|>",
  32. "read_url": "<|read_url|>",
  33. }
  34. VALID_TASKS = set(DS_TASK_SP_TOKENS.keys())
  35. # ============================================================
  36. # Templates
  37. # ============================================================
  38. system_msg_template: str = "{content}"
  39. user_msg_template: str = "{content}"
  40. latest_reminder_msg_template: str = "{content}"
  41. assistant_msg_template: str = "{reasoning}{content}{tool_calls}" + eos_token
  42. assistant_msg_wo_eos_template: str = "{reasoning}{content}{tool_calls}"
  43. thinking_template: str = "{reasoning}"
  44. response_format_template: str = (
  45. "## Response Format:\n\nYou MUST strictly adhere to the following schema to reply:\n{schema}"
  46. )
  47. tool_call_template: str = (
  48. "<{dsml_token}invoke name=\"{name}\">\n{arguments}\n</{dsml_token}invoke>"
  49. )
  50. tool_calls_template = (
  51. "<{dsml_token}{tc_block_name}>\n{tool_calls}\n</{dsml_token}{tc_block_name}>"
  52. )
  53. tool_calls_block_name: str = "tool_calls"
  54. tool_output_template: str = (
  55. "<tool_result>{content}</tool_result>"
  56. )
  57. REASONING_EFFORT_PROMPTS = {
  58. "low": "",
  59. "high": (
  60. "Reasoning Effort: Absolute maximum with no shortcuts permitted.\n"
  61. "You MUST be very thorough in your thinking and comprehensively decompose the problem to resolve the root cause, rigorously stress-testing your logic against all potential paths, edge cases, and adversarial scenarios.\n"
  62. "Explicitly write out your entire deliberation process, documenting every intermediate step, considered alternative, and rejected hypothesis to ensure absolutely no assumption is left unchecked.\n\n"
  63. ),
  64. "max": (
  65. "Reasoning Effort: Beyond maximum — exhaustive, relentless, and uncompromising.\n"
  66. "You MUST reason with the utmost depth and rigor, leaving absolutely nothing to chance: exhaustively decompose the problem into its most fundamental components, trace every causal chain to its root, and resolve the underlying cause rather than any surface symptom.\n"
  67. "Do not stop reasoning until you have independently verified the solution from multiple angles and are certain that no assumption remains unchecked and no error remains undiscovered.\n\n"
  68. ),
  69. }
  70. TOOLS_TEMPLATE = """## Tools
  71. You have access to a set of tools to help answer the user's question. You can invoke tools by writing a "<{dsml_token}tool_calls>" block like the following:
  72. <{dsml_token}tool_calls>
  73. <{dsml_token}invoke name="$TOOL_NAME">
  74. <{dsml_token}parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</{dsml_token}parameter>
  75. ...
  76. </{dsml_token}invoke>
  77. <{dsml_token}invoke name="$TOOL_NAME2">
  78. ...
  79. </{dsml_token}invoke>
  80. </{dsml_token}tool_calls>
  81. String parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`.
  82. If thinking_mode is enabled (triggered by {thinking_start_token}), you MUST output your complete reasoning inside {thinking_start_token}...{thinking_end_token} BEFORE any tool calls or final response.
  83. Otherwise, output directly after {thinking_end_token} with tool calls or final response.
  84. ### Available Tool Schemas
  85. {tool_schemas}
  86. You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.
  87. """
  88. # ============================================================
  89. # Utility Functions
  90. # ============================================================
  91. def to_json(value: Any) -> str:
  92. """Serialize a value to JSON string."""
  93. try:
  94. return json.dumps(value, ensure_ascii=False)
  95. except Exception:
  96. return json.dumps(value, ensure_ascii=True)
  97. def tools_from_openai_format(tools):
  98. """Extract function definitions from OpenAI-format tool list."""
  99. return [tool["function"] for tool in tools]
  100. def tool_calls_from_openai_format(tool_calls):
  101. """Convert OpenAI-format tool calls to internal format."""
  102. return [
  103. {
  104. "name": tool_call["function"]["name"],
  105. "arguments": tool_call["function"]["arguments"],
  106. }
  107. for tool_call in tool_calls
  108. ]
  109. def tool_calls_to_openai_format(tool_calls):
  110. """Convert internal tool calls to OpenAI format."""
  111. return [
  112. {
  113. "type": "function",
  114. "function": {
  115. "name": tool_call["name"],
  116. "arguments": tool_call["arguments"],
  117. }
  118. }
  119. for tool_call in tool_calls
  120. ]
  121. def encode_arguments_to_dsml(tool_call: Dict[str, Any]) -> str:
  122. """
  123. Encode tool call arguments into DSML parameter format.
  124. Args:
  125. tool_call: Dict with "name" and "arguments" keys.
  126. Returns:
  127. DSML-formatted parameter string.
  128. """
  129. p_dsml_template = '<{dsml_token}parameter name="{key}" string="{is_str}">{value}</{dsml_token}parameter>'
  130. P_dsml_strs = []
  131. if isinstance(tool_call["arguments"], str):
  132. arguments = json.loads(tool_call["arguments"])
  133. else:
  134. arguments = tool_call["arguments"]
  135. for k, v in arguments.items():
  136. p_dsml_str = p_dsml_template.format(
  137. dsml_token=dsml_token,
  138. key=k,
  139. is_str="true" if isinstance(v, str) else "false",
  140. value=v if isinstance(v, str) else to_json(v),
  141. )
  142. P_dsml_strs.append(p_dsml_str)
  143. return "\n".join(P_dsml_strs)
  144. def decode_dsml_to_arguments(tool_name: str, tool_args: Dict[str, Tuple[str, str]]) -> Dict[str, str]:
  145. """
  146. Decode DSML parameters back to a tool call dict.
  147. Args:
  148. tool_name: Name of the tool.
  149. tool_args: Dict mapping param_name -> (value, is_string_flag).
  150. Returns:
  151. Dict with "name" and "arguments" (JSON string) keys.
  152. """
  153. def _decode_value(key: str, value: str, string: str):
  154. if string == "true":
  155. value = to_json(value)
  156. return f"{to_json(key)}: {value}"
  157. tool_args_json = "{" + ", ".join([_decode_value(k, v, string=is_str) for k, (v, is_str) in tool_args.items()]) + "}"
  158. return dict(name=tool_name, arguments=tool_args_json)
  159. def render_tools(tools: List[Dict[str, Union[str, Dict[str, Any]]]]) -> str:
  160. """
  161. Render tool schemas into the system prompt format.
  162. Args:
  163. tools: List of tool schema dicts (each with name, description, parameters).
  164. Returns:
  165. Formatted tools section string.
  166. """
  167. tools_json = [to_json(t) for t in tools]
  168. return TOOLS_TEMPLATE.format(
  169. tool_schemas="\n".join(tools_json),
  170. dsml_token=dsml_token,
  171. thinking_start_token=thinking_start_token,
  172. thinking_end_token=thinking_end_token,
  173. )
  174. def find_last_user_index(messages: List[Dict[str, Any]]) -> int:
  175. """Find the index of the last user/developer message."""
  176. last_user_index = -1
  177. for idx in range(len(messages) - 1, -1, -1):
  178. if messages[idx].get("role") in ["user", "developer"]:
  179. last_user_index = idx
  180. break
  181. return last_user_index
  182. # ============================================================
  183. # Message Rendering
  184. # ============================================================
  185. def render_message(index: int, messages: List[Dict[str, Any]], thinking_mode: str, drop_thinking: bool = True, reasoning_effort: Optional[str] = None) -> str:
  186. """
  187. Render a single message at the given index into its encoded string form.
  188. This is the core function that converts each message in the conversation
  189. into the DeepSeek-V4 format.
  190. Args:
  191. index: Index of the message to render.
  192. messages: Full list of messages in the conversation.
  193. thinking_mode: Either "chat" or "thinking".
  194. drop_thinking: Whether to drop reasoning content from earlier turns.
  195. reasoning_effort: Optional reasoning effort level ("max", "high", or None).
  196. Returns:
  197. Encoded string for this message.
  198. """
  199. assert 0 <= index < len(messages)
  200. assert thinking_mode in ["chat", "thinking"], f"Invalid thinking_mode `{thinking_mode}`"
  201. prompt = ""
  202. msg = messages[index]
  203. last_user_idx = find_last_user_index(messages)
  204. role = msg.get("role")
  205. content = msg.get("content")
  206. tools = msg.get("tools")
  207. response_format = msg.get("response_format")
  208. tool_calls = msg.get("tool_calls")
  209. reasoning = msg.get("reasoning")
  210. wo_eos = msg.get("wo_eos", False)
  211. if tools:
  212. tools = tools_from_openai_format(tools)
  213. if tool_calls:
  214. tool_calls = tool_calls_from_openai_format(tool_calls)
  215. # Reasoning effort prefix (only at index 0 in thinking mode; "low" adds nothing)
  216. assert reasoning_effort is None or reasoning_effort in REASONING_EFFORT_PROMPTS, \
  217. f"Invalid reasoning effort: {reasoning_effort}"
  218. if index == 0 and thinking_mode == "thinking" and reasoning_effort in ("high", "max"):
  219. prompt += REASONING_EFFORT_PROMPTS[reasoning_effort]
  220. if role == "system":
  221. prompt += system_msg_template.format(content=content or "")
  222. if tools:
  223. prompt += "\n\n" + render_tools(tools)
  224. if response_format:
  225. prompt += "\n\n" + response_format_template.format(schema=to_json(response_format))
  226. elif role == "developer":
  227. assert content, f"Invalid message for role `{role}`: {msg}"
  228. content_developer = USER_SP_TOKEN
  229. content_developer += content
  230. if tools:
  231. content_developer += "\n\n" + render_tools(tools)
  232. if response_format:
  233. content_developer += "\n\n" + response_format_template.format(schema=to_json(response_format))
  234. prompt += user_msg_template.format(content=content_developer)
  235. elif role == "user":
  236. prompt += USER_SP_TOKEN
  237. # Handle content blocks (tool results mixed with text)
  238. content_blocks = msg.get("content_blocks")
  239. if content_blocks:
  240. parts = []
  241. for block in content_blocks:
  242. block_type = block.get("type")
  243. if block_type == "text":
  244. parts.append(block.get("text", ""))
  245. elif block_type == "tool_result":
  246. tool_content = block.get("content", "")
  247. if isinstance(tool_content, list):
  248. text_parts = []
  249. for b in tool_content:
  250. if b.get("type") == "text":
  251. text_parts.append(b.get("text", ""))
  252. else:
  253. text_parts.append(f"[Unsupported {b.get('type')}]")
  254. tool_content = "\n\n".join(text_parts)
  255. parts.append(tool_output_template.format(content=tool_content))
  256. else:
  257. parts.append(f"[Unsupported {block_type}]")
  258. prompt += "\n\n".join(parts)
  259. else:
  260. prompt += content or ""
  261. elif role == "latest_reminder":
  262. prompt += LATEST_REMINDER_SP_TOKEN + latest_reminder_msg_template.format(content=content)
  263. elif role == "tool":
  264. raise NotImplementedError("deepseek_v4 merges tool messages into user; please preprocess with merge_tool_messages()")
  265. elif role == "assistant":
  266. thinking_part = ""
  267. tc_content = ""
  268. if tool_calls:
  269. tc_list = [
  270. tool_call_template.format(
  271. dsml_token=dsml_token,
  272. name=tc.get("name"),
  273. arguments=encode_arguments_to_dsml(tc)
  274. )
  275. for tc in tool_calls
  276. ]
  277. tc_content += '\n\n' + tool_calls_template.format(
  278. dsml_token=dsml_token,
  279. tool_calls="\n".join(tc_list),
  280. tc_block_name=tool_calls_block_name,
  281. )
  282. summary_content = content or ""
  283. reasoning = reasoning or ""
  284. # Check if previous message has a task - if so, this is a task output (no thinking)
  285. prev_has_task = index - 1 >= 0 and messages[index - 1].get("task") is not None
  286. if thinking_mode == "thinking" and not prev_has_task:
  287. if not drop_thinking or index > last_user_idx:
  288. thinking_part = thinking_template.format(reasoning=reasoning) + thinking_end_token
  289. else:
  290. thinking_part = ""
  291. if wo_eos:
  292. prompt += assistant_msg_wo_eos_template.format(
  293. reasoning=thinking_part,
  294. content=summary_content,
  295. tool_calls=tc_content,
  296. )
  297. else:
  298. prompt += assistant_msg_template.format(
  299. reasoning=thinking_part,
  300. content=summary_content,
  301. tool_calls=tc_content,
  302. )
  303. else:
  304. raise NotImplementedError(f"Unknown role: {role}")
  305. # Append transition tokens based on what follows
  306. if index + 1 < len(messages) and messages[index + 1].get("role") not in ["assistant", "latest_reminder"]:
  307. return prompt
  308. task = messages[index].get("task")
  309. if task is not None:
  310. # Task special token for internal classification tasks
  311. assert task in VALID_TASKS, f"Invalid task: '{task}'. Valid tasks are: {list(VALID_TASKS)}"
  312. task_sp_token = DS_TASK_SP_TOKENS[task]
  313. if task != "action":
  314. # Non-action tasks: append task sp token directly after the message
  315. prompt += task_sp_token
  316. else:
  317. # Action task: append Assistant + thinking token + action sp token
  318. prompt += ASSISTANT_SP_TOKEN
  319. prompt += thinking_end_token if thinking_mode != "thinking" else thinking_start_token
  320. prompt += task_sp_token
  321. elif messages[index].get("role") in ["user", "developer"]:
  322. # Normal generation: append Assistant + thinking token
  323. prompt += ASSISTANT_SP_TOKEN
  324. if not drop_thinking and thinking_mode == "thinking":
  325. prompt += thinking_start_token
  326. elif drop_thinking and thinking_mode == "thinking" and index >= last_user_idx:
  327. prompt += thinking_start_token
  328. else:
  329. prompt += thinking_end_token
  330. return prompt
  331. # ============================================================
  332. # Preprocessing
  333. # ============================================================
  334. def merge_tool_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
  335. """
  336. Merge tool messages into the preceding user message using content_blocks format.
  337. DeepSeek-V4 does not have a standalone "tool" role; instead, tool results
  338. are encoded as <tool_result> blocks within user messages.
  339. This function converts a standard OpenAI-format conversation (with separate
  340. "tool" role messages) into V4 format where tool results are merged into
  341. user messages.
  342. Args:
  343. messages: List of message dicts in OpenAI format.
  344. Returns:
  345. Processed message list with tool messages merged into user messages.
  346. """
  347. merged: List[Dict[str, Any]] = []
  348. for msg in messages:
  349. msg = copy.deepcopy(msg)
  350. role = msg.get("role")
  351. if role == "tool":
  352. # Convert tool message to a user message with tool_result block
  353. tool_block = {
  354. "type": "tool_result",
  355. "tool_use_id": msg.get("tool_call_id", ""),
  356. "content": msg.get("content", ""),
  357. }
  358. # Merge into previous message if it's already a user (merged tool)
  359. if merged and merged[-1].get("role") == "user" and "content_blocks" in merged[-1]:
  360. merged[-1]["content_blocks"].append(tool_block)
  361. else:
  362. merged.append({
  363. "role": "user",
  364. "content_blocks": [tool_block],
  365. })
  366. elif role == "user":
  367. text_block = {"type": "text", "text": msg.get("content", "")}
  368. if merged and merged[-1].get("role") == "user" and "content_blocks" in merged[-1] and merged[-1].get("task") is None:
  369. merged[-1]["content_blocks"].append(text_block)
  370. else:
  371. new_msg = {
  372. "role": "user",
  373. "content": msg.get("content", ""),
  374. "content_blocks": [text_block],
  375. }
  376. # Preserve extra fields (task, wo_eos, mask, etc.)
  377. for key in ("task", "wo_eos", "mask"):
  378. if key in msg:
  379. new_msg[key] = msg[key]
  380. merged.append(new_msg)
  381. else:
  382. merged.append(msg)
  383. return merged
  384. def sort_tool_results_by_call_order(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
  385. """
  386. Sort tool_result blocks within user messages by the order of tool_calls
  387. in the preceding assistant message.
  388. Args:
  389. messages: Preprocessed message list (after merge_tool_messages).
  390. Returns:
  391. Message list with sorted tool result blocks.
  392. """
  393. last_tool_call_order: Dict[str, int] = {}
  394. for msg in messages:
  395. role = msg.get("role")
  396. if role == "assistant" and msg.get("tool_calls"):
  397. last_tool_call_order = {}
  398. for idx, tc in enumerate(msg["tool_calls"]):
  399. tc_id = tc.get("id") or tc.get("function", {}).get("id", "")
  400. if tc_id:
  401. last_tool_call_order[tc_id] = idx
  402. elif role == "user" and msg.get("content_blocks"):
  403. tool_blocks = [b for b in msg["content_blocks"] if b.get("type") == "tool_result"]
  404. if len(tool_blocks) > 1 and last_tool_call_order:
  405. sorted_blocks = sorted(
  406. tool_blocks,
  407. key=lambda b: last_tool_call_order.get(b.get("tool_use_id", ""), 0)
  408. )
  409. sorted_idx = 0
  410. new_blocks = []
  411. for block in msg["content_blocks"]:
  412. if block.get("type") == "tool_result":
  413. new_blocks.append(sorted_blocks[sorted_idx])
  414. sorted_idx += 1
  415. else:
  416. new_blocks.append(block)
  417. msg["content_blocks"] = new_blocks
  418. return messages
  419. # ============================================================
  420. # Main Encoding Function
  421. # ============================================================
  422. def encode_messages(
  423. messages: List[Dict[str, Any]],
  424. thinking_mode: str,
  425. context: Optional[List[Dict[str, Any]]] = None,
  426. drop_thinking: bool = True,
  427. add_default_bos_token: bool = True,
  428. reasoning_effort: Optional[str] = None,
  429. ) -> str:
  430. """
  431. Encode a list of messages into the DeepSeek-V4 prompt format.
  432. This is the main entry point for encoding conversations. It handles:
  433. - BOS token insertion
  434. - Thinking mode with optional reasoning content dropping
  435. - Tool message merging into user messages
  436. - Multi-turn conversation context
  437. Args:
  438. messages: List of message dicts to encode.
  439. thinking_mode: Either "chat" or "thinking".
  440. context: Optional preceding context messages (already encoded prefix).
  441. drop_thinking: If True, drop reasoning from earlier assistant turns
  442. (only keep reasoning for messages after the last user message).
  443. add_default_bos_token: Whether to prepend BOS token at conversation start.
  444. reasoning_effort: Optional reasoning effort level ("max", "high", or None).
  445. Returns:
  446. The encoded prompt string.
  447. """
  448. context = context if context else []
  449. # Preprocess: merge tool messages and sort tool results
  450. messages = merge_tool_messages(messages)
  451. messages = sort_tool_results_by_call_order(context + messages)[len(context):]
  452. if context:
  453. context = merge_tool_messages(context)
  454. context = sort_tool_results_by_call_order(context)
  455. full_messages = context + messages
  456. prompt = bos_token if add_default_bos_token and len(context) == 0 else ""
  457. # Resolve drop_thinking: if any message has tools defined, don't drop thinking
  458. effective_drop_thinking = drop_thinking
  459. if any(m.get("tools") for m in full_messages):
  460. effective_drop_thinking = False
  461. if thinking_mode == "thinking" and effective_drop_thinking:
  462. full_messages = _drop_thinking_messages(full_messages)
  463. # After dropping, recalculate how many messages to render
  464. # (context may have shrunk too)
  465. num_to_render = len(full_messages) - len(_drop_thinking_messages(context))
  466. context_len = len(full_messages) - num_to_render
  467. else:
  468. num_to_render = len(messages)
  469. context_len = len(context)
  470. for idx in range(num_to_render):
  471. prompt += render_message(
  472. idx + context_len,
  473. full_messages,
  474. thinking_mode=thinking_mode,
  475. drop_thinking=effective_drop_thinking,
  476. reasoning_effort=reasoning_effort,
  477. )
  478. return prompt
  479. def _drop_thinking_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
  480. """
  481. Drop reasoning and non-essential messages before the last user message.
  482. Behavior:
  483. - Messages with role in ["user", "system", "tool", "latest_reminder"] are always kept.
  484. - Messages at or after the last user index are always kept.
  485. - Assistant messages before the last user get reasoning removed.
  486. - Developer messages before the last user are dropped entirely.
  487. """
  488. last_user_idx = find_last_user_index(messages)
  489. result = []
  490. keep_roles = {"user", "system", "tool", "latest_reminder", "direct_search_results"}
  491. for idx, msg in enumerate(messages):
  492. role = msg.get("role")
  493. if role in keep_roles or idx >= last_user_idx:
  494. result.append(msg)
  495. elif role == "assistant":
  496. msg = copy.copy(msg)
  497. msg.pop("reasoning", None)
  498. result.append(msg)
  499. # developer and other roles before last_user_idx are dropped
  500. return result
  501. # ============================================================
  502. # Parsing (Decoding model output)
  503. # ============================================================
  504. def _read_until_stop(index: int, text: str, stop: List[str]) -> Tuple[int, str, Optional[str]]:
  505. """
  506. Read text from index until one of the stop strings is found.
  507. Returns:
  508. Tuple of (new_index, content_before_stop, matched_stop_string_or_None).
  509. """
  510. min_pos = len(text)
  511. matched_stop = None
  512. for s in stop:
  513. pos = text.find(s, index)
  514. if pos != -1 and pos < min_pos:
  515. min_pos = pos
  516. matched_stop = s
  517. if matched_stop:
  518. content = text[index:min_pos]
  519. return min_pos + len(matched_stop), content, matched_stop
  520. else:
  521. content = text[index:]
  522. return len(text), content, None
  523. def parse_tool_calls(index: int, text: str) -> Tuple[int, Optional[str], List[Dict[str, str]]]:
  524. """
  525. Parse DSML tool calls from text starting at the given index.
  526. Args:
  527. index: Starting position in text.
  528. text: The full text to parse.
  529. Returns:
  530. Tuple of (new_index, last_stop_token, list_of_tool_call_dicts).
  531. Each tool call dict has "name" and "arguments" keys.
  532. """
  533. tool_calls: List[Dict[str, Any]] = []
  534. stop_token = None
  535. tool_calls_end_token = f"</{dsml_token}{tool_calls_block_name}>"
  536. while index < len(text):
  537. index, content_before, stop_token = _read_until_stop(index, text, [f"<{dsml_token}invoke", tool_calls_end_token])
  538. if content_before != ">\n":
  539. raise ValueError(f"Tool call format error: expected '>\\n' but got '{content_before}'")
  540. if stop_token == tool_calls_end_token:
  541. break
  542. if stop_token is None:
  543. raise ValueError("Missing special token in tool calls")
  544. index, tool_name_content, stop_token = _read_until_stop(index, text, [f"<{dsml_token}parameter", f"</{dsml_token}invoke"])
  545. p_tool_name = re.findall(r'^\s*name="(.*?)">\n$', tool_name_content, flags=re.DOTALL)
  546. if len(p_tool_name) != 1:
  547. raise ValueError(f"Tool name format error: '{tool_name_content}'")
  548. tool_name = p_tool_name[0]
  549. tool_args: Dict[str, Tuple[str, str]] = {}
  550. while stop_token == f"<{dsml_token}parameter":
  551. index, param_content, stop_token = _read_until_stop(index, text, [f"/{dsml_token}parameter"])
  552. param_kv = re.findall(r'^ name="(.*?)" string="(true|false)">(.*?)<$', param_content, flags=re.DOTALL)
  553. if len(param_kv) != 1:
  554. raise ValueError(f"Parameter format error: '{param_content}'")
  555. param_name, string, param_value = param_kv[0]
  556. if param_name in tool_args:
  557. raise ValueError(f"Duplicate parameter name: '{param_name}'")
  558. tool_args[param_name] = (param_value, string)
  559. index, content, stop_token = _read_until_stop(index, text, [f"<{dsml_token}parameter", f"</{dsml_token}invoke"])
  560. if content != ">\n":
  561. raise ValueError(f"Parameter format error: expected '>\\n' but got '{content}'")
  562. tool_call = decode_dsml_to_arguments(tool_name=tool_name, tool_args=tool_args)
  563. tool_calls.append(tool_call)
  564. return index, stop_token, tool_calls
  565. def parse_message_from_completion_text(text: str, thinking_mode: str) -> Dict[str, Any]:
  566. """
  567. Parse a model completion text into a structured assistant message.
  568. This function takes the raw text output from the model (a single assistant turn)
  569. and extracts:
  570. - reasoning (thinking block)
  571. - content (summary/response)
  572. - tool_calls (if any)
  573. NOTE: This function is designed to parse only correctly formatted strings and
  574. will raise ValueError for malformed output.
  575. Args:
  576. text: The raw completion text (including EOS token).
  577. thinking_mode: Either "chat" or "thinking".
  578. Returns:
  579. Dict with keys: "role", "content", "reasoning", "tool_calls".
  580. tool_calls are in OpenAI format.
  581. """
  582. summary_content, reasoning = "", ""
  583. tool_calls: List[Dict[str, str]] = []
  584. index, stop_token = 0, None
  585. tool_calls_start_token = f"\n\n<{dsml_token}{tool_calls_block_name}"
  586. is_thinking = thinking_mode == "thinking"
  587. is_tool_calling = False
  588. if is_thinking:
  589. index, content_delta, stop_token = _read_until_stop(index, text, [thinking_end_token, tool_calls_start_token])
  590. reasoning = content_delta
  591. if stop_token != thinking_end_token:
  592. raise ValueError("Invalid thinking format: missing </think>")
  593. index, content_delta, stop_token = _read_until_stop(index, text, [eos_token, tool_calls_start_token])
  594. summary_content = content_delta
  595. if stop_token == tool_calls_start_token:
  596. is_tool_calling = True
  597. else:
  598. if stop_token != eos_token:
  599. raise ValueError("Invalid format: missing EOS token")
  600. if is_tool_calling:
  601. index, stop_token, tool_calls = parse_tool_calls(index, text)
  602. index, tool_ends_text, stop_token = _read_until_stop(index, text, [eos_token])
  603. if tool_ends_text:
  604. raise ValueError("Unexpected content after tool calls")
  605. if len(text) != index or stop_token not in [eos_token, None]:
  606. raise ValueError("Unexpected content at end")
  607. for sp_token in [bos_token, eos_token, thinking_start_token, thinking_end_token, dsml_token]:
  608. if sp_token in summary_content or sp_token in reasoning:
  609. raise ValueError(f"Unexpected special token '{sp_token}' in content")
  610. return {
  611. "role": "assistant",
  612. "content": summary_content,
  613. "reasoning": reasoning,
  614. "tool_calls": tool_calls_to_openai_format(tool_calls)
  615. }
  616. # fmt: on