初始化提交
This commit is contained in:
lvyulong
505
test/services/test_llm.py
Normal file
505
test/services/test_llm.py
Normal file
@@ -0,0 +1,505 @@
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
from pydantic import ValidationError
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).parent.parent.parent))
|
||||
|
||||
from app.config import config
|
||||
from app.models.schema import VideoScriptRequest
|
||||
from app.services import llm
|
||||
|
||||
|
||||
class TestScriptPromptOptions(unittest.TestCase):
|
||||
def test_build_script_prompt_appends_advanced_requirements(self):
|
||||
"""
|
||||
高级文案要求只作为附加约束，不替换默认系统提示词。
|
||||
这样普通用户不配置时仍然走稳定默认规则，高级用户也能细化风格。
|
||||
"""
|
||||
prompt = llm.build_script_prompt(
|
||||
video_subject="咖啡",
|
||||
language="zh-CN",
|
||||
paragraph_number=3,
|
||||
video_script_prompt="语气轻松，面向程序员",
|
||||
)
|
||||
|
||||
self.assertIn("# Role: Video Script Generator", prompt)
|
||||
self.assertIn("- video subject: 咖啡", prompt)
|
||||
self.assertIn("- number of paragraphs: 3", prompt)
|
||||
self.assertIn("- language: zh-CN", prompt)
|
||||
self.assertIn("# Additional User Requirements:", prompt)
|
||||
self.assertIn("语气轻松，面向程序员", prompt)
|
||||
|
||||
def test_custom_system_prompt_keeps_runtime_context(self):
|
||||
"""
|
||||
自定义 system prompt 会替换默认脚本规则，但视频主题、语言、段落数
|
||||
仍由服务层统一追加，避免高级用户漏写必要上下文。
|
||||
"""
|
||||
prompt = llm.build_script_prompt(
|
||||
video_subject="露营",
|
||||
language="en",
|
||||
paragraph_number=2,
|
||||
custom_system_prompt="Only write cinematic narration.",
|
||||
)
|
||||
|
||||
self.assertNotIn("# Role: Video Script Generator", prompt)
|
||||
self.assertIn("Only write cinematic narration.", prompt)
|
||||
self.assertIn("- video subject: 露营", prompt)
|
||||
self.assertIn("- number of paragraphs: 2", prompt)
|
||||
self.assertIn("- language: en", prompt)
|
||||
|
||||
def test_generate_script_sends_custom_prompt_to_llm(self):
|
||||
captured = {}
|
||||
|
||||
def fake_generate_response(prompt):
|
||||
captured["prompt"] = prompt
|
||||
return "第一段。\n\n第二段。"
|
||||
|
||||
with patch.object(llm, "_generate_response", side_effect=fake_generate_response):
|
||||
result = llm.generate_script(
|
||||
video_subject="咖啡",
|
||||
language="zh-CN",
|
||||
paragraph_number=2,
|
||||
video_script_prompt="开头更有悬念",
|
||||
)
|
||||
|
||||
self.assertEqual(result, "第一段。\n\n第二段。")
|
||||
self.assertIn("- number of paragraphs: 2", captured["prompt"])
|
||||
self.assertIn("开头更有悬念", captured["prompt"])
|
||||
|
||||
def test_video_script_request_rejects_invalid_advanced_options(self):
|
||||
"""
|
||||
API 请求模型需要限制高级 prompt 参数，避免外部调用绕过 WebUI
|
||||
传入异常段落数或超长提示词，导致模型成本和结果不可控。
|
||||
"""
|
||||
with self.assertRaises(ValidationError):
|
||||
VideoScriptRequest(video_subject="咖啡", paragraph_number=0)
|
||||
|
||||
with self.assertRaises(ValidationError):
|
||||
VideoScriptRequest(
|
||||
video_subject="咖啡",
|
||||
video_script_prompt="x" * (llm.MAX_SCRIPT_PROMPT_LENGTH + 1),
|
||||
)
|
||||
|
||||
|
||||
class TestLiteLLMProvider(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.original_app_config = dict(config.app)
|
||||
|
||||
def tearDown(self):
|
||||
config.app.clear()
|
||||
config.app.update(self.original_app_config)
|
||||
|
||||
def _use_litellm_provider(self, model_name="openai/gpt-4o-mini"):
|
||||
config.app["llm_provider"] = "litellm"
|
||||
config.app["litellm_model_name"] = model_name
|
||||
|
||||
def test_litellm_provider_returns_normalized_text(self):
|
||||
"""
|
||||
验证 LiteLLM provider 的主路径不依赖真实网络和私有 API key。
|
||||
|
||||
这里用 fake module 注入 `sys.modules`，直接覆盖动态 import 的
|
||||
`litellm.completion()`，确保测试稳定覆盖 `_generate_response()` 里的
|
||||
litellm 分支。
|
||||
"""
|
||||
self._use_litellm_provider()
|
||||
|
||||
fake_litellm = types.SimpleNamespace()
|
||||
|
||||
def _completion(**kwargs):
|
||||
self.assertEqual(kwargs["model"], "openai/gpt-4o-mini")
|
||||
self.assertEqual(
|
||||
kwargs["messages"], [{"role": "user", "content": "Say hello"}]
|
||||
)
|
||||
self.assertTrue(kwargs["drop_params"])
|
||||
message = types.SimpleNamespace(content="hello\nworld")
|
||||
choice = types.SimpleNamespace(message=message)
|
||||
return types.SimpleNamespace(choices=[choice])
|
||||
|
||||
fake_litellm.completion = _completion
|
||||
|
||||
with patch.dict(sys.modules, {"litellm": fake_litellm}):
|
||||
result = llm._generate_response("Say hello")
|
||||
|
||||
self.assertEqual(result, "helloworld")
|
||||
|
||||
def test_litellm_provider_requires_model_name(self):
|
||||
self._use_litellm_provider(model_name="")
|
||||
|
||||
result = llm._generate_response("test")
|
||||
|
||||
self.assertIn("Error:", result)
|
||||
self.assertIn("model_name is not set", result)
|
||||
|
||||
def test_litellm_provider_handles_empty_response(self):
|
||||
self._use_litellm_provider()
|
||||
|
||||
fake_litellm = types.SimpleNamespace(
|
||||
completion=lambda **kwargs: types.SimpleNamespace(choices=[])
|
||||
)
|
||||
|
||||
with patch.dict(sys.modules, {"litellm": fake_litellm}):
|
||||
result = llm._generate_response("test")
|
||||
|
||||
self.assertIn("Error:", result)
|
||||
self.assertIn("returned empty response", result)
|
||||
|
||||
def test_litellm_provider_handles_empty_message(self):
|
||||
"""
|
||||
某些 OpenAI-compatible 网关在内容过滤或安全拦截时会返回
|
||||
HTTP 200，但 `choices[0].message` 为 None。这里必须返回
|
||||
可诊断的错误，而不是抛出 AttributeError。
|
||||
"""
|
||||
self._use_litellm_provider()
|
||||
|
||||
fake_litellm = types.SimpleNamespace(
|
||||
completion=lambda **kwargs: types.SimpleNamespace(
|
||||
choices=[types.SimpleNamespace(message=None)]
|
||||
)
|
||||
)
|
||||
|
||||
with patch.dict(sys.modules, {"litellm": fake_litellm}):
|
||||
result = llm._generate_response("test")
|
||||
|
||||
self.assertIn("Error:", result)
|
||||
self.assertIn("returned empty message", result)
|
||||
|
||||
def test_openai_provider_still_uses_existing_path(self):
|
||||
config.app["llm_provider"] = "openai"
|
||||
config.app["openai_api_key"] = ""
|
||||
config.app["openai_base_url"] = "https://api.openai.com/v1"
|
||||
config.app["openai_model_name"] = "gpt-4o-mini"
|
||||
|
||||
result = llm._generate_response("test")
|
||||
|
||||
self.assertIn("Error:", result)
|
||||
self.assertIn("api_key is not set", result)
|
||||
self.assertNotIn("litellm", result.lower())
|
||||
|
||||
def test_grok_provider_still_uses_existing_path(self):
|
||||
config.app["llm_provider"] = "grok"
|
||||
config.app["grok_api_key"] = ""
|
||||
config.app["grok_base_url"] = "https://api.x.ai/v1"
|
||||
config.app["grok_model_name"] = "grok-4.3"
|
||||
|
||||
result = llm._generate_response("test")
|
||||
|
||||
self.assertIn("Error:", result)
|
||||
self.assertIn("api_key is not set", result)
|
||||
self.assertNotIn("litellm", result.lower())
|
||||
|
||||
def _use_ollama_provider(self, base_url=""):
|
||||
config.app["llm_provider"] = "ollama"
|
||||
config.app["ollama_api_key"] = ""
|
||||
config.app["ollama_base_url"] = base_url
|
||||
config.app["ollama_model_name"] = "llama3"
|
||||
|
||||
def _assert_ollama_base_url(self, expected_base_url: str):
|
||||
class FakeCompletions:
|
||||
def create(self, **kwargs):
|
||||
self.kwargs = kwargs
|
||||
message = types.SimpleNamespace(content="hello\nollama")
|
||||
choice = types.SimpleNamespace(message=message)
|
||||
return types.SimpleNamespace(choices=[choice])
|
||||
|
||||
fake_completions = FakeCompletions()
|
||||
fake_client = types.SimpleNamespace(
|
||||
chat=types.SimpleNamespace(completions=fake_completions)
|
||||
)
|
||||
|
||||
with (
|
||||
patch.object(llm, "OpenAI", return_value=fake_client) as openai_client,
|
||||
patch.object(llm, "ChatCompletion", types.SimpleNamespace),
|
||||
):
|
||||
result = llm._generate_response("Say hello")
|
||||
|
||||
openai_client.assert_called_once_with(
|
||||
api_key="ollama",
|
||||
base_url=expected_base_url,
|
||||
)
|
||||
self.assertEqual(
|
||||
fake_completions.kwargs,
|
||||
{
|
||||
"model": "llama3",
|
||||
"messages": [{"role": "user", "content": "Say hello"}],
|
||||
},
|
||||
)
|
||||
self.assertEqual(result, "helloollama")
|
||||
|
||||
def test_ollama_default_base_url_uses_localhost_outside_container(self):
|
||||
"""
|
||||
普通本机运行时，Ollama 默认仍然使用 localhost，避免影响已有用户。
|
||||
"""
|
||||
self._use_ollama_provider()
|
||||
|
||||
with patch.object(config, "is_running_in_container", return_value=False):
|
||||
self._assert_ollama_base_url("http://localhost:11434/v1")
|
||||
|
||||
def test_ollama_default_base_url_uses_host_gateway_inside_container(self):
|
||||
"""
|
||||
容器内运行时，localhost 指向容器自身；默认改为 host.docker.internal，
|
||||
方便 Docker Desktop 用户访问宿主机上的 Ollama。
|
||||
"""
|
||||
self._use_ollama_provider()
|
||||
|
||||
with (
|
||||
patch.object(config, "is_running_in_container", return_value=True),
|
||||
patch.object(config, "_can_resolve_hostname", return_value=True),
|
||||
):
|
||||
self._assert_ollama_base_url("http://host.docker.internal:11434/v1")
|
||||
|
||||
def test_ollama_default_base_url_falls_back_to_container_gateway(self):
|
||||
"""
|
||||
原生 Linux Docker 里不一定能解析 host.docker.internal。此时使用容器
|
||||
默认网关作为兜底地址，比直接返回不可解析的 hostname 更稳。
|
||||
"""
|
||||
self._use_ollama_provider()
|
||||
|
||||
with (
|
||||
patch.object(config, "is_running_in_container", return_value=True),
|
||||
patch.object(config, "_can_resolve_hostname", return_value=False),
|
||||
patch.object(config, "get_container_default_gateway_ip", return_value="172.17.0.1"),
|
||||
):
|
||||
self._assert_ollama_base_url("http://172.17.0.1:11434/v1")
|
||||
|
||||
def test_ollama_explicit_base_url_takes_precedence(self):
|
||||
"""
|
||||
用户手动配置的 ollama_base_url 优先级最高，不受容器检测影响。
|
||||
"""
|
||||
self._use_ollama_provider(base_url="http://ollama:11434/v1")
|
||||
|
||||
with patch.object(config, "is_running_in_container", return_value=True):
|
||||
self._assert_ollama_base_url("http://ollama:11434/v1")
|
||||
|
||||
def test_mimo_provider_uses_openai_compatible_client(self):
|
||||
"""
|
||||
MiMo 官方接口兼容 OpenAI Chat Completions 协议。这里用 fake OpenAI
|
||||
client 验证 provider 会使用 MiMo 独立配置和默认 base_url，不依赖
|
||||
真实网络或私有 API Key。
|
||||
"""
|
||||
config.app["llm_provider"] = "mimo"
|
||||
config.app["mimo_api_key"] = "mimo-key"
|
||||
config.app["mimo_base_url"] = ""
|
||||
config.app["mimo_model_name"] = ""
|
||||
|
||||
class FakeCompletions:
|
||||
def create(self, **kwargs):
|
||||
self.kwargs = kwargs
|
||||
message = types.SimpleNamespace(content="hello\nmimo")
|
||||
choice = types.SimpleNamespace(message=message)
|
||||
return types.SimpleNamespace(choices=[choice])
|
||||
|
||||
fake_completions = FakeCompletions()
|
||||
fake_client = types.SimpleNamespace(
|
||||
chat=types.SimpleNamespace(completions=fake_completions)
|
||||
)
|
||||
|
||||
with (
|
||||
patch.object(llm, "OpenAI", return_value=fake_client) as openai_client,
|
||||
patch.object(llm, "ChatCompletion", types.SimpleNamespace),
|
||||
):
|
||||
result = llm._generate_response("Say hello")
|
||||
|
||||
openai_client.assert_called_once_with(
|
||||
api_key="mimo-key",
|
||||
base_url="https://api.xiaomimimo.com/v1",
|
||||
)
|
||||
self.assertEqual(
|
||||
fake_completions.kwargs,
|
||||
{
|
||||
"model": "mimo-v2.5-pro",
|
||||
"messages": [{"role": "user", "content": "Say hello"}],
|
||||
},
|
||||
)
|
||||
self.assertEqual(result, "hellomimo")
|
||||
|
||||
def test_azure_provider_uses_azure_client_directly(self):
|
||||
"""
|
||||
Azure OpenAI 的鉴权、endpoint 和 api-version 都由 AzureOpenAI 客户端处理。
|
||||
这个测试覆盖 issue #892：azure 分支必须直接调用 AzureOpenAI 创建的客户端，
|
||||
不能继续落入普通 OpenAI-compatible 分支，否则会丢失 Azure 专用请求配置。
|
||||
"""
|
||||
config.app["llm_provider"] = "azure"
|
||||
config.app["azure_api_key"] = "azure-key"
|
||||
config.app["azure_base_url"] = "https://example.openai.azure.com"
|
||||
config.app["azure_model_name"] = "gpt-4o-mini"
|
||||
config.app["azure_api_version"] = "2024-02-15-preview"
|
||||
|
||||
class FakeCompletions:
|
||||
def create(self, **kwargs):
|
||||
self.kwargs = kwargs
|
||||
message = types.SimpleNamespace(content="hello\nazure")
|
||||
choice = types.SimpleNamespace(message=message)
|
||||
return types.SimpleNamespace(choices=[choice])
|
||||
|
||||
fake_completions = FakeCompletions()
|
||||
fake_client = types.SimpleNamespace(
|
||||
chat=types.SimpleNamespace(completions=fake_completions)
|
||||
)
|
||||
|
||||
with (
|
||||
patch.object(llm, "AzureOpenAI", return_value=fake_client) as azure_client,
|
||||
patch.object(llm, "OpenAI") as openai_client,
|
||||
patch.object(llm, "ChatCompletion", types.SimpleNamespace),
|
||||
):
|
||||
result = llm._generate_response("Say hello")
|
||||
|
||||
azure_client.assert_called_once_with(
|
||||
api_key="azure-key",
|
||||
api_version="2024-02-15-preview",
|
||||
azure_endpoint="https://example.openai.azure.com",
|
||||
)
|
||||
openai_client.assert_not_called()
|
||||
self.assertEqual(
|
||||
fake_completions.kwargs,
|
||||
{
|
||||
"model": "gpt-4o-mini",
|
||||
"messages": [{"role": "user", "content": "Say hello"}],
|
||||
},
|
||||
)
|
||||
self.assertEqual(result, "helloazure")
|
||||
|
||||
def test_g4f_provider_requires_explicit_opt_in(self):
|
||||
"""
|
||||
g4f 存在供应链和稳定性风险，不能因为用户把 provider 写成 g4f
|
||||
就默认加载第三方包并访问逆向接口，必须显式启用。
|
||||
"""
|
||||
config.app["llm_provider"] = "g4f"
|
||||
config.app["enable_g4f"] = False
|
||||
|
||||
result = llm._generate_response("test")
|
||||
|
||||
self.assertIn("Error:", result)
|
||||
self.assertIn("g4f provider is disabled", result)
|
||||
|
||||
def test_g4f_provider_uses_lazy_import_after_opt_in(self):
|
||||
config.app["llm_provider"] = "g4f"
|
||||
config.app["enable_g4f"] = True
|
||||
config.app["g4f_model_name"] = "gpt-3.5-turbo"
|
||||
|
||||
fake_g4f = types.SimpleNamespace()
|
||||
fake_g4f.ChatCompletion = types.SimpleNamespace(
|
||||
create=lambda **kwargs: "hello from g4f"
|
||||
)
|
||||
|
||||
with patch.dict(sys.modules, {"g4f": fake_g4f}):
|
||||
result = llm._generate_response("test")
|
||||
|
||||
self.assertEqual(result, "hello from g4f")
|
||||
|
||||
def test_g4f_provider_reports_missing_optional_dependency(self):
|
||||
config.app["llm_provider"] = "g4f"
|
||||
config.app["enable_g4f"] = True
|
||||
config.app["g4f_model_name"] = "gpt-3.5-turbo"
|
||||
|
||||
with patch.dict(sys.modules, {"g4f": None}):
|
||||
result = llm._generate_response("test")
|
||||
|
||||
self.assertIn("Error:", result)
|
||||
self.assertIn("g4f package is not installed by default", result)
|
||||
|
||||
|
||||
class TestRuntimeEnvironmentDetection(unittest.TestCase):
|
||||
def test_container_detection_ignores_plain_linux_cgroup_file(self):
|
||||
"""
|
||||
普通 Linux 也有 /proc/1/cgroup，不能因为文件存在就判定为容器。
|
||||
"""
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
cgroup_path = Path(tmp_dir) / "cgroup"
|
||||
cgroup_path.write_text("0::/init.scope\n", encoding="utf-8")
|
||||
|
||||
self.assertFalse(
|
||||
config.is_running_in_container(
|
||||
dockerenv_path=str(Path(tmp_dir) / "missing-dockerenv"),
|
||||
containerenv_path=str(Path(tmp_dir) / "missing-containerenv"),
|
||||
cgroup_path=str(cgroup_path),
|
||||
)
|
||||
)
|
||||
|
||||
def test_container_detection_accepts_dockerenv_marker(self):
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
dockerenv_path = Path(tmp_dir) / ".dockerenv"
|
||||
dockerenv_path.write_text("", encoding="utf-8")
|
||||
|
||||
self.assertTrue(
|
||||
config.is_running_in_container(
|
||||
dockerenv_path=str(dockerenv_path),
|
||||
containerenv_path=str(Path(tmp_dir) / "missing-containerenv"),
|
||||
cgroup_path=str(Path(tmp_dir) / "missing-cgroup"),
|
||||
)
|
||||
)
|
||||
|
||||
def test_container_detection_accepts_cgroup_container_marker(self):
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
cgroup_path = Path(tmp_dir) / "cgroup"
|
||||
cgroup_path.write_text(
|
||||
"0::/system.slice/docker-abcdef.scope\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
self.assertTrue(
|
||||
config.is_running_in_container(
|
||||
dockerenv_path=str(Path(tmp_dir) / "missing-dockerenv"),
|
||||
containerenv_path=str(Path(tmp_dir) / "missing-containerenv"),
|
||||
cgroup_path=str(cgroup_path),
|
||||
)
|
||||
)
|
||||
|
||||
def test_container_gateway_ip_decodes_default_route(self):
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
route_path = Path(tmp_dir) / "route"
|
||||
route_path.write_text(
|
||||
"Iface\tDestination\tGateway\tFlags\tRefCnt\tUse\tMetric\tMask\tMTU\tWindow\tIRTT\n"
|
||||
"eth0\t00000000\t010011AC\t0003\t0\t0\t0\t00000000\t0\t0\t0\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
config.get_container_default_gateway_ip(str(route_path)),
|
||||
"172.17.0.1",
|
||||
)
|
||||
|
||||
def test_container_gateway_ip_ignores_missing_default_route(self):
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
route_path = Path(tmp_dir) / "route"
|
||||
route_path.write_text(
|
||||
"Iface\tDestination\tGateway\tFlags\tRefCnt\tUse\tMetric\tMask\tMTU\tWindow\tIRTT\n"
|
||||
"eth0\t0011AC0A\t00000000\t0001\t0\t0\t0\t00FFFFFF\t0\t0\t0\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
self.assertEqual(config.get_container_default_gateway_ip(str(route_path)), "")
|
||||
|
||||
|
||||
FOUNDRY_KEY = os.environ.get("ANTHROPIC_FOUNDRY_API_KEY", "")
|
||||
FOUNDRY_BASE = "https://amanrai-test-resource.services.ai.azure.com/anthropic"
|
||||
FOUNDRY_MODEL = "azure_ai/claude-sonnet-4-6"
|
||||
|
||||
|
||||
@unittest.skipUnless(FOUNDRY_KEY, "ANTHROPIC_FOUNDRY_API_KEY not set")
|
||||
class TestLiteLLMLiveIntegration(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.original_app_config = dict(config.app)
|
||||
config.app["llm_provider"] = "litellm"
|
||||
config.app["litellm_model_name"] = FOUNDRY_MODEL
|
||||
os.environ["AZURE_AI_API_KEY"] = FOUNDRY_KEY
|
||||
os.environ["AZURE_AI_API_BASE"] = FOUNDRY_BASE
|
||||
|
||||
def tearDown(self):
|
||||
config.app.clear()
|
||||
config.app.update(self.original_app_config)
|
||||
|
||||
def test_live_litellm_completion(self):
|
||||
result = llm._generate_response("What is 2+2? Reply with just the number.")
|
||||
|
||||
self.assertNotIn("Error:", result)
|
||||
self.assertIn("4", result)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user