diff --git a/litellm/router_rust/steps/callbacks.py b/litellm/router_rust/steps/callbacks.py index cd2d69d13a8..48cb998df9b 100644 --- a/litellm/router_rust/steps/callbacks.py +++ b/litellm/router_rust/steps/callbacks.py @@ -452,6 +452,10 @@ class OptionalChecks: *self._policies_not_set_globally(), ) + def has_filters(self) -> bool: + """Whether ``filter_deployments`` would run any callback: without one it keeps every deployment.""" + return any(filters_deployments(callback) for callback in self.routing_callbacks()) + async def filter_deployments( self, model: str, @@ -464,12 +468,10 @@ class OptionalChecks: the Rust core takes only if its later filters keep it; a filter's error is raised for the Rust core to decide on.""" remaining: list[object] = list(deployments) # mutable-ok: filter contract # rebind-ok: narrowed - base: Final[object] = CustomLogger.async_filter_deployments # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # untyped deployment dicts request_kwargs: Final = dict(request) callbacks: Final = tuple(self.routing_callbacks()) for callback in callbacks: - # The base filter returns its input unchanged: only overriding filters are run. - if not isinstance(callback, CustomLogger) or type(callback).async_filter_deployments is base: # pyright: ignore[reportUnknownMemberType] # see above + if not isinstance(callback, CustomLogger) or not filters_deployments(callback): continue if isinstance(callback, PromptCachingDeploymentCheck): continue @@ -572,14 +574,22 @@ def _ids_of(value: object) -> Sequence[object]: return value if isinstance(value, list | tuple) else () # pyright: ignore[reportUnknownVariableType] # the step's id list +def filters_deployments(callback: object) -> bool: + """A routing callback that overrides ``async_filter_deployments``: the base filter returns its + input unchanged, so only overriding filters are run.""" + base: Final[object] = CustomLogger.async_filter_deployments # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # untyped deployment dicts + return isinstance(callback, CustomLogger) and type(callback).async_filter_deployments is not base # pyright: ignore[reportUnknownMemberType] # see above + + def filtered_ids(remaining: Iterable[object], candidates: Sequence[str]) -> list[str]: # mutable-ok: a JSON list """The ids of the deployments the filters left, in their order, among the candidates.""" - ids: Final[list[str]] = [] # mutable-ok: built here + allowed: Final = frozenset(candidates) + ids: Final[dict[str, None]] = {} # mutable-ok: built here, insertion-ordered for deployment in remaining: found = (as_str_mapping((as_str_mapping(deployment) or {}).get("model_info")) or {}).get("id") - if isinstance(found, str) and found in candidates and found not in ids: - ids.append(found) - return ids + if isinstance(found, str) and found in allowed: + ids.setdefault(found) + return list(ids) class RustOptionalChecks: @@ -698,6 +708,8 @@ class RustOptionalChecks: from litellm.router_rust.core import attempt_failure candidates: Final = [str(id) for id in step.get("candidates") or ()] # pyright: ignore[reportGeneralTypeIssues, reportUnknownVariableType, reportUnknownArgumentType] # the step's id list + if not self._opt.has_filters(): + return (list(dict.fromkeys(candidates)), None), None records: Final = self._records_of(candidates) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownVariableType] # RustCore's deployment dicts narrowed: Final = _request_with_failed_ids(request, step) try: diff --git a/litellm/rust_bridge/host/model_capabilities.py b/litellm/rust_bridge/host/model_capabilities.py index 39eb0b08aa2..efb8e10022a 100644 --- a/litellm/rust_bridge/host/model_capabilities.py +++ b/litellm/rust_bridge/host/model_capabilities.py @@ -1,6 +1,27 @@ from __future__ import annotations +from collections.abc import Mapping +from functools import lru_cache +from typing import Final, TypeAlias + import litellm +from litellm.litellm_core_utils.core_helpers import as_str_mapping + + +class _Held: + __slots__ = ("value",) + + def __init__(self, value: object) -> None: + self.value = value + + def __hash__(self) -> int: + return id(self.value) + + def __eq__(self, other: object) -> bool: + return isinstance(other, _Held) and other.value is self.value + + +_Fingerprint: TypeAlias = tuple[str, str | None, _Held, int, int, tuple[_Held, ...]] def _resolved_provider(model: str, custom_llm_provider: str | None) -> tuple[str, str]: @@ -11,11 +32,50 @@ def _resolved_provider(model: str, custom_llm_provider: str | None) -> tuple[str return resolved_model, provider +def _fingerprint(model: str, custom_llm_provider: str | None, resolved_model: str, provider: str) -> _Fingerprint: + """Changes whenever an answer could: ``register_model`` bumps the generation, and replacing + any entry the probes look up (``litellm.model_cost[key] = ...``) changes that entry's identity. + The key holds the entries it names, so a cached answer's entries are never freed and their + addresses never reused.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + from litellm.utils import get_model_cost_mutation_generation + + model_cost: Final = as_str_mapping(litellm.model_cost) or {} # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # the model map is an untyped dict of rows + keys: Final = dict.fromkeys( + ( + *AnthropicModelInfo._model_map_lookup_candidates(model), # pyright: ignore[reportPrivateUsage] # the keys the capability walk reads + resolved_model, + f"{provider}/{resolved_model}", + ) + ) + return ( + model, + custom_llm_provider, + _Held(model_cost), + get_model_cost_mutation_generation(), + len(model_cost), + tuple(_Held(model_cost.get(key)) for key in keys), + ) + + def anthropic_model_capabilities(model: str, custom_llm_provider: str | None) -> dict[str, object]: + """The capability probes the Rust messages route reads, once per model-map state: each probe + walks the model map (``get_model_info`` plus generalization matching), and the route asks on + every request.""" + resolved_model, provider = _resolved_provider(model, custom_llm_provider) + cached: Final = _capabilities( + _fingerprint(model, custom_llm_provider, resolved_model, provider), resolved_model, provider + ) + tiers: Final = as_str_mapping(cached["effort_tiers"]) + return {**cached, "effort_tiers": dict(tiers) if tiers is not None else cached["effort_tiers"]} + + +@lru_cache(maxsize=256) +def _capabilities(fingerprint: _Fingerprint, resolved_model: str, provider: str) -> Mapping[str, object]: from litellm.llms.anthropic.chat.transformation import AnthropicConfig from litellm.llms.anthropic.common_utils import AnthropicModelInfo - resolved_model, provider = _resolved_provider(model, custom_llm_provider) + model: Final = fingerprint[0] def supports(flag: str) -> bool: return AnthropicModelInfo._supports_model_capability(model, flag, provider) # pyright: ignore[reportPrivateUsage] # same probes the Python transform runs; forking them would drift diff --git a/tests/unit/rust_bridge/host/test_model_capabilities.py b/tests/unit/rust_bridge/host/test_model_capabilities.py index dffa1b2cde3..501730431a7 100644 --- a/tests/unit/rust_bridge/host/test_model_capabilities.py +++ b/tests/unit/rust_bridge/host/test_model_capabilities.py @@ -86,3 +86,72 @@ def test_capability_source_observes_runtime_registration_changes(monkeypatch: py def test_default_thinking_capability_is_projected(model: str, expected: bool) -> None: capabilities: Final = anthropic_model_capabilities(model, "anthropic") assert capabilities["thinking_on_by_default"] is expected + + +def test_capabilities_are_probed_once_per_model_map_state(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + model: Final = "claude-test-probe-count" + _flag_model(monkeypatch, model, supports_reasoning=True) + probes: list[str] = [] + probe: Final = AnthropicModelInfo._supports_model_capability + + def counted(name: str, key: str, provider: str) -> bool: + probes.append(key) + return probe(name, key, provider) + + monkeypatch.setattr(AnthropicModelInfo, "_supports_model_capability", staticmethod(counted)) + + first: Final = anthropic_model_capabilities(model, "anthropic") + probed: Final = len(probes) + second: Final = anthropic_model_capabilities(model, "anthropic") + + assert probed > 0 + assert len(probes) == probed + assert second == first + + +def test_a_cached_answer_is_a_fresh_copy(monkeypatch: pytest.MonkeyPatch) -> None: + model: Final = "claude-test-fresh-copy" + _flag_model(monkeypatch, model, supports_xhigh_reasoning_effort=True) + first: Final = anthropic_model_capabilities(model, "anthropic") + tiers: Final = first["effort_tiers"] + assert isinstance(tiers, dict) + first["supports_reasoning"] = "changed" + tiers["xhigh"] = "changed" + + second: Final = anthropic_model_capabilities(model, "anthropic") + + assert second["supports_reasoning"] is False + assert second["effort_tiers"] == { + **dict.fromkeys(("minimal", "low", "medium", "high", "max"), False), + "xhigh": True, + } + + +def test_register_model_is_observed_after_a_cached_answer() -> None: + model: Final = "claude-test-registered-later" + try: + litellm.register_model({model: {"litellm_provider": "anthropic", "mode": "chat", "supports_reasoning": False}}) + before: Final = anthropic_model_capabilities(model, "anthropic") + litellm.register_model({model: {"litellm_provider": "anthropic", "mode": "chat", "supports_reasoning": True}}) + after: Final = anthropic_model_capabilities(model, "anthropic") + finally: + litellm.model_cost.pop(model, None) + + assert before["supports_reasoning"] is False + assert after["supports_reasoning"] is True + + +def test_the_cache_key_holds_the_entries_it_names(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.rust_bridge.host.model_capabilities import _fingerprint # pyright: ignore[reportPrivateUsage] + + model: Final = "claude-test-held-entry" + _flag_model(monkeypatch, model, supports_reasoning=True) + entry: Final = litellm.model_cost[model] # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the model map is an untyped dict of rows + held: Final = _fingerprint(model, "anthropic", model, "anthropic") + _flag_model(monkeypatch, model, supports_reasoning=True) + same_content: Final = _fingerprint(model, "anthropic", model, "anthropic") + + assert any(h.value is entry for h in held[5]) + assert held != same_content diff --git a/tests/unit/test_router/core/test_callback_filters_step.py b/tests/unit/test_router/core/test_callback_filters_step.py new file mode 100644 index 00000000000..ee1253f900a --- /dev/null +++ b/tests/unit/test_router/core/test_callback_filters_step.py @@ -0,0 +1,84 @@ +""" +The ``callback_filters`` host step (``RustCallbackChecks._answer_callback_filters``, +litellm/router_rust/steps/callbacks.py): a pick reads the candidates' deployments back only when a +routing callback filters them, and the kept ids come back in order, once each. +""" + +from collections.abc import Mapping +from typing import Final + +import pytest + +from litellm.integrations.custom_logger import CustomLogger +from litellm.router_rust.steps.callbacks import filtered_ids, filters_deployments +from tests.unit.test_router.core.registry_helpers import MakeCore +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment + +GROUP: Final = [deployment("g", "a"), deployment("g", "b"), deployment("g", "c")] + + +class Narrow(CustomLogger): + """Keeps the deployments other than ``dropped``.""" + + def __init__(self, dropped: str) -> None: + super().__init__() + self.dropped = dropped + + async def async_filter_deployments( + self, + model: str, + healthy_deployments: list, + messages: object, + request_kwargs: object = None, + parent_otel_span: object = None, + ) -> list: + return [d for d in healthy_deployments if d["model_info"]["id"] != self.dropped] + + +class Observer(CustomLogger): + """A callback that does not override ``async_filter_deployments``.""" + + +def _entry(deployment_id: object) -> Mapping[str, object]: + return {"model_info": {"id": deployment_id}} + + +def test_filtered_ids_keep_the_filter_order_once_each_among_the_candidates() -> None: + remaining: Final = [_entry("c"), _entry("x"), _entry("a"), _entry("c"), _entry(7), {"model_info": None}, "junk"] + + assert filtered_ids(remaining, ["a", "b", "c"]) == ["c", "a"] + + +def test_only_an_overriding_filter_counts_as_a_filter() -> None: + assert filters_deployments(Narrow(dropped="a")) + assert not filters_deployments(Observer()) + assert not filters_deployments(object()) + + +@pytest.mark.asyncio +async def test_a_pick_without_filters_reads_no_deployment_back( + make_core: MakeCore, host: ScriptedHost, monkeypatch: pytest.MonkeyPatch +) -> None: + host.lists["callbacks"].append(Observer()) + core: Final = make_core(model_list=list(GROUP)) + + def refuse(_self: object, _ids: object) -> list[dict[str, object]]: # mutable-ok: the patched signature + raise AssertionError("no filter runs, so no deployment is read back") + + monkeypatch.setattr(type(core), "_records_of", refuse) + + picked: Final = {(await core.select("g"))["model_info"]["id"] for _ in range(30)} # pyright: ignore[reportIndexIssue] + + assert picked == {"a", "b", "c"} + + +@pytest.mark.asyncio +async def test_a_pick_with_a_filter_reads_the_candidates_and_honours_it( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.lists["callbacks"].append(Narrow(dropped="b")) + core: Final = make_core(model_list=list(GROUP)) + + picked: Final = {(await core.select("g"))["model_info"]["id"] for _ in range(30)} # pyright: ignore[reportIndexIssue] + + assert picked == {"a", "c"}