test_api_utils.py

# SPDX-License-Identifier: Apache-2.0
"""
Tests for API utility functions.

Tests clean_output_text, is_mllm_model, and extract_multimodal_content
from vllm_mlx/api/utils.py. No MLX dependency.
"""

from vllm_mlx.api.models import ContentPart, ImageUrl, Message
from vllm_mlx.api.utils import (
    MLLM_PATTERNS,
    SPECIAL_TOKENS_PATTERN,
    _content_to_text,
    clean_output_text,
    extract_multimodal_content,
    is_mllm_model,
    is_vlm_model,
)


class TestCleanOutputText:
    """Tests for clean_output_text function."""

    def test_empty_string(self):
        assert clean_output_text("") == ""

    def test_none_returns_none(self):
        assert clean_output_text(None) is None

    def test_plain_text_unchanged(self):
        assert clean_output_text("Hello world") == "Hello world"

    def test_removes_im_end(self):
        assert clean_output_text("Hello<|im_end|>") == "Hello"

    def test_removes_im_start(self):
        assert clean_output_text("<|im_start|>Hello") == "Hello"

    def test_removes_endoftext(self):
        assert clean_output_text("Hello<|endoftext|>") == "Hello"

    def test_removes_eot_id(self):
        assert clean_output_text("Hello<|eot_id|>") == "Hello"

    def test_removes_end_token(self):
        assert clean_output_text("Hello<|end|>") == "Hello"

    def test_removes_start_header_id(self):
        result = clean_output_text("<|start_header_id|>assistant<|end_header_id|>Hello")
        assert "<|start_header_id|>" not in result
        assert "<|end_header_id|>" not in result

    def test_removes_s_tags(self):
        assert clean_output_text("<s>Hello</s>") == "Hello"

    def test_removes_pad_tokens(self):
        assert clean_output_text("[PAD]Hello[PAD]") == "Hello"

    def test_removes_sep_cls(self):
        assert clean_output_text("[CLS]Hello[SEP]") == "Hello"

    def test_removes_multiple_special_tokens(self):
        text = "<|im_start|>assistant\nHello world<|im_end|><|endoftext|>"
        result = clean_output_text(text)
        assert result == "assistant\nHello world"

    def test_preserves_think_tags(self):
        text = "<think>Let me think about this.</think>The answer is 42."
        result = clean_output_text(text)
        assert "<think>" in result
        assert "</think>" in result
        assert "The answer is 42." in result

    def test_adds_missing_opening_think_tag(self):
        text = "Some thinking content.</think>The answer is 42."
        result = clean_output_text(text)
        assert result.startswith("<think>")
        assert "</think>" in result

    def test_no_extra_think_tag_when_already_present(self):
        text = "<think>Thinking.</think>Answer."
        result = clean_output_text(text)
        assert result.count("<think>") == 1

    def test_strips_whitespace(self):
        assert clean_output_text("  Hello  ") == "Hello"

    def test_combined_special_tokens_and_think(self):
        text = "<|im_start|><think>I need to think.</think>42<|im_end|>"
        result = clean_output_text(text)
        assert "<think>" in result
        assert "</think>" in result
        assert "42" in result
        assert "<|im_start|>" not in result


class TestSpecialTokensPattern:
    """Tests for the special tokens regex pattern."""

    def test_matches_all_expected_tokens(self):
        tokens = [
            "<|im_end|>",
            "<|im_start|>",
            "<|endoftext|>",
            "<|end|>",
            "<|eot_id|>",
            "<|start_header_id|>",
            "<|end_header_id|>",
            "</s>",
            "<s>",
            "<pad>",
            "[PAD]",
            "[SEP]",
            "[CLS]",
        ]
        for token in tokens:
            assert SPECIAL_TOKENS_PATTERN.search(token) is not None, (
                f"Pattern should match {token}"
            )

    def test_does_not_match_think_tags(self):
        assert SPECIAL_TOKENS_PATTERN.search("<think>") is None
        assert SPECIAL_TOKENS_PATTERN.search("</think>") is None

    def test_does_not_match_normal_text(self):
        assert SPECIAL_TOKENS_PATTERN.search("Hello world") is None


