chopratejas--headroom/tests/test_backend_anyllm.py

from __future__ import annotations

from types import SimpleNamespace

import pytest

from headroom.backends import anyllm
from headroom.backends.base import BackendResponse, StreamEvent


class FakeAsyncStream:
    def __init__(self, items) -> None:  # noqa: ANN001
        self._items = list(items)

    def __aiter__(self):
        self._iter = iter(self._items)
        return self

    async def __anext__(self):
        try:
            return next(self._iter)
        except StopIteration as exc:
            raise StopAsyncIteration from exc


class FakeAnyLLMInstance:
    def __init__(self) -> None:
        self.calls: list[dict[str, object]] = []
        self.response = None
        self.raise_error: Exception | None = None

    async def acompletion(self, **kwargs):  # noqa: ANN003
        self.calls.append(kwargs)
        if self.raise_error is not None:
            raise self.raise_error
        return self.response


def make_backend(
    monkeypatch: pytest.MonkeyPatch, provider: str = "groq"
) -> tuple[anyllm.AnyLLMBackend, FakeAnyLLMInstance]:
    fake_instance = FakeAnyLLMInstance()

    class FakeAnyLLM:
        @staticmethod
        def create(requested_provider: str, **kwargs):  # noqa: ANN003
            assert requested_provider == provider
            return fake_instance

    monkeypatch.setattr(anyllm, "ANYLLM_AVAILABLE", True)
    monkeypatch.setattr(anyllm, "AnyLLM", FakeAnyLLM)
    return anyllm.AnyLLMBackend(provider=provider.upper()), fake_instance


def test_init_forwards_api_base_and_api_key(monkeypatch: pytest.MonkeyPatch) -> None:
    """Regression for #942: custom api_base/api_key must reach AnyLLM.create."""
    fake_instance = FakeAnyLLMInstance()
    create_calls: list[dict[str, object]] = []

    class FakeAnyLLM:
        @staticmethod
        def create(requested_provider: str, **kwargs):  # noqa: ANN003
            create_calls.append({"provider": requested_provider, **kwargs})
            return fake_instance

    monkeypatch.setattr(anyllm, "ANYLLM_AVAILABLE", True)
    monkeypatch.setattr(anyllm, "AnyLLM", FakeAnyLLM)

    backend = anyllm.AnyLLMBackend(
        provider="openai",
        api_key="sk-custom",
        api_base="https://custom-provider.example/v1",
    )

    assert backend.api_base == "https://custom-provider.example/v1"
    assert create_calls == [
        {
            "provider": "openai",
            "api_key": "sk-custom",
            "api_base": "https://custom-provider.example/v1",
        }
    ]


def test_init_omits_unset_api_base_and_api_key(monkeypatch: pytest.MonkeyPatch) -> None:
    """Unset overrides must not be forwarded, preserving provider env defaults."""
    fake_instance = FakeAnyLLMInstance()
    create_calls: list[dict[str, object]] = []

    class FakeAnyLLM:
        @staticmethod
        def create(requested_provider: str, **kwargs):  # noqa: ANN003
            create_calls.append({"provider": requested_provider, **kwargs})
            return fake_instance

    monkeypatch.setattr(anyllm, "ANYLLM_AVAILABLE", True)
    monkeypatch.setattr(anyllm, "AnyLLM", FakeAnyLLM)

    anyllm.AnyLLMBackend(provider="openai")

    assert create_calls == [{"provider": "openai"}]


def test_init_treats_empty_overrides_as_unset(monkeypatch: pytest.MonkeyPatch) -> None:
    """Empty-string api_base/api_key must not be forwarded (env var set to "")."""
    fake_instance = FakeAnyLLMInstance()
    create_calls: list[dict[str, object]] = []

    class FakeAnyLLM:
        @staticmethod
        def create(requested_provider: str, **kwargs):  # noqa: ANN003
            create_calls.append({"provider": requested_provider, **kwargs})
            return fake_instance

    monkeypatch.setattr(anyllm, "ANYLLM_AVAILABLE", True)
    monkeypatch.setattr(anyllm, "AnyLLM", FakeAnyLLM)

    backend = anyllm.AnyLLMBackend(provider="openai", api_key="", api_base="")

    assert backend.api_base is None
    assert backend.api_key is None
    assert create_calls == [{"provider": "openai"}]


