11fbb591c3
LLMClient 改为供应商链模式: - 主供应商 + GAOKAO_LLM_FALLBACK_MODELS/PROVIDERS/API_KEYS/BASE_URLS - 按顺序尝试, 第一个成功即返回 - 全部失败时抛出聚合错误(含供应商数量和最后错误) 配置层: - Settings 新增 4 个 fallback 字段 - load_settings() 读取对应环境变量 新增测试(4 个): - test_fallback_config_parsed: 验证 fallback 链解析正确 - test_fallback_no_config: 无 fallback 时只有主供应商 - test_fallback_falls_through_to_second_provider: 主失败自动切换 - test_fallback_all_fail_raises: 全部失败抛聚合错误 部署示例: .env.docker.example / .env.payment.example 补充注释掉的 fallback 模板 验证: data/llm/tests 16 passed, admin/tests 40 passed
239 lines
8.6 KiB
Python
239 lines
8.6 KiB
Python
"""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"}])
|