import ast import re import unittest from pathlib import Path from typing import ClassVar from jinja2 import Environment, StrictUndefined, select_autoescape from utils.cache.files import PROJECT_ROOT from utils.cache.yaml import load_yaml _AI_VARS = Path(PROJECT_ROOT) / "group_vars" / "all" / "16_ai.yml" def _ansible_bool(value): if isinstance(value, str): return value.strip().lower() in ("true", "yes", "on", "1") return bool(value) def _stub_lookup(preload_models, lmstudio_models): def lookup(kind, role, path, *args): if (kind, path) == ("config", "services.ollama.preload_models"): return preload_models if (kind, path) == ("config", "services.lmstudio.preload_models"): return lmstudio_models raise AssertionError(f"unexpected lookup({kind!r}, {role!r}, {path!r})") return lookup class TestLitellmChatModel(unittest.TestCase): """The gateway only publishes a model it has a backend for (roles/svc-ai-litellm/templates/config.yaml.j2). These expressions decide what consumers are told to ask for, so naming an unpublished model makes every prompt fail with a 400. """ env: ClassVar[Environment] source: ClassVar[dict] @classmethod def setUpClass(cls): cls.env = Environment(undefined=StrictUndefined, autoescape=select_autoescape()) cls.env.tests["search"] = lambda value, pattern: bool( re.search(pattern, str(value)) ) cls.env.filters["bool"] = _ansible_bool cls.source = load_yaml(_AI_VARS) def _render( self, name, *, roles, api_key, preload_models=(), lmstudio_models=(), remote_aliases=None, ): return ( self.env.from_string(self.source[name]) .render( LITELLM_BACKEND_ROLES=list(roles), LITELLM_OLLAMA_BACKEND=str("svc-ai-ollama" in roles), LITELLM_LMSTUDIO_BACKEND=str("svc-ai-lmstudio" in roles), AI_REMOTE_ALIASES=( list(remote_aliases) if remote_aliases is not None else (["openrouter/auto"] if api_key else []) ), lookup=_stub_lookup( [{"alias": alias, "name": alias} for alias in preload_models], [ {"alias": alias, "name": f"{alias}-gguf"} for alias in lmstudio_models ], ), ) .strip() ) def _both(self, **kwargs): return ( self._render("LITELLM_CHAT_MODEL", **kwargs), self._render("LITELLM_CHAT_MODEL_SERVED", **kwargs), ) def test_no_backend_and_no_key_names_no_model(self): model, served = self._both(roles=["web-app-mattermost"], api_key="") self.assertEqual(model, "") self.assertEqual(served, "False") def test_no_backend_with_a_key_uses_openrouter(self): model, served = self._both(roles=["web-app-mattermost"], api_key="sk-test") self.assertEqual(model, "openrouter/auto") self.assertEqual(served, "True") def test_a_declared_mock_answers_when_no_backend_preloaded_anything(self): model = self._render( "LITELLM_CHAT_MODEL", roles=["svc-ai-ollama"], api_key="", preload_models=[], remote_aliases=["mock/deterministic"], ) self.assertEqual( model, "mock/deterministic", "a CI round that deploys the backend without preloading must fall to the mock", ) def test_a_preloaded_model_wins_over_the_mock(self): model = self._render( "LITELLM_CHAT_MODEL", roles=["svc-ai-ollama"], api_key="", preload_models=["qwen2.5:0.5b"], remote_aliases=["mock/deterministic"], ) self.assertEqual( model, "qwen2.5:0.5b", "the round that loads a real model must prove that model, not the mock", ) def test_without_a_mock_the_chain_is_unchanged(self): self.assertEqual( self._render( "LITELLM_CHAT_MODEL", roles=["web-app-mattermost"], api_key="sk-test" ), "openrouter/auto", "production declares no mock, so it must resolve exactly as before", ) def test_lmstudio_wins_over_the_openrouter_fallback(self): model, served = self._both( roles=["svc-ai-lmstudio"], api_key="", lmstudio_models=["qwen2.5:0.5b"] ) self.assertEqual(model, "qwen2.5:0.5b") self.assertEqual(served, "True") def test_lmstudio_without_a_preloaded_model_serves_nothing(self): served = self._render( "LITELLM_CHAT_MODEL_SERVED", roles=["svc-ai-lmstudio"], api_key="" ) self.assertEqual(served, "False") def test_both_backends_name_the_same_model(self): alias = "qwen2.5:0.5b" with_ollama = self._render( "LITELLM_CHAT_MODEL", roles=["svc-ai-ollama"], api_key="", preload_models=[alias], ) with_lmstudio = self._render( "LITELLM_CHAT_MODEL", roles=["svc-ai-lmstudio"], api_key="", lmstudio_models=[alias], ) self.assertEqual(with_ollama, with_lmstudio) self.assertEqual(with_ollama, alias) def test_ollama_serves_its_first_non_embedding_model(self): model, served = self._both( roles=["svc-ai-ollama"], api_key="", preload_models=["nomic-embed-text", "llama3.2"], ) self.assertEqual(model, "llama3.2") self.assertEqual(served, "True") def test_ollama_with_only_embedding_models_serves_nothing(self): served = self._render( "LITELLM_CHAT_MODEL_SERVED", roles=["svc-ai-ollama"], api_key="", preload_models=["nomic-embed-text"], ) self.assertEqual(served, "False") def test_the_model_is_named_exactly_when_one_is_served(self): cases = ( (["web-app-mattermost"], "", (), ()), (["web-app-mattermost"], "sk-test", (), ()), (["svc-ai-lmstudio"], "", (), ()), (["svc-ai-lmstudio"], "", (), ("qwen2.5:0.5b",)), (["svc-ai-ollama"], "", ("llama3.2",), ()), ) for roles, api_key, preload, lmstudio in cases: with self.subTest(roles=roles, api_key=bool(api_key), lmstudio=lmstudio): model, served = self._both( roles=roles, api_key=api_key, preload_models=preload, lmstudio_models=lmstudio, ) self.assertEqual(bool(model), served == "True") class TestAiRemoteAliases(unittest.TestCase): models: ClassVar[list] template: ClassVar[str] @classmethod def setUpClass(cls): cls.models = [ {"alias": "openrouter/auto", "provider": "openrouter"}, {"alias": "openrouter/ternary-bonsai-2-27b", "provider": "openrouter"}, {"alias": "openai/gpt-4o-mini", "provider": "openai"}, ] cls.template = load_yaml(_AI_VARS)["AI_REMOTE_ALIASES"] def _render(self, enabled): models = self.models def lookup(kind, *terms): if kind == "config": self.assertEqual( terms, ("svc-ai-litellm", "services.litellm.remote_models") ) return models if kind == "api_enabled": return terms[0] in enabled raise AssertionError(f"unexpected lookup({kind!r}, {terms!r})") env = Environment(undefined=StrictUndefined, autoescape=False) # noqa: S701 - renders a Python list literal, not markup return ast.literal_eval( env.from_string(self.template).render( lookup=lookup, AI_MOCK_PROVIDER="mock" ) ) def test_every_model_of_an_enabled_provider_is_listed(self): self.assertEqual( self._render({"openrouter"}), ["openrouter/auto", "openrouter/ternary-bonsai-2-27b"], ) def test_no_enabled_provider_lists_nothing(self): self.assertEqual(self._render(set()), []) class TestAiAgentModel(unittest.TestCase): """Every agent the broker starts asks the gateway for this alias, and roles/svc-ai-agent-broker/tasks/00_core.yml stops the deploy when it is empty. A deployment that carries neither a mock nor a backend nor a provider key resolves to '', so the two sources are pinned together. """ template: ClassVar[str] @classmethod def setUpClass(cls): cls.template = load_yaml(_AI_VARS)["AI_AGENT_MODEL"] def _render(self, *, mock_aliases, chat_model): env = Environment(undefined=StrictUndefined, autoescape=False) # noqa: S701 - renders a bare alias, not markup return env.from_string(self.template).render( AI_MOCK_ALIASES=list(mock_aliases), LITELLM_CHAT_MODEL=chat_model ) def test_a_declared_mock_is_preferred_over_the_served_model(self): self.assertEqual( self._render( mock_aliases=["mock/deterministic"], chat_model="qwen2.5:0.5b" ), "mock/deterministic", ) def test_without_a_mock_the_served_model_answers(self): self.assertEqual( self._render(mock_aliases=[], chat_model="qwen2.5:0.5b"), "qwen2.5:0.5b" ) def test_neither_a_mock_nor_a_served_model_names_nothing(self): self.assertEqual( self._render(mock_aliases=[], chat_model=""), "", "a broker deployed without a backend role and without a provider " "key has no model to start its agents on, which 00_core.yml turns " "into a stopped deploy rather than agents that all answer 400", ) if __name__ == "__main__": unittest.main()