diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index 8ff0798cafa..5015795657f 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -2009,6 +2009,7 @@ class RustCore( def retire(self) -> None: """``_core_discard``: only the router's global registrations go; it keeps serving calls.""" untrack_live_registry(self._hosted) + self._hosted.retire_shared_capabilities() self._retire_callbacks() if self._hosted.strategies is not None: self._hosted.strategies.retire() diff --git a/litellm/router_rust/surface/registry.py b/litellm/router_rust/surface/registry.py index 0fca3e781d9..e48569c4c9d 100644 --- a/litellm/router_rust/surface/registry.py +++ b/litellm/router_rust/surface/registry.py @@ -42,6 +42,7 @@ router and every ``registry()`` query dispatches the way it does on the Python c import copy import dataclasses +import itertools import json import os import re @@ -102,6 +103,7 @@ from litellm.router_rust.invoke.shaping import ROUTER_CLIENTS_TARGET from litellm.router_rust.surface.indexes import HostedIndexes from litellm.router_rust.surface.model_info import HostModelInfo from litellm.router_rust.surface.model_list import HostedModelList +from litellm.router_rust.surface.replay import live_registries from litellm.router_strategy.auto_router.auto_router_model_naming import classify_strategy_router_model from litellm.router_utils.deployment_numbers import refuse_unservable_numbers, refuse_wide_numbers from litellm.router_utils.endpoint_eligibility import warn_on_unenforced_endpoint_flag @@ -121,6 +123,7 @@ from litellm.types.router import ( ) from litellm.types.utils import ( DEPLOYMENT_SCOPED_TOKEN_LIMIT_FIELDS, + SHARED_ONLY_WHEN_TRUE_CAPABILITY_FIELDS, CustomPricingLiteLLMParams, LlmProviders, LlmProvidersSet, @@ -133,6 +136,7 @@ from litellm.utils import ( get_potential_model_names, get_secret, lookup_model_info_for_registration, + model_cost_registration_key, model_info_with_pricing_overrides, ) @@ -222,6 +226,17 @@ class _Build: # step is answered as invalid and the build ends, then this is raised. failure: Exception | None = None + def reached_entries(self) -> tuple[UntypedEntry, ...]: + return tuple( + {"model_name": entry.model_name, "litellm_params": entry.litellm_params, "model_info": entry.model_info} + for entry in self.entries.values() + ) + + +def _creating(model_list: Sequence[object]) -> _Build: + raw: Final[list[UntypedEntry]] = copy.deepcopy(list(model_list)) # pyright: ignore[reportAssignmentType] # model_list entries # mutable-ok: a deep copy the build consumes + return _Build(path="create", raw=raw) + def model_info_is_active_for_environment(model_info: Mapping[str, object] | None) -> bool: """The environment-gating rule (``model_info_is_active_for_environment``): a deployment whose @@ -551,6 +566,95 @@ def register_deployment_pricing(host: "PythonHost", deployment: Deployment) -> N ) +def _entry_section(entry: Mapping[str, object], section: str) -> Mapping[str, object]: + value: Final = entry.get(section) + return _MODEL_INFO_ADAPTER.validate_python(value) if isinstance(value, Mapping) else {} + + +def _entry_backend(entry: Mapping[str, object]) -> tuple[str, str | None] | None: + params: Final = _entry_section(entry, "litellm_params") + model: Final = params.get("model") + provider: Final = params.get("custom_llm_provider") + return (model, provider if isinstance(provider, str) else None) if isinstance(model, str) else None + + +def _entry_backend_keys(entry: Mapping[str, object]) -> frozenset[str]: + backend: Final = _entry_backend(entry) + return frozenset() if backend is None else frozenset(_backend_cost_map_keys(*backend)) + + +def _entry_row_keys(entry: Mapping[str, object]) -> frozenset[str]: + return frozenset(model_cost_registration_key(key) for key in _entry_backend_keys(entry)) + + +def _shared_capabilities_set_true(entry: Mapping[str, object]) -> frozenset[str]: + model_info: Final = _entry_section(entry, "model_info") + return frozenset(field for field in SHARED_ONLY_WHEN_TRUE_CAPABILITY_FIELDS if model_info.get(field) is True) + + +def _distinct_withdrawals(entries: Iterable[Mapping[str, object]]) -> tuple[Mapping[str, object], ...]: + return tuple( + {(_entry_backend_keys(entry), _shared_capabilities_set_true(entry)): entry for entry in entries}.values() + ) + + +def _row_sets_true(host: "PythonHost", key: str, fields: frozenset[str]) -> bool: + row: Final[object] = host.dynamic().model_cost.get(model_cost_registration_key(key)) + return isinstance(row, dict) and any( + _MODEL_INFO_ADAPTER.validate_python(row).get(field) is True for field in fields + ) + + +def _catalog_sets_true(catalog: Mapping[str, Mapping[str, object]], catalog_keys: tuple[str, ...], field: str) -> bool: + entry: Final = next((catalog[key] for key in catalog_keys if key in catalog), None) + return entry is not None and entry.get(field) is True + + +def withdraw_shared_capabilities( + host: "PythonHost", + removed: Mapping[str, object], + replacement: Mapping[str, object] | None, + served: Iterable[Mapping[str, object]], +) -> None: + """For each ``SHARED_ONLY_WHEN_TRUE_CAPABILITY_FIELDS`` flag ``removed`` set ``True``, set the shared + ``litellm.model_cost`` row(s) ``removed`` registers onto that still hold it ``True`` to whether ``replacement``, + a ``served`` deployment registering onto one of them, or the catalog still sets it ``True``.""" + backend: Final = _entry_backend(removed) + set_true: Final = _shared_capabilities_set_true(removed) + if backend is None or not set_true: + return + row_keys: Final = _entry_row_keys(removed) + kept: Final = ( + _shared_capabilities_set_true(replacement) + if replacement is not None and _entry_row_keys(replacement) == row_keys + else frozenset[str]() + ) + withdrawn: Final = set_true - kept + held: Final = tuple(key for key in _entry_backend_keys(removed) if _row_sets_true(host, key, withdrawn)) + if not held: + return + candidates: Final = tuple( + entry for entry in served if not withdrawn.isdisjoint(_shared_capabilities_set_true(entry)) + ) + rows_by_backend: Final = { + keys: frozenset(model_cost_registration_key(key) for key in keys) + for keys in frozenset(_entry_backend_keys(entry) for entry in candidates) + } + siblings_set_true: Final = frozenset( + itertools.chain.from_iterable( + _shared_capabilities_set_true(entry) & withdrawn + for entry in candidates + if not rows_by_backend[_entry_backend_keys(entry)].isdisjoint(row_keys) + ) + ) + catalog: Final = catalog_with_runtime_registrations() + catalog_keys: Final = _catalog_backend_keys(*backend, _backend_cost_map_keys(*backend)) + restored: Final = { + field: field in siblings_set_true or _catalog_sets_true(catalog, catalog_keys, field) for field in withdrawn + } + host.invoke("register_model")(model_cost={key: restored for key in held}, persist_across_reloads=False) + + # -- the registry -------------------------------------------------------------------------------- @@ -850,15 +954,24 @@ class HostRegistry(HostModelInfo): """Replace every deployment (``set_model_list``). All or nothing: when an entry raises, the deployments served before the call are put back and the error is raised.""" previous: Final = [copy.deepcopy(dict(entry)) for entry in self._core_model_list] + attempt: Final = _creating(model_list) try: - self._replace_model_list(model_list) + self._replace_model_list(attempt) except Exception: - self._replace_model_list(previous) + self._replace_model_list(_creating(previous)) + self._withdraw_shared_capabilities(attempt.reached_entries(), self._served_by_live_routers()) raise + self._withdraw_shared_capabilities(previous, self._served_by_live_routers()) - def _replace_model_list(self, model_list: Sequence[object]) -> None: + def _withdraw_shared_capabilities( + self, entries: Iterable[Mapping[str, object]], served: Sequence[Mapping[str, object]] + ) -> None: + for entry in _distinct_withdrawals(entries): + withdraw_shared_capabilities(self._host, entry, None, served) + + def _replace_model_list(self, build: _Build) -> None: self._discovered_model_info_cache.flush_cache() - original_model_list: Final[list[UntypedEntry]] = copy.deepcopy(list(model_list)) # pyright: ignore[reportAssignmentType] # model_list entries # mutable-ok: a deep copy the build consumes + original_model_list: Final = build.raw if self.strategies is not None: self.strategies.reset() self._provider_unresolved_deployments = () @@ -879,12 +992,7 @@ class HostRegistry(HostModelInfo): # router-core walks the list (one deployment per organization, secrets resolved first, # ids assigned), asking this registry about each deployment in turn; it starts the list # and its id, name and team indexes over, and takes the given names as `model_names`. - self._run_build( - _Build(path="create", raw=original_model_list), - original_model_list, - resolve_secrets=True, - install="replace", - ) + self._run_build(build, original_model_list, resolve_secrets=True, install="replace") if self.strategies is not None: self.strategies.finalize() self._on_change() @@ -1312,8 +1420,11 @@ class HostRegistry(HostModelInfo): with self._deployment_index_lock: removal_idx: Final = self.model_id_to_deployment_index_map.get(model_id) if removal_idx is not None: - _ = self._core_model_list[removal_idx] # pyright: ignore[reportArgumentType] # raises IndexError for a position past the list, as the pop does + removed: Final = self._core_model_list[removal_idx] # pyright: ignore[reportArgumentType] # raises IndexError for a position past the list, as the pop does self._remove(model_id) + withdraw_shared_capabilities( + self._host, removed, deployment.to_json(exclude_none=True), self._served_by_live_routers() + ) if self.strategies is not None: self.strategies.release(on_router) self._on_change() @@ -1490,6 +1601,18 @@ class HostRegistry(HostModelInfo): self._registry_version += 1 self._on_change() + def _served_by_other_routers(self) -> tuple[UntypedEntry, ...]: + others: Final = (tuple(registry._core_model_list) for registry in live_registries() if registry is not self) + return tuple(itertools.chain.from_iterable(others)) + + def _served_by_live_routers(self) -> tuple[UntypedEntry, ...]: + return (*self._core_model_list, *self._served_by_other_routers()) + + def retire_shared_capabilities(self) -> None: + """Take back the ``SHARED_ONLY_WHEN_TRUE_CAPABILITY_FIELDS`` flags this router's deployments put on shared + ``litellm.model_cost`` rows that no other live router's deployment still sets (``retire``).""" + self._withdraw_shared_capabilities(self._core_model_list, self._served_by_other_routers()) + def _core_delete_deployment(self, id: str) -> UntypedEntry | None: """Remove one deployment and return its stored entry (``_delete_deployment``).""" try: @@ -1499,6 +1622,7 @@ class HostRegistry(HostModelInfo): return None item: Final = self._core_model_list[deployment_idx] # pyright: ignore[reportArgumentType] # raises IndexError for a position past the list, as the pop does self._remove(id) + withdraw_shared_capabilities(self._host, item, None, self._served_by_live_routers()) if self.checks is not None: self.checks.unregister_deployment_budget(id) if self.strategies is not None: diff --git a/litellm/router_rust/surface/replay.py b/litellm/router_rust/surface/replay.py index 6c770e07883..cd78b1ec76e 100644 --- a/litellm/router_rust/surface/replay.py +++ b/litellm/router_rust/surface/replay.py @@ -22,6 +22,10 @@ def untrack_live_registry(registry: HostRegistry) -> None: _live_registries.discard(registry) +def live_registries() -> tuple[HostRegistry, ...]: + return tuple(_live_registries) + + def _replay_live_registries() -> None: for registry in tuple(_live_registries): registry.replay_model_cost_registrations() diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 3ddeb118c33..ffc9e915ea4 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -4344,6 +4344,14 @@ SHARED_BACKEND_MODEL_INFO_FIELDS: Final[frozenset[str]] = ( - DEPLOYMENT_SCOPED_TOKEN_COUNT_FIELDS ) +SHARED_ONLY_WHEN_TRUE_CAPABILITY_FIELDS: Final[frozenset[str]] = frozenset({"supports_box_questions"}) + + +def _belongs_on_shared_backend_key(field: str, value: object) -> bool: + if field in SHARED_ONLY_WHEN_TRUE_CAPABILITY_FIELDS: + return value is True + return field in SHARED_BACKEND_MODEL_INFO_FIELDS + def shared_backend_model_info(model_info: dict[str, Any]) -> dict[str, Any]: """Return only the fields safe to register under a shared ``{provider}/{model}`` @@ -4352,8 +4360,13 @@ def shared_backend_model_info(model_info: dict[str, Any]) -> dict[str, Any]: ``off_peak_pricing``. Per-deployment metadata (``id``, ``access_via_team_ids``, arbitrary custom keys) never belongs on the shared key; it stays under the deployment's unique model id. + + A ``SHARED_ONLY_WHEN_TRUE_CAPABILITY_FIELDS`` flag reaches the shared key only as + ``True``. Its reader checks the deployment's own flag first, so a ``False`` or + non-boolean value already applies to its deployment, and on the shared key it would + only take the capability away from same-model deployments that set no flag. """ - return {k: v for k, v in model_info.items() if k in SHARED_BACKEND_MODEL_INFO_FIELDS} + return {k: v for k, v in model_info.items() if _belongs_on_shared_backend_key(k, v)} ABOVE_THRESHOLD_COST_KEY_PATTERN: Final = re.compile(r"_above_\d+k?_tokens(?:_batches)?$") diff --git a/litellm/utils.py b/litellm/utils.py index 7f7c6ffbc4d..ca916e049d1 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -4147,6 +4147,12 @@ def model_cost_registration_row(key: str, provider: object) -> Mapping[str, obje return _builtin_cost_map().get(base.key) or base.base +def model_cost_registration_key(key: str) -> str: + """The ``litellm.model_cost`` row a ``register_model`` write of ``key`` lands on (``typesafe/typesafe/jev-1.13.0`` + lands on ``typesafe/jev-1.13.0``).""" + return registration_base(key, {}, custom_llm_provider=None).key + + def registration_base(key: str, value: Mapping[str, object], custom_llm_provider: str | None) -> Landing: provider: Final = str(value.get("litellm_provider") or "") if _skips_builtin_registration_lookup(key, provider): diff --git a/tests/integration/providers/test_decisions_box_deployment_flag_wire.py b/tests/integration/providers/test_decisions_box_deployment_flag_wire.py index 70be286569a..6dc2d3c2809 100644 --- a/tests/integration/providers/test_decisions_box_deployment_flag_wire.py +++ b/tests/integration/providers/test_decisions_box_deployment_flag_wire.py @@ -354,6 +354,73 @@ async def test_a_router_deployment_flagged_false_refuses_boxes_even_when_its_bac ] * 2 +def _catalog_member(box_questions: bool) -> _Member: + backend: Final = next( + key + for key in tuple(litellm.model_cost) + if key.startswith("typesafe/") and litellm.supports_box_questions(model=key) is box_questions + ) + return _Member( + "typesafe", + backend, + "/v1/systemone", + backend.removeprefix("typesafe/"), + headers=MappingProxyType({_IMAGES_HEADER: "true"}), + ) + + +@pytest.mark.parametrize("sibling_flag", (False, "true"), ids=("false", "string")) +def test_a_later_same_model_deployment_flag_never_takes_catalog_box_questions_from_an_unflagged_one( + gateway: Gateway, sibling_flag: JsonValue +) -> None: + member: Final = _catalog_member(box_questions=True) + with wire_server(member.reply) as wire, gateway.scenario() as scenario: + unflagged: Final = member.deployment(scenario, wire.url, None) + member.deployment(scenario, wire.url, {"supports_box_questions": sibling_flag}) + answered: Final = _decide(gateway, unflagged) + assert answered.status_code == 200, answered.text + assert answered.json()["answers"] == _expected_answers(_QUESTIONS), answered.text + (call,) = wire.drain() + assert _JSON_OBJECT.validate_json(call.body) == member.upstream_body(_QUESTIONS) + + +def test_an_unflagged_deployment_stops_sending_box_questions_once_its_last_true_sibling_is_deleted( + gateway: Gateway, +) -> None: + member: Final = _catalog_member(box_questions=False) + with wire_server(member.reply) as wire, gateway.scenario() as scenario: + unflagged: Final = member.deployment(scenario, wire.url, None) + with gateway.scenario() as sibling: + member.deployment(sibling, wire.url, {"supports_box_questions": True}) + borrowed: Final = _decide(gateway, unflagged, (_BOX_QUESTION,)) + assert borrowed.status_code == 200, borrowed.text + assert borrowed.json()["answers"] == [_BOX_ANSWER], borrowed.text + refused: Final = eventually( + lambda: tuple(_decide(gateway, unflagged, (_BOX_QUESTION,)) for _ in range(8)), + lambda responses: all(response.status_code == 400 for response in responses), + seconds=3 * PROXY_CONFIG_RELOAD_INTERVAL_SECONDS, + ) + for response in refused: + _assert_refused_box(response) + received: Final = wire.drain() + assert received, "the borrowed box question never reached the upstream" + assert all(_JSON_OBJECT.validate_json(call.body) == member.upstream_body((_BOX_QUESTION,)) for call in received) + + +@pytest.mark.parametrize("value", ("true", 1), ids=("string", "int")) +def test_a_non_boolean_deployment_flag_keeps_the_catalog_box_questions_of_its_backend( + gateway: Gateway, value: JsonValue +) -> None: + member: Final = _catalog_member(box_questions=True) + with wire_server(member.reply) as wire, gateway.scenario() as scenario: + model: Final = member.deployment(scenario, wire.url, {"supports_box_questions": value}) + answered: Final = _decide(gateway, model) + assert answered.status_code == 200, answered.text + assert answered.json()["answers"] == _expected_answers(_QUESTIONS), answered.text + (call,) = wire.drain() + assert _JSON_OBJECT.validate_json(call.body) == member.upstream_body(_QUESTIONS) + + @pytest.mark.parametrize( "value", ("true", 1, [True], "", "x" * 5000, None), ids=("string", "int", "list", "empty", "5kb", "null") ) diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index f5c71668e63..8d15ba02a38 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -8,6 +8,7 @@ should still use the built-in pricing. """ import asyncio +import contextlib import copy import itertools import json @@ -18,6 +19,7 @@ import queue import re import sys import threading +from collections.abc import Iterator from pathlib import Path from types import MappingProxyType from typing import Final, cast @@ -32,8 +34,10 @@ from litellm import Router from litellm.constants import DEFAULT_MAX_LRU_CACHE_SIZE from litellm.litellm_core_utils.ptu_pricing import ptu_config_error from litellm.litellm_core_utils.llm_cost_calc.utils import SERVICE_TIER_COST_KEY_SUFFIXES +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy.utils import create_model_info_response +from litellm.router_api import ChangeSet from litellm.types.router import ( Deployment, DeploymentTypedDict, @@ -634,6 +638,286 @@ def test_pdf_opt_out_on_one_deployment_does_not_block_its_sibling(respx_mock: re _restore_model_cost_entries(model_keys) +@pytest.mark.parametrize( + ("catalog_flag", "sibling_flag", "shared_flag"), + ((True, False, True), (True, "true", True), (None, True, True)), + ids=("false-keeps-catalog-true", "string-keeps-catalog-true", "true-reaches-shared"), +) +def test_a_sibling_box_questions_flag_reaches_the_shared_key_only_as_true( + catalog_flag: bool | None, sibling_flag: bool | str, shared_flag: bool +) -> None: + backend_model: Final = "perplexity/box-isolation-check" + catalog_row: Final = { + "litellm_provider": "perplexity", + "mode": "evaluation", + **({} if catalog_flag is None else {"supports_box_questions": catalog_flag}), + } + model_keys: Final = { + key: copy.deepcopy(litellm.model_cost.get(key)) for key in ("unflagged-id", "sibling-id", backend_model) + } + litellm.model_cost[backend_model] = catalog_row + publish(litellm.model_cost) + try: + Router( + model_list=[ + { + "model_name": "box-isolation", + "litellm_params": {"model": backend_model}, + "model_info": {"id": "unflagged-id"}, + }, + { + "model_name": "box-isolation", + "litellm_params": {"model": backend_model}, + "model_info": {"id": "sibling-id", "supports_box_questions": sibling_flag}, + }, + ], + ) + + assert litellm.supports_box_questions(model=backend_model) is shared_flag + assert litellm.model_cost["sibling-id"]["supports_box_questions"] == sibling_flag + finally: + _restore_model_cost_entries(model_keys) + + +def _box_flag_deployment( + backend_model: str, model_id: str, flag: bool | None, custom_llm_provider: str | None = None +) -> DeploymentTypedDict: + return { + "model_name": "box-withdraw", + "litellm_params": { + "model": backend_model, + **({} if custom_llm_provider is None else {"custom_llm_provider": custom_llm_provider}), + }, + "model_info": {"id": model_id, **({} if flag is None else {"supports_box_questions": flag})}, + } + + +@contextlib.contextmanager +def _live_box_router(backend_model: str, deployments: tuple[tuple[str, bool | None], ...]) -> Iterator[Router]: + router: Final = Router( + model_list=[_box_flag_deployment(backend_model, model_id, flag) for model_id, flag in deployments] + ) + try: + yield router + finally: + router.retire() + + +@pytest.mark.parametrize( + ("catalog_box", "true_sibling_stays", "change", "shared_flag"), + ( + (False, False, "update", False), + (False, False, "delete", False), + (False, False, "replace", False), + (False, True, "update", True), + (True, False, "update", True), + ), + ids=( + "update-to-false-clears", + "delete-clears", + "replace-clears", + "another-true-sibling-keeps", + "catalog-true-keeps", + ), +) +def test_a_withdrawn_box_questions_true_leaves_the_shared_key_once_no_sibling_sets_it( + catalog_box: bool, true_sibling_stays: bool, change: str, shared_flag: bool +) -> None: + backend_model: Final = ( + next( + key + for key, entry in GetModelCostMap.loaded_model_cost_map().items() + if entry.get("supports_box_questions") is True + ) + if catalog_box + else "perplexity/box-withdraw-check" + ) + deployments: Final = ( + ("unflagged-id", None), + ("sibling-id", True), + *((("other-true-id", True),) if true_sibling_stays else ()), + ) + model_keys: Final = { + key: copy.deepcopy(litellm.model_cost.get(key)) + for key in (*(model_id for model_id, _ in deployments), backend_model) + } + try: + with _live_box_router(backend_model, deployments) as router: + assert litellm.supports_box_questions(model=backend_model) is True + if change == "replace": + router.set_model_list( + [ + _box_flag_deployment(backend_model, model_id, flag) + for model_id, flag in deployments + if model_id != "sibling-id" + ] + ) + else: + _ = router.apply( + ChangeSet(upsert=[Deployment(**_box_flag_deployment(backend_model, "sibling-id", False))]) + if change == "update" + else ChangeSet(delete=["sibling-id"]) + ) + + assert litellm.supports_box_questions(model=backend_model) is shared_flag + finally: + _restore_model_cost_entries(model_keys) + + +def test_a_failed_set_model_list_takes_back_the_box_questions_true_its_aborted_build_registered() -> None: + backend_model: Final = "perplexity/box-withdraw-on-rollback" + model_keys: Final = { + key: copy.deepcopy(litellm.model_cost.get(key)) + for key in ("unflagged-id", "false-id", "aborted-true-id", "unservable-id", backend_model) + } + unservable: Final = { + "model_name": "box-withdraw", + "litellm_params": {"model": backend_model, "input_cost_per_token": float("inf")}, + "model_info": {"id": "unservable-id"}, + } + try: + with _live_box_router(backend_model, (("unflagged-id", None), ("false-id", False))) as router: + with pytest.raises(ValueError, match="cannot be served"): + router.set_model_list( + [ + _box_flag_deployment(backend_model, "unflagged-id", None), + _box_flag_deployment(backend_model, "aborted-true-id", True), + unservable, + ] + ) + + assert sorted( + deployment["model_info"]["id"] for deployment in router.registry().deployments("box-withdraw") + ) == ["false-id", "unflagged-id"] + assert litellm.supports_box_questions(model=backend_model) is False + finally: + _restore_model_cost_entries(model_keys) + + +def test_a_failed_set_model_list_writes_no_row_for_a_box_questions_true_it_never_built() -> None: + backend_model: Final = "perplexity/box-withdraw-on-rollback-served" + never_built_model: Final = "perplexity/box-withdraw-never-built" + model_keys: Final = { + key: copy.deepcopy(litellm.model_cost.get(key)) + for key in ("unflagged-id", "unservable-id", "never-built-id", backend_model, never_built_model) + } + unservable: Final = { + "model_name": "box-withdraw", + "litellm_params": {"model": backend_model, "input_cost_per_token": float("inf")}, + "model_info": {"id": "unservable-id"}, + } + try: + with _live_box_router(backend_model, (("unflagged-id", None),)) as router: + with pytest.raises(ValueError, match="cannot be served"): + router.set_model_list([unservable, _box_flag_deployment(never_built_model, "never-built-id", True)]) + + assert never_built_model not in litellm.model_cost + finally: + _restore_model_cost_entries(model_keys) + + +def test_a_failed_set_model_list_takes_back_a_box_questions_true_whose_model_comes_from_the_environment( + monkeypatch: pytest.MonkeyPatch, +) -> None: + backend_model: Final = "perplexity/box-withdraw-on-rollback-env" + monkeypatch.setenv("BOX_WITHDRAW_ROLLBACK_MODEL", backend_model) + model_keys: Final = { + key: copy.deepcopy(litellm.model_cost.get(key)) + for key in ("unflagged-id", "env-true-id", "unservable-id", backend_model) + } + unservable: Final = { + "model_name": "box-withdraw", + "litellm_params": {"model": backend_model, "input_cost_per_token": float("inf")}, + "model_info": {"id": "unservable-id"}, + } + try: + with _live_box_router(backend_model, (("unflagged-id", None),)) as router: + with pytest.raises(ValueError, match="cannot be served"): + router.set_model_list( + [ + _box_flag_deployment(backend_model, "unflagged-id", None), + _box_flag_deployment("os.environ/BOX_WITHDRAW_ROLLBACK_MODEL", "env-true-id", True), + unservable, + ] + ) + + assert litellm.supports_box_questions(model=backend_model) is False + finally: + _restore_model_cost_entries(model_keys) + + +def test_a_box_questions_true_another_live_router_serves_stays_on_the_shared_key() -> None: + backend_model: Final = "perplexity/box-withdraw-across-routers" + model_keys: Final = { + key: copy.deepcopy(litellm.model_cost.get(key)) + for key in ("unflagged-id", "sibling-id", "other-router-id", backend_model) + } + try: + with ( + _live_box_router(backend_model, (("unflagged-id", None), ("sibling-id", True))) as router, + _live_box_router(backend_model, (("other-router-id", True),)) as other_router, + ): + _ = router.apply(ChangeSet(delete=["sibling-id"])) + assert litellm.supports_box_questions(model=backend_model) is True + + _ = other_router.apply(ChangeSet(delete=["other-router-id"])) + assert litellm.supports_box_questions(model=backend_model) is False + finally: + _restore_model_cost_entries(model_keys) + + +@pytest.mark.parametrize("deleted_first", ("ui-true-id", "config-true-id")) +def test_a_box_questions_true_stays_while_a_sibling_spelled_with_its_provider_still_sets_it(deleted_first: str) -> None: + backend_model: Final = next( + key + for key, entry in GetModelCostMap.loaded_model_cost_map().items() + if entry.get("litellm_provider") == "typesafe" and entry.get("supports_box_questions") is not True + ) + model_keys: Final = { + key: copy.deepcopy(litellm.model_cost.get(key)) + for key in ("unflagged-id", "config-true-id", "ui-true-id", backend_model) + } + try: + with _live_box_router(backend_model, (("unflagged-id", None), ("config-true-id", True))) as router: + _ = router.apply( + ChangeSet( + upsert=[ + Deployment( + **_box_flag_deployment(backend_model, "ui-true-id", True, custom_llm_provider="typesafe") + ) + ] + ) + ) + + _ = router.apply(ChangeSet(delete=[deleted_first])) + assert litellm.supports_box_questions(model=backend_model) is True + + _ = router.apply(ChangeSet(delete=[({"ui-true-id", "config-true-id"} - {deleted_first}).pop()])) + assert litellm.supports_box_questions(model=backend_model) is False + finally: + _restore_model_cost_entries(model_keys) + + +def test_a_retired_router_takes_back_a_box_questions_true_no_live_router_still_sets() -> None: + backend_model: Final = "perplexity/box-withdraw-on-retire" + model_keys: Final = { + key: copy.deepcopy(litellm.model_cost.get(key)) + for key in ("unflagged-id", "first-true-id", "second-true-id", backend_model) + } + try: + with _live_box_router(backend_model, (("unflagged-id", None),)): + first: Final = Router(model_list=[_box_flag_deployment(backend_model, "first-true-id", True)]) + second: Final = Router(model_list=[_box_flag_deployment(backend_model, "second-true-id", True)]) + + first.retire() + assert litellm.supports_box_questions(model=backend_model) is True + + second.retire() + assert litellm.supports_box_questions(model=backend_model) is False + finally: + _restore_model_cost_entries(model_keys) + + def test_token_count_settings_stay_on_the_deployment_that_set_them() -> None: """A deployment's ``token_count_model`` and framing settings must not land on the shared ``{provider}/{model}`` key, or a sibling deployment of that model would count with them."""