def make_choice(
    content: str = "hello", finish_reason: str = "stop", tool_calls=None, index: int = 0
):
    return SimpleNamespace(
        index=index,
        finish_reason=finish_reason,
        message=SimpleNamespace(role="assistant", content=content, tool_calls=tool_calls),
    )


def make_response(*choices, usage=None):
    return SimpleNamespace(
        id="resp_123",
        created=123456,
        choices=list(choices),
        usage=usage,
    )


def make_tool_call(tool_id: str, name: str, arguments):
    return SimpleNamespace(id=tool_id, function=SimpleNamespace(name=name, arguments=arguments))


def test_init_raises_without_anyllm() -> None:
    original_available = anyllm.ANYLLM_AVAILABLE
    try:
        anyllm.ANYLLM_AVAILABLE = False
        with pytest.raises(ImportError):
            anyllm.AnyLLMBackend()
    finally:
        anyllm.ANYLLM_AVAILABLE = original_available


def test_init_name_and_basic_methods(monkeypatch: pytest.MonkeyPatch) -> None:
    backend, instance = make_backend(monkeypatch, provider="groq")

    assert backend.provider == "groq"
    assert backend.name == "anyllm-groq"
    assert backend.map_model_id("claude-3-5") == "claude-3-5"
    assert backend.supports_model("anything") is True
    assert backend.llm is instance


def test_convert_content_blocks_and_messages(monkeypatch: pytest.MonkeyPatch) -> None:
    backend, _instance = make_backend(monkeypatch)

    assert backend._convert_content_blocks([{"type": "text", "text": "hello"}]) == "hello"
    assert backend._convert_content_blocks(
        [
            {"type": "text", "text": "caption"},
            {
                "type": "image",
                "source": {"type": "base64", "media_type": "image/jpeg", "data": "abc"},
            },
            {"type": "image", "source": {"type": "url", "url": "https://example.com/img.png"}},
        ]
    ) == [
        {"type": "text", "text": "caption"},
        {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,abc"}},
        {"type": "image_url", "image_url": {"url": "https://example.com/img.png"}},
    ]
    assert backend._convert_content_blocks([{"type": "tool_use", "id": "ignored"}]) == ""

    converted = backend._convert_messages(
        [
            {"role": "user", "content": "plain text"},
            {
                "role": "assistant",
                "content": [{"type": "text", "text": "a"}, {"type": "text", "text": "b"}],
            },
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": "look"},
                    {"type": "image", "source": {"type": "url", "url": "https://example.com"}},
                ],
            },
            {"role": "user", "content": 123},
        ]
    )

    assert converted == [
        {"role": "user", "content": "plain text"},
        {"role": "assistant", "content": "a\nb"},
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "look"},
                {"type": "image_url", "image_url": {"url": "https://example.com"}},
            ],
        },
    ]


def test_to_anthropic_response_maps_tool_calls_and_usage(monkeypatch: pytest.MonkeyPatch) -> None:
    backend, _instance = make_backend(monkeypatch)
    response = make_response(
        make_choice(
            content="hello",
            finish_reason="tool_calls",
            tool_calls=[
                make_tool_call("tc1", "memory_save", '{"content":"python"}'),
                make_tool_call("tc2", "memory_search", {"query": "python"}),
            ],
        ),
        usage=SimpleNamespace(prompt_tokens=12, completion_tokens=7),
    )

    converted = backend._to_anthropic_response(response, "claude-sonnet")

    assert converted["type"] == "message"
    assert converted["role"] == "assistant"
    assert converted["model"] == "claude-sonnet"
    assert converted["stop_reason"] == "tool_use"
    assert converted["usage"] == {"input_tokens": 12, "output_tokens": 7}
    assert converted["content"][0] == {"type": "text", "text": "hello"}
    assert converted["content"][1]["input"] == {"content": "python"}
    assert converted["content"][2]["input"] == {"query": "python"}


