hunjing/backend/tests/ai_provider/routing/test_multimodal.py

from __future__ import annotations

from app.ai_provider.routing.multimodal import MultiModalDecision, MultiModalRouter


class TestMultiModalDecision:
    def test_enum_values(self):
        assert MultiModalDecision.PLAIN_TEXT == "plain_text"
        assert MultiModalDecision.VISION == "vision"
        assert MultiModalDecision.THINKING == "thinking"
        assert MultiModalDecision.TOOL_USE == "tool_use"
        assert MultiModalDecision.STRUCTURED_OUTPUT == "structured_output"

    def test_enum_is_str_enum(self):
        assert issubclass(MultiModalDecision, str)


class TestMultiModalRouter:
    def test_plain_text_message(self):
        messages = [{"role": "user", "content": "Hello, world!"}]
        decisions = MultiModalRouter.decide_from_messages(messages)
        assert decisions == {MultiModalDecision.PLAIN_TEXT}

    def test_plain_text_multiple_messages(self):
        messages = [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": "What is Python?"},
            {"role": "assistant", "content": "Python is a programming language."},
        ]
        decisions = MultiModalRouter.decide_from_messages(messages)
        assert decisions == {MultiModalDecision.PLAIN_TEXT}

    def test_vision_with_openai_image_url_block(self):
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": "What's in this image?"},
                    {"type": "image_url", "image_url": {"url": "https://example.com/img.png"}},
                ],
            }
        ]
        decisions = MultiModalRouter.decide_from_messages(messages)
        assert MultiModalDecision.VISION in decisions

    def test_vision_with_image_type_block(self):
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": "Describe this"},
                    {
                        "type": "image",
                        "source": {"type": "base64", "media_type": "image/png", "data": "..."},
                    },
                ],
            }
        ]
        decisions = MultiModalRouter.decide_from_messages(messages)
        assert MultiModalDecision.VISION in decisions

    def test_plain_text_with_single_content_block(self):
        """A single text-only content block should not trigger VISION."""
        messages = [
            {
                "role": "user",
                "content": [{"type": "text", "text": "Hello!"}],
            }
        ]
        decisions = MultiModalRouter.decide_from_messages(messages)
        assert MultiModalDecision.VISION not in decisions

    def test_tool_use_from_assistant_tool_calls(self):
        messages = [
            {"role": "user", "content": "What's the weather?"},
            {
                "role": "assistant",
                "content": "",
                "tool_calls": [
                    {"id": "call_1", "function": {"name": "get_weather", "arguments": "{}"}}
                ],
            },
        ]
        decisions = MultiModalRouter.decide_from_messages(messages)
        assert MultiModalDecision.TOOL_USE in decisions

    def test_tool_use_from_tool_role(self):
        messages = [
            {"role": "user", "content": "What's the weather?"},
            {
                "role": "assistant",
                "content": "",
                "tool_calls": [{"id": "call_1", "name": "get_weather"}],
            },
            {"role": "tool", "content": "25°C", "tool_call_id": "call_1"},
        ]
        decisions = MultiModalRouter.decide_from_messages(messages)
        assert MultiModalDecision.TOOL_USE in decisions

    def test_thinking_from_reasoning_field(self):
        messages = [
            {"role": "user", "content": "Think step by step"},
            {"role": "assistant", "content": "Final answer", "reasoning": "Let me think..."},
        ]
        decisions = MultiModalRouter.decide_from_messages(messages)
        assert MultiModalDecision.THINKING in decisions

    def test_empty_messages(self):
        decisions = MultiModalRouter.decide_from_messages([])
        assert decisions == {MultiModalDecision.PLAIN_TEXT}

    def test_multiple_capabilities_combined(self):
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": "Analyze this image step by step"},
                    {"type": "image_url", "image_url": {"url": "https://example.com/img.png"}},
                ],
            },
            {
                "role": "assistant",
                "content": "Let me think about this image...",
                "reasoning": "First, I see an image. Let me analyze it carefully.",
                "tool_calls": [
                    {"id": "call_1", "function": {"name": "analyze_image", "arguments": "{}"}}
                ],
            },
            {"role": "tool", "content": "Analysis complete", "tool_call_id": "call_1"},
        ]
        decisions = MultiModalRouter.decide_from_messages(messages)
        assert MultiModalDecision.VISION in decisions
        assert MultiModalDecision.THINKING in decisions
        assert MultiModalDecision.TOOL_USE in decisions

    def test_decide_with_empty_content(self):
        messages = [{"role": "user", "content": ""}]
        decisions = MultiModalRouter.decide_from_messages(messages)
        assert decisions == {MultiModalDecision.PLAIN_TEXT}