scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-19 00:01:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 88

fix(router): reject non-finite selective head outputs

Browse source
This commit is contained in:
Tin Chi Lo 2026-09-15 13:35:44 -07:00
parent 52295c09fa
commit d076743146
2 changed files with 19 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
litellm/router_strategy/complexity_router/selective_policy.py
View file

@ -76,10 +76,13 @@ class SelectiveHead(BaseModel):
return self
def linear(self, values: Mapping[str, float]) -> tuple[float, ...]:
return tuple(
outputs: Final = tuple(
sum(coefficient * values.get(feature, 0.0) for feature, coefficient in zip(self.features, row)) + intercept
for row, intercept in zip(self.coefficients, self.intercept)
)
if not all(math.isfinite(value) for value in outputs):
raise OverflowError("Selective head produced a non-finite value")
return outputs
def probabilities(self, values: Mapping[str, float]) -> tuple[float, ...]:
if self.constant is not None:

15
tests/test_litellm/router_strategy/test_selective_policy.py
View file

@ -72,6 +72,21 @@ def test_overflow_falls_back_to_the_capable_model() -> None:
assert not policy.evaluate({}, 0.5, 0.9).use_efficient
@pytest.mark.parametrize("coefficient", [1e308, -1e308])
def test_overflowing_binary_logits_cannot_hide_as_equal_probabilities(coefficient: float) -> None:
head: Final = SelectiveHead(
features=("logit_p",), coefficients=((coefficient,),), intercept=(0.0,), classes=(0, 1)
)
policy: Final = SelectivePolicy(
version="test",
feature_schema="v2-v1",
target="scalar_calibration",
threshold=0.05,
heads=(head, head),
)
assert not policy.evaluate({}, 0.99999, 0.99999).use_efficient
def test_per_model_policy_uses_both_fitted_heads() -> None:
efficient: Final = SelectiveHead(
features=("p_e",), coefficients=((0.0,),), intercept=(math.log(4.0),), classes=(0, 1)