9a3926544a
Replace transformers.pipeline('zero-shot-classification', bart-large-mnli)
with AutoModel/AutoTokenizer + mean-pool + L2-normalize + cosine
similarity + softmax (BAAI/bge-large-en-v1.5). One forward pass for the
input, cost independent of category count.
Preserves the exact public interface (ensure_available, classify_zero_shot)
— dispatcher.py needs zero changes.
Also updates LocalEncoderConfig's default model and class docstring in
config.py, and rewrites test_local_encoder.py with 13 offline tests
mocking both transformers and torch.
This is Phase 1 of Wave 5.1 in plans/token-waste-waves.md — implementation
only. Phase 2 (re-tuning confidence_threshold against real traffic) and
Phase 3 (delta-detector) are not included.
179 lines
6.7 KiB
Python
179 lines
6.7 KiB
Python
"""classifier.mode: local_llm / cloud_llm / local_encoder.
|
|
|
|
Every existing deployment is on "local_llm" implicitly today; these tests
|
|
pin that the new field defaults to it (so nothing changes for anyone who
|
|
never touches the new keys) and that the two cross-field validators reject
|
|
the ambiguous or incomplete cloud_llm/local_encoder configurations the same
|
|
way ``tool_use_category_is_a_real_category`` already rejects a dangling
|
|
category reference elsewhere in this file's config.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
import yaml
|
|
from pydantic import ValidationError
|
|
|
|
from config import RouterConfig
|
|
|
|
ROOT = Path(__file__).resolve().parent.parent
|
|
|
|
|
|
@pytest.fixture
|
|
def raw() -> dict:
|
|
with open(ROOT / "config" / "config.yaml") as fh:
|
|
return yaml.safe_load(fh)
|
|
|
|
|
|
def test_mode_defaults_to_local_llm(raw):
|
|
"""The shipped config.yaml writes it explicitly, but a deployment that
|
|
doesn't should still get today's behavior, unchanged."""
|
|
cfg = copy.deepcopy(raw)
|
|
del cfg["classifier"]["mode"]
|
|
loaded = RouterConfig(**cfg)
|
|
assert loaded.classifier.mode == "local_llm"
|
|
|
|
|
|
def test_shipped_config_loads_with_mode_local_llm(raw):
|
|
loaded = RouterConfig(**raw)
|
|
assert loaded.classifier.mode == "local_llm"
|
|
assert loaded.classifier.cloud_primary is None
|
|
assert loaded.classifier.cloud_primary_auto is False
|
|
assert loaded.classifier.encoder is None
|
|
|
|
|
|
def test_unknown_mode_value_is_rejected(raw):
|
|
cfg = copy.deepcopy(raw)
|
|
cfg["classifier"]["mode"] = "cloud_encoder_llm_hybrid" # not a real mode
|
|
with pytest.raises(ValidationError, match="literal_error"):
|
|
# match on the error TYPE, not just any ValidationError: on a tree
|
|
# that doesn't know about classifier.mode at all yet, setting it
|
|
# also raises ValidationError (extra_forbidden), which would let
|
|
# this test pass for the wrong reason on the pre-change code.
|
|
RouterConfig(**cfg)
|
|
|
|
|
|
# --- cloud_llm: exactly one of cloud_primary / cloud_primary_auto ---------
|
|
|
|
|
|
def test_cloud_llm_with_neither_primary_set_is_rejected(raw):
|
|
cfg = copy.deepcopy(raw)
|
|
cfg["classifier"]["mode"] = "cloud_llm"
|
|
with pytest.raises(ValueError, match="neither classifier.cloud_primary"):
|
|
RouterConfig(**cfg)
|
|
|
|
|
|
def test_cloud_llm_with_both_primaries_set_is_rejected(raw):
|
|
cfg = copy.deepcopy(raw)
|
|
cfg["classifier"]["mode"] = "cloud_llm"
|
|
cfg["classifier"]["cloud_primary"] = {
|
|
"base_url": "https://api.neuralwatt.com/v1",
|
|
"model": "deepseek-v4-flash",
|
|
}
|
|
cfg["classifier"]["cloud_primary_auto"] = True
|
|
with pytest.raises(ValueError, match="BOTH classifier.cloud_primary"):
|
|
RouterConfig(**cfg)
|
|
|
|
|
|
def test_cloud_llm_with_pinned_primary_loads(raw):
|
|
cfg = copy.deepcopy(raw)
|
|
cfg["classifier"]["mode"] = "cloud_llm"
|
|
cfg["classifier"]["cloud_primary"] = {
|
|
"base_url": "https://api.neuralwatt.com/v1",
|
|
"model": "deepseek-v4-flash",
|
|
"api_key_env": "NEURALWATT_API_KEY",
|
|
}
|
|
loaded = RouterConfig(**cfg)
|
|
assert loaded.classifier.cloud_primary.model == "deepseek-v4-flash"
|
|
assert loaded.classifier.cloud_primary_auto is False
|
|
|
|
|
|
def test_cloud_llm_with_auto_loads(raw):
|
|
cfg = copy.deepcopy(raw)
|
|
cfg["classifier"]["mode"] = "cloud_llm"
|
|
cfg["classifier"]["cloud_primary_auto"] = True
|
|
loaded = RouterConfig(**cfg)
|
|
assert loaded.classifier.cloud_primary_auto is True
|
|
assert loaded.classifier.cloud_primary is None
|
|
|
|
|
|
def test_local_llm_mode_does_not_require_either_primary(raw):
|
|
"""The validator is scoped to mode == cloud_llm -- it must not fire for
|
|
the default mode just because cloud_primary/auto are both unset."""
|
|
cfg = copy.deepcopy(raw)
|
|
assert cfg["classifier"]["mode"] == "local_llm"
|
|
RouterConfig(**cfg) # must not raise
|
|
|
|
|
|
# --- local_encoder: needs the encoder block ------------------------------
|
|
|
|
|
|
def test_local_encoder_without_encoder_block_is_rejected(raw):
|
|
cfg = copy.deepcopy(raw)
|
|
cfg["classifier"]["mode"] = "local_encoder"
|
|
with pytest.raises(ValueError, match="classifier.encoder is not"):
|
|
RouterConfig(**cfg)
|
|
|
|
|
|
def test_local_encoder_with_empty_encoder_block_loads_with_defaults(raw):
|
|
cfg = copy.deepcopy(raw)
|
|
cfg["classifier"]["mode"] = "local_encoder"
|
|
cfg["classifier"]["encoder"] = {}
|
|
loaded = RouterConfig(**cfg)
|
|
assert loaded.classifier.encoder.model == "BAAI/bge-large-en-v1.5"
|
|
assert loaded.classifier.encoder.device == "cpu"
|
|
assert loaded.classifier.encoder.confidence_threshold == 0.5
|
|
|
|
|
|
def test_local_encoder_rejects_unknown_device(raw):
|
|
cfg = copy.deepcopy(raw)
|
|
cfg["classifier"]["mode"] = "local_encoder"
|
|
cfg["classifier"]["encoder"] = {"device": "tpu"}
|
|
with pytest.raises(ValidationError, match="literal_error"):
|
|
# Same reasoning as test_unknown_mode_value_is_rejected: on a tree
|
|
# without classifier.encoder at all, this also raises (extra_
|
|
# forbidden on the block itself), which would pass for the wrong
|
|
# reason before LocalEncoderConfig's device field exists.
|
|
RouterConfig(**cfg)
|
|
|
|
|
|
def test_local_encoder_rejects_percent_style_confidence_threshold(raw):
|
|
"""classify_zero_shot returns a 0.0-1.0 probability, so a percent-style
|
|
value (e.g. 80 meaning "80%") must be rejected -- otherwise no real
|
|
confidence score can ever clear the threshold and every classification
|
|
silently fails. Caught live 2026-09-06 via the admin UI taking a raw
|
|
number with no conversion."""
|
|
cfg = copy.deepcopy(raw)
|
|
cfg["classifier"]["mode"] = "local_encoder"
|
|
cfg["classifier"]["encoder"] = {"confidence_threshold": 80}
|
|
with pytest.raises(ValueError, match=r"must be in \[0.0, 1.0\]"):
|
|
RouterConfig(**cfg)
|
|
|
|
|
|
def test_local_encoder_rejects_negative_confidence_threshold(raw):
|
|
cfg = copy.deepcopy(raw)
|
|
cfg["classifier"]["mode"] = "local_encoder"
|
|
cfg["classifier"]["encoder"] = {"confidence_threshold": -0.1}
|
|
with pytest.raises(ValueError, match=r"must be in \[0.0, 1.0\]"):
|
|
RouterConfig(**cfg)
|
|
|
|
|
|
def test_local_encoder_accepts_confidence_threshold_at_bounds(raw):
|
|
cfg = copy.deepcopy(raw)
|
|
cfg["classifier"]["mode"] = "local_encoder"
|
|
cfg["classifier"]["encoder"] = {"confidence_threshold": 0.0}
|
|
assert RouterConfig(**cfg).classifier.encoder.confidence_threshold == 0.0
|
|
cfg["classifier"]["encoder"] = {"confidence_threshold": 1.0}
|
|
assert RouterConfig(**cfg).classifier.encoder.confidence_threshold == 1.0
|
|
|
|
|
|
def test_local_encoder_unaffected_by_the_cloud_llm_validator(raw):
|
|
"""A local_encoder config leaving cloud_primary/auto both unset must not
|
|
trip the cloud_llm validator -- it's scoped to mode == cloud_llm."""
|
|
cfg = copy.deepcopy(raw)
|
|
cfg["classifier"]["mode"] = "local_encoder"
|
|
cfg["classifier"]["encoder"] = {}
|
|
RouterConfig(**cfg) # must not raise
|