vllm/tests/parser/test_harmony.py

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

import json
from collections.abc import Sequence

import pytest
from openai_harmony import (
    Conversation,
    Message,
    RenderConversationConfig,
    Role,
)
from transformers import AutoTokenizer

from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
from vllm.entrypoints.openai.engine.protocol import FunctionCall
from vllm.entrypoints.openai.parser.harmony_utils import (
    get_encoding,
)
from vllm.parser.harmony import HarmonyParser
from vllm.parser.parser_manager import ParserManager

REASONING_MODEL_NAME = "openai/gpt-oss-20b"


@pytest.fixture(scope="module")
def gpt_oss_tokenizer():
    return AutoTokenizer.from_pretrained(REASONING_MODEL_NAME)


@pytest.fixture
def harmony_parser(gpt_oss_tokenizer):
    parser_cls = ParserManager.get_parser(
        tool_parser_name="openai",
        reasoning_parser_name="openai_gptoss",
        enable_auto_tools=True,
        model_name=REASONING_MODEL_NAME,
        is_harmony=True,
    )
    assert parser_cls is HarmonyParser
    return parser_cls(gpt_oss_tokenizer)


@pytest.fixture
def chat_request():
    return ChatCompletionRequest(
        model="openai/gpt-oss-20b",
        messages=[{"role": "user", "content": "Hello"}],
    )


def encode_output(harmony_str: str) -> list[int]:
    return get_encoding().encode(harmony_str, allowed_special="all")


def assistant(content: str, channel: str) -> Message:
    return Message.from_role_and_content(Role.ASSISTANT, content).with_channel(channel)


def tool_call(
    recipient: str,
    content: str,
    channel: str = "commentary",
    content_type: str | None = "json",
) -> Message:
    message = assistant(content, channel).with_recipient(recipient)
    return message if content_type is None else message.with_content_type(content_type)


def get_model_output_tokens(
    prompt_messages: Sequence[Message],
    response_messages: Sequence[Message],
) -> list[int]:
    enc = get_encoding()
    # Keep analysis messages when synthesizing model-output-only token sequences
    # for parser tests; the default render path drops them after a later final turn.
    config = RenderConversationConfig(auto_drop_analysis=False)
    prompt_ids = enc.render_conversation_for_completion(
        Conversation.from_messages(list(prompt_messages)),
        Role.ASSISTANT,
        config=config,
    )
    full_ids = enc.render_conversation_for_completion(
        Conversation.from_messages([*prompt_messages, *response_messages]),
        Role.ASSISTANT,
        config=config,
    )
    assert full_ids[: len(prompt_ids)] == prompt_ids
    return full_ids[len(prompt_ids) :]


def get_text(msg: Message) -> str:
    return msg.content[0].text if msg.content else ""


def visible_segments(result) -> list[tuple[str | None, str | None, str]]:
    return [
        (segment.channel, segment.recipient, segment.delta)
        for segment in result.segments
        if not segment.is_boundary and segment.delta
    ]


def tool_call_tuples(tool_calls: list[FunctionCall] | None) -> list[tuple[str, str]]:
    return [] if tool_calls is None else [(tc.name, tc.arguments) for tc in tool_calls]