@pytest.mark.asyncio
async def test_send_message_builds_anthropic_response(monkeypatch: pytest.MonkeyPatch) -> None:
    backend, instance = make_backend(monkeypatch)
    instance.response = make_response(
        make_choice("done", "stop"),
        usage=SimpleNamespace(prompt_tokens=4, completion_tokens=6),
    )

    result = await backend.send_message(
        {
            "model": "claude-3-7-sonnet",
            "messages": [{"role": "user", "content": [{"type": "text", "text": "hello"}]}],
            "system": [{"text": "system rule"}, "extra"],
            "max_tokens": 200,
            "temperature": 0.3,
            "top_p": 0.8,
            "stop_sequences": ["END"],
            "tools": [{"name": "t"}],
            "tool_choice": {"type": "auto"},
        },
        {},
    )

    assert isinstance(result, BackendResponse)
    assert result.status_code == 200
    assert result.headers == {"content-type": "application/json"}
    assert result.body["content"][0]["text"] == "done"
    assert instance.calls[0]["messages"][0] == {"role": "system", "content": "system rule extra"}
    assert instance.calls[0]["stop"] == ["END"]


@pytest.mark.asyncio
async def test_send_message_returns_error_response(monkeypatch: pytest.MonkeyPatch) -> None:
    backend, instance = make_backend(monkeypatch)
    instance.raise_error = RuntimeError("authentication api_key missing")

    result = await backend.send_message({"messages": []}, {})

    assert result.status_code == 401
    assert result.body["error"]["type"] == "authentication_error"
    assert result.error == "authentication api_key missing"


@pytest.mark.asyncio
async def test_stream_message_yields_events_and_error(monkeypatch: pytest.MonkeyPatch) -> None:
    backend, instance = make_backend(monkeypatch)
    instance.response = FakeAsyncStream(
        [
            SimpleNamespace(choices=[SimpleNamespace(delta=SimpleNamespace(content="hel"))]),
            SimpleNamespace(choices=[SimpleNamespace(delta=SimpleNamespace(content="lo"))]),
            SimpleNamespace(choices=[]),
        ]
    )

    events = [
        event
        async for event in backend.stream_message(
            {"model": "claude", "messages": [], "system": "sys"}, {}
        )
    ]

    assert [event.event_type for event in events] == [
        "message_start",
        "content_block_start",
        "content_block_delta",
        "content_block_delta",
        "content_block_stop",
        "message_delta",
        "message_stop",
    ]
    assert events[0].data["message"]["model"] == "claude"
    assert events[5].data["usage"] == {"output_tokens": 2}
    assert instance.calls[0]["stream"] is True
    assert instance.calls[0]["messages"][0] == {"role": "system", "content": "sys"}

    backend_error, instance_error = make_backend(monkeypatch, provider="openai")
    instance_error.raise_error = RuntimeError("stream broke")
    error_events = [event async for event in backend_error.stream_message({"messages": []}, {})]
    assert error_events[-1].event_type == "error"
    assert error_events[-1].data["error"]["message"] == "stream broke"


