137 lines
4.8 KiB
Python
137 lines
4.8 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""本地 70B 推理后端测试（issue #44）。
|
||
|
||
覆盖：
|
||
1. 参数化：endpoint / model / timeout / max_tokens / temperature / echo_context；
|
||
2. dry-run（未配置 endpoint）：占位输出 + 来源回显，与旧 LocalBackend 兼容；
|
||
3. OpenAI 兼容调用：mock /v1/chat/completions 响应 → 提取 answer；
|
||
4. 响应格式异常 → RuntimeError；
|
||
5. 与 LLMGateway 组合：本地后端承载敏感内容（数据不出厂）。
|
||
"""
|
||
import os
|
||
import sys
|
||
import unittest
|
||
from unittest import mock
|
||
|
||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||
import _bootstrap # noqa: F401
|
||
|
||
from llm_gateway.backends import Local70BBackend # noqa: E402
|
||
from llm_gateway.gateway import LLMGateway # noqa: E402
|
||
from llm_gateway.prompts import PromptRegistry # noqa: E402
|
||
|
||
PROMPTS_CONFIG = os.path.join(
|
||
os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
|
||
"config", "prompts.template.yaml",
|
||
)
|
||
|
||
|
||
def make_gateway(**kwargs) -> LLMGateway:
|
||
"""构建带提示词版本库的网关（默认本地 70B 后端）。"""
|
||
return LLMGateway(
|
||
local=Local70BBackend(),
|
||
prompts=PromptRegistry.from_template_config(PROMPTS_CONFIG),
|
||
**kwargs,
|
||
)
|
||
|
||
CONFIG = os.path.join(
|
||
os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
|
||
"config", "local70b.template.yaml",
|
||
)
|
||
|
||
|
||
def load_config():
|
||
import yaml
|
||
with open(CONFIG, "r", encoding="utf-8") as fh:
|
||
return yaml.safe_load(fh) or {}
|
||
|
||
|
||
class TestParameterization(unittest.TestCase):
|
||
"""参数化提取与配置资产。"""
|
||
|
||
def test_defaults(self):
|
||
b = Local70BBackend()
|
||
self.assertEqual(b.endpoint, "")
|
||
self.assertEqual(b.model, "iaop-local-70b")
|
||
self.assertEqual(b.max_tokens, 1024)
|
||
self.assertEqual(b.temperature, 0.1)
|
||
|
||
def test_config_asset_parses(self):
|
||
cfg = load_config()["local70b"]
|
||
b = Local70BBackend(**{k: v for k, v in cfg.items()})
|
||
self.assertEqual(b.endpoint, "http://10.20.0.30:8000/v1")
|
||
self.assertEqual(b.model, "iaop-local-70b")
|
||
|
||
def test_name(self):
|
||
self.assertEqual(Local70BBackend().name, "local-70b")
|
||
|
||
|
||
class TestDryRun(unittest.TestCase):
|
||
"""未配置 endpoint：占位 + 来源回显（兼容旧行为）。"""
|
||
|
||
def setUp(self):
|
||
self.b = Local70BBackend() # endpoint 默认空
|
||
|
||
def test_dry_run_with_context(self):
|
||
out = self.b.generate("请解释炉温报警", ["SOP-CL-001", "工艺规范"])
|
||
self.assertIn("[本地70B占位]", out)
|
||
self.assertIn("[来源: SOP-CL-001]", out)
|
||
|
||
def test_dry_run_no_echo(self):
|
||
b = Local70BBackend(echo_context=False)
|
||
out = b.generate("hi", ["s1"])
|
||
self.assertNotIn("[来源", out)
|
||
|
||
def test_health_dry_run(self):
|
||
health = self.b.health()
|
||
self.assertEqual(health["status"], "dry-run")
|
||
|
||
|
||
class TestOpenAICompat(unittest.TestCase):
|
||
"""OpenAI 兼容 /v1/chat/completions 调用。"""
|
||
|
||
def setUp(self):
|
||
self.b = Local70BBackend(endpoint="http://local:8000/v1")
|
||
|
||
def test_generate_extracts_answer(self):
|
||
fake = {"choices": [{"message": {"content": "炉温偏高，建议降氯气流量"}}]}
|
||
with mock.patch.object(self.b, "_post_json", return_value=fake) as post:
|
||
out = self.b.generate("炉温异常", ["SOP-CL-001"])
|
||
post.assert_called_once()
|
||
path, payload = post.call_args[0]
|
||
self.assertEqual(path, "/v1/chat/completions")
|
||
self.assertEqual(payload["model"], "iaop-local-70b")
|
||
# system 提示注入 RAG 引用（溯源）
|
||
self.assertIn("SOP-CL-001", payload["messages"][0]["content"])
|
||
self.assertEqual(out, "炉温偏高，建议降氯气流量")
|
||
|
||
def test_bad_response_raises(self):
|
||
with mock.patch.object(self.b, "_post_json", return_value={"choices": []}):
|
||
with self.assertRaises(RuntimeError):
|
||
self.b.generate("x", [])
|
||
|
||
def test_health_ok(self):
|
||
with mock.patch("urllib.request.urlopen") as urlopen:
|
||
resp = mock.MagicMock()
|
||
resp.status = 200
|
||
urlopen.return_value.__enter__ = mock.MagicMock(return_value=resp)
|
||
urlopen.return_value.__exit__ = mock.MagicMock(return_value=False)
|
||
health = self.b.health()
|
||
self.assertEqual(health["status"], "ok")
|
||
self.assertEqual(health["backend"], "local-70b")
|
||
|
||
|
||
class TestGatewayIntegration(unittest.TestCase):
|
||
"""与 LLMGateway 组合：本地后端承载敏感内容（数据不出厂）。"""
|
||
|
||
def test_gateway_with_local70b(self):
|
||
gw = make_gateway()
|
||
result = gw.ask("炉温是多少", rag_context=["工艺规范"])
|
||
self.assertIn("本地70B", result.answer)
|
||
# 敏感内容路由本地（CLF 工艺参数 → local）
|
||
self.assertEqual(result.route.target, "local")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|