class TestParse:
    # Rendered conversation outputs.

    def test_reasoning_only(self, harmony_parser, chat_request):
        prompt = [Message.from_role_and_content(Role.USER, "Why?")]
        response = [assistant("This is reasoning", "analysis")]

        reasoning, content, tool_calls = harmony_parser.parse(
            "",
            chat_request,
            model_output_token_ids=get_model_output_tokens(prompt, response),
        )

        assert reasoning == "This is reasoning"
        assert content is None
        assert tool_calls is None

    def test_content_only(self, harmony_parser, chat_request):
        prompt = [Message.from_role_and_content(Role.USER, "Hello")]
        response = [assistant("This is a test", "final")]

        reasoning, content, tool_calls = harmony_parser.parse(
            "",
            chat_request,
            model_output_token_ids=get_model_output_tokens(prompt, response),
        )

        assert reasoning is None
        assert content == "This is a test"
        assert tool_calls is None

    def test_reasoning_and_content(self, harmony_parser, chat_request):
        prompt = [Message.from_role_and_content(Role.USER, "What is 2+2?")]
        response = [
            assistant("I should think first.", "analysis"),
            assistant("The answer is 4.", "final"),
        ]

        reasoning, content, tool_calls = harmony_parser.parse(
            "",
            chat_request,
            model_output_token_ids=get_model_output_tokens(prompt, response),
        )

        assert reasoning == "I should think first."
        assert content == "The answer is 4."
        assert tool_calls is None

    @pytest.mark.parametrize(
        "tool_args",
        [
            '{"location": "Tokyo"}',
            '{\n"location": "Tokyo"\n}',
        ],
    )
    @pytest.mark.parametrize("tool_channel", ["commentary", "analysis"])
    def test_single_tool_call(
        self, harmony_parser, chat_request, tool_args, tool_channel
    ):
        prompt = [
            Message.from_role_and_content(Role.USER, "What is the weather in Tokyo?")
        ]
        response = [tool_call("functions.get_current_weather", tool_args, tool_channel)]

        reasoning, content, tool_calls = harmony_parser.parse(
            "",
            chat_request,
            model_output_token_ids=get_model_output_tokens(prompt, response),
        )

        assert reasoning is None
        assert content is None
        assert tool_call_tuples(tool_calls) == [
            ("get_current_weather", json.dumps({"location": "Tokyo"}))
        ]

    def test_multiple_tool_calls_varied_formats(self, harmony_parser, chat_request):
        prompt = [
            Message.from_role_and_content(
                Role.USER, "What is the weather in Tokyo based on where I'm at?"
            )
        ]
        response = [
            tool_call("functions.get_current_weather", '{"location": "Tokyo"}'),
            tool_call("functions.get_user_location", '{"location": "Tokyo"}'),
            tool_call(
                "functions.no_content_type",
                '{"location": "Tokyo"}',
                content_type=None,
            ),
            tool_call("functions.not_json_no_content_type", "foo", content_type=None),
            tool_call("functions.empty_args", "{}"),
            tool_call("functions.no_args", ""),
        ]

        _, content, tool_calls = harmony_parser.parse(
            "",
            chat_request,
            model_output_token_ids=get_model_output_tokens(prompt, response),
        )

        assert content is None
        assert tool_call_tuples(tool_calls) == [
            ("get_current_weather", json.dumps({"location": "Tokyo"})),
            ("get_user_location", json.dumps({"location": "Tokyo"})),
            ("no_content_type", json.dumps({"location": "Tokyo"})),
            ("not_json_no_content_type", "foo"),
            ("empty_args", json.dumps({})),
            ("no_args", ""),
        ]

    def test_tool_call_bare_recipient(self, harmony_parser, chat_request):
        prompt = [Message.from_role_and_content(Role.USER, "Weather?")]
        response = [tool_call("get_current_weather", '{"location": "Tokyo"}')]

        _, _, tool_calls = harmony_parser.parse(
            "",
            chat_request,
            model_output_token_ids=get_model_output_tokens(prompt, response),
        )

        assert tool_call_tuples(tool_calls) == [
            ("get_current_weather", json.dumps({"location": "Tokyo"}))
        ]

    def test_multiple_tool_calls_bare_recipients(self, harmony_parser, chat_request):
        prompt = [Message.from_role_and_content(Role.USER, "Use both tools.")]
        response = [
            tool_call("get_current_weather", '{"location": "Tokyo"}'),
            tool_call("get_user_location", "{}"),
        ]

        _, _, tool_calls = harmony_parser.parse(
            "",
            chat_request,
            model_output_token_ids=get_model_output_tokens(prompt, response),
        )

        assert tool_call_tuples(tool_calls) == [
            ("get_current_weather", json.dumps({"location": "Tokyo"})),
            ("get_user_location", json.dumps({})),
        ]

    def test_assistant_recipient_not_tool(self, harmony_parser, chat_request):
        prompt = [Message.from_role_and_content(Role.USER, "Hello")]
        response = [
            tool_call("assistant", "Some tool response", content_type=None),
            assistant("Here is the answer", "final"),
        ]

        reasoning, content, tool_calls = harmony_parser.parse(
            "",
            chat_request,
            model_output_token_ids=get_model_output_tokens(prompt, response),
        )

        assert reasoning is None
        assert content == "Here is the answer"
        assert tool_calls is None

    def test_tool_call_dotted_name(self, harmony_parser, chat_request):
        prompt = [Message.from_role_and_content(Role.USER, "Compute 2+3")]
        response = [tool_call("math.sum", '{"a": 2, "b": 3}')]

        _, _, tool_calls = harmony_parser.parse(
            "",
            chat_request,
            model_output_token_ids=get_model_output_tokens(prompt, response),
        )

        assert tool_call_tuples(tool_calls) == [
            ("math.sum", json.dumps({"a": 2, "b": 3}))
        ]

    def test_tool_calls_with_final_content(self, harmony_parser, chat_request):
        prompt = [Message.from_role_and_content(Role.USER, "What is the weather?")]
        response = [
            assistant("User asked about the weather.", "analysis"),
            tool_call("functions.get_current_weather", '{"location": "Tokyo"}'),
            assistant("This tool call will get the weather.", "final"),
        ]

        reasoning, content, tool_calls = harmony_parser.parse(
            "",
            chat_request,
            model_output_token_ids=get_model_output_tokens(prompt, response),
        )

        assert reasoning == "User asked about the weather."
        assert content == "This tool call will get the weather."
        assert tool_call_tuples(tool_calls) == [
            ("get_current_weather", json.dumps({"location": "Tokyo"}))
        ]

    # Raw/truncated Harmony output streams.

    def test_interrupted_first_message(self, harmony_parser, chat_request):
        reasoning, content, tool_calls = harmony_parser.parse(
            "",
            chat_request,
            model_output_token_ids=encode_output(
                "<|channel|>final<|message|>I'm in the middle of answering"
            ),
        )

        assert reasoning is None
        assert content == "I'm in the middle of answering"
        assert tool_calls is None

    def test_interrupted_reasoning_first_message(self, harmony_parser, chat_request):
        reasoning, content, tool_calls = harmony_parser.parse(
            "",
            chat_request,
            model_output_token_ids=encode_output(
                "<|channel|>analysis<|message|>I'm in the middle of thinking"
            ),
        )

        assert reasoning == "I'm in the middle of thinking"
        assert content is None
        assert tool_calls is None

    def test_truncated_output(self, harmony_parser, chat_request):
        reasoning, content, tool_calls = harmony_parser.parse(
            "",
            chat_request,
            model_output_token_ids=encode_output(
                "<|channel|>analysis<|message|>I'm thinking.<|end|>"
                "<|start|>assistant<|channel|>final<|message|>"
                "I'm in the middle of answering"
            ),
        )

        assert reasoning == "I'm thinking."
        assert content == "I'm in the middle of answering"
        assert tool_calls is None

    @pytest.mark.parametrize(
        ("harmony_str", "expected_content"),
        [
            (
                "<|channel|>commentary<|message|>I'll search for that",
                "I'll search for that",
            ),
            (
                "<|channel|>commentary<|message|>Let me look that up.<|end|>"
                "<|start|>assistant<|channel|>final<|message|>The answer is 42.<|end|>",
                "Let me look that up.\nThe answer is 42.",
            ),
        ],
    )
    def test_commentary_preambles(
        self,
        harmony_parser,
        chat_request,
        harmony_str,
        expected_content,
    ):
        reasoning, content, tool_calls = harmony_parser.parse(
            "",
            chat_request,
            model_output_token_ids=encode_output(harmony_str),
        )

        assert reasoning is None
        assert content == expected_content
        assert tool_calls is None

    def test_commentary_with_recipient_excluded(self, harmony_parser, chat_request):
        reasoning, content, tool_calls = harmony_parser.parse(
            "",
            chat_request,
            model_output_token_ids=encode_output(
                "<|channel|>commentary"
                "<|message|>Let me check the weather.<|end|>"
                "<|start|>assistant to=functions.get_weather"
                "<|channel|>commentary"
                '<|message|>{"location": "SF"}<|end|>'
            ),
        )

        assert reasoning is None
        assert content == "Let me check the weather."
        assert tool_call_tuples(tool_calls) == [
            ("get_weather", json.dumps({"location": "SF"}))
        ]


