deepseek-ai/DeepSeek-V4-Flash-Vision-Exp
928925k
1"""2DeepSeek-V4 Text and Vision Encoding3 4A self-contained implementation for encoding/decoding DeepSeek-V4 chat messages5with tool calling, thinking mode, quick instruction tasks, and image content blocks.6"""7 8from typing import Any, Dict, List, Union, Optional, Tuple9import copy10import json11import re12 13# ============================================================14# Special Tokens15# ============================================================16 17bos_token: str = "<｜begin▁of▁sentence｜>"18eos_token: str = "<｜end▁of▁sentence｜>"19thinking_start_token: str = "<think>"20thinking_end_token: str = "</think>"21dsml_token: str = "｜DSML｜"22 23USER_SP_TOKEN = "<｜User｜>"24ASSISTANT_SP_TOKEN = "<｜Assistant｜>"25LATEST_REMINDER_SP_TOKEN = "<｜latest_reminder｜>"26IMAGE_PLACEHOLDER = "<｜deepseek_image｜>"27IMAGE_TAG_PATTERN = re.compile(r"<image>(.*?)</image>", re.DOTALL)28 29# Task special tokens for internal classification tasks30DS_TASK_SP_TOKENS = {31 "action": "<｜action｜>",32 "query": "<｜query｜>",33 "authority": "<｜authority｜>",34 "domain": "<｜domain｜>",35 "title": "<｜title｜>",36 "read_url": "<｜read_url｜>",37}38VALID_TASKS = set(DS_TASK_SP_TOKENS.keys())39 40# ============================================================41# Templates42# ============================================================43 44system_msg_template: str = "{content}"45user_msg_template: str = "{content}"46latest_reminder_msg_template: str = "{content}"47assistant_msg_template: str = "{reasoning}{content}{tool_calls}" + eos_token48assistant_msg_wo_eos_template: str = "{reasoning}{content}{tool_calls}"49thinking_template: str = "{reasoning_content}"50 51response_format_template: str = (52 "## Response Format:\n\nYou MUST strictly adhere to the following schema to reply:\n{schema}"53)54tool_call_template: str = (55 "<{dsml_token}invoke name=\"{name}\">\n{arguments}\n</{dsml_token}invoke>"56)57tool_calls_template = (58 "<{dsml_token}{tc_block_name}>\n{tool_calls}\n</{dsml_token}{tc_block_name}>"59)60tool_calls_block_name: str = "tool_calls"61 62tool_output_template: str = (63 "<tool_result>{content}</tool_result>"64)65 66# Reasoning effort levels. In thinking mode, the prompt for the selected level is67# prepended at the very beginning of the conversation. `low` is the default and68# adds nothing.69REASONING_EFFORT_PROMPTS: Dict[str, str] = {70 "low": "",71 "high": (72 "Reasoning Effort: Absolute maximum with no shortcuts permitted.\n"73 "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"74 "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"75 ),76 "max": (77 "Reasoning Effort: Beyond maximum — exhaustive, relentless, and uncompromising.\n"78 "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"79 "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"80 ),81}82DEFAULT_REASONING_EFFORT = "low"83 84TOOLS_TEMPLATE = """## Tools85 86You 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:87 88<{dsml_token}tool_calls>89<{dsml_token}invoke name="$TOOL_NAME">90<{dsml_token}parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</{dsml_token}parameter>91...92</{dsml_token}invoke>93<{dsml_token}invoke name="$TOOL_NAME2">94...95</{dsml_token}invoke>96</{dsml_token}tool_calls>97 98String 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"`.99 100If 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.101 102Otherwise, output directly after {thinking_end_token} with tool calls or final response.103 104### Available Tool Schemas105 106{tool_schemas}107 108You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.109"""110 111# ============================================================112# Utility Functions113# ============================================================114 115def to_json(value: Any) -> str:116 """Serialize a value to JSON string."""117 try:118 return json.dumps(value, ensure_ascii=False)119 except:120 return json.dumps(value, ensure_ascii=True)121 122 123def tools_from_openai_format(tools):124 """Extract function definitions from OpenAI-format tool list."""125 return [tool["function"] for tool in tools]126 127 128def tool_calls_from_openai_format(tool_calls):129 """Convert OpenAI-format tool calls to internal format."""130 return [131 {132 "name": tool_call["function"]["name"],133 "arguments": tool_call["function"]["arguments"],134 }135 for tool_call in tool_calls136 ]137 138 139def tool_calls_to_openai_format(tool_calls):140 """Convert internal tool calls to OpenAI format."""141 return [142 {143 "type": "function",144 "function": {145 "name": tool_call["name"],146 "arguments": tool_call["arguments"],147 }148 }149 for tool_call in tool_calls150 ]151 152 153def encode_arguments_to_dsml(tool_call: Dict[str, str]) -> str:154 """155 Encode tool call arguments into DSML parameter format.156 157 Args:158 tool_call: Dict with "name" and "arguments" (JSON string) keys.159 160 Returns:161 DSML-formatted parameter string.162 """163 p_dsml_template = '<{dsml_token}parameter name="{key}" string="{is_str}">{value}</{dsml_token}parameter>'164 P_dsml_strs = []165 166 try:167 arguments = json.loads(tool_call["arguments"])168 except Exception as err:169 arguments = {"arguments": tool_call["arguments"]}170 171 for k, v in arguments.items():172 p_dsml_str = p_dsml_template.format(173 dsml_token=dsml_token,174 key=k,175 is_str="true" if isinstance(v, str) else "false",176 value=v if isinstance(v, str) else to_json(v),177 )178 P_dsml_strs.append(p_dsml_str)179 180 return "\n".join(P_dsml_strs)181 182 183def decode_dsml_to_arguments(tool_name: str, tool_args: Dict[str, Tuple[str, str]]) -> Dict[str, str]:184 """185 Decode DSML parameters back to a tool call dict.186 187 Args:188 tool_name: Name of the tool.189 tool_args: Dict mapping param_name -> (value, is_string_flag).190 191 Returns:192 Dict with "name" and "arguments" (JSON string) keys.193 """194 def _decode_value(key: str, value: str, string: str):195 if string == "true":196 value = to_json(value)197 return f"{to_json(key)}: {value}"198 199 tool_args_json = "{" + ", ".join([_decode_value(k, v, string=is_str) for k, (v, is_str) in tool_args.items()]) + "}"200 return dict(name=tool_name, arguments=tool_args_json)201 202 203def render_tools(tools: List[Dict[str, Union[str, Dict[str, Any]]]]) -> str:204 """205 Render tool schemas into the system prompt format.206 207 Args:208 tools: List of tool schema dicts (each with name, description, parameters).209 210 Returns:211 Formatted tools section string.212 """213 tools_json = [to_json(t) for t in tools]214 215 return TOOLS_TEMPLATE.format(216 tool_schemas="\n".join(tools_json),217 dsml_token=dsml_token,218 thinking_start_token=thinking_start_token,219 thinking_end_token=thinking_end_token,220 )221 222 223def find_last_user_index(messages: List[Dict[str, Any]]) -> int:224 """Find the index of the last user/developer message."""225 last_user_index = -1226 for idx in range(len(messages) - 1, -1, -1):227 if messages[idx].get("role") in ["user", "developer"]:228 last_user_index = idx229 break230 return last_user_index231 232 233# ============================================================234# Message Rendering235# ============================================================236 237def render_message(index: int, messages: List[Dict[str, Any]], thinking_mode: str, drop_thinking: bool = True, reasoning_effort: Optional[str] = None) -> str:238 """239 Render a single message at the given index into its encoded string form.240 241 This is the core function that converts each message in the conversation242 into the DeepSeek-V4 format.243 244 Args:245 index: Index of the message to render.246 messages: Full list of messages in the conversation.247 thinking_mode: Either "chat" or "thinking".248 drop_thinking: Whether to drop reasoning content from earlier turns.249 reasoning_effort: Reasoning effort level, one of "low", "high", "max".250 None is treated as "low".251 252 Returns:253 Encoded string for this message.254 """255 assert 0 <= index < len(messages)256 assert thinking_mode in ["chat", "thinking"], f"Invalid thinking_mode `{thinking_mode}`"257 258 prompt = ""259 msg = messages[index]260 last_user_idx = find_last_user_index(messages)261 262 role = msg.get("role")263 content = msg.get("content")264 tools = msg.get("tools")265 response_format = msg.get("response_format")266 tool_calls = msg.get("tool_calls")267 reasoning_content = msg.get("reasoning_content")268 wo_eos = msg.get("wo_eos", False)269 270 if tools:271 tools = tools_from_openai_format(tools)272 if tool_calls:273 tool_calls = tool_calls_from_openai_format(tool_calls)274 275 # Reasoning effort prefix (only at index 0 in thinking mode; "low" adds nothing)276 reasoning_effort = reasoning_effort or DEFAULT_REASONING_EFFORT277 assert reasoning_effort in REASONING_EFFORT_PROMPTS, \278 f"Invalid reasoning effort: {reasoning_effort}, expected one of {list(REASONING_EFFORT_PROMPTS)}"279 if index == 0 and thinking_mode == "thinking":280 prompt += REASONING_EFFORT_PROMPTS[reasoning_effort]281 282 if role == "system":283 prompt += system_msg_template.format(content=content or "")284 if tools:285 prompt += "\n\n" + render_tools(tools)286 if response_format:287 prompt += "\n\n" + response_format_template.format(schema=to_json(response_format))288 289 elif role == "developer":290 assert content, f"Invalid message for role `{role}`: {msg}"291 292 content_developer = USER_SP_TOKEN293 content_developer += content294 295 if tools:296 content_developer += "\n\n" + render_tools(tools)297 if response_format:298 content_developer += "\n\n" + response_format_template.format(schema=to_json(response_format))299 300 prompt += user_msg_template.format(content=content_developer)301 302 elif role == "user":303 prompt += USER_SP_TOKEN304 305 # Handle content blocks (tool results mixed with text)306 content_blocks = msg.get("content_blocks")307 if content_blocks:308 parts = []309 for block in content_blocks:310 block_type = block.get("type")311 if block_type == "text":312 parts.append(block.get("text", ""))313 elif block_type == "tool_result":314 tool_content = block.get("content", "")315 if isinstance(tool_content, list):316 text_parts = []317 for b in tool_content:318 if b.get("type") == "text":319 text_parts.append(b.get("text", ""))320 else:321 text_parts.append(f"[Unsupported {b.get('type')}]")322 tool_content = "\n\n".join(text_parts)323 parts.append(tool_output_template.format(content=tool_content))324 else:325 parts.append(f"[Unsupported {block_type}]")326 prompt += "\n\n".join(parts)327 else:328 prompt += content or ""329 330 elif role == "latest_reminder":331 prompt += LATEST_REMINDER_SP_TOKEN + latest_reminder_msg_template.format(content=content)332 333 elif role == "tool":334 raise NotImplementedError("deepseek_v4 merges tool messages into user; please preprocess with merge_tool_messages()")335 336 elif role == "assistant":337 thinking_part = ""338 tc_content = ""339 340 if tool_calls:341 tc_list = [342 tool_call_template.format(343 dsml_token=dsml_token,344 name=tc.get("name"),345 arguments=encode_arguments_to_dsml(tc)346 )347 for tc in tool_calls348 ]349 tc_content += '\n\n' + tool_calls_template.format(350 dsml_token=dsml_token,351 tool_calls="\n".join(tc_list),352 tc_block_name=tool_calls_block_name,353 )354 355 summary_content = content or ""356 rc = reasoning_content or ""357 358 # Check if previous message has a task - if so, this is a task output (no thinking)359 prev_has_task = index - 1 >= 0 and messages[index - 1].get("task") is not None360 361 if thinking_mode == "thinking" and not prev_has_task:362 if not drop_thinking or index > last_user_idx:363 thinking_part = thinking_template.format(reasoning_content=rc) + thinking_end_token364 else:365 thinking_part = ""366 367 if wo_eos:368 prompt += assistant_msg_wo_eos_template.format(369 reasoning=thinking_part,370 content=summary_content,371 tool_calls=tc_content,372 )373 else:374 prompt += assistant_msg_template.format(375 reasoning=thinking_part,376 content=summary_content,377 tool_calls=tc_content,378 )379 else:380 raise NotImplementedError(f"Unknown role: {role}")381 382 # Append transition tokens based on what follows383 if index + 1 < len(messages) and messages[index + 1].get("role") not in ["assistant", "latest_reminder"]:384 return prompt385 386 task = messages[index].get("task")387 if task is not None:388 # Task special token for internal classification tasks389 assert task in VALID_TASKS, f"Invalid task: '{task}'. Valid tasks are: {list(VALID_TASKS)}"390 task_sp_token = DS_TASK_SP_TOKENS[task]391 392 if task != "action":393 # Non-action tasks: append task sp token directly after the message394 prompt += task_sp_token395 else:396 # Action task: append Assistant + thinking token + action sp token397 prompt += ASSISTANT_SP_TOKEN398 prompt += thinking_end_token if thinking_mode != "thinking" else thinking_start_token399 prompt += task_sp_token400 401 elif messages[index].get("role") in ["user", "developer"]:402 # Normal generation: append Assistant + thinking token403 prompt += ASSISTANT_SP_TOKEN404 if not drop_thinking and thinking_mode == "thinking":405 prompt += thinking_start_token406 elif drop_thinking and thinking_mode == "thinking" and index >= last_user_idx:407 prompt += thinking_start_token408 else:409 prompt += thinking_end_token410 411 return prompt412 413 414# ============================================================415# Preprocessing416# ============================================================417 418def merge_tool_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:419 """420 Merge tool messages into the preceding user message using content_blocks format.421 422 DeepSeek-V4 does not have a standalone "tool" role; instead, tool results423 are encoded as <tool_result> blocks within user messages.424 425 This function converts a standard OpenAI-format conversation (with separate426 "tool" role messages) into V4 format where tool results are merged into427 user messages.428 429 Args:430 messages: List of message dicts in OpenAI format.431 432 Returns:433 Processed message list with tool messages merged into user messages.434 """435 merged: List[Dict[str, Any]] = []436 437 for msg in messages:438 msg = copy.deepcopy(msg)439 role = msg.get("role")440 441 if role == "tool":442 # Convert tool message to a user message with tool_result block443 tool_block = {444 "type": "tool_result",445 "tool_use_id": msg.get("tool_call_id", ""),446 "content": msg.get("content", ""),447 }448 # Merge into previous message if it's already a user (merged tool)449 if merged and merged[-1].get("role") == "user" and "content_blocks" in merged[-1]:450 merged[-1]["content_blocks"].append(tool_block)451 else:452 merged.append({453 "role": "user",454 "content_blocks": [tool_block],455 })456 elif role == "user":457 content_blocks = msg.get("content_blocks")458 if content_blocks is None:459 content_blocks = [{"type": "text", "text": msg.get("content", "")}]460 if merged and merged[-1].get("role") == "user" and "content_blocks" in merged[-1] and merged[-1].get("task") is None:461 merged[-1]["content_blocks"].extend(content_blocks)462 else:463 # Preserve structured content and all message-level metadata.464 new_msg = msg465 new_msg["content_blocks"] = content_blocks466 merged.append(new_msg)467 else:468 merged.append(msg)469 470 return merged471 472 473def sort_tool_results_by_call_order(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:474 """475 Sort tool_result blocks within user messages by the order of tool_calls476 in the preceding assistant message.477 478 Args:479 messages: Preprocessed message list (after merge_tool_messages).480 481 Returns:482 Message list with sorted tool result blocks.483 """484 last_tool_call_order: Dict[str, int] = {}485 486 for msg in messages:487 role = msg.get("role")488 if role == "assistant" and msg.get("tool_calls"):489 last_tool_call_order = {}490 for idx, tc in enumerate(msg["tool_calls"]):491 tc_id = tc.get("id") or tc.get("function", {}).get("id", "")492 if tc_id:493 last_tool_call_order[tc_id] = idx494 495 elif role == "user" and msg.get("content_blocks"):496 tool_blocks = [b for b in msg["content_blocks"] if b.get("type") == "tool_result"]497 if len(tool_blocks) > 1 and last_tool_call_order:498 sorted_blocks = sorted(499 tool_blocks,500 key=lambda b: last_tool_call_order.get(b.get("tool_use_id", ""), 0)501 )502 sorted_idx = 0503 new_blocks = []504 for block in msg["content_blocks"]:505 if block.get("type") == "tool_result":506 new_blocks.append(sorted_blocks[sorted_idx])507 sorted_idx += 1508 else:509 new_blocks.append(block)510 msg["content_blocks"] = new_blocks511 512 return messages513 514 515# ============================================================516# Main Encoding Function517# ============================================================518 519def _encode_messages_text(520 messages: List[Dict[str, Any]],521 thinking_mode: str,522 context: Optional[List[Dict[str, Any]]] = None,523 drop_thinking: bool = True,524 add_default_bos_token: bool = True,525 reasoning_effort: Optional[str] = None,526) -> str:527 """528 Encode a list of messages into the DeepSeek-V4 prompt format.529 530 This is the main entry point for encoding conversations. It handles:531 - BOS token insertion532 - Thinking mode with optional reasoning content dropping533 - Tool message merging into user messages534 - Multi-turn conversation context535 536 Args:537 messages: List of message dicts to encode.538 thinking_mode: Either "chat" or "thinking".539 context: Optional preceding context messages (already encoded prefix).540 drop_thinking: If True, drop reasoning_content from earlier assistant turns541 (only keep reasoning for messages after the last user message).542 add_default_bos_token: Whether to prepend BOS token at conversation start.543 reasoning_effort: Reasoning effort level, one of "low", "high", "max".544 Only takes effect in thinking mode. None is treated as "low".545 546 Returns:547 The encoded prompt string.548 """549 context = context if context else []550 551 # Preprocess: merge tool messages and sort tool results552 messages = merge_tool_messages(messages)553 messages = sort_tool_results_by_call_order(context + messages)[len(context):]554 if context:555 context = merge_tool_messages(context)556 context = sort_tool_results_by_call_order(context)557 558 full_messages = context + messages559 560 prompt = bos_token if add_default_bos_token and len(context) == 0 else ""561 562 # Resolve drop_thinking: if any message has tools defined, don't drop thinking563 effective_drop_thinking = drop_thinking564 if any(m.get("tools") for m in full_messages):565 effective_drop_thinking = False566 567 if thinking_mode == "thinking" and effective_drop_thinking:568 full_messages = _drop_thinking_messages(full_messages)569 # After dropping, recalculate how many messages to render570 # (context may have shrunk too)571 num_to_render = len(full_messages) - len(_drop_thinking_messages(context))572 context_len = len(full_messages) - num_to_render573 else:574 num_to_render = len(messages)575 context_len = len(context)576 577 for idx in range(num_to_render):578 prompt += render_message(579 idx + context_len,580 full_messages,581 thinking_mode=thinking_mode,582 drop_thinking=effective_drop_thinking,583 reasoning_effort=reasoning_effort,584 )585 586 return prompt587 588 589def _drop_thinking_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:590 """591 Drop reasoning_content and non-essential messages before the last user message.592 593 Behavior:594 - Messages with role in ["user", "system", "tool", "latest_reminder"] are always kept.595 - Messages at or after the last user index are always kept.596 - Assistant messages before the last user get reasoning_content removed.597 - Developer messages before the last user are dropped entirely.598 """599 last_user_idx = find_last_user_index(messages)600 result = []601 keep_roles = {"user", "system", "tool", "latest_reminder", "direct_search_results"}602 603 for idx, msg in enumerate(messages):604 role = msg.get("role")605 if role in keep_roles or idx >= last_user_idx:606 result.append(msg)607 elif role == "assistant":608 msg = copy.copy(msg)609 msg.pop("reasoning_content", None)610 result.append(msg)611 # developer and other roles before last_user_idx are dropped612 613 return result614 615 616# ============================================================617# Vision Message Preprocessing618# ============================================================619 620def parse_tagged_text(text: str) -> Union[str, List[Dict[str, Any]]]:621 """Convert ``<image>path</image>`` text into standard content blocks."""622 matches = list(IMAGE_TAG_PATTERN.finditer(text))623 remaining = IMAGE_TAG_PATTERN.sub("", text)624 if "<image>" in remaining or "</image>" in remaining:625 raise ValueError("Malformed <image>path</image> tag")626 if not matches:627 return text628 629 blocks: List[Dict[str, Any]] = []630 cursor = 0631 for match in matches:632 if match.start() > cursor:633 blocks.append({"type": "text", "text": text[cursor:match.start()]})634 path = match.group(1)635 if not path:636 raise ValueError("Image path must not be empty")637 blocks.append({638 "type": "image_url",639 "image_url": {"url": path},640 })641 cursor = match.end()642 if cursor < len(text):643 blocks.append({"type": "text", "text": text[cursor:]})644 return blocks645 646 647def _is_image_block(block: Dict[str, Any]) -> bool:648 """Return whether a content block is an OpenAI/Anthropic/internal image."""649 return isinstance(block, dict) and block.get("type") in ("image", "image_url")650 651 652def _extract_image(block: Dict[str, Any]) -> Dict[str, Any]:653 """Normalize a supported image block into an internal image record."""654 record: Dict[str, Any] = {"type": "image"}655 if block.get("type") == "image_url":656 image_url = block.get("image_url")657 if isinstance(image_url, str):658 record["url"] = image_url659 else:660 record["url"] = (image_url or {}).get("url", "")661 else:662 for key in ("source", "url", "data"):663 if key in block:664 record[key] = block[key]665 if not any(record.get(key) for key in ("source", "url", "data")):666 raise ValueError("Image block does not contain a valid source")667 return record668 669 670def _process_image_blocks(671 blocks: List[Any], image_placeholder: str = IMAGE_PLACEHOLDER672) -> Tuple[List[Any], List[Dict[str, Any]]]:673 """Replace image blocks and collect their records in one ordered traversal."""674 new_blocks: List[Any] = []675 images: List[Dict[str, Any]] = []676 for block in blocks:677 if not isinstance(block, dict):678 new_blocks.append(block)679 continue680 if _is_image_block(block):681 new_blocks.append({"type": "text", "text": image_placeholder})682 images.append(_extract_image(block))683 elif block.get("type") == "tool_result" and isinstance(block.get("content"), list):684 block = copy.copy(block)685 block["content"], nested_images = _process_image_blocks(686 block["content"], image_placeholder)687 new_blocks.append(block)688 images.extend(nested_images)689 elif block.get("type") == "text":690 text = block.get("text") or ""691 if IMAGE_PLACEHOLDER in text:692 raise ValueError(693 f"Text block contains image placeholder '{IMAGE_PLACEHOLDER}': "694 f"'{text[:100]}'. Images should be separate content blocks."695 )696 new_blocks.append(block)697 else:698 new_blocks.append(block)699 return new_blocks, images700 701 702def _validate_no_image_sp_tokens(msg: Dict[str, Any]) -> None:703 """Reject user-supplied image placeholder tokens in textual fields."""704 content = msg.get("content")705 if isinstance(content, str) and IMAGE_PLACEHOLDER in content:706 raise ValueError(707 f"Message content contains image special token '{IMAGE_PLACEHOLDER}'. "708 "Images should be provided as image content blocks."709 )710 reasoning_content = msg.get("reasoning_content")711 if isinstance(reasoning_content, str) and IMAGE_PLACEHOLDER in reasoning_content:712 raise ValueError(713 f"reasoning_content contains image special token '{IMAGE_PLACEHOLDER}'"714 )715 716 717def process_image_messages(718 messages: List[Dict[str, Any]],719) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:720 """Normalize image blocks and return their records in prompt order."""721 processed: List[Dict[str, Any]] = []722 images: List[Dict[str, Any]] = []723 for msg in messages:724 msg = copy.deepcopy(msg)725 _validate_no_image_sp_tokens(msg)726 727 if isinstance(msg.get("content"), list) and "content_blocks" not in msg:728 msg["content_blocks"] = msg.pop("content")729 730 if msg.get("content_blocks"):731 msg["content_blocks"], message_images = _process_image_blocks(732 msg["content_blocks"])733 images.extend(message_images)734 if not isinstance(msg.get("content"), str):735 texts = [736 block.get("text", "")737 for block in msg["content_blocks"]738 if isinstance(block, dict) and block.get("type") == "text"739 ]740 msg["content"] = "\n\n".join(texts)741 742 processed.append(msg)743 return processed, images744 745 746def encode_messages(747 messages: List[Dict[str, Any]],748 thinking_mode: str,749 context: Optional[List[Dict[str, Any]]] = None,750 drop_thinking: bool = True,751 add_default_bos_token: bool = True,752 reasoning_effort: Optional[str] = None,753 return_multi_modal_data: bool = False,754) -> Any:755 """Encode text or multimodal messages through one canonical public entrypoint.756 757 Text-only calls preserve the original string-returning API. When758 return_multi_modal_data is true, the result is ``(prompt, media_data)``.759 """760 context = context or []761 processed_context, _ = process_image_messages(context) if context else ([], [])762 processed_messages, images = process_image_messages(messages)763 prompt = _encode_messages_text(764 processed_messages,765 thinking_mode=thinking_mode,766 context=processed_context if processed_context else None,767 drop_thinking=drop_thinking,768 add_default_bos_token=add_default_bos_token,769 reasoning_effort=reasoning_effort,770 )771 if return_multi_modal_data:772 return prompt, {"images": images}773 return prompt774 775 776def load_cases(input_file: str) -> List[Dict[str, Any]]:777 """Load one or more OpenAI-format conversation cases from JSON."""778 with open(input_file) as file:779 data = json.load(file)780 if isinstance(data, dict):781 data = [data]782 elif data and isinstance(data[0], dict) and "role" in data[0]:783 data = [{"messages": data}]784 785 cases = []786 for case in data:787 messages = copy.deepcopy(case["messages"])788 if "tools" in case:789 if not messages:790 raise ValueError("A case with tools must contain at least one message")791 messages[0]["tools"] = case["tools"]792 cases.append({793 "messages": messages,794 "context": case.get("context"),795 "thinking_mode": case.get("thinking_mode"),796 "reasoning_effort": case.get("reasoning_effort"),797 })798 return cases799 800 801def encode_case(802 case: Dict[str, Any], thinking_mode: str803) -> Tuple[str, List[Dict[str, Any]]]:804 """Encode one JSON case and return its current-turn image records."""805 prompt, media_data = encode_messages(806 case["messages"],807 thinking_mode=case.get("thinking_mode") or thinking_mode,808 context=case.get("context"),809 reasoning_effort=case.get("reasoning_effort"),810 return_multi_modal_data=True,811 )812 return prompt, media_data["images"]813 814 815# ============================================================816# Parsing (Decoding model output)817# ============================================================818 819def _read_until_stop(index: int, text: str, stop: List[str]) -> Tuple[int, str, Optional[str]]:820 """821 Read text from index until one of the stop strings is found.822 823 Returns:824 Tuple of (new_index, content_before_stop, matched_stop_string_or_None).825 """826 min_pos = len(text)827 matched_stop = None828 829 for s in stop:830 pos = text.find(s, index)831 if pos != -1 and pos < min_pos:832 min_pos = pos833 matched_stop = s834 835 if matched_stop:836 content = text[index:min_pos]837 return min_pos + len(matched_stop), content, matched_stop838 else:839 content = text[index:]840 return len(text), content, None841 842 843def parse_tool_calls(index: int, text: str) -> Tuple[int, Optional[str], List[Dict[str, str]]]:844 """845 Parse DSML tool calls from text starting at the given index.846 847 Args:848 index: Starting position in text.849 text: The full text to parse.850 851 Returns:852 Tuple of (new_index, last_stop_token, list_of_tool_call_dicts).853 Each tool call dict has "name" and "arguments" keys.854 """855 tool_calls: List[Dict[str, Any]] = []856 stop_token = None857 tool_calls_end_token = f"</{dsml_token}{tool_calls_block_name}>"858 859 while index < len(text):860 index, _, stop_token = _read_until_stop(index, text, [f"<{dsml_token}invoke", tool_calls_end_token])861 if _ != ">\n":862 raise ValueError(f"Tool call format error: expected '>\\n' but got '{_}'")863 864 if stop_token == tool_calls_end_token:865 break866 867 if stop_token is None:868 raise ValueError("Missing special token in tool calls")869 870 index, tool_name_content, stop_token = _read_until_stop(index, text, [f"<{dsml_token}parameter", f"</{dsml_token}invoke"])871 872 p_tool_name = re.findall(r'^\s*name="(.*?)">\n$', tool_name_content, flags=re.DOTALL)873 if len(p_tool_name) != 1:874 raise ValueError(f"Tool name format error: '{tool_name_content}'")875 tool_name = p_tool_name[0]876 877 tool_args: Dict[str, Tuple[str, str]] = {}878 while stop_token == f"<{dsml_token}parameter":879 index, param_content, stop_token = _read_until_stop(index, text, [f"/{dsml_token}parameter"])880 881 param_kv = re.findall(r'^ name="(.*?)" string="(true|false)">(.*?)<$', param_content, flags=re.DOTALL)882 if len(param_kv) != 1:883 raise ValueError(f"Parameter format error: '{param_content}'")884 param_name, string, param_value = param_kv[0]885 886 if param_name in tool_args:887 raise ValueError(f"Duplicate parameter name: '{param_name}'")888 tool_args[param_name] = (param_value, string)889 890 index, content, stop_token = _read_until_stop(index, text, [f"<{dsml_token}parameter", f"</{dsml_token}invoke"])891 if content != ">\n":892 raise ValueError(f"Parameter format error: expected '>\\n' but got '{content}'")893 894 tool_call = decode_dsml_to_arguments(tool_name=tool_name, tool_args=tool_args)895 tool_calls.append(tool_call)896 897 return index, stop_token, tool_calls898 899 900def parse_message_from_completion_text(text: str, thinking_mode: str) -> Dict[str, Any]:901 """902 Parse a model completion text into a structured assistant message.903 904 This function takes the raw text output from the model (a single assistant turn)905 and extracts:906 - reasoning_content (thinking block)907 - content (summary/response)908 - tool_calls (if any)909 910 NOTE: This function is designed to parse only correctly formatted strings and911 will raise ValueError for malformed output.912 913 Args:914 text: The raw completion text (including EOS token).915 thinking_mode: Either "chat" or "thinking".916 917 Returns:918 Dict with keys: "role", "content", "reasoning_content", "tool_calls".919 tool_calls are in OpenAI format.920 """921 summary_content, reasoning_content, tool_calls = "", "", []922 index, stop_token = 0, None923 tool_calls_start_token = f"\n\n<{dsml_token}{tool_calls_block_name}"924 925 is_thinking = thinking_mode == "thinking"926 is_tool_calling = False927 928 if is_thinking:929 index, content_delta, stop_token = _read_until_stop(index, text, [thinking_end_token, tool_calls_start_token])930 reasoning_content = content_delta931 assert stop_token == thinking_end_token, "Invalid thinking format: missing </think>"932 933 index, content_delta, stop_token = _read_until_stop(index, text, [eos_token, tool_calls_start_token])934 summary_content = content_delta935 if stop_token == tool_calls_start_token:936 is_tool_calling = True937 else:938 assert stop_token == eos_token, "Invalid format: missing EOS token"939 940 if is_tool_calling:941 index, stop_token, tool_calls = parse_tool_calls(index, text)942 943 index, tool_ends_text, stop_token = _read_until_stop(index, text, [eos_token])944 assert not tool_ends_text, "Unexpected content after tool calls"945 946 assert len(text) == index and stop_token in [eos_token, None], "Unexpected content at end"947 948 for sp_token in [bos_token, eos_token, thinking_start_token, thinking_end_token, dsml_token]:949 assert sp_token not in summary_content and sp_token not in reasoning_content, \950 f"Unexpected special token '{sp_token}' in content"951 952 return {953 "role": "assistant",954 "content": summary_content,955 "reasoning_content": reasoning_content,956 "tool_calls": tool_calls_to_openai_format(tool_calls)957 }958