@pytest.mark.asyncio
async def test_send_openai_message_maps_choices_and_tool_calls(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    backend, instance = make_backend(monkeypatch)
    instance.response = make_response(
        make_choice(
            content="answer",
            finish_reason="stop",
            tool_calls=[
                make_tool_call("tc1", "memory_search", '{"query":"python"}'),
                SimpleNamespace(id="tc2", function=None),
            ],
            index=0,
        ),
        usage=SimpleNamespace(prompt_tokens=2, completion_tokens=3, total_tokens=5),
    )

    result = await backend.send_openai_message(
        {
            "model": "gpt-4o",
            "messages": [{"role": "user", "content": "hi"}],
            "max_tokens": 50,
            "temperature": 0.2,
            "top_p": 0.9,
            "stop": ["END"],
            "tools": [{"name": "memory"}],
            "tool_choice": "auto",
            "response_format": {"type": "json_object"},
            "seed": 1,
            "n": 2,
        },
        {},
    )

    assert result.status_code == 200
    assert result.body["object"] == "chat.completion"
    assert (
        result.body["choices"][0]["message"]["tool_calls"][0]["function"]["name"] == "memory_search"
    )
    assert result.body["choices"][0]["message"]["tool_calls"][1] == {
        "id": "tc2",
        "type": "function",
    }
    assert result.body["usage"] == {"prompt_tokens": 2, "completion_tokens": 3, "total_tokens": 5}


@pytest.mark.asyncio
async def test_send_openai_message_returns_error_response(monkeypatch: pytest.MonkeyPatch) -> None:
    backend, instance = make_backend(monkeypatch)
    instance.raise_error = RuntimeError("model not found")

    result = await backend.send_openai_message({"messages": []}, {})

    assert result.status_code == 404
    assert result.body["error"]["type"] == "model_not_found"


@pytest.mark.asyncio
async def test_stream_openai_message_yields_sse_chunks_and_done(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    backend, instance = make_backend(monkeypatch)
    instance.response = FakeAsyncStream(
        [
            SimpleNamespace(
                model_dump=lambda **kwargs: {
                    "id": "chunk1",
                    "choices": [{"delta": {"content": "a"}}],
                }
            ),
            SimpleNamespace(
                model_dump=lambda **kwargs: {
                    "id": "chunk2",
                    "choices": [{"delta": {"content": "b"}}],
                }
            ),
        ]
    )

    chunks = [
        chunk
        async for chunk in backend.stream_openai_message(
            {
                "messages": [{"role": "user", "content": "hi"}],
                "stream_options": {"include_usage": True},
            },
            {},
        )
    ]

    assert chunks[0].startswith("data: {")
    assert chunks[-1] == "data: [DONE]\n\n"
    assert instance.calls[0]["stream"] is True
    assert instance.calls[0]["stream_options"] == {"include_usage": True}

    backend_error, instance_error = make_backend(monkeypatch, provider="anthropic")
    instance_error.raise_error = RuntimeError("rate limit hit")
    error_chunks = [
        chunk async for chunk in backend_error.stream_openai_message({"messages": []}, {})
    ]
    assert '"backend_error"' in error_chunks[0]
    assert error_chunks[-1] == "data: [DONE]\n\n"


def test_error_response_classifies_common_failures(monkeypatch: pytest.MonkeyPatch) -> None:
    backend, _instance = make_backend(monkeypatch)

    auth = backend._error_response(RuntimeError("authentication api key missing"))
    rate = backend._error_response(RuntimeError("rate limit exceeded"), openai_format=True)
    model = backend._error_response(RuntimeError("model not found"), openai_format=True)
    generic = backend._error_response(RuntimeError("other error"))

    assert auth.status_code == 401
    assert auth.body["error"]["type"] == "authentication_error"
    assert rate.status_code == 429
    assert rate.body["error"]["type"] == "rate_limit_exceeded"
    assert model.status_code == 404
    assert model.body["error"]["type"] == "model_not_found"
    assert generic.status_code == 500
    assert generic.body["error"]["type"] == "api_error"


@pytest.mark.asyncio
async def test_close_is_noop(monkeypatch: pytest.MonkeyPatch) -> None:
    backend, _instance = make_backend(monkeypatch)
    assert await backend.close() is None
    assert isinstance(StreamEvent(event_type="message_start", data={}), StreamEvent)