class TestProcessChunk:
    def test_empty(self, harmony_parser):
        result = harmony_parser.process_chunk([])
        assert result.segments == []
        assert result.reasoning_token_count == 0

    def test_single_channel(self, harmony_parser):
        result = harmony_parser.process_chunk(
            encode_output("<|channel|>final<|message|>Hello")
        )

        assert visible_segments(result) == [("final", None, "Hello")]

    def test_cross_channel(self, harmony_parser):
        result = harmony_parser.process_chunk(
            encode_output(
                "<|channel|>analysis<|message|>Think<|end|>"
                "<|start|>assistant<|channel|>final<|message|>Answer"
            )
        )

        assert visible_segments(result) == [
            ("analysis", None, "Think"),
            ("final", None, "Answer"),
        ]

    def test_boundary_detection(self, harmony_parser):
        result = harmony_parser.process_chunk(
            encode_output("<|channel|>final<|message|>Done<|end|>")
        )

        boundary_segments = [
            segment for segment in result.segments if segment.is_boundary
        ]
        assert len(boundary_segments) == 1
        assert boundary_segments[0].completed_message is not None
        assert boundary_segments[0].completed_message.channel == "final"
        assert get_text(boundary_segments[0].completed_message) == "Done"

    def test_multi_boundary(self, harmony_parser):
        result = harmony_parser.process_chunk(
            encode_output(
                "<|channel|>analysis<|message|>One<|end|>"
                "<|start|>assistant<|channel|>final<|message|>Two<|end|>"
            )
        )

        boundary_segments = [
            segment for segment in result.segments if segment.is_boundary
        ]
        assert [
            get_text(segment.completed_message) for segment in boundary_segments
        ] == [
            "One",
            "Two",
        ]