class TestIsMllmModel:
    """Tests for is_mllm_model function."""

    def test_qwen_vl_models(self):
        assert is_mllm_model("mlx-community/Qwen3-VL-4B-Instruct-3bit") is True
        assert is_mllm_model("mlx-community/Qwen2-VL-7B-Instruct-4bit") is True

    def test_llava_models(self):
        assert is_mllm_model("mlx-community/llava-1.5-7b-4bit") is True
        assert is_mllm_model("mlx-community/LLaVA-NeXT-7b") is True

    def test_idefics_models(self):
        assert is_mllm_model("mlx-community/Idefics3-8B-Llama3-4bit") is True
        assert is_mllm_model("mlx-community/idefics2-8b-4bit") is True

    def test_paligemma_models(self):
        assert is_mllm_model("mlx-community/paligemma2-3b-mix-224-4bit") is True
        assert is_mllm_model("mlx-community/PaliGemma-3b-mix") is True

    def test_gemma3_models(self):
        assert is_mllm_model("mlx-community/gemma-3-12b-it-4bit") is True
        assert is_mllm_model("mlx-community/gemma3-4b-it-4bit") is True

    def test_medgemma_models(self):
        assert is_mllm_model("mlx-community/MedGemma-4b-it-4bit") is True
        assert is_mllm_model("mlx-community/medgemma-4b") is True

    def test_pixtral_models(self):
        assert is_mllm_model("mlx-community/pixtral-12b-4bit") is True
        assert is_mllm_model("mlx-community/Pixtral-12b-8bit") is True

    def test_molmo_models(self):
        assert is_mllm_model("mlx-community/Molmo-7B-D-0924-4bit") is True
        assert is_mllm_model("mlx-community/molmo-7b") is True

    def test_phi3_vision(self):
        assert is_mllm_model("mlx-community/phi3-vision-128k") is True
        assert is_mllm_model("mlx-community/phi-3-vision-128k-instruct-4bit") is True

    def test_cogvlm(self):
        assert is_mllm_model("mlx-community/CogVLM-chat-hf") is True
        assert is_mllm_model("mlx-community/cogvlm-chat-hf") is True

    def test_internvl(self):
        assert is_mllm_model("mlx-community/InternVL2-8B") is True

    def test_deepseek_vl(self):
        assert is_mllm_model("mlx-community/deepseek-vl-7b-chat-4bit") is True
        assert is_mllm_model("mlx-community/DeepSeek-VL2-small-4bit") is True

    def test_non_mllm_models(self):
        assert is_mllm_model("mlx-community/Llama-3.2-3B-Instruct-4bit") is False
        assert is_mllm_model("mlx-community/Qwen3-8B-4bit") is False
        assert is_mllm_model("mlx-community/Mistral-7B-Instruct-v0.3-4bit") is False
        assert is_mllm_model("mlx-community/DeepSeek-R1-Distill-Qwen-7B") is False

    def test_case_insensitive(self):
        assert is_mllm_model("LLAVA-7B") is True
        assert is_mllm_model("pixtral-12b") is True

    def test_backwards_compatibility_alias(self):
        assert is_vlm_model is is_mllm_model

    def test_all_patterns_defined(self):
        assert len(MLLM_PATTERNS) > 20


