"""Contract tests for BAML client selection (issue 46). These tests assert that: 1. Each BAML function declares the correct client in its .baml source. 2. Dev mode (localhost) resolves to the default BAML clients (MlxPrecise/MlxCreative). 3. Prod mode (HF Router base_url) resolves to the "HF" client override. 4. Explicit env override takes precedence over base_url inference. These are *static* contracts on the .baml sources plus *runtime* contracts on the _resolve_client_name helper — fast, no network, no model loading. """ from __future__ import annotations import re from pathlib import Path import pytest from densefeed.backends.baml_backend import BAMLBackend, _resolve_client_name from densefeed.config import DenseFeedConfig BAML_DIR = Path(__file__).resolve().parent.parent / "baml_src" # --------------------------------------------------------------------------- # Mapping: function name -> expected default client in the .baml source # --------------------------------------------------------------------------- EXPECTED_CLIENT_BINDINGS: dict[str, str] = { "FilterItems": "MlxPrecise", "ScoreItems": "MlxPrecise", "GenerateOutline": "MlxCreative", "GenerateScriptSegment": "MlxCreative", } def _extract_function_client(baml_text: str, function_name: str) -> str | None: """Parse the `client ` line from a BAML function definition.""" pattern = rf"function\s+{re.escape(function_name)}\s*\([^)]*\)\s*[^{{]*\{{\s*client\s+(\w+)" m = re.search(pattern, baml_text) return m.group(1) if m else None # -- Static contract: .baml files declare the right default client ---------- @pytest.mark.parametrize( "func_name, expected_client", list(EXPECTED_CLIENT_BINDINGS.items()), ids=list(EXPECTED_CLIENT_BINDINGS.keys()), ) def test_baml_function_declares_correct_client(func_name: str, expected_client: str): """Each BAML function must bind to its intended default client.""" # Scan all .baml files for the function definition for baml_file in sorted(BAML_DIR.glob("*.baml")): text = baml_file.read_text() client = _extract_function_client(text, func_name) if client is not None: assert client == expected_client, ( f"{func_name} in {baml_file.name} declares client '{client}', " f"expected '{expected_client}'" ) return pytest.fail(f"BAML function '{func_name}' not found in {BAML_DIR}") def test_all_rank_functions_use_precise_client(): """Filter + Score must use the deterministic MlxPrecise client (temp 0.2).""" rank_functions = ["FilterItems", "ScoreItems"] for func_name in rank_functions: for baml_file in sorted(BAML_DIR.glob("*.baml")): client = _extract_function_client(baml_file.read_text(), func_name) if client is not None: assert client == "MlxPrecise", ( f"{func_name} must use MlxPrecise for deterministic ranking, " f"got '{client}' in {baml_file.name}" ) def test_all_script_functions_use_creative_client(): """Outline + Script generation must use the creative MlxCreative client (temp 0.7).""" script_functions = ["GenerateOutline", "GenerateScriptSegment"] for func_name in script_functions: for baml_file in sorted(BAML_DIR.glob("*.baml")): client = _extract_function_client(baml_file.read_text(), func_name) if client is not None: assert client == "MlxCreative", ( f"{func_name} must use MlxCreative for creative generation, " f"got '{client}' in {baml_file.name}" ) def test_no_baml_function_uses_hf_client_by_default(): """No function should default to the HF Router client — that is runtime-only.""" hf_functions = [] for baml_file in sorted(BAML_DIR.glob("*.baml")): text = baml_file.read_text() for func_name in EXPECTED_CLIENT_BINDINGS: client = _extract_function_client(text, func_name) if client == "HF": hf_functions.append(f"{func_name} in {baml_file.name}") assert not hf_functions, ( f"Functions defaulting to HF client (should be runtime override only): " f"{', '.join(hf_functions)}" ) # -- BAML client definitions exist and are distinct ------------------------- def test_clients_baml_defines_mlx_precise(): """clients.baml must define the MlxPrecise client for ranking.""" text = (BAML_DIR / "clients.baml").read_text() assert "client MlxPrecise" in text assert "localhost:8081" in text assert "temperature 0.2" in text def test_clients_baml_defines_mlx_creative(): """clients.baml must define the MlxCreative client for script generation.""" text = (BAML_DIR / "clients.baml").read_text() assert "client MlxCreative" in text assert "localhost:8080" in text assert "temperature 0.7" in text def test_clients_baml_defines_hf_router(): """clients.baml must define the HF client for production override.""" text = (BAML_DIR / "clients.baml").read_text() assert "client HF" in text assert "router.huggingface.co" in text # -- Runtime contract: _resolve_client_name --------------------------------- def test_dev_mode_resolves_to_baml_defaults(): """Local base_url -> empty string, meaning .baml file defaults apply.""" cfg = DenseFeedConfig() assert "localhost" in cfg.llm.base_url assert _resolve_client_name(cfg) == "" def test_prod_mode_resolves_to_hf(): """HF Router base_url -> 'HF' client override.""" cfg = DenseFeedConfig() cfg.llm.base_url = "https://router.huggingface.co/v1" assert _resolve_client_name(cfg) == "HF" def test_env_override_forces_client(monkeypatch): """DENSEFEED_BAML_CLIENT env var overrides base_url logic.""" monkeypatch.setenv("DENSEFEED_BAML_CLIENT", "MlxPrecise") # Even with HF base_url, explicit override wins cfg = DenseFeedConfig() cfg.llm.base_url = "https://router.huggingface.co/v1" assert _resolve_client_name(cfg) == "MlxPrecise" def test_dev_backend_uses_default_baml_client(): """BAMLBackend with localhost base_url uses the global `b` (no override).""" from baml_client import b cfg = DenseFeedConfig() backend = BAMLBackend(cfg) assert backend._b is b, "Dev mode must use default BAML client (MlxPrecise/MlxCreative)" def test_prod_backend_overrides_to_hf(): """BAMLBackend with HF base_url creates a with_options client.""" from baml_client import b cfg = DenseFeedConfig() cfg.llm.base_url = "https://router.huggingface.co/v1" backend = BAMLBackend(cfg) assert backend._b is not b, "Prod mode must override to HF client" # -- Factory contract: backend selection ------------------------------------ def test_baml_backend_type_for_baml_config(): """config.llm.backend='baml' -> BAMLBackend (uses BAML client routing).""" from densefeed.backends import create_backend cfg = DenseFeedConfig() cfg.llm.backend = "baml" assert isinstance(create_backend(cfg), BAMLBackend) def test_job_backend_type_for_job_config(monkeypatch): """config.llm.backend='job' -> JobBackend (bypasses BAML client entirely).""" from densefeed.backends import create_backend from densefeed.backends.job_backend import JobBackend monkeypatch.setattr(JobBackend, "_load_model", lambda self: None) cfg = DenseFeedConfig() cfg.llm.backend = "job" assert isinstance(create_backend(cfg), JobBackend) # -- ZeroGPU client contract (T04) ------------------------------------------- def test_zerogpu_rank_client_resolves(): """b.with_options(client='ZeroGPURank') must produce a valid BAML client.""" from baml_client import b opts = b.with_options(client="ZeroGPURank") assert opts is not None assert opts is not b, "with_options must return a new client instance" def test_zerogpu_script_client_resolves(): """b.with_options(client='ZeroGPUScript') must produce a valid BAML client.""" from baml_client import b opts = b.with_options(client="ZeroGPUScript") assert opts is not None assert opts is not b, "with_options must return a new client instance" def test_clients_baml_defines_zerogpu_rank(): """clients.baml must define the ZeroGPURank client for ZeroGPU ranking.""" text = (BAML_DIR / "clients.baml").read_text() assert "client ZeroGPURank" in text assert "openbmb/MiniCPM5-1B" in text def test_clients_baml_defines_zerogpu_script(): """clients.baml must define the ZeroGPUScript client for ZeroGPU script generation.""" text = (BAML_DIR / "clients.baml").read_text() assert "client ZeroGPUScript" in text assert "google/gemma-4-12b-it" in text