tgit/tests/llm/test_providers.py

import json
from typing import Any

import httpx
import pytest
import respx

from app.config import LLMSettings
from app.llm.factory import make_provider
from app.llm.mock import MockProvider, MockRule
from app.llm.openai_compat import OpenAICompatProvider
from app.llm.parsing import parse_json_object
from app.llm.provider import ChatMessage, JsonSchema, LLMError

BASE = "https://llm.example.test/v1"
SCHEMA = JsonSchema("verdict", {"type": "object", "properties": {"verdict": {"type": "string"}}})


def settings(**kw: Any) -> LLMSettings:
    base = {"provider": "openai_compat", "base_url": BASE, "model": "m-1", "api_key": "k", "max_attempts": 3}
    return LLMSettings(**(base | kw))


def completion(content: str, prompt: int = 5, completion_tokens: int = 3) -> dict[str, Any]:
    return {
        "choices": [{"message": {"role": "assistant", "content": content}}],
        "usage": {"prompt_tokens": prompt, "completion_tokens": completion_tokens},
    }


class Sleeps:
    def __init__(self) -> None:
        self.delays: list[float] = []

    async def __call__(self, delay: float) -> None:
        self.delays.append(delay)


@respx.mock
async def test_plain_completion_sends_messages_and_counts_tokens() -> None:
    route = respx.post(f"{BASE}/chat/completions").respond(json=completion("привет"))
    provider = OpenAICompatProvider(settings())

    result = await provider.complete([ChatMessage("system", "s"), ChatMessage("user", "u")])
    await provider.aclose()

    sent = json.loads(route.calls.last.request.content)
    assert route.calls.last.request.headers["Authorization"] == "Bearer k"
    assert sent["model"] == "m-1" and sent["temperature"] == 0.0
    assert sent["messages"] == [{"role": "system", "content": "s"}, {"role": "user", "content": "u"}]
    assert "response_format" not in sent
    assert (result.text, result.data, result.tokens) == ("привет", None, 8)


@respx.mock
async def test_schema_requests_structured_output_and_parses_it() -> None:
    route = respx.post(f"{BASE}/chat/completions").respond(json=completion('```json\n{"verdict": "safe"}\n```'))
    provider = OpenAICompatProvider(settings())

    result = await provider.complete([ChatMessage("user", "x")], SCHEMA)

    fmt = json.loads(route.calls.last.request.content)["response_format"]
    assert fmt == {"type": "json_schema", "json_schema": {"name": "verdict", "schema": SCHEMA.schema, "strict": True}}
    assert result.data == {"verdict": "safe"}


@respx.mock
async def test_invalid_json_gives_no_data() -> None:
    respx.post(f"{BASE}/chat/completions").respond(json=completion("not json"))

    result = await OpenAICompatProvider(settings()).complete([ChatMessage("user", "x")], SCHEMA)

    assert result.data is None and result.text == "not json"


@respx.mock
async def test_retries_rate_limit_and_network_errors() -> None:
    sleeps = Sleeps()
    respx.post(f"{BASE}/chat/completions").mock(
        side_effect=[httpx.Response(429), httpx.ConnectError("x"), httpx.Response(200, json=completion("ok"))]
    )

    result = await OpenAICompatProvider(settings(), sleep=sleeps).complete([ChatMessage("user", "x")])

    assert result.text == "ok" and sleeps.delays == [1.0, 2.0]


@respx.mock
async def test_client_error_is_not_retried() -> None:
    route = respx.post(f"{BASE}/chat/completions").respond(401, text="bad key")

    with pytest.raises(LLMError, match="401"):
        await OpenAICompatProvider(settings()).complete([ChatMessage("user", "x")])
    assert route.call_count == 1


@respx.mock
async def test_unreachable_after_all_attempts() -> None:
    respx.post(f"{BASE}/chat/completions").mock(side_effect=httpx.ConnectError("down"))

    with pytest.raises(LLMError, match="unreachable"):
        await OpenAICompatProvider(settings(), sleep=Sleeps()).complete([ChatMessage("user", "x")])


@respx.mock
async def test_unexpected_shape_is_an_error() -> None:
    respx.post(f"{BASE}/chat/completions").respond(json={"choices": []})

    with pytest.raises(LLMError, match="shape"):
        await OpenAICompatProvider(settings()).complete([ChatMessage("user", "x")])


def test_openai_compat_requires_url_and_model() -> None:
    with pytest.raises(LLMError):
        OpenAICompatProvider(settings(base_url=""))


async def test_mock_rules_and_fallbacks() -> None:
    mock = MockProvider([MockRule("взлом", {"verdict": "attack"}), MockRule("привет", "здравствуй")])

    assert (await mock.complete([ChatMessage("user", "попытка взлома")], SCHEMA)).data == {"verdict": "attack"}
    assert (await mock.complete([ChatMessage("user", "привет")])).text == "здравствуй"
    assert (await mock.complete([ChatMessage("user", "что-то")])).text == "mock: что-то"
    assert (await mock.complete([ChatMessage("user", "что-то")], SCHEMA)).data == {}
    assert len(mock.calls) == 4
    await mock.aclose()


def test_factory_picks_provider() -> None:
    assert isinstance(make_provider(LLMSettings(provider="mock")), MockProvider)
    assert isinstance(make_provider(settings()), OpenAICompatProvider)


@pytest.mark.parametrize(
    ("text", "expected"),
    [('{"a": 1}', {"a": 1}), ('```json\n{"a": 1}\n```', {"a": 1}), ("[1]", None), ("nope", None)],
)
def test_parse_json_object(text: str, expected: dict[str, Any] | None) -> None:
    assert parse_json_object(text) == expected


def test_openrouter_key_name_is_accepted(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.delenv("LLM_API_KEY", raising=False)
    monkeypatch.setenv("OPENROUTER_API_KEY", "or-key")

    assert LLMSettings(_env_file=None).api_key.get_secret_value() == "or-key"  # type: ignore[call-arg]

    monkeypatch.setenv("LLM_API_KEY", "llm-key")  # the project's own name wins
    assert LLMSettings(_env_file=None).api_key.get_secret_value() == "llm-key"  # type: ignore[call-arg]


async def test_requests_go_through_the_proxy() -> None:
    seen: list[str] = []

    class Proxy(httpx.AsyncBaseTransport):
        """Stands in for the HTTP proxy: httpx routes requests to the transport mounted for the proxy."""

        async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
            seen.append(str(request.url))
            return httpx.Response(200, json=completion("через прокси"))

    provider = OpenAICompatProvider(settings(proxy_url="http://user:pass@proxy.example.test:3128"))
    mounts = provider._http._mounts  # httpx keeps the proxy transport here, no public accessor
    assert mounts, "no proxy transport mounted"
    provider._http._mounts = {pattern: Proxy() for pattern in mounts}

    result = await provider.complete([ChatMessage("user", "x")])
    await provider.aclose()

    assert result.text == "через прокси" and seen == [f"{BASE}/chat/completions"]


def test_proxy_is_a_secret() -> None:
    s = settings(proxy_url="http://user:pass@proxy.example.test:3128")
    assert "pass" not in repr(s)