506 lines
19 KiB
Python
506 lines
19 KiB
Python
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()
|