diff --git a/litellm/proxy/auth/legacy_empty_models_check.py b/litellm/proxy/auth/legacy_empty_models_check.py new file mode 100644 index 00000000000..d1ecd899b46 --- /dev/null +++ b/litellm/proxy/auth/legacy_empty_models_check.py @@ -0,0 +1,198 @@ +import asyncio +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final, Protocol + +from litellm._logging import verbose_proxy_logger +from litellm.constants import UI_SESSION_TOKEN_TEAM_ID +from litellm.proxy._types import LiteLLMRoutes, SpecialModelNames +from litellm.repositories.prisma_protocols import CountTable + +DEPRECATION_DISCUSSION_URL: Final = "https://github.com/BerriAI/litellm/discussions/32090" + + +class CountTableRepository(Protocol): + @property + def table(self) -> CountTable: ... + + +@dataclass(frozen=True, slots=True) +class LegacyEmptyModelsCounts: + teamless_keys_with_empty_models: int + teamless_keys_with_all_team_models: int + team_keys_with_empty_models: int + teams_with_empty_models: int + keys_with_access_groups_and_empty_models: int + teams_with_access_groups_and_empty_models: int + + +def _model_calling_keys(where: Mapping[str, object]) -> Mapping[str, object]: + return { + "AND": [ + where, + { + "NOT": [ + {"allowed_routes": {"equals": [LiteLLMRoutes.management_routes.name]}}, + {"allowed_routes": {"equals": [LiteLLMRoutes.info_routes.name]}}, + ] + }, + ] + } + + +async def count_legacy_empty_models_rows(keys: CountTable, teams: CountTable) -> LegacyEmptyModelsCounts: + ( + teamless_empty, + teamless_all_team, + team_keys_empty, + teams_empty, + keys_with_groups_empty, + teams_with_groups_empty, + ) = await asyncio.gather( + keys.count( + where=_model_calling_keys( + {"team_id": None, "models": {"isEmpty": True}, "access_group_ids": {"isEmpty": True}} + ) + ), + keys.count( + where=_model_calling_keys( + { + "team_id": None, + "models": {"has": SpecialModelNames.all_team_models.value}, + "access_group_ids": {"isEmpty": True}, + } + ) + ), + keys.count( + where=_model_calling_keys( + { + "AND": [{"team_id": {"not": None}}, {"team_id": {"not": UI_SESSION_TOKEN_TEAM_ID}}], + "models": {"isEmpty": True}, + "access_group_ids": {"isEmpty": True}, + } + ) + ), + teams.count(where={"models": {"isEmpty": True}, "access_group_ids": {"isEmpty": True}}), + keys.count( + where=_model_calling_keys( + { + "OR": [{"team_id": None}, {"team_id": {"not": UI_SESSION_TOKEN_TEAM_ID}}], + "models": {"isEmpty": True}, + "NOT": {"access_group_ids": {"isEmpty": True}}, + } + ) + ), + teams.count(where={"models": {"isEmpty": True}, "NOT": {"access_group_ids": {"isEmpty": True}}}), + ) + return LegacyEmptyModelsCounts( + teamless_keys_with_empty_models=teamless_empty, + teamless_keys_with_all_team_models=teamless_all_team, + team_keys_with_empty_models=team_keys_empty, + teams_with_empty_models=teams_empty, + keys_with_access_groups_and_empty_models=keys_with_groups_empty, + teams_with_access_groups_and_empty_models=teams_with_groups_empty, + ) + + +def _counted(count: int, singular: str, plural: str) -> str: + return f"{count} {singular if count == 1 else plural}" + + +def legacy_empty_models_warning(counts: LegacyEmptyModelsCounts) -> str | None: + findings: Final = tuple( + f"{_counted(count, singular, plural)}, which means {meaning}." + for count, singular, plural, meaning in ( + ( + counts.teamless_keys_with_empty_models, + "key without a team has an empty models list", + "keys without a team have an empty models list", + "every proxy model", + ), + ( + counts.teamless_keys_with_all_team_models, + "key without a team lists all-team-models", + "keys without a team list all-team-models", + "every proxy model", + ), + ( + counts.team_keys_with_empty_models, + "team key has an empty models list", + "team keys have an empty models list", + "all of the team's models", + ), + ( + counts.teams_with_empty_models, + "team has an empty models list", + "teams have an empty models list", + "every proxy model", + ), + ) + if count > 0 + ) + access_group_findings: Final = tuple( + f"{_counted(count, singular, plural)}." + for count, singular, plural in ( + ( + counts.keys_with_access_groups_and_empty_models, + "key with access groups has an empty models list", + "keys with access groups have an empty models list", + ), + ( + counts.teams_with_access_groups_and_empty_models, + "team with access groups has an empty models list", + "teams with access groups have an empty models list", + ), + ) + if count > 0 + ) + if not findings and not access_group_findings: + return None + all_proxy_models: Final = SpecialModelNames.all_proxy_models.value + all_team_models: Final = SpecialModelNames.all_team_models.value + keep_access_advice: Final = ( + ( + f"To keep today's access, set models to {all_proxy_models} or a list of models on each one, " + f"or to {all_team_models} on keys that belong to a team, before you upgrade.", + ) + if findings + else () + ) + teamless_all_team_models_advice: Final = ( + ( + f"Keys without a team that list {all_team_models} have no team to inherit from, " + f"so set each one to {all_proxy_models} or attach it to a team.", + ) + if counts.teamless_keys_with_all_team_models > 0 + else () + ) + access_group_advice: Final = ( + ("Check which models each key and team with access groups should reach before you upgrade.",) + if access_group_findings + else () + ) + return " ".join( + ( + "Deprecated model access defaults are still in use.", + *findings, + *access_group_findings, + "These counts include blocked and expired keys and blocked teams.", + "LiteLLM 2.0.0 removes these defaults.", + *keep_access_advice, + *teamless_all_team_models_advice, + *access_group_advice, + "Keys and teams created without a models list still get an empty one until then, " + "so these counts can grow back.", + f"See {DEPRECATION_DISCUSSION_URL}", + ) + ) + + +async def warn_about_legacy_empty_models(keys: CountTableRepository, teams: CountTableRepository) -> None: + try: + counts: Final = await count_legacy_empty_models_rows(keys=keys.table, teams=teams.table) + except Exception as e: # noqa: BLE001 # a startup check must never take the proxy down on a DB error + verbose_proxy_logger.warning("Skipped the deprecated model access defaults check: %s", e) + return + message: Final = legacy_empty_models_warning(counts) + if message is not None: + verbose_proxy_logger.warning(message) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3bc7d5a1b7d..e6447754436 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -455,6 +455,7 @@ from litellm.proxy.auth.auth_utils import ( from litellm.proxy.auth.fallback_budget import RouterFallbackBudgetCheck, router_fallback_budget_check from litellm.proxy.auth.fallback_model_access import router_fallback_access_check from litellm.proxy.auth.handle_jwt import JWTHandler +from litellm.proxy.auth.legacy_empty_models_check import warn_about_legacy_empty_models from litellm.proxy.auth.litellm_license import AUTO_ROUTER_LICENSE_REMEDY, LicenseCheck from litellm.proxy.auth.login_throttle import ( LoginThrottle, @@ -1795,6 +1796,14 @@ async def _proxy_lifespan(app: FastAPI, resources: AsyncExitStack) -> AsyncGener await asyncio.sleep(5) await resources.enter_async_context(database_runtime.background_task(_run_agent_grant_id_migration())) + await resources.enter_async_context( + database_runtime.background_task( + warn_about_legacy_empty_models( + keys=VerificationTokenRepository(prisma_client), + teams=TeamRepository(prisma_client), + ) + ) + ) ## A coordination_redis block saved from the admin UI lives in the database, ## which is only reachable once the prisma client exists. Apply it here, before diff --git a/tests/integration/authorization/test_legacy_empty_models_boot_warning.py b/tests/integration/authorization/test_legacy_empty_models_boot_warning.py new file mode 100644 index 00000000000..a131f978423 --- /dev/null +++ b/tests/integration/authorization/test_legacy_empty_models_boot_warning.py @@ -0,0 +1,92 @@ +"""A proxy booting against keys and teams that still rely on the deprecated empty models defaults logs how many. + +The seeding proxy owns a fresh database, creates keys and teams through the management API in every shape the +deprecation covers plus controls that must not count, and a second proxy booted against that database is the boot +whose log an admin reads before upgrading. +""" + +import uuid +from pathlib import Path +from typing import Final + +import psycopg + +from tests.integration._support.client import JSON_OBJECT, Gateway, eventually, string_value +from tests.integration._support.process import OwnedState, owned_proxy_process + +LISTED_MODEL: Final = "openai/gpt-5.4-mini" +WARNING_HEADER: Final = "Deprecated model access defaults are still in use." + + +def _team(gateway: Gateway, models: list[str], access_group_ids: tuple[str, ...] = ()) -> str: + created: Final = gateway.post( + "/team/new", + { + "team_alias": f"empty-models-{uuid.uuid4().hex}", + "models": models, + "access_group_ids": list(access_group_ids), + }, + ) + return string_value(created["team_id"]) + + +def _access_group(gateway: Gateway) -> str: + created: Final = gateway.request( + "POST", + "/v1/access_group", + {"access_group_name": f"empty-models-{uuid.uuid4().hex}", "access_model_names": [LISTED_MODEL]}, + ) + assert created.status_code == 201, created.text + return string_value(JSON_OBJECT.validate_json(created.content)["access_group_id"]) + + +def _seed(gateway: Gateway, database_url: str) -> None: + empty_team: Final = _team(gateway, []) + listed_team: Final = _team(gateway, [LISTED_MODEL]) + access_group: Final = _access_group(gateway) + _team(gateway, [], access_group_ids=(access_group,)) + keys: Final = ( + {"models": []}, + {"models": []}, + {"models": ["all-team-models"]}, + {"models": [], "team_id": empty_team}, + {"models": [], "team_id": listed_team}, + {"models": [], "team_id": listed_team}, + {"models": ["all-team-models"], "team_id": listed_team}, + {"models": [LISTED_MODEL]}, + {"models": ["all-proxy-models"]}, + {"models": [], "access_group_ids": [access_group]}, + {"models": ["all-team-models"], "access_group_ids": [access_group]}, + {"models": [], "team_id": listed_team, "access_group_ids": [access_group]}, + {"models": [], "key_type": "management"}, + {"models": [], "key_type": "read_only"}, + {"models": [], "team_id": listed_team, "key_type": "management"}, + ) + for body in keys: + gateway.post("/key/generate", body) + with psycopg.connect(database_url, autocommit=True) as connection: + connection.execute( + 'INSERT INTO "LiteLLM_VerificationToken" (token, team_id, models) VALUES (%s, %s, %s)', + (f"ui-session-{uuid.uuid4().hex}", "litellm-dashboard", []), + ) + + +def test_boot_logs_the_keys_and_teams_on_the_deprecated_empty_models_defaults( + gateway: Gateway, tmp_path: Path +) -> None: + with owned_proxy_process(gateway, tmp_path / "seed", {}, state=OwnedState.PRIVATE) as seeding: + database_url: Final = seeding.gateway.database_url + assert database_url is not None + _seed(seeding.gateway, database_url) + with owned_proxy_process( + gateway, tmp_path / "reboot", {"DATABASE_URL": database_url}, database_setup=() + ) as rebooted: + log: Final = eventually(lambda: rebooted.log.read_text(), lambda text: WARNING_HEADER in text, seconds=60) + + assert "2 keys without a team have an empty models list, which means every proxy model." in log + assert "1 key without a team lists all-team-models, which means every proxy model." in log + assert "3 team keys have an empty models list, which means all of the team's models." in log + assert "1 team has an empty models list, which means every proxy model." in log + assert "2 keys with access groups have an empty models list." in log + assert "1 team with access groups has an empty models list." in log + assert "Keys without a team that list all-team-models have no team to inherit from" in log diff --git a/tests/unit/proxy/auth/test_legacy_empty_models_check.py b/tests/unit/proxy/auth/test_legacy_empty_models_check.py new file mode 100644 index 00000000000..437298af18d --- /dev/null +++ b/tests/unit/proxy/auth/test_legacy_empty_models_check.py @@ -0,0 +1,203 @@ +import logging +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from typing import Final + +import pytest + +from litellm._logging import verbose_proxy_logger +from litellm.proxy.auth.legacy_empty_models_check import ( + DEPRECATION_DISCUSSION_URL, + CountTableRepository, + LegacyEmptyModelsCounts, + legacy_empty_models_warning, + warn_about_legacy_empty_models, +) + +NOTHING_ON_THE_DEFAULTS: Final = LegacyEmptyModelsCounts( + teamless_keys_with_empty_models=0, + teamless_keys_with_all_team_models=0, + team_keys_with_empty_models=0, + teams_with_empty_models=0, + keys_with_access_groups_and_empty_models=0, + teams_with_access_groups_and_empty_models=0, +) + + +def test_no_warning_when_nothing_relies_on_the_deprecated_defaults() -> None: + assert legacy_empty_models_warning(NOTHING_ON_THE_DEFAULTS) is None + + +@pytest.mark.parametrize( + ("counts", "reported", "left_out"), + [ + ( + LegacyEmptyModelsCounts(3, 0, 0, 0, 0, 0), + "3 keys without a team have an empty models list, which means every proxy model.", + ("all-team-models, which", "team key", "team has", "teams have", "access groups"), + ), + ( + LegacyEmptyModelsCounts(0, 1, 0, 0, 0, 0), + "1 key without a team lists all-team-models, which means every proxy model.", + ("empty models list", "team key", "access groups"), + ), + ( + LegacyEmptyModelsCounts(0, 0, 4, 0, 0, 0), + "4 team keys have an empty models list, which means all of the team's models.", + ("without a team", "teams have", "access groups"), + ), + ( + LegacyEmptyModelsCounts(0, 0, 0, 1, 0, 0), + "1 team has an empty models list, which means every proxy model.", + ("without a team", "team key", "access groups"), + ), + ( + LegacyEmptyModelsCounts(0, 0, 0, 0, 2, 0), + "2 keys with access groups have an empty models list.", + ( + "which means", + "without a team", + "team key", + "teams with access groups have", + "team with access groups has", + ), + ), + ( + LegacyEmptyModelsCounts(0, 0, 0, 0, 0, 1), + "1 team with access groups has an empty models list.", + ("which means", "without a team", "team key", "keys with access groups have", "key with access groups has"), + ), + ], +) +def test_warning_reports_only_the_categories_in_use( + counts: LegacyEmptyModelsCounts, reported: str, left_out: tuple[str, ...] +) -> None: + message: Final = legacy_empty_models_warning(counts) + + assert message is not None + assert reported in message + assert all(phrase not in message for phrase in left_out), message + assert DEPRECATION_DISCUSSION_URL in message + + +KEEP_ACCESS_ADVICE: Final = "To keep today's access, set models to all-proxy-models" +TEAMLESS_ALL_TEAM_MODELS_ADVICE: Final = "Keys without a team that list all-team-models have no team to inherit from" +ACCESS_GROUP_ADVICE: Final = "Check which models each key and team with access groups should reach" + + +@pytest.mark.parametrize( + ("counts", "advice"), + [ + (LegacyEmptyModelsCounts(3, 0, 0, 0, 0, 0), (KEEP_ACCESS_ADVICE,)), + (LegacyEmptyModelsCounts(0, 0, 4, 0, 0, 0), (KEEP_ACCESS_ADVICE,)), + (LegacyEmptyModelsCounts(0, 0, 0, 1, 0, 0), (KEEP_ACCESS_ADVICE,)), + (LegacyEmptyModelsCounts(0, 1, 0, 0, 0, 0), (KEEP_ACCESS_ADVICE, TEAMLESS_ALL_TEAM_MODELS_ADVICE)), + (LegacyEmptyModelsCounts(0, 0, 0, 0, 2, 0), (ACCESS_GROUP_ADVICE,)), + (LegacyEmptyModelsCounts(0, 0, 0, 0, 0, 1), (ACCESS_GROUP_ADVICE,)), + ], +) +def test_warning_gives_only_the_advice_for_the_categories_in_use( + counts: LegacyEmptyModelsCounts, advice: tuple[str, ...] +) -> None: + message: Final = legacy_empty_models_warning(counts) + + assert message is not None + assert { + phrase + for phrase in (KEEP_ACCESS_ADVICE, TEAMLESS_ALL_TEAM_MODELS_ADVICE, ACCESS_GROUP_ADVICE) + if phrase in message + } == set(advice) + + +def test_warning_reports_every_category_in_use_at_once() -> None: + message: Final = legacy_empty_models_warning(LegacyEmptyModelsCounts(2, 5, 7, 11, 13, 17)) + + assert message is not None + assert "2 keys without a team have an empty models list" in message + assert "5 keys without a team list all-team-models" in message + assert "7 team keys have an empty models list" in message + assert "11 teams have an empty models list" in message + assert "13 keys with access groups have an empty models list" in message + assert "17 teams with access groups have an empty models list" in message + assert all( + phrase in message for phrase in (KEEP_ACCESS_ADVICE, TEAMLESS_ALL_TEAM_MODELS_ADVICE, ACCESS_GROUP_ADVICE) + ) + + +class _RecordingHandler(logging.Handler): + def __init__(self) -> None: + super().__init__(level=logging.WARNING) + self.messages: list[str] = [] + + def emit(self, record: logging.LogRecord) -> None: + self.messages.append(record.getMessage()) + + +@pytest.fixture +def proxy_warnings() -> Iterator[list[str]]: + handler: Final = _RecordingHandler() + verbose_proxy_logger.addHandler(handler) + try: + yield handler.messages + finally: + verbose_proxy_logger.removeHandler(handler) + + +class _FixedCountTable: + def __init__(self, count: int) -> None: + self._count = count + + async def count(self, *, where: Mapping[str, object]) -> int: + return self._count + + +class _UnreachableTable: + async def count(self, *, where: Mapping[str, object]) -> int: + raise ConnectionError("database is down") + + +@dataclass(frozen=True, slots=True) +class _Repository: + table: _FixedCountTable | _UnreachableTable + + +class _UnconnectedRepository: + @property + def table(self) -> _FixedCountTable: + raise RuntimeError("No DB Connected") + + +def _repository(table: _FixedCountTable | _UnreachableTable) -> CountTableRepository: + return _Repository(table=table) + + +@pytest.mark.asyncio +async def test_boot_check_logs_the_warning_when_rows_rely_on_the_defaults(proxy_warnings: list[str]) -> None: + await warn_about_legacy_empty_models(keys=_repository(_FixedCountTable(2)), teams=_repository(_FixedCountTable(1))) + + assert len(proxy_warnings) == 1 + assert proxy_warnings[0].startswith("Deprecated model access defaults are still in use.") + assert "1 team has an empty models list" in proxy_warnings[0] + + +@pytest.mark.asyncio +async def test_boot_check_logs_nothing_when_no_rows_rely_on_the_defaults(proxy_warnings: list[str]) -> None: + await warn_about_legacy_empty_models(keys=_repository(_FixedCountTable(0)), teams=_repository(_FixedCountTable(0))) + + assert proxy_warnings == [] + + +@pytest.mark.asyncio +async def test_a_database_error_skips_the_boot_check_with_a_warning(proxy_warnings: list[str]) -> None: + await warn_about_legacy_empty_models(keys=_repository(_UnreachableTable()), teams=_repository(_UnreachableTable())) + + assert proxy_warnings == ["Skipped the deprecated model access defaults check: database is down"] + + +@pytest.mark.asyncio +async def test_a_repository_without_a_connected_client_skips_the_boot_check_with_a_warning( + proxy_warnings: list[str], +) -> None: + await warn_about_legacy_empty_models(keys=_UnconnectedRepository(), teams=_repository(_FixedCountTable(1))) + + assert proxy_warnings == ["Skipped the deprecated model access defaults check: No DB Connected"]