class TestExtractMultimodalContent:
    """Tests for extract_multimodal_content function."""

    def test_simple_text_messages(self):
        messages = [
            Message(role="system", content="You are helpful."),
            Message(role="user", content="Hello"),
        ]
        processed, images, videos = extract_multimodal_content(messages)

        assert len(processed) == 2
        assert processed[0] == {"role": "system", "content": "You are helpful."}
        assert processed[1] == {"role": "user", "content": "Hello"}
        assert images == []
        assert videos == []

    def test_none_content(self):
        messages = [Message(role="assistant", content=None)]
        processed, images, videos = extract_multimodal_content(messages)
        assert processed[0] == {"role": "assistant", "content": ""}

    def test_multimodal_with_image_url(self):
        messages = [
            Message(
                role="user",
                content=[
                    ContentPart(type="text", text="What is this?"),
                    ContentPart(
                        type="image_url",
                        image_url=ImageUrl(url="https://example.com/img.png"),
                    ),
                ],
            )
        ]
        processed, images, videos = extract_multimodal_content(messages)

        assert len(processed) == 1
        assert processed[0]["content"] == "What is this?"
        assert images == ["https://example.com/img.png"]
        assert videos == []

    def test_multimodal_with_dict_image_url(self):
        messages = [
            Message(
                role="user",
                content=[
                    {"type": "text", "text": "Describe this"},
                    {
                        "type": "image_url",
                        "image_url": {"url": "data:image/png;base64,abc"},
                    },
                ],
            )
        ]
        processed, images, videos = extract_multimodal_content(messages)
        assert images == ["data:image/png;base64,abc"]

    def test_multimodal_with_string_image_url(self):
        messages = [
            Message(
                role="user",
                content=[
                    {"type": "text", "text": "Look"},
                    {"type": "image_url", "image_url": "https://example.com/img.png"},
                ],
            )
        ]
        processed, images, videos = extract_multimodal_content(messages)
        assert images == ["https://example.com/img.png"]

    def test_multimodal_with_video(self):
        messages = [
            Message(
                role="user",
                content=[
                    {"type": "text", "text": "What happens?"},
                    {"type": "video", "video": "/path/to/video.mp4"},
                ],
            )
        ]
        processed, images, videos = extract_multimodal_content(messages)
        assert videos == ["/path/to/video.mp4"]

    def test_multimodal_with_video_url(self):
        messages = [
            Message(
                role="user",
                content=[
                    {"type": "text", "text": "Describe"},
                    {
                        "type": "video_url",
                        "video_url": {"url": "https://example.com/v.mp4"},
                    },
                ],
            )
        ]
        processed, images, videos = extract_multimodal_content(messages)
        assert videos == ["https://example.com/v.mp4"]

    def test_multimodal_with_string_video_url(self):
        messages = [
            Message(
                role="user",
                content=[
                    {"type": "text", "text": "Look"},
                    {"type": "video_url", "video_url": "https://example.com/v.mp4"},
                ],
            )
        ]
        processed, images, videos = extract_multimodal_content(messages)
        assert videos == ["https://example.com/v.mp4"]

    def test_multiple_images(self):
        messages = [
            Message(
                role="user",
                content=[
                    {"type": "text", "text": "Compare these"},
                    {"type": "image_url", "image_url": {"url": "img1.png"}},
                    {"type": "image_url", "image_url": {"url": "img2.png"}},
                ],
            )
        ]
        processed, images, videos = extract_multimodal_content(messages)
        assert len(images) == 2

    def test_tool_response_message(self):
        messages = [
            Message(role="tool", content="72F and sunny", tool_call_id="call_1")
        ]
        processed, images, videos = extract_multimodal_content(messages)
        assert processed[0]["role"] == "user"
        assert "Tool Result" in processed[0]["content"]
        assert "call_1" in processed[0]["content"]

    def test_tool_response_preserve_native(self):
        messages = [
            Message(role="tool", content="72F and sunny", tool_call_id="call_1")
        ]
        processed, images, videos = extract_multimodal_content(
            messages, preserve_native_format=True
        )
        assert processed[0]["role"] == "tool"
        assert processed[0]["tool_call_id"] == "call_1"
        assert processed[0]["content"] == "72F and sunny"

    def test_assistant_with_tool_calls(self):
        messages = [
            Message(
                role="assistant",
                content=None,
                tool_calls=[
                    {
                        "id": "call_1",
                        "type": "function",
                        "function": {
                            "name": "get_weather",
                            "arguments": '{"city": "NYC"}',
                        },
                    }
                ],
            )
        ]
        processed, images, videos = extract_multimodal_content(messages)
        assert processed[0]["role"] == "assistant"
        assert "get_weather" in processed[0]["content"]

    def test_assistant_with_tool_calls_preserve_native(self):
        messages = [
            Message(
                role="assistant",
                content="Let me check.",
                tool_calls=[
                    {
                        "id": "call_1",
                        "type": "function",
                        "function": {
                            "name": "get_weather",
                            "arguments": '{"city": "NYC"}',
                        },
                    }
                ],
            )
        ]
        processed, images, videos = extract_multimodal_content(
            messages, preserve_native_format=True
        )
        assert processed[0]["role"] == "assistant"
        assert processed[0]["content"] == "Let me check."
        assert "tool_calls" in processed[0]
        assert len(processed[0]["tool_calls"]) == 1

    def test_dict_messages(self):
        messages = [
            Message(role="user", content="Hello"),
        ]
        # Also test with raw dicts (the function handles both)
        processed, images, videos = extract_multimodal_content(messages)
        assert processed[0]["content"] == "Hello"

    def test_image_type_content_with_raw_dicts(self):
        # type="image" path handles raw dict content (not Pydantic ContentPart)
        # Pass a raw dict message to avoid Pydantic stripping unknown fields
        raw_messages = [
            type(
                "Msg",
                (),
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": "What is this?"},
                        {"type": "image", "image": "https://example.com/img.png"},
                    ],
                    "tool_calls": None,
                    "tool_call_id": None,
                },
            )()
        ]
        processed, images, videos = extract_multimodal_content(raw_messages)
        assert images == ["https://example.com/img.png"]

    def test_empty_messages(self):
        processed, images, videos = extract_multimodal_content([])
        assert processed == []
        assert images == []
        assert videos == []

    def test_tool_response_none_content(self):
        messages = [Message(role="tool", content=None, tool_call_id="call_1")]
        processed, images, videos = extract_multimodal_content(messages)
        assert processed[0]["role"] == "user"
        assert "call_1" in processed[0]["content"]

    def test_multiple_text_parts_combined(self):
        messages = [
            Message(
                role="user",
                content=[
                    {"type": "text", "text": "First part."},
                    {"type": "text", "text": "Second part."},
                ],
            )
        ]
        processed, images, videos = extract_multimodal_content(messages)
        assert "First part." in processed[0]["content"]
        assert "Second part." in processed[0]["content"]

    def test_assistant_tool_calls_with_list_content(self):
        """Regression test for issue #61: list content + tool_calls causes TypeError."""
        messages = [
            Message(
                role="assistant",
                content=[ContentPart(type="text", text="Let me check.")],
                tool_calls=[
                    {
                        "id": "call_1",
                        "type": "function",
                        "function": {
                            "name": "get_weather",
                            "arguments": '{"city": "Prague"}',
                        },
                    }
                ],
            )
        ]
        result, images, videos = extract_multimodal_content(messages)
        assert isinstance(result[0]["content"], str)
        assert "Let me check." in result[0]["content"]
        assert "get_weather" in result[0]["content"]

    def test_assistant_tool_calls_with_list_content_native(self):
        """Regression test for issue #61: list content + tool_calls with native format."""
        messages = [
            Message(
                role="assistant",
                content=[ContentPart(type="text", text="Checking now.")],
                tool_calls=[
                    {
                        "id": "call_2",
                        "type": "function",
                        "function": {
                            "name": "search",
                            "arguments": '{"q": "test"}',
                        },
                    }
                ],
            )
        ]
        result, images, videos = extract_multimodal_content(
            messages, preserve_native_format=True
        )
        assert isinstance(result[0]["content"], str)
        assert "Checking now." in result[0]["content"]
        assert "tool_calls" in result[0]


