tests/test_llm.py

"""LLM client wire-protocol tests with a mocked httpx transport."""

from __future__ import annotations

import json

import httpx
import pytest

from app.llm import (LLMClient, LLMError, _chat_completions_url, _messages_url,
                     extract_json_object, user_turn, assistant_turn)

PNG = b"\x89PNG\r\n\x1a\nfake"


def make_settings(make_settings, **overrides):
    defaults = {"LLM_MAX_RETRIES": 0, "LLM_TIMEOUT_S": 5}
    defaults.update(overrides)
    return make_settings(**defaults)


def anthropic_ok(text: str) -> httpx.Response:
    return httpx.Response(200, json={
        "content": [{"type": "text", "text": text}],
        "stop_reason": "end_turn",
    })


def openai_ok(text: str) -> httpx.Response:
    return httpx.Response(200, json={
        "choices": [{"message": {"content": text}, "finish_reason": "stop"}],
    })


class TestUrlJoin:
    def test_base_without_v1(self):
        assert _messages_url("https://api.anthropic.com") == \
            "https://api.anthropic.com/v1/messages"
        assert _messages_url("https://openrouter.ai/api") == \
            "https://openrouter.ai/api/v1/messages"

    def test_base_with_v1(self):
        assert _messages_url("https://x.test/v1") == "https://x.test/v1/messages"
        assert _messages_url("https://x.test/v1/") == "https://x.test/v1/messages"

    def test_chat_completions_join(self):
        assert _chat_completions_url("https://openrouter.ai/api") == \
            "https://openrouter.ai/api/v1/chat/completions"
        assert _chat_completions_url("https://x.test/v1") == \
            "https://x.test/v1/chat/completions"


class TestAnthropicCall:
    @pytest.mark.asyncio
    async def test_happy_path_headers_and_body(self, make_settings):
        seen = {}

        def handler(request: httpx.Request) -> httpx.Response:
            seen["url"] = str(request.url)
            seen["x-api-key"] = request.headers.get("x-api-key")
            seen["authorization"] = request.headers.get("authorization")
            seen["anthropic-version"] = request.headers.get("anthropic-version")
            seen["body"] = json.loads(request.content)
            return anthropic_ok('{"ok": true}')

        settings = make_settings()
        transport = httpx.MockTransport(handler)
        async with httpx.AsyncClient(transport=transport) as http:
            client = LLMClient(settings, http=http)
            reply = await client.complete_vision(
                system="sys", messages=[user_turn("hello", image_png=PNG)])

        assert reply.text == '{"ok": true}'
        assert reply.style == "anthropic"
        assert seen["url"] == "https://llm.test/api/v1/messages"
        assert seen["x-api-key"] == "test-key-not-real"
        assert seen["authorization"] == "Bearer test-key-not-real"
        assert seen["anthropic-version"] == "2023-06-01"
        assert seen["body"]["model"] == "test-model"
        assert seen["body"]["system"] == "sys"
        content = seen["body"]["messages"][0]["content"]
        assert content[0]["type"] == "image"
        assert content[0]["source"]["media_type"] == "image/png"
        assert content[1] == {"type": "text", "text": "hello"}

    @pytest.mark.asyncio
    async def test_no_retry_on_401(self, make_settings):
        calls = []

        def handler(request):
            calls.append(request)
            return httpx.Response(401, json={"error": {"message": "bad key"}})

        settings = make_settings(LLM_MAX_RETRIES=3)
        transport = httpx.MockTransport(handler)
        async with httpx.AsyncClient(transport=transport) as http:
            client = LLMClient(settings, http=http)
            with pytest.raises(LLMError) as exc:
                await client.complete_vision(system="s", messages=[user_turn("hi")])
        assert exc.value.code == "llm_rejected"
        assert len(calls) == 1  # never retried
        assert "test-key-not-real" not in str(exc.value)  # key never leaks

    @pytest.mark.asyncio
    async def test_retry_on_429_then_success(self, make_settings, monkeypatch):
        calls = []

        async def no_sleep(_):
            return None

        monkeypatch.setattr("app.llm.asyncio.sleep", no_sleep)

        def handler(request):
            calls.append(request)
            if len(calls) < 3:
                return httpx.Response(429, headers={"Retry-After": "0"})
            return anthropic_ok("done")

        settings = make_settings(LLM_MAX_RETRIES=3)
        transport = httpx.MockTransport(handler)
        async with httpx.AsyncClient(transport=transport) as http:
            client = LLMClient(settings, http=http)
            reply = await client.complete_vision(system="s", messages=[user_turn("hi")])
        assert reply.text == "done"
        assert len(calls) == 3

    @pytest.mark.asyncio
    async def test_retry_on_non_json_success_then_success(
        self, make_settings, monkeypatch
    ):
        calls = []

        async def no_sleep(_):
            return None

        monkeypatch.setattr("app.llm.asyncio.sleep", no_sleep)

        def handler(request):
            calls.append(request)
            if len(calls) < 3:
                return httpx.Response(
                    200,
                    text="<html>temporary upstream gateway response</html>",
                    headers={"content-type": "text/html"},
                )
            return anthropic_ok("done")

        settings = make_settings(LLM_MAX_RETRIES=2)
        transport = httpx.MockTransport(handler)
        async with httpx.AsyncClient(transport=transport) as http:
            client = LLMClient(settings, http=http)
            reply = await client.complete_vision(
                system="s", messages=[user_turn("hi")]
            )
        assert reply.text == "done"
        assert len(calls) == 3

    @pytest.mark.asyncio
    async def test_truncation_raises(self, make_settings):
        def handler(request):
            return httpx.Response(200, json={
                "content": [{"type": "text", "text": "partial"}],
                "stop_reason": "max_tokens",
            })

        settings = make_settings()
        transport = httpx.MockTransport(handler)
        async with httpx.AsyncClient(transport=transport) as http:
            client = LLMClient(settings, http=http)
            with pytest.raises(LLMError) as exc:
                await client.complete_vision(system="s", messages=[user_turn("hi")])
        assert exc.value.code == "llm_truncated"


