117 lines
4.1 KiB
Python
117 lines
4.1 KiB
Python
import os
|
|
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_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_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, "high")
|
|
self.assertEqual(caps.max_output, 8192)
|
|
self.assertEqual(caps.output_cny_per_mtoken, 2.0)
|
|
self.assertEqual(caps.cache_hit_cny_per_mtoken, 0.02)
|
|
self.assertEqual(caps.thinking_transport, "extra_body")
|
|
self.assertEqual(caps.reasoning_replay, "tool_calls")
|
|
|
|
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.assertFalse(glm.thinking_enabled)
|
|
self.assertEqual(glm.thinking_transport, "extra_body")
|
|
self.assertTrue(doubao.thinking_enabled)
|
|
self.assertEqual(doubao.thinking_transport, "extra_body")
|
|
self.assertEqual(doubao.default_reasoning_effort, "")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|