From 1035971a0536d0d91933df1b053e831b75d5fa3f Mon Sep 17 00:00:00 2001 From: tavren27 <348969968+tavren27@users.noreply.github.com> Date: Tue, 10 Nov 2026 09:26:46 -0800 Subject: [PATCH] fix(decisions): keep box questions on a deployment when a same-model sibling turns them off (#50724) * fix(decisions): keep a same-model deployment's false box-question flag off the shared cost-map row * fix(router): take a withdrawn supports_box_questions true off the shared row on update and delete A deployment's true reached the shared backend row, but a later false no longer overwrites it, so updating or deleting that deployment left box questions on for its unflagged siblings until restart. The row now gets the catalog's value back unless another served deployment of the model still sets true. * fix(router): count every live router and set_model_list when withdrawing a shared box-questions true * fix(router): match box-questions siblings by their cost-map row and withdraw on retire A sibling spelled with its provider registers onto the same litellm.model_cost row, so the withdrawal now compares rows instead of backend keys. Retiring a router takes back the true its deployments put on rows no other live router still sets. Withdrawals are deduped per backend so a large set_model_list stays cheap. * fix(router): take back a box-questions true an aborted set_model_list registered A set_model_list that fails partway rolls back the deployment list, but a true the aborted build already wrote to the shared cost-map row stayed there, so unflagged same-model deployments kept sending box questions. The rollback now withdraws the attempted entries too A withdrawal now only rewrites rows that still hold true, so an entry the failed build never reached gets no row written for it, even on a model the catalog does not list The delete wire test now uses a decisions-mode catalog backend whose catalog lists no box questions, since an unflagged deployment only reads the shared row for a decisions-mode model * fix(router): take back an env-named box-questions true a failed set_model_list registered The rollback took flags back from the raw model_list entries, so a deployment whose model came from os.environ/ was looked up under the reference instead of the model it registered onto, and its true stayed on the shared row It now takes them back from the entries the failed build reached, whose os.environ/ values are already read --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/router_rust/core.py | 1 + litellm/router_rust/surface/registry.py | 146 ++++++++- litellm/router_rust/surface/replay.py | 4 + litellm/types/utils.py | 15 +- litellm/utils.py | 6 + ...test_decisions_box_deployment_flag_wire.py | 67 +++++ .../unit/test_router_model_cost_isolation.py | 284 ++++++++++++++++++ 7 files changed, 511 insertions(+), 12 deletions(-) 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."""