class TestContentToText:
    """Tests for the _content_to_text helper."""

    def test_none(self):
        assert _content_to_text(None) == ""

    def test_string(self):
        assert _content_to_text("hello") == "hello"

    def test_empty_string(self):
        assert _content_to_text("") == ""

    def test_list_of_content_parts(self):
        parts = [
            ContentPart(type="text", text="Hello"),
            ContentPart(type="text", text="World"),
        ]
        assert _content_to_text(parts) == "Hello\nWorld"

    def test_list_of_dicts(self):
        parts = [
            {"type": "text", "text": "foo"},
            {"type": "image_url", "image_url": "http://img"},
        ]
        assert _content_to_text(parts) == "foo"

    def test_list_with_no_text_parts(self):
        parts = [{"type": "image_url", "image_url": "http://img"}]
        assert _content_to_text(parts) == ""

    def test_empty_list(self):
        assert _content_to_text([]) == ""


class TestGptOssSpecialTokens:
    """Tests for GPT-OSS channel token handling in utils."""

    def test_pattern_matches_channel_token(self):
        assert SPECIAL_TOKENS_PATTERN.search("<|channel|>") is not None

    def test_pattern_matches_message_token(self):
        assert SPECIAL_TOKENS_PATTERN.search("<|message|>") is not None

    def test_pattern_matches_start_token(self):
        assert SPECIAL_TOKENS_PATTERN.search("<|start|>") is not None

    def test_pattern_matches_return_token(self):
        assert SPECIAL_TOKENS_PATTERN.search("<|return|>") is not None

    def test_pattern_matches_call_token(self):
        assert SPECIAL_TOKENS_PATTERN.search("<|call|>") is not None

    def test_clean_output_extracts_final_channel(self):
        text = (
            "<|channel|>analysis<|message|>Thinking about it"
            "<|start|>assistant<|channel|>final<|message|>The answer is 42<|return|>"
        )
        result = clean_output_text(text)
        assert result == "The answer is 42"
        assert "<|" not in result

    def test_clean_output_final_only(self):
        text = "<|channel|>final<|message|>Just the answer<|return|>"
        result = clean_output_text(text)
        assert result == "Just the answer"

    def test_clean_output_strips_return_token(self):
        text = "<|channel|>final<|message|>Hello world<|return|>"
        result = clean_output_text(text)
        assert "<|return|>" not in result
        assert result == "Hello world"

    def test_clean_output_no_channel_tokens_passthrough(self):
        text = "Normal text without any channel tokens."
        result = clean_output_text(text)
        assert result == text

    def test_pattern_matches_constrain_token(self):
        assert SPECIAL_TOKENS_PATTERN.search("<|constrain|>") is not None

    def test_clean_output_constrain_format(self):
        """Should extract final content from extended constrain format."""
        text = (
            "<|channel|>analysis<|message|>Thinking"
            "<|end|><|channel|>final <|constrain|>JSON<|message|>"
            '{"hello":"world"}<|return|>'
        )
        result = clean_output_text(text)
        assert result == '{"hello":"world"}'
        assert "<|constrain|>" not in result
        assert "<|channel|>" not in result

    def test_clean_output_constrain_final_only(self):
        """Should handle constrain format with only final channel."""
        text = '<|channel|>final <|constrain|>JSON<|message|>{"key":"value"}<|return|>'
        result = clean_output_text(text)
        assert result == '{"key":"value"}'

    def test_clean_output_no_final_strips_constrain(self):
        """When no final channel found, constrain tokens should be stripped."""
        text = "<|channel|>analysis<|message|>Just thinking <|constrain|>something"
        result = clean_output_text(text)
        assert "<|constrain|>" not in result