"""LLM 模块测试。""" from __future__ import annotations import pytest from unittest.mock import patch, MagicMock from dataclasses import dataclass from data.llm.client import LLMClient, LLMResponse, LLMError from data.llm.prompts import ( build_audit_prompt, build_cwb_prompt, build_full_plan_prompt, ) @dataclass class MockSettings: llm_provider: str = "none" llm_api_key: str = "" llm_base_url: str = "https://dashscope.aliyuncs.com/compatible-mode/v1" llm_model: str = "qwen-plus" llm_timeout_seconds: int = 60 llm_max_tokens: int = 4096 llm_fallback_models: str = "" llm_fallback_providers: str = "" llm_fallback_api_keys: str = "" llm_fallback_base_urls: str = "" class TestLLMClient: def test_not_configured_when_provider_none(self): client = LLMClient(MockSettings(llm_provider="none")) assert not client.is_configured def test_not_configured_when_no_api_key(self): client = LLMClient(MockSettings(llm_provider="openai", llm_api_key="")) assert not client.is_configured def test_configured_when_provider_and_key_set(self): client = LLMClient(MockSettings(llm_provider="openai", llm_api_key="sk-test")) assert client.is_configured def test_chat_raises_when_not_configured(self): client = LLMClient(MockSettings(llm_provider="none")) with pytest.raises(LLMError, match="LLM 未配置"): client.chat([{"role": "user", "content": "test"}]) @patch("data.llm.client.urllib.request.urlopen") def test_chat_success(self, mock_urlopen): mock_resp = MagicMock() mock_resp.read.return_value = b'{"choices":[{"message":{"content":"test response"}}],"usage":{"total_tokens":10},"model":"qwen-plus"}' mock_resp.__enter__ = MagicMock(return_value=mock_resp) mock_resp.__exit__ = MagicMock(return_value=False) mock_urlopen.return_value = mock_resp client = LLMClient( MockSettings(llm_provider="dashscope", llm_api_key="sk-test") ) result = client.chat([{"role": "user", "content": "hello"}]) assert isinstance(result, LLMResponse) assert result.content == "test response" assert result.model == "qwen-plus" assert result.usage["total_tokens"] == 10 @patch("data.llm.client.urllib.request.urlopen") def test_chat_http_error(self, mock_urlopen): import urllib.error import io error_fp = io.BytesIO(b'{"error":"bad key"}') mock_urlopen.side_effect = urllib.error.HTTPError( "http://test", 401, "Unauthorized", {}, error_fp ) client = LLMClient(MockSettings(llm_provider="openai", llm_api_key="bad")) with pytest.raises(LLMError, match="HTTP 401"): client.chat([{"role": "user", "content": "test"}]) def test_chat_with_system(self): """chat_with_system 构建正确的消息结构。""" with patch.object(LLMClient, "chat") as mock_chat: mock_chat.return_value = LLMResponse(content="ok") client = LLMClient( MockSettings(llm_provider="openai", llm_api_key="sk-test") ) client.chat_with_system("you are helpful", "hello") call_args = mock_chat.call_args messages = call_args[0][0] assert len(messages) == 2 assert messages[0]["role"] == "system" assert messages[0]["content"] == "you are helpful" assert messages[1]["role"] == "user" assert messages[1]["content"] == "hello" class TestPrompts: def test_audit_prompt_structure(self): system, user = build_audit_prompt( province="湖南", score=578, rank=12000, subjects=["物理", "化学", "生物"], existing_plan="已有一版方案", ) assert "志愿填报顾问" in system assert "湖南" in user assert "578" in user assert "物理" in user assert "JSON" in user assert "risk_level" in user def test_audit_prompt_with_crowd_db(self): recs = [ {"name": "湖南大学", "major": "计算机"}, {"name": "中南大学", "major": "软件工程"}, ] system, user = build_audit_prompt( province="湖南", score=578, rank=12000, subjects=["物理"], existing_plan="test", crowd_db_recs=recs, ) assert "湖南大学" in user assert "中南大学" in user def test_audit_prompt_minimal(self): system, user = build_audit_prompt( province="广东", score=None, rank=None, subjects=[], existing_plan="", ) assert "未提供" in user def test_cwb_prompt_structure(self): system, user = build_cwb_prompt( province="湖南", score=578, rank=12000, subjects=["物理", "化学", "生物"], target_cities=["长沙", "深圳"], ) assert "冲稳保" in system assert "长沙" in user assert "598" in user # score + 20 assert "558" in user # score - 20 def test_full_plan_prompt_structure(self): system, user = build_full_plan_prompt( province="湖南", score=578, rank=12000, subjects=["物理", "化学", "生物"], target_majors=["计算机科学", "人工智能"], family_background="家长希望省内优先", ) assert "完整" in system assert "计算机科学" in user assert "家长希望省内优先" in user assert "volunteers" in user assert "至少 8 条" in user class TestFallback: def test_fallback_config_parsed(self): """fallback 配置被正确解析为供应商链。""" settings = MockSettings( llm_provider="dashscope", llm_api_key="key-main", llm_model="qwen-plus", llm_fallback_models="gpt-4o-mini,deepseek-chat", llm_fallback_providers="openai,deepseek", llm_fallback_api_keys="key-openai,key-deepseek", llm_fallback_base_urls="https://api.openai.com/v1,https://api.deepseek.com/v1", ) client = LLMClient(settings) assert client.provider_count == 3 assert client._providers[0]["model"] == "qwen-plus" assert client._providers[1]["model"] == "gpt-4o-mini" assert client._providers[2]["model"] == "deepseek-chat" assert client._providers[1]["api_key"] == "key-openai" assert client._providers[2]["provider"] == "deepseek" def test_fallback_no_config(self): """无 fallback 配置时只有主供应商。""" client = LLMClient(MockSettings(llm_provider="openai", llm_api_key="sk-test")) assert client.provider_count == 1 def test_fallback_falls_through_to_second_provider(self): """主供应商失败时，自动尝试第二个供应商。""" settings = MockSettings( llm_provider="dashscope", llm_api_key="key-main", llm_model="qwen-plus", llm_fallback_models="gpt-4o-mini", llm_fallback_providers="openai", llm_fallback_api_keys="key-openai", llm_fallback_base_urls="https://api.openai.com/v1", ) client = LLMClient(settings) call_count = {"n": 0} def mock_call_side_effect(**kwargs): call_count["n"] += 1 if call_count["n"] == 1: raise LLMError("first provider failed") return LLMResponse(content="fallback success", model="gpt-4o-mini") with patch.object( client, "_call_single_provider", side_effect=mock_call_side_effect ): result = client.chat([{"role": "user", "content": "test"}]) assert result.content == "fallback success" assert call_count["n"] == 2 def test_fallback_all_fail_raises(self): """全部供应商都失败时抛出聚合错误。""" settings = MockSettings( llm_provider="dashscope", llm_api_key="key-main", llm_model="qwen-plus", llm_fallback_models="gpt-4o-mini", llm_fallback_providers="openai", llm_fallback_api_keys="key-openai", ) client = LLMClient(settings) with patch.object( client, "_call_single_provider", side_effect=LLMError("all fail") ): with pytest.raises(LLMError, match="全部 2 个"): client.chat([{"role": "user", "content": "test"}])