# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

import regex as re

from vllm.tool_parsers.deepseekv32_tool_parser import DeepSeekV32ToolParser


class DeepSeekV4ToolParser(DeepSeekV32ToolParser):
    """
    DeepSeek V4 DSML tool parser.

    V4 keeps the V3.2 DSML invoke/parameter grammar, but wraps tool calls in
    ``<｜DSML｜tool_calls>`` instead of ``<｜DSML｜function_calls>``.
    """

    tool_call_start_token: str = "<｜DSML｜tool_calls>"
    tool_call_end_token: str = "</｜DSML｜tool_calls>"

    def __init__(self, tokenizer, tools=None):
        super().__init__(tokenizer, tools)
        # The V3.2 __init__ overwrites the class-level start/end tokens
        # with instance attributes set to "function_calls". Re-override
        # them to the V4 "tool_calls" tags so _extract_content(),
        # _extract_invoke_regions(), etc. match the correct tags.
        self.tool_call_start_token = "<｜DSML｜tool_calls>"
        self.tool_call_end_token = "</｜DSML｜tool_calls>"
        self.tool_calls_start_token = self.tool_call_start_token
        # Recompile regexes with the V4 tool_calls tags instead of
        # the V3.2 function_calls tags that the base class hardcodes.
        self.tool_call_complete_regex = re.compile(
            r"<｜DSML｜tool_calls>(.*?)</｜DSML｜tool_calls>", re.DOTALL
        )