class TestOpenAIFallback:
    @pytest.mark.asyncio
    async def test_404_falls_back_to_chat_completions(self, make_settings):
        urls = []

        def handler(request: httpx.Request) -> httpx.Response:
            urls.append(str(request.url))
            if request.url.path.endswith("/v1/messages"):
                return httpx.Response(404, json={"error": "not found"})
            return openai_ok("fallback-ok")

        settings = make_settings()  # LLM_API_STYLE default auto
        transport = httpx.MockTransport(handler)
        async with httpx.AsyncClient(transport=transport) as http:
            client = LLMClient(settings, http=http)
            reply = await client.complete_vision(
                system="sys", messages=[user_turn("hello", image_png=PNG)])
        assert reply.text == "fallback-ok"
        assert reply.style == "openai"
        assert urls == [
            "https://llm.test/api/v1/messages",
            "https://llm.test/api/v1/chat/completions",
        ]

    @pytest.mark.asyncio
    async def test_openai_body_shape(self, make_settings):
        seen = {}

        def handler(request):
            seen["body"] = json.loads(request.content)
            return openai_ok("ok")

        settings = make_settings(LLM_API_STYLE="openai")
        transport = httpx.MockTransport(handler)
        async with httpx.AsyncClient(transport=transport) as http:
            client = LLMClient(settings, http=http)
            await client.complete_vision(
                system="SYS", messages=[
                    user_turn("first", image_png=PNG),
                    assistant_turn("reply-one"),
                    user_turn("second"),
                ])
        body = seen["body"]
        assert body["messages"][0] == {"role": "system", "content": "SYS"}
        user1 = body["messages"][1]
        assert user1["role"] == "user"
        types = [part["type"] for part in user1["content"]]
        assert types == ["image_url", "text"]
        image_part = user1["content"][0]
        assert image_part["image_url"]["url"].startswith("data:image/png;base64,")
        assert user1["content"][1] == {"type": "text", "text": "first"}
        assert body["messages"][2] == {"role": "assistant", "content": "reply-one"}
        assert body["messages"][3] == {"role": "user", "content": "second"}


class TestJsonExtraction:
    def test_plain(self):
        assert extract_json_object('{"a": 1}') == {"a": 1}

    def test_fenced(self):
        assert extract_json_object('```json\n{"a": 2}\n```') == {"a": 2}

    def test_prose_wrapped(self):
        assert extract_json_object('Here you go:\n{"a": 3}\nHope that helps!') == {"a": 3}

    def test_malformed(self):
        with pytest.raises(LLMError) as exc:
            extract_json_object("no json at all")
        assert exc.value.code == "llm_bad_json"

    def test_non_object(self):
        with pytest.raises(LLMError):
            extract_json_object("[1, 2, 3]")