from clients import art_client_from_env, boss_client_from_env, card_client_from_env from art import DiffusersImageClient, LazyArtClient from boss import NemotronBossClient from generator import LlamaCppCardClient, MiniCPMCardClient from local_llm import LocalCompletionClient, LocalJsonChatClient, nemotron_prompt class FakeChat: # Complete one fake prompt. def complete(self, system: str, user: str) -> str: return "{}" class FakeArtBackend: # Return a deterministic image URI. def create_art(self, prompt: str) -> str: return f"uri:{prompt}" # Verify card client is disabled without an endpoint. def test_card_client_from_env_disabled(monkeypatch) -> None: monkeypatch.delenv("TABRAS_CARD_ENDPOINT", raising=False) assert card_client_from_env() is None # Verify card client reads endpoint environment. def test_card_client_from_env(monkeypatch) -> None: monkeypatch.setenv("TABRAS_CARD_ENDPOINT", "http://cards") monkeypatch.setenv("TABRAS_CARD_MODEL", "mini") client = card_client_from_env() assert isinstance(client, MiniCPMCardClient) assert client.chat.endpoint == "http://cards" assert client.chat.model == "mini" # Verify card client can use the direct MiniCPM Transformers backend. def test_card_client_from_env_transformers(monkeypatch) -> None: monkeypatch.setenv("TABRAS_CARD_BACKEND", "transformers") monkeypatch.setenv("TABRAS_CARD_MODEL", "mini-local") monkeypatch.setattr("clients.MiniCPMTransformersChatClient.load", lambda model_path, **kwargs: FakeChat()) client = card_client_from_env() assert isinstance(client, MiniCPMCardClient) assert isinstance(client.chat, FakeChat) # Verify card client can use a llama.cpp MiniCPM backend. def test_card_client_from_env_llamacpp(monkeypatch) -> None: monkeypatch.setenv("TABRAS_CARD_BACKEND", "llamacpp") monkeypatch.setenv("TABRAS_CARD_ENDPOINT", "http://cards") monkeypatch.setenv("TABRAS_CARD_MODEL", "mini-q4") client = card_client_from_env() assert isinstance(client, LlamaCppCardClient) assert isinstance(client.chat, LocalJsonChatClient) assert client.chat.endpoint == "http://cards" assert client.chat.model == "mini-q4" assert client.chat.max_tokens == 320 # Verify boss client is disabled unless enabled. def test_boss_client_from_env_disabled(monkeypatch) -> None: monkeypatch.delenv("TABRAS_AI_BOSS", raising=False) assert boss_client_from_env() is None # Verify boss client reads endpoint environment. def test_boss_client_from_env(monkeypatch) -> None: monkeypatch.setenv("TABRAS_AI_BOSS", "1") monkeypatch.setenv("TABRAS_BOSS_ENDPOINT", "http://boss") monkeypatch.setenv("TABRAS_BOSS_MODEL", "nemotron") client = boss_client_from_env() assert isinstance(client, NemotronBossClient) assert client.chat.endpoint == "http://boss" assert client.chat.model == "nemotron" # Verify boss client can use a Nemotron formatted completion endpoint. def test_boss_client_from_env_completion(monkeypatch) -> None: monkeypatch.setenv("TABRAS_AI_BOSS", "1") monkeypatch.setenv("TABRAS_BOSS_BACKEND", "completion") monkeypatch.setenv("TABRAS_BOSS_ENDPOINT", "http://boss-completion") client = boss_client_from_env() assert isinstance(client, NemotronBossClient) assert isinstance(client.chat, LocalCompletionClient) assert client.chat.endpoint == "http://boss-completion" assert client.chat.prompt_template is nemotron_prompt # Verify boss client can use the direct Nemotron Transformers backend. def test_boss_client_from_env_transformers(monkeypatch) -> None: monkeypatch.setenv("TABRAS_AI_BOSS", "1") monkeypatch.setenv("TABRAS_BOSS_BACKEND", "transformers") monkeypatch.setenv("TABRAS_BOSS_MODEL", "nemotron-local") monkeypatch.setattr("clients.NemotronTransformersChatClient.load", lambda *_args, **_kwargs: FakeChat()) client = boss_client_from_env() assert isinstance(client, NemotronBossClient) assert isinstance(client.chat, FakeChat) # Verify boss client can use a quantized MLX backend. def test_boss_client_from_env_mlx(monkeypatch) -> None: monkeypatch.setenv("TABRAS_AI_BOSS", "1") monkeypatch.setenv("TABRAS_BOSS_BACKEND", "mlx") monkeypatch.setenv("TABRAS_BOSS_MODEL", "nemotron-mlx") monkeypatch.setattr("clients.MLXChatClient.load", lambda *_args, **_kwargs: FakeChat()) client = boss_client_from_env() assert isinstance(client, NemotronBossClient) assert isinstance(client.chat, FakeChat) # Verify art client is disabled unless explicitly enabled. def test_art_client_from_env_disabled(monkeypatch) -> None: monkeypatch.delenv("TABRAS_ART_BACKEND", raising=False) monkeypatch.setattr("clients._art_client_cache", None) assert art_client_from_env() is None # Verify art client loads and caches the SD-Turbo diffusers backend. def test_art_client_from_env_diffusers(monkeypatch) -> None: monkeypatch.setenv("TABRAS_ART_BACKEND", "diffusers") monkeypatch.setenv("TABRAS_ART_MODEL", "sdxl-local") monkeypatch.setenv("TABRAS_ART_STEPS", "1") monkeypatch.setenv("TABRAS_ART_WIDTH", "128") monkeypatch.setenv("TABRAS_ART_HEIGHT", "96") monkeypatch.setattr("clients._art_client_cache", None) calls = [] def fake_load(model_id: str, **kwargs): calls.append(model_id) assert model_id == "sdxl-local" assert kwargs["steps"] == 1 assert kwargs["width"] == 128 assert kwargs["height"] == 96 return FakeArtBackend() monkeypatch.setattr(DiffusersImageClient, "load", staticmethod(fake_load)) client = art_client_from_env() assert isinstance(client, LazyArtClient) assert art_client_from_env() is client assert calls == [] assert client.create_art("spell") == "uri:spell" assert calls == ["sdxl-local"] # Verify art client defaults use the lazy SDXL-Turbo backend. def test_art_client_from_env_fast_defaults(monkeypatch) -> None: monkeypatch.setenv("TABRAS_ART_BACKEND", "diffusers") monkeypatch.delenv("TABRAS_ART_MODEL", raising=False) monkeypatch.delenv("TABRAS_ART_STEPS", raising=False) monkeypatch.delenv("TABRAS_ART_WIDTH", raising=False) monkeypatch.delenv("TABRAS_ART_HEIGHT", raising=False) monkeypatch.setattr("clients._art_client_cache", None) def fake_load(model_id: str, **kwargs): assert model_id == "stabilityai/sdxl-turbo" assert kwargs["steps"] == 1 assert kwargs["width"] == 512 assert kwargs["height"] == 320 return FakeArtBackend() monkeypatch.setattr(DiffusersImageClient, "load", staticmethod(fake_load)) assert art_client_from_env().create_art("spark") == "uri:spark"