tests/llm/test_router.py

from __future__ import annotations

import pytest

from app.llm.base import ChatResult, ProviderError, RateLimitError
from app.llm.router import LLMRouter


class FakeProvider:
    def __init__(self, result: ChatResult | None = None, raise_exc: Exception | None = None):
        self.result = result
        self.raise_exc = raise_exc
        self.calls = 0
        self.last_model: str | None = None

    async def chat(self, messages, tools, model, temperature=0.3):
        self.calls += 1
        self.last_model = model
        if self.raise_exc is not None:
            raise self.raise_exc
        assert self.result is not None
        return self.result


def _models():
    return {
        "groq": {"small": "gs", "large": "gl"},
        "cloudflare": {"small": "cs", "large": "cl"},
    }


async def test_router_routes_tier_to_correct_model():
    groq = FakeProvider(result=ChatResult(content="ok", tool_calls=[], finish_reason="stop"))
    router = LLMRouter(providers={"groq": groq}, order=["groq"], models=_models())
    await router.chat(messages=[{"role": "user", "content": "x"}], tools=[], tier="small")
    assert groq.last_model == "gs"
    await router.chat(messages=[{"role": "user", "content": "x"}], tools=[], tier="large")
    assert groq.last_model == "gl"


async def test_router_failover_on_ratelimit():
    groq = FakeProvider(raise_exc=RateLimitError())
    cf = FakeProvider(result=ChatResult(content="ok", tool_calls=[], finish_reason="stop"))
    router = LLMRouter(
        providers={"groq": groq, "cloudflare": cf},
        order=["groq", "cloudflare"],
        models=_models(),
    )
    r = await router.chat(messages=[{"role": "user", "content": "x"}], tools=[], tier="small")
    assert r.content == "ok"
    assert groq.calls == 1
    assert cf.calls == 1
    assert cf.last_model == "cs"


def test_build_router_multiple_groq_keys_and_cloudflare():
    from types import SimpleNamespace

    from app.llm.router import build_router_from_settings

    s = SimpleNamespace(
        groq_api_key="k1",
        groq_api_keys=["k2", "k3", "k1"],  # k1 deduped
        groq_base_url="https://api.groq.com/openai/v1",
        model_small="s",
        model_large="l",
        cloudflare_account_id="acct",
        cloudflare_api_token="tok",
        cf_model_small="cs",
        cf_model_large="cl",
    )
    router = build_router_from_settings(s)
    assert router.order == ["groq0", "groq1", "groq2", "cloudflare"]  # 3 keys + CF
    assert len(router.providers) == 4


def test_build_router_with_three_providers():
    from types import SimpleNamespace

    from app.llm.router import build_router_from_settings

    s = SimpleNamespace(
        groq_api_key="k1", groq_api_keys=[], groq_base_url="u",
        model_small="s", model_large="l",
        cloudflare_account_id="a", cloudflare_api_token="t", cf_model_small="cs", cf_model_large="cl",
        cerebras_api_key="cb", cerebras_base_url="https://api.cerebras.ai/v1", cerebras_model="gpt-oss-120b",
    )
    router = build_router_from_settings(s)
    assert router.order == ["groq0", "cloudflare", "cerebras"]  # 3 independent providers
    assert len(router.providers) == 3


async def test_router_tags_provider_name():
    groq = FakeProvider(result=ChatResult(content="ok", tool_calls=[], finish_reason="stop"))
    router = LLMRouter(providers={"groq0": groq}, order=["groq0"], models={"groq0": {"large": "l"}})
    r = await router.chat(messages=[{"role": "user", "content": "x"}], tools=[], tier="large")
    assert r.provider == "groq0"


async def test_router_raises_when_all_fail():
    groq = FakeProvider(raise_exc=ProviderError("g"))
    cf = FakeProvider(raise_exc=ProviderError("c"))
    router = LLMRouter(
        providers={"groq": groq, "cloudflare": cf},
        order=["groq", "cloudflare"],
        models=_models(),
    )
    with pytest.raises(ProviderError):
        await router.chat(messages=[{"role": "user", "content": "x"}], tools=[], tier="large")