from __future__ import annotations

import unittest

from utils.api_client import DeepSeekClient


class APIClientMultimodalSanitizeTest(unittest.TestCase):
    def test_text_only_model_strips_image_parts_to_text(self):
        client = DeepSeekClient(web_mode=True)
        client.model_key = "DeepSeek-V4-Pro High"  # custom_models.json: multimodal=none

        messages = [
            {
                "role": "tool",
                "name": "mcp__playwright__browser_take_screenshot",
                "content": [
                    {"type": "text", "text": "截图成功"},
                    {"type": "image_url", "image_url": {"url": "data:image/png;base64,abc"}},
                ],
            }
        ]

        sanitized = client._sanitize_messages_for_model_capability(messages)
        self.assertEqual(len(sanitized), 1)
        self.assertIsInstance(sanitized[0].get("content"), str, sanitized[0])
        self.assertIn("截图成功", sanitized[0].get("content") or "", sanitized[0])

    def test_multimodal_model_keeps_image_parts(self):
        client = DeepSeekClient(web_mode=True)
        client.model_key = "qwen3-vl-plus"  # 内置多模态模型

        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": "看图"},
                    {"type": "image_url", "image_url": {"url": "data:image/png;base64,abc"}},
                ],
            }
        ]

        sanitized = client._sanitize_messages_for_model_capability(messages)
        self.assertEqual(len(sanitized), 1)
        self.assertIsInstance(sanitized[0].get("content"), list, sanitized[0])
        self.assertEqual(
            [part.get("type") for part in sanitized[0].get("content") or []],
            ["text", "image_url"],
            sanitized[0],
        )


if __name__ == "__main__":
    unittest.main()