iAOP/core/llm-gateway/tests/test_local70b.py

# -*- 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()