vLLM/vllm/tool_parsers/cohere_command_tool_parser.py

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

from collections.abc import Sequence

try:
    from cohere_melody import PyFilter, PyFilterOptions
except ImportError as e:
    raise ImportError(
        "The Cohere tool parser requires the `cohere_melody` "
        "package, which is not installed. Install it with:\n"
        "    pip install 'cohere-melody>=0.11.1'"
    ) from e

from vllm.entrypoints.generate.base.protocol import (
    DeltaFunctionCall,
    DeltaMessage,
    DeltaToolCall,
    ExtractedToolCallInformation,
    FunctionCall,
    ToolCall,
)
from vllm.entrypoints.openai.chat_completion.protocol import (
    ChatCompletionRequest,
)
from vllm.entrypoints.openai.responses.protocol import (
    ResponsesRequest,
)
from vllm.tokenizers import TokenizerLike
from vllm.tool_parsers import ToolParser
from vllm.tool_parsers.utils import Tool


class BaseCohereCommandToolParser(ToolParser):
    def __init__(
        self,
        tokenizer: TokenizerLike,
        streaming_opts: PyFilterOptions,
        unary_opts: PyFilterOptions,
    ):
        super().__init__(tokenizer)
        self.melody_streaming = PyFilter(streaming_opts)
        self.melody_unary = PyFilter(unary_opts)
        # Melody can emit the tool-call id before the function name. Keep it
        # until the first real name delta so clients receive both together.
        self._pending_streaming_tool_call_ids: dict[int, str] = {}

    def adjust_request(
        self, request: ChatCompletionRequest | ResponsesRequest
    ) -> ChatCompletionRequest | ResponsesRequest:
        request = super().adjust_request(request)
        request.skip_special_tokens = False
        return request

    def extract_tool_calls_streaming(
        self,
        previous_text: str,
        current_text: str,
        delta_text: str,
        previous_token_ids: Sequence[int],
        current_token_ids: Sequence[int],
        delta_token_ids: Sequence[int],
        request: ChatCompletionRequest,
    ) -> DeltaMessage | None:
        r = self.melody_streaming.write_decoded(delta_text)
        if r.content is not None:
            return DeltaMessage(content=r.content)
        if r.reasoning is not None:
            return DeltaMessage(reasoning=r.reasoning)
        if r.tool_calls:
            tool_calls: list[DeltaToolCall] = []
            for tc in r.tool_calls:
                if tc.id:
                    self._pending_streaming_tool_call_ids[tc.index] = tc.id
                name = tc.name or None
                arguments = tc.arguments or None
                # Empty strings are placeholders in Melody's streaming output;
                # omit them from OpenAI-compatible deltas instead of sending
                # invalid tool names or empty argument fragments.
                if name is None and arguments is None:
                    continue

                function_kwargs = {}
                if name is not None:
                    function_kwargs["name"] = name
                if arguments is not None:
                    function_kwargs["arguments"] = arguments
                tool_call_kwargs = {
                    "index": tc.index,
                    "function": DeltaFunctionCall(**function_kwargs),
                }
                if name is not None:
                    tool_call_id = tc.id or self._pending_streaming_tool_call_ids.pop(
                        tc.index, None
                    )
                    if tool_call_id is not None:
                        tool_call_kwargs["id"] = tool_call_id
                    tool_call_kwargs["type"] = "function"
                tool_calls.append(DeltaToolCall(**tool_call_kwargs))
            if tool_calls:
                return DeltaMessage(tool_calls=tool_calls)
        return None

    def extract_tool_calls(
        self,
        model_output: str,
        request: ChatCompletionRequest,
    ) -> ExtractedToolCallInformation:
        result = self.melody_unary.process_full_text(model_output)
        tool_calls = [
            ToolCall(
                id=tc.id,
                type="function",
                function=FunctionCall(name=tc.name, arguments=tc.arguments),
            )
            for tc in result.tool_calls
        ]
        return ExtractedToolCallInformation(
            tools_called=len(tool_calls) > 0,
            tool_calls=tool_calls,
            content=result.content,
        )


class CohereCommand3ToolParser(BaseCohereCommandToolParser):
    def __init__(
        self,
        tokenizer: TokenizerLike,
        tools: list[Tool] | None = None,
    ):
        super().__init__(
            tokenizer,
            streaming_opts=PyFilterOptions().cmd3().start_in_answer(),
            unary_opts=PyFilterOptions().cmd3().start_in_answer(),
        )


class CohereCommand4ToolParser(BaseCohereCommandToolParser):
    def __init__(
        self,
        tokenizer: TokenizerLike,
        tools: list[Tool] | None = None,
    ):
        super().__init__(
            tokenizer,
            streaming_opts=PyFilterOptions().cmd4().start_in_answer(),
            unary_opts=PyFilterOptions().cmd4().start_in_answer(),
        )