group_fqcd_jr/tests/unit/service/test_model_gateway.py

from types import SimpleNamespace

import httpx
import pytest

from app.core.errors import RecoverableAgentError
from app.service.model_gateway import (
    ModelDispatchService,
    ModelGenerationService,
    OpenAICompatibleGateway,
)


class Gateway:
    def __init__(self) -> None:
        self.calls: list[str] = []

    async def generate(self, *, endpoint_code: str, prompt: str, timeout_ms: int) -> str:
        self.calls.append(endpoint_code)
        if endpoint_code == "primary":
            raise TimeoutError("timeout")
        return "fallback answer"


class Endpoint:
    def __init__(self, code: str) -> None:
        self.endpoint_code = code
        self.timeout_ms = 1000


@pytest.mark.asyncio
async def test_model_dispatch_uses_declared_fallback_order() -> None:
    gateway = Gateway()
    result = await ModelDispatchService(gateway).generate(
        [Endpoint("primary"), Endpoint("fallback")], "hello"
    )
    assert result.text == "fallback answer"
    assert result.degraded is True
    assert gateway.calls == ["primary", "fallback"]


@pytest.mark.asyncio
async def test_openai_compatible_gateway_uses_secret_ref_and_parses_text(monkeypatch) -> None:
    monkeypatch.setenv("TEST_MODEL_KEY", "secret-value")
    requests: list[httpx.Request] = []

    async def handler(request: httpx.Request) -> httpx.Response:
        requests.append(request)
        return httpx.Response(200, json={"choices": [{"message": {"content": "answer"}}]})

    client = httpx.AsyncClient(transport=httpx.MockTransport(handler))
    endpoint = SimpleNamespace(endpoint_code="primary", base_url="https://model.test/v1",
                               model_name="test-model", secret_ref="env:TEST_MODEL_KEY")
    gateway = OpenAICompatibleGateway({"primary": endpoint}, client=client)
    try:
        assert await gateway.generate(
            endpoint_code="primary", prompt="hello", timeout_ms=1000
        ) == "answer"
        assert str(requests[0].url) == "https://model.test/v1/chat/completions"
        assert requests[0].headers["authorization"] == "Bearer secret-value"
        assert requests[0].content.find(b"secret-value") == -1
    finally:
        await client.aclose()


@pytest.mark.asyncio
async def test_gateway_missing_secret_and_empty_route_fail_closed(monkeypatch) -> None:
    monkeypatch.delenv("MISSING_MODEL_KEY", raising=False)
    endpoint = SimpleNamespace(endpoint_code="primary", base_url="https://model.test",
                               model_name="test-model", secret_ref="env:MISSING_MODEL_KEY")
    with pytest.raises(RecoverableAgentError, match="密钥未配置"):
        await OpenAICompatibleGateway({"primary": endpoint}).generate(
            endpoint_code="primary", prompt="hello", timeout_ms=1000)
    with pytest.raises(RecoverableAgentError, match="没有可用"):
        await ModelGenerationService(ModelDispatchService(Gateway())).generate([], "hello")


# ---------------------------------------------------------------------------
# `DatabaseModelEndpointResolver` 按 task_type 过滤能力（补上被漏掉的契约兑现）
# ---------------------------------------------------------------------------


class _FakeScalarSession:
    """只实现 `scalars()`：解析器只用它取 active 端点列表。"""

    def __init__(self, rows: list[object]) -> None:
        self._rows = rows

    async def scalars(self, _statement: object) -> list[object]:
        return self._rows

    async def __aenter__(self) -> "_FakeScalarSession":
        return self

    async def __aexit__(self, *args: object) -> None:
        return None


def _endpoint(code: str, capabilities: list[str] | None) -> SimpleNamespace:
    return SimpleNamespace(endpoint_code=code, capabilities=capabilities)


@pytest.mark.asyncio
@pytest.mark.parametrize(
    ("task_type", "expected"),
    [
        # 向量化只能走声明 embedding 的端点：交给 chat 端点会打到 /chat/completions 上。
        ("embedding", ["embedding-primary"]),
        # 生成与意图分类都走 chat 端点：交给 embedding 端点会 404。
        ("chat", ["chat-primary"]),
        ("intent_classification", ["chat-primary"]),
        # 未知任务类型无法判断该要哪种能力 → 不过滤（返回全部 active）。
        ("something_new", ["embedding-primary", "chat-primary"]),
    ],
)
async def test_resolve_filters_endpoints_by_task_type_capability(
    monkeypatch: pytest.MonkeyPatch, task_type: str, expected: list[str]
) -> None:
    """必须按 `task_type` 过滤能力。

    不过滤的后果：`ModelDispatchService` 的 `generate`/`embed` 只取前
    `max(1, max_attempts)`（默认 2）个端点，错叫一个就吃掉一次尝试机会 ——
    端点一多会直接耗尽尝试而失败（原实现 `del agent_type, task_type` 即此缺陷）。
    """
    from app.service import model_gateway

    rows = [_endpoint("embedding-primary", ["embedding"]), _endpoint("chat-primary", ["chat"])]
    monkeypatch.setattr(model_gateway, "SessionFactory", lambda: _FakeScalarSession(rows))

    resolved = await model_gateway.DatabaseModelEndpointResolver().resolve(
        agent_type="customer_service", task_type=task_type
    )
    assert [e.endpoint_code for e in resolved] == expected


@pytest.mark.asyncio
async def test_resolve_skips_endpoints_without_declared_capabilities(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """`capabilities` 为 NULL/空 的端点不得被任何筛选选中（不能裸奔到错误的网关方法上）。"""
    from app.service import model_gateway

    rows = [_endpoint("broken", None), _endpoint("empty", []), _endpoint("chat-primary", ["chat"])]
    monkeypatch.setattr(model_gateway, "SessionFactory", lambda: _FakeScalarSession(rows))

    resolved = await model_gateway.DatabaseModelEndpointResolver().resolve(
        agent_type="customer_service", task_type="chat"
    )
    assert [e.endpoint_code for e in resolved] == ["chat-primary"]