dieya_risk_prediction_backend/tests/test_llm_protocol.py

import json

import httpx
import pytest
from app.core.config import get_settings
from app.services.llm_client import LLMClient


@pytest.mark.asyncio
async def test_openai_compatible_semantic_and_stream_protocol(monkeypatch: pytest.MonkeyPatch) -> None:
    get_settings.cache_clear()
    monkeypatch.setenv("RISK_LLM_API_KEY", "test-key")
    monkeypatch.setenv("RISK_LLM_BASE_URL", "https://llm.test/v1")
    monkeypatch.setenv("RISK_LLM_MODEL", "test-model")
    calls = 0

    async def handler(request: httpx.Request) -> httpx.Response:
        nonlocal calls
        calls += 1
        if calls == 1:
            request_body = json.loads(request.content)
            assert request_body["response_format"] == {"type": "json_object"}
            assert "json" in request_body["messages"][0]["content"]
            return httpx.Response(
                200,
                json={
                    "choices": [
                        {
                            "message": {
                                "content": json.dumps(
                                    {"intent": "assessment", "category": "洗发水", "confidence": 0.92},
                                    ensure_ascii=False,
                                )
                            }
                        }
                    ]
                },
            )
        return httpx.Response(
            200,
            headers={"content-type": "text/event-stream"},
            content=(
                'data: {"choices":[{"delta":{"content":"第一段"}}]}\n\n'
                'data: {"choices":[{"delta":{"content":"第二段"}}]}\n\n'
                "data: [DONE]\n\n"
            ).encode(),
        )

    client = LLMClient(transport=httpx.MockTransport(handler))
    semantic = await client.extract_semantics("评估洗发水新品")
    chunks = [chunk async for chunk in client.stream_reply([{"role": "user", "content": "继续分析"}])]

    assert semantic["intent"] == "assessment"
    assert semantic["confidence"] == 0.92
    assert "".join(chunks) == "第一段第二段"
    assert calls == 2
    get_settings.cache_clear()


@pytest.mark.asyncio
async def test_semantic_request_fails_over_to_deepseek_responses(monkeypatch: pytest.MonkeyPatch) -> None:
    get_settings.cache_clear()
    monkeypatch.setenv("RISK_LLM_API_KEY", "primary-key")
    monkeypatch.setenv("RISK_LLM_BASE_URL", "https://primary.test/v1")
    monkeypatch.setenv("RISK_LLM_MODEL", "primary-model")
    monkeypatch.setenv("RISK_LLM_FALLBACK_API_KEY", "fallback-key")
    monkeypatch.setenv("RISK_LLM_FALLBACK_BASE_URL", "https://fallback.test")
    monkeypatch.setenv("RISK_LLM_FALLBACK_MODEL", "deepseek-v4-pro")
    monkeypatch.setenv("RISK_LLM_FALLBACK_API_MODE", "responses")
    calls: list[str] = []

    async def handler(request: httpx.Request) -> httpx.Response:
        calls.append(str(request.url))
        if request.url.path.endswith("/chat/completions"):
            return httpx.Response(503, json={"error": {"message": "primary unavailable"}})
        body = json.loads(request.content)
        assert body["model"] == "deepseek-v4-pro"
        assert body["input"][-1]["content"] == "评估洗发水新品"
        return httpx.Response(
            200,
            json={
                "output_text": json.dumps(
                    {"intent": "assessment", "risk_requested": True, "confidence": 0.8},
                    ensure_ascii=False,
                )
            },
        )

    client = LLMClient(transport=httpx.MockTransport(handler))
    result = await client.extract_semantics("评估洗发水新品")

    assert result["intent"] == "assessment"
    assert calls == [
        "https://primary.test/v1/chat/completions",
        "https://fallback.test/responses",
    ]
    get_settings.cache_clear()


@pytest.mark.asyncio
async def test_stream_request_fails_over_before_primary_emits_text(monkeypatch: pytest.MonkeyPatch) -> None:
    get_settings.cache_clear()
    monkeypatch.setenv("RISK_LLM_API_KEY", "primary-key")
    monkeypatch.setenv("RISK_LLM_BASE_URL", "https://primary.test/v1")
    monkeypatch.setenv("RISK_LLM_FALLBACK_API_KEY", "fallback-key")
    monkeypatch.setenv("RISK_LLM_FALLBACK_BASE_URL", "https://fallback.test")
    monkeypatch.setenv("RISK_LLM_FALLBACK_MODEL", "deepseek-v4-pro")
    monkeypatch.setenv("RISK_LLM_FALLBACK_API_MODE", "responses")
    calls: list[str] = []

    async def handler(request: httpx.Request) -> httpx.Response:
        calls.append(str(request.url))
        if request.url.path.endswith("/chat/completions"):
            raise httpx.ConnectError("primary connection failed", request=request)
        return httpx.Response(
            200,
            headers={"content-type": "text/event-stream"},
            content=(
                'data: {"type":"response.output_text.delta","delta":"兜底"}\n\n'
                'data: {"type":"response.output_text.delta","delta":"回复"}\n\n'
                "data: [DONE]\n\n"
            ).encode(),
        )

    client = LLMClient(transport=httpx.MockTransport(handler))
    chunks = [chunk async for chunk in client.stream_reply([{"role": "user", "content": "继续分析"}])]

    assert "".join(chunks) == "兜底回复"
    assert calls == [
        "https://primary.test/v1/chat/completions",
        "https://fallback.test/responses",
    ]
    get_settings.cache_clear()