tests/test_baml_client_selection.py

"""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 <Name>` 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<llm> 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<llm> 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<llm> 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<llm> 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<llm> ZeroGPUScript" in text
    assert "google/gemma-4-12b-it" in text