76 lines
2.9 KiB
Python
76 lines
2.9 KiB
Python
|
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")
|