zcbot/tests/test_llm_kwargs.py

import os
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch

from core.capabilities import ModelCapabilities
from core.llm import LLM


class LLMKwargsTests(unittest.TestCase):
    def _llm(
        self, *, family: str, thinking_enabled: bool, thinking_transport: str
    ) -> LLM:
        caps = ModelCapabilities(
            family=family,
            model_id=f"{family}/model",
            api_key_env="TEST_LLM_API_KEY",
            thinking_enabled=thinking_enabled,
            thinking_transport=thinking_transport,
            optimal_temperature=0.3,
        )
        with patch.dict(os.environ, {"TEST_LLM_API_KEY": "test-key"}):
            return LLM(caps)

    def test_deepseek_explicitly_enables_thinking_and_sets_effort(self) -> None:
        llm = self._llm(
            family="deepseek_v4", thinking_enabled=True, thinking_transport="extra_body"
        )

        kwargs = llm._build_kwargs(
            [{"role": "user", "content": "hello"}], None, None, "high"
        )

        self.assertNotIn("reasoning_effort", kwargs)
        self.assertEqual(
            kwargs["extra_body"],
            {"thinking": {"type": "enabled"}, "reasoning_effort": "high"},
        )

    def test_deepseek_explicitly_disables_thinking_without_effort(self) -> None:
        llm = self._llm(
            family="deepseek_v4", thinking_enabled=False, thinking_transport="extra_body"
        )

        kwargs = llm._build_kwargs(
            [{"role": "user", "content": "hello"}], None, None, "high"
        )

        self.assertNotIn("reasoning_effort", kwargs)
        self.assertEqual(
            kwargs["extra_body"], {"thinking": {"type": "disabled"}}
        )

    def test_other_openai_compatible_provider_gets_no_thinking_body(self) -> None:
        llm = self._llm(
            family="unifyllm", thinking_enabled=False, thinking_transport="none"
        )

        kwargs = llm._build_kwargs(
            [{"role": "user", "content": "hello"}], None, None, None
        )

        self.assertNotIn("extra_body", kwargs)

    def test_extra_body_transport_enables_thinking_without_effort(self) -> None:
        llm = self._llm(
            family="doubao", thinking_enabled=True, thinking_transport="extra_body"
        )

        kwargs = llm._build_kwargs(
            [{"role": "user", "content": "hello"}], None, None, None
        )

        self.assertEqual(
            kwargs["extra_body"], {"thinking": {"type": "enabled"}}
        )

    def test_glm_preserved_thinking_is_nested_in_thinking_body(self) -> None:
        caps = ModelCapabilities(
            family="glm", variant="flash53", model_id="zai/glm-5.3-flash",
            api_key_env="TEST_LLM_API_KEY", thinking_enabled=True,
            thinking_transport="extra_body", thinking_clear=False,
            default_reasoning_effort="high",
        )
        with patch.dict(os.environ, {"TEST_LLM_API_KEY": "test-key"}):
            kwargs = LLM(caps)._build_kwargs(
                [{"role": "user", "content": "hello"}], None, None, "high"
            )
        self.assertEqual(kwargs["extra_body"], {
            "thinking": {"type": "enabled", "clear_thinking": False},
            "reasoning_effort": "high",
        })

    def test_invalid_transport_is_rejected_before_request(self) -> None:
        llm = self._llm(
            family="test", thinking_enabled=True, thinking_transport="unknown"
        )

        with self.assertRaisesRegex(ValueError, "thinking_transport"):
            llm._build_kwargs(
                [{"role": "user", "content": "hello"}], None, None, "high"
            )

    def test_auto_is_rejected_at_provider_boundary(self) -> None:
        llm = self._llm(
            family="deepseek_v4", thinking_enabled=True, thinking_transport="extra_body"
        )
        with self.assertRaisesRegex(ValueError, "auto"):
            llm._build_kwargs(
                [{"role": "user", "content": "hello"}], None, None, "auto"
            )

    def test_chat_accepts_a_request_specific_short_timeout(self) -> None:
        llm = self._llm(
            family="deepseek_v4", thinking_enabled=True, thinking_transport="extra_body"
        )
        with patch("core.llm.litellm.completion", return_value=object()) as completion:
            llm.chat(
                [{"role": "user", "content": "route"}],
                reasoning_effort="low",
                max_retries=1,
                timeout_s=8,
            )
        self.assertEqual(completion.call_args.kwargs["timeout"], 8.0)

    def test_flash_profile_matches_0731_capabilities(self) -> None:
        caps = ModelCapabilities.load(
            "deepseek_v4.flash", Path(__file__).resolve().parents[1] / "config" / "models"
        )

        self.assertTrue(caps.thinking_enabled)
        self.assertEqual(caps.reasoning_effort_levels, ["low", "high", "max"])
        self.assertEqual(caps.default_reasoning_effort, "auto")
        self.assertEqual(caps.max_output, 8192)
        self.assertEqual(caps.output_cny_per_mtoken, 4.752)
        self.assertEqual(caps.cache_hit_cny_per_mtoken, 0.0504)
        self.assertEqual(
            caps.pricing["periods"][0]["revision"], "deepseek-v4-20260816"
        )
        self.assertEqual(caps.thinking_transport, "extra_body")
        self.assertEqual(caps.reasoning_replay, "tool_turn")

        pro = ModelCapabilities.load(
            "deepseek_v4.pro", Path(__file__).resolve().parents[1] / "config" / "models"
        )
        self.assertEqual(pro.default_reasoning_effort, "medium")

    def test_auto_profile_requires_low_and_high_levels(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            path = Path(tmp) / "test.yaml"
            path.write_text(
                "family: test\nvariants:\n  bad:\n"
                "    thinking_enabled: true\n"
                "    thinking_transport: extra_body\n"
                "    reasoning_effort_levels: [low, max]\n"
                "    default_reasoning_effort: auto\n",
                encoding="utf-8",
            )
            with self.assertRaisesRegex(ValueError, "low/high"):
                ModelCapabilities.load("test.bad", Path(tmp))

    def test_other_controllable_profiles_declare_transport(self) -> None:
        models_dir = Path(__file__).resolve().parents[1] / "config" / "models"

        glm = ModelCapabilities.load("glm.pro52", models_dir)
        doubao = ModelCapabilities.load("doubao.turbo", models_dir)

        self.assertEqual(glm.profile, "glm.flash53")
        self.assertTrue(glm.thinking_enabled)
        self.assertEqual(glm.thinking_transport, "extra_body")
        self.assertFalse(glm.thinking_clear)
        self.assertEqual(glm.reasoning_replay, "conversation")
        self.assertTrue(glm.native_image_input)
        self.assertTrue(doubao.thinking_enabled)
        self.assertEqual(doubao.thinking_transport, "extra_body")
        self.assertEqual(doubao.default_reasoning_effort, "")


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