185 lines
7.0 KiB
Python
185 lines
7.0 KiB
Python
"""LLM 弹幕回复单元测试：协议适配、过滤、限流、上下文、重试。"""
|
|
|
|
import asyncio
|
|
import unittest
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
from app.llm_client import (
|
|
LLMError,
|
|
LLMReplyEngine,
|
|
OpenAICompatProvider,
|
|
OllamaNativeProvider,
|
|
AnthropicProvider,
|
|
create_provider,
|
|
)
|
|
|
|
|
|
def _make_engine(**overrides) -> LLMReplyEngine:
|
|
cfg = {
|
|
"enabled": True,
|
|
"provider": "openai",
|
|
"base_url": "https://api.example.com/v1",
|
|
"api_key": "sk-test",
|
|
"model": "test-model",
|
|
"system_prompt": "你是直播间助手",
|
|
"max_reply_chars": 60,
|
|
"replies_per_minute": 6,
|
|
"per_user_cooldown_sec": 30,
|
|
"context_turns": 4,
|
|
"min_text_len": 2,
|
|
"blacklist_uids": [666],
|
|
"ignore_prefixes": ["//"],
|
|
}
|
|
cfg.update(overrides)
|
|
return LLMReplyEngine(cfg)
|
|
|
|
|
|
class _FakeProvider:
|
|
"""可注入的假 provider：返回固定回复，可记录调用。"""
|
|
|
|
name = "fake"
|
|
model = "fake-model"
|
|
|
|
def __init__(self, reply: str = "你好呀~", delay: float = 0.0, fail_times: int = 0):
|
|
self.reply = reply
|
|
self.delay = delay
|
|
self.fail_times = fail_times
|
|
self.calls: list[list[dict]] = []
|
|
|
|
async def chat(self, messages, *, temperature=0.8, max_tokens=256):
|
|
self.calls.append(messages)
|
|
if self.delay:
|
|
await asyncio.sleep(self.delay)
|
|
if self.fail_times > 0:
|
|
self.fail_times -= 1
|
|
raise LLMError("模拟失败")
|
|
return self.reply
|
|
|
|
|
|
def _install(engine: LLMReplyEngine, provider: _FakeProvider):
|
|
engine._provider = provider
|
|
engine._cfg_fingerprint = engine._provider_fingerprint()
|
|
return provider
|
|
|
|
|
|
class TestProviders(unittest.TestCase):
|
|
def test_create_openai(self):
|
|
p = create_provider({"provider": "openai", "base_url": "https://x/v1",
|
|
"api_key": "k", "model": "m"})
|
|
self.assertIsInstance(p, OpenAICompatProvider)
|
|
self.assertEqual(p.base_url, "https://x/v1")
|
|
|
|
def test_create_ollama(self):
|
|
p = create_provider({"provider": "ollama", "base_url": "http://localhost:11434",
|
|
"model": "qwen3"})
|
|
self.assertIsInstance(p, OllamaNativeProvider)
|
|
|
|
def test_create_anthropic(self):
|
|
p = create_provider({"provider": "anthropic", "base_url": "https://x",
|
|
"api_key": "k", "model": "claude"})
|
|
self.assertIsInstance(p, AnthropicProvider)
|
|
|
|
def test_unknown_provider(self):
|
|
with self.assertRaises(LLMError):
|
|
create_provider({"provider": "nope", "base_url": "https://x"})
|
|
|
|
def test_missing_base_url(self):
|
|
with self.assertRaises(LLMError):
|
|
create_provider({"provider": "openai"})
|
|
|
|
|
|
class TestFiltering(unittest.IsolatedAsyncioTestCase):
|
|
async def test_disabled_engine(self):
|
|
engine = _make_engine(enabled=False)
|
|
self.assertIsNone(await engine.generate_reply("你好", 1, "路人"))
|
|
|
|
async def test_too_short(self):
|
|
engine = _make_engine()
|
|
self.assertIsNone(engine.should_reply("哈", 1, "路人"))
|
|
|
|
async def test_ignore_prefix(self):
|
|
engine = _make_engine()
|
|
self.assertIsNone(engine.should_reply("//悄悄话", 1, "路人"))
|
|
|
|
async def test_blacklist(self):
|
|
engine = _make_engine()
|
|
self.assertIsNone(engine.should_reply("你好呀", 666, "黑名单"))
|
|
|
|
async def test_command_like_text_passes_through(self):
|
|
"""指令拦截在 CommandHandler 层完成；引擎本身不拦截普通文本。"""
|
|
engine = _make_engine()
|
|
self.assertIsNotNone(engine.should_reply("上号", 1, "路人"))
|
|
|
|
async def test_per_user_cooldown(self):
|
|
engine = _make_engine()
|
|
p = _install(engine, _FakeProvider())
|
|
first = await engine.generate_reply("在吗", 1, "小明")
|
|
second = await engine.generate_reply("还在吗", 1, "小明")
|
|
self.assertTrue(first)
|
|
self.assertIsNone(second) # 冷却期内被过滤
|
|
self.assertEqual(len(p.calls), 1)
|
|
|
|
async def test_zero_cooldown_disabled(self):
|
|
"""per_user_cooldown_sec=0 表示禁用冷却，不应被吞成默认 30。"""
|
|
engine = _make_engine(per_user_cooldown_sec=0)
|
|
_install(engine, _FakeProvider())
|
|
self.assertTrue(await engine.generate_reply("第一条", 1, "甲"))
|
|
self.assertTrue(await engine.generate_reply("第二条", 1, "甲"))
|
|
|
|
async def test_global_rate_limit(self):
|
|
engine = _make_engine(replies_per_minute=1, per_user_cooldown_sec=0)
|
|
_install(engine, _FakeProvider())
|
|
self.assertTrue(await engine.generate_reply("第一条", 1, "甲"))
|
|
self.assertIsNone(engine.should_reply("第二条", 2, "乙"))
|
|
|
|
async def test_busy_skip(self):
|
|
"""并发保护：在途请求未完成时新请求直接放弃。"""
|
|
engine = _make_engine(per_user_cooldown_sec=0)
|
|
_install(engine, _FakeProvider(delay=0.2))
|
|
t1 = asyncio.create_task(engine.generate_reply("第一条", 1, "甲"))
|
|
await asyncio.sleep(0.05)
|
|
t2 = asyncio.create_task(engine.generate_reply("第二条", 2, "乙"))
|
|
self.assertTrue(await t1)
|
|
self.assertIsNone(await t2)
|
|
|
|
async def test_retry_then_success(self):
|
|
engine = _make_engine(retry_attempts=2, per_user_cooldown_sec=0)
|
|
_install(engine, _FakeProvider(fail_times=1))
|
|
with patch("app.llm_client.asyncio.sleep", new=AsyncMock()):
|
|
reply = await engine.generate_reply("在吗", 1, "小明")
|
|
self.assertEqual(reply, "你好呀~")
|
|
|
|
async def test_sanitize_truncates_and_strips_markdown(self):
|
|
engine = _make_engine(max_reply_chars=10)
|
|
# 去掉 ** 后共 14 字，截断到 10 字
|
|
self.assertEqual(engine._sanitize("**很长**的一段话超过十个字了哦朋友"),
|
|
"很长的一段话超过十个…")
|
|
|
|
|
|
class TestContext(unittest.IsolatedAsyncioTestCase):
|
|
async def test_context_included_and_capped(self):
|
|
engine = _make_engine(context_turns=2, per_user_cooldown_sec=0)
|
|
p = _install(engine, _FakeProvider())
|
|
for i in range(5):
|
|
await engine.generate_reply(f"消息{i}", 1, "小明")
|
|
# 最后一次调用：system(1) + 最多 2*2 条上下文 + 当前 1 条
|
|
msgs = p.calls[-1]
|
|
user_msgs = [m for m in msgs if m["role"] == "user"]
|
|
self.assertLessEqual(len(msgs), 1 + 4 + 1)
|
|
self.assertEqual(user_msgs[-1]["content"], "小明: 消息4")
|
|
|
|
async def test_system_prompt_present(self):
|
|
engine = _make_engine()
|
|
p = _install(engine, _FakeProvider())
|
|
await engine.generate_reply("在吗", 1, "小明")
|
|
first = p.calls[0][0]
|
|
self.assertEqual(first["role"], "system")
|
|
# 人设在前，情绪标签指令在后（桌宠联动需要）
|
|
self.assertIn("你是直播间助手", first["content"])
|
|
self.assertIn("情绪标签", first["content"])
|
|
self.assertIn("[开心]", first["content"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|