DealDocumentScreening/tests/unit/test_llm_yandex_gpt.py
2026-08-17 20:49:29 +03:00

294 lines
10 KiB
Python

"""YandexGPT adapter unit tests (respx-mocked).
Covers: 200 happy path, 429→fallback, invalid-JSON→repair→success,
auth/config errors, and the provider returning a merged AnalysisResult.
No network.
"""
from __future__ import annotations
import json
import httpx
import pytest
import respx
from contract_check.core.analysis.report_schema import ReportPayload
from contract_check.core.llm.yandex_gpt import (
LLMConfigError,
LLMError,
LLMQuotaError,
LLMUnavailableError,
YandexGPTProvider,
)
HOST = "https://llm.test"
URL = f"{HOST}/foundationModels/v1/completion"
FOLDER_ID = "b1folder"
MODEL = "yandexgpt-lite"
FALLBACK_MODEL = "yandexgpt"
VALID_PAYLOAD = ReportPayload(
findings=[
{
"checklist_id": "penalties",
"severity": "high",
"quote": "Штраф 0,5% за каждый день просрочки",
"section_ref": "п. 6.3",
"risk": "Высокая неустойка",
"recommendation": "Ограничить cap",
}
]
)
def _completion_body(content: str, model: str = MODEL) -> dict[str, object]:
return {
"result": {
"alternatives": [{"message": {"role": "assistant", "text": content}}],
"usage": {
"inputTextTokens": 100,
"completionTokens": 20,
"totalTokens": 120,
},
"modelVersion": model,
}
}
def _provider(fallback: str | None = FALLBACK_MODEL) -> YandexGPTProvider:
return YandexGPTProvider(
api_key="key",
folder_id=FOLDER_ID,
model=MODEL,
fallback_model=fallback,
base_url=HOST,
completion_path="/foundationModels/v1/completion",
max_concurrency=1,
chunk_size=10000,
)
@pytest.mark.asyncio
async def test_happy_path_returns_findings() -> None:
async with _provider() as provider:
with respx.mock(base_url=HOST) as mock:
route = mock.post("/foundationModels/v1/completion")
route.mock(
return_value=httpx.Response(
200, json=_completion_body(VALID_PAYLOAD.model_dump_json())
)
)
result = await provider.analyze("договор " * 200)
assert len(result.findings) == 1
assert result.findings[0].checklist_id == "penalties"
assert result.findings[0].severity == "high"
assert result.prompt_tokens == 100
assert result.eval_tokens == 20
assert result.fell_back is False
assert result.model_used == MODEL
@pytest.mark.asyncio
async def test_429_falls_back_to_secondary_model() -> None:
async with _provider() as provider:
with respx.mock(base_url=HOST) as mock:
route = mock.post("/foundationModels/v1/completion")
route.mock(
side_effect=[
httpx.Response(429, json={"error": "quota"}),
httpx.Response(
200,
json=_completion_body(
VALID_PAYLOAD.model_dump_json(), model=FALLBACK_MODEL
),
),
]
)
result = await provider.analyze("договор " * 200)
assert result.fell_back is True
assert FALLBACK_MODEL in result.models_used
@pytest.mark.asyncio
async def test_invalid_json_repaired_on_second_attempt() -> None:
async with _provider() as provider:
with respx.mock(base_url=HOST) as mock:
route = mock.post("/foundationModels/v1/completion")
route.mock(
side_effect=[
httpx.Response(200, json=_completion_body("not valid json {")),
httpx.Response(200, json=_completion_body(VALID_PAYLOAD.model_dump_json())),
]
)
result = await provider.analyze("договор " * 200)
assert result.repaired is True
assert len(result.findings) == 1
@pytest.mark.asyncio
async def test_model_aliases_validate_without_repair() -> None:
"""Models often emit risk_type/description and omit recommendation."""
aliased = json.dumps(
{
"findings": [
{
"risk_type": "penalties",
"severity": "medium",
"description": "Высокая неустойка.",
"quote": "Штраф 0,5% за каждый день просрочки",
"section_ref": "п. 6.3",
}
]
}
)
async with _provider() as provider:
with respx.mock(base_url=HOST) as mock:
route = mock.post("/foundationModels/v1/completion")
route.mock(return_value=httpx.Response(200, json=_completion_body(aliased)))
result = await provider.analyze("договор " * 200)
assert route.call_count == 1
assert result.repaired is False
assert result.findings[0].checklist_id == "penalties"
assert result.findings[0].risk == "Высокая неустойка."
assert result.findings[0].recommendation == ""
@pytest.mark.asyncio
async def test_429_with_no_fallback_raises_quota() -> None:
provider = YandexGPTProvider(
api_key="key",
folder_id=FOLDER_ID,
model=MODEL,
fallback_model=None,
base_url=HOST,
completion_path="/foundationModels/v1/completion",
)
async with provider:
with respx.mock(base_url=HOST) as mock:
mock.post("/foundationModels/v1/completion").mock(
return_value=httpx.Response(429, json={"error": "quota"})
)
with pytest.raises(LLMQuotaError):
await provider.analyze("договор " * 200)
@pytest.mark.asyncio
async def test_invalid_json_after_repair_raises_llm_error() -> None:
async with _provider() as provider:
with respx.mock(base_url=HOST) as mock:
mock.post("/foundationModels/v1/completion").mock(
return_value=httpx.Response(200, json=_completion_body("still not json"))
)
with pytest.raises(LLMError):
await provider.analyze("договор " * 200)
@pytest.mark.asyncio
async def test_short_text_single_chunk_one_request() -> None:
async with _provider() as provider:
with respx.mock(base_url=HOST) as mock:
route = mock.post("/foundationModels/v1/completion")
route.mock(
return_value=httpx.Response(
200, json=_completion_body(VALID_PAYLOAD.model_dump_json())
)
)
await provider.analyze("короткий договор")
assert route.call_count == 1
@pytest.mark.asyncio
async def test_multi_chunk_merges_findings() -> None:
findings = json.loads(VALID_PAYLOAD.model_dump_json())["findings"]
async with _provider() as provider:
with respx.mock(base_url=HOST) as mock:
mock.post("/foundationModels/v1/completion").mock(
return_value=httpx.Response(
200,
json=_completion_body(json.dumps({"findings": findings + findings})),
)
)
result = await provider.analyze("договор " * 3000, checklist="x")
assert len(result.findings) == 1
@pytest.mark.asyncio
async def test_401_raises_config_error() -> None:
async with _provider() as provider:
with respx.mock(base_url=HOST) as mock:
route = mock.post("/foundationModels/v1/completion")
route.mock(return_value=httpx.Response(401, json={"error": "unauthorized"}))
with pytest.raises(LLMConfigError):
await provider.analyze("договор " * 200)
assert route.call_count == 1
@pytest.mark.asyncio
async def test_400_raises_config_error() -> None:
async with _provider() as provider:
with respx.mock(base_url=HOST) as mock:
route = mock.post("/foundationModels/v1/completion")
route.mock(return_value=httpx.Response(400, json={"error": "bad request"}))
with pytest.raises(LLMConfigError):
await provider.analyze("договор " * 200)
assert route.call_count == 1
@pytest.mark.asyncio
async def test_404_raises_config_error() -> None:
async with _provider() as provider:
with respx.mock(base_url=HOST) as mock:
route = mock.post("/foundationModels/v1/completion")
route.mock(return_value=httpx.Response(404, text="Not Found"))
with pytest.raises(LLMConfigError):
await provider.analyze("договор " * 200)
assert route.call_count == 1
@pytest.mark.asyncio
async def test_connect_error_raises_config_error_immediately() -> None:
async with _provider() as provider:
with respx.mock(base_url=HOST) as mock:
route = mock.post("/foundationModels/v1/completion")
route.mock(side_effect=httpx.ConnectError("Connection refused"))
with pytest.raises(LLMConfigError):
await provider.analyze("договор " * 200)
assert route.call_count == 1
@pytest.mark.asyncio
async def test_503_retries_then_raises_unavailable() -> None:
async with _provider() as provider:
with respx.mock(base_url=HOST) as mock:
route = mock.post("/foundationModels/v1/completion")
route.mock(return_value=httpx.Response(503, text="Unavailable"))
with pytest.raises(LLMUnavailableError):
await provider.analyze("договор " * 200)
assert route.call_count == 3
@pytest.mark.asyncio
async def test_request_payload_contains_model_uri_and_json_schema() -> None:
async with _provider(fallback=None) as provider:
with respx.mock(base_url=HOST) as mock:
route = mock.post("/foundationModels/v1/completion")
route.mock(
return_value=httpx.Response(
200, json=_completion_body(VALID_PAYLOAD.model_dump_json())
)
)
await provider.analyze("короткий договор", checklist="test checklist")
request = route.calls[0].request
body = json.loads(request.content)
assert body["modelUri"] == f"gpt://{FOLDER_ID}/{MODEL}"
assert body["completionOptions"]["responseFormat"]["type"] == "JSON_OBJECT"
assert "schema" in body["completionOptions"]["responseFormat"]["json_schema"]
assert body["messages"][0]["role"] == "system"
assert "test checklist" in body["messages"][0]["text"]
assert request.headers["authorization"] == "Api-Key key"
assert request.headers["x-folder-id"] == FOLDER_ID