DealDocumentScreening/tests/unit/test_llm_prescreen_extraction.py

"""extract_prescreen provider tests (respx-mocked, both adapters).

Covers (plan Phase 2): valid dict passthrough, prescreen_max_chars input cap,
malformed JSON → repair once → success, and malformed after repair → LLMError.
"""

from __future__ import annotations

import json
from typing import Any

import httpx
import pytest
import respx

from contract_check.core.llm.ollama_cloud import LLMError, OllamaCloudProvider
from contract_check.core.llm.yandex_gpt import YandexGPTProvider

OLLAMA_HOST = "https://ollama.test"
OLLAMA_URL = f"{OLLAMA_HOST}/api/chat"
YANDEX_HOST = "https://llm.test"
YANDEX_URL = f"{YANDEX_HOST}/foundationModels/v1/completion"

VALID_PRESCREEN: dict[str, Any] = {
    "contract_type": "supply",
    "party_a": "ООО «Альфа»",
    "party_b": "ООО «Бета»",
    "total_amount": 1250000.0,
    "currency": "RUB",
    "start_date": "2025-09-01",
    "end_date": "2026-08-31",
    "has_penalty_clause": True,
    "has_termination_clause": False,
    "has_arbitration": True,
}


def _ollama_body(content: str) -> dict[str, object]:
    return {"message": {"content": content}, "prompt_eval_count": 5, "eval_count": 7}


def _yandex_body(content: str) -> dict[str, object]:
    return {
        "result": {
            "alternatives": [{"message": {"role": "assistant", "text": content}}],
            "usage": {"inputTextTokens": 5, "completionTokens": 7, "totalTokens": 12},
        }
    }


# ── Ollama Cloud ─────────────────────────────────────────────────────────────


async def test_ollama_extract_prescreen_returns_dict() -> None:
    provider = OllamaCloudProvider(host=OLLAMA_HOST, api_key="key", model="m")
    async with provider:
        with respx.mock(base_url=OLLAMA_HOST) as mock:
            route = mock.post(OLLAMA_URL).mock(
                return_value=httpx.Response(200, json=_ollama_body(json.dumps(VALID_PRESCREEN)))
            )
            result = await provider.extract_prescreen("Договор поставки...")
    assert route.call_count == 1
    assert result == VALID_PRESCREEN


async def test_ollama_extract_prescreen_truncates_input() -> None:
    provider = OllamaCloudProvider(
        host=OLLAMA_HOST, api_key="key", model="m", prescreen_max_chars=100
    )
    async with provider:
        with respx.mock(base_url=OLLAMA_HOST) as mock:
            route = mock.post(OLLAMA_URL).mock(
                return_value=httpx.Response(200, json=_ollama_body(json.dumps(VALID_PRESCREEN)))
            )
            await provider.extract_prescreen("договор " * 500)
            body = json.loads(route.calls.last.request.read())
    user_msg = body["messages"][1]["content"]
    assert len(user_msg) == 100
    assert body["messages"][0]["role"] == "system"
    assert "метаданных" in body["messages"][0]["content"]


async def test_ollama_extract_prescreen_repairs_invalid_json() -> None:
    provider = OllamaCloudProvider(host=OLLAMA_HOST, api_key="key", model="m")
    async with provider:
        with respx.mock(base_url=OLLAMA_HOST) as mock:
            route = mock.post(OLLAMA_URL).mock(
                side_effect=[
                    httpx.Response(200, json=_ollama_body("not json {")),
                    httpx.Response(200, json=_ollama_body(json.dumps(VALID_PRESCREEN))),
                ]
            )
            result = await provider.extract_prescreen("Договор...")
    assert route.call_count == 2
    assert result == VALID_PRESCREEN


async def test_ollama_extract_prescreen_invalid_after_repair_raises() -> None:
    provider = OllamaCloudProvider(host=OLLAMA_HOST, api_key="key", model="m")
    async with provider:
        with respx.mock(base_url=OLLAMA_HOST) as mock:
            mock.post(OLLAMA_URL).mock(
                return_value=httpx.Response(200, json=_ollama_body("still not json"))
            )
            with pytest.raises(LLMError):
                await provider.extract_prescreen("Договор...")


async def test_ollama_extract_prescreen_ignores_extra_fields() -> None:
    """Hallucinated fields are dropped by the wire model (extra=ignore)."""
    payload = {**VALID_PRESCREEN, "risk_score": "high", "summary": "безопасен"}
    provider = OllamaCloudProvider(host=OLLAMA_HOST, api_key="key", model="m")
    async with provider:
        with respx.mock(base_url=OLLAMA_HOST) as mock:
            mock.post(OLLAMA_URL).mock(
                return_value=httpx.Response(200, json=_ollama_body(json.dumps(payload)))
            )
            result = await provider.extract_prescreen("Договор...")
    assert "risk_score" not in result
    assert result["contract_type"] == "supply"


# ── YandexGPT ────────────────────────────────────────────────────────────────


async def test_yandex_extract_prescreen_returns_dict() -> None:
    provider = YandexGPTProvider(api_key="key", folder_id="b1", base_url=YANDEX_HOST)
    async with provider:
        with respx.mock(base_url=YANDEX_HOST) as mock:
            route = mock.post(YANDEX_URL).mock(
                return_value=httpx.Response(200, json=_yandex_body(json.dumps(VALID_PRESCREEN)))
            )
            result = await provider.extract_prescreen("Договор поставки...")
    assert route.call_count == 1
    assert result == VALID_PRESCREEN


async def test_yandex_extract_prescreen_truncates_input() -> None:
    provider = YandexGPTProvider(
        api_key="key", folder_id="b1", base_url=YANDEX_HOST, prescreen_max_chars=100
    )
    async with provider:
        with respx.mock(base_url=YANDEX_HOST) as mock:
            route = mock.post(YANDEX_URL).mock(
                return_value=httpx.Response(200, json=_yandex_body(json.dumps(VALID_PRESCREEN)))
            )
            await provider.extract_prescreen("договор " * 500)
            body = json.loads(route.calls.last.request.read())
    assert len(body["messages"][1]["text"]) == 100


async def test_yandex_extract_prescreen_repairs_invalid_json() -> None:
    provider = YandexGPTProvider(api_key="key", folder_id="b1", base_url=YANDEX_HOST)
    async with provider:
        with respx.mock(base_url=YANDEX_HOST) as mock:
            route = mock.post(YANDEX_URL).mock(
                side_effect=[
                    httpx.Response(200, json=_yandex_body("{broken")),
                    httpx.Response(200, json=_yandex_body(json.dumps(VALID_PRESCREEN))),
                ]
            )
            result = await provider.extract_prescreen("Договор...")
    assert route.call_count == 2
    assert result == VALID_PRESCREEN