mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-20 00:11:50 +00:00
fix(proxy): rewrite every model allowlist in one statement on rename
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
ryan
parent
de0047c802
commit
1d50d1ad3b
2 changed files with 89 additions and 72 deletions
|
|
@ -8,37 +8,36 @@ alone denies the new name while the old entry grants a name nothing serves any m
|
|||
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.management_helpers.access_group_model_sync import RawExecutor, raw_executor, still_backed
|
||||
from litellm.proxy.management_helpers.access_group_model_sync import raw_executor, still_backed
|
||||
from litellm.router import Router
|
||||
|
||||
|
||||
class _TouchedRow(BaseModel):
|
||||
kind: str
|
||||
object_id: str
|
||||
team_alias: str | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _AllowlistTable:
|
||||
kind: str
|
||||
table: str
|
||||
returning: str
|
||||
id_column: str
|
||||
cache_keys: Callable[[_TouchedRow], tuple[str, ...]]
|
||||
alias_column: str | None = None
|
||||
|
||||
def replace_sql(self) -> str:
|
||||
def update_cte(self, set_clause: str, where_clause: str) -> str:
|
||||
alias: Final = f'"{self.alias_column}"' if self.alias_column else "NULL::text"
|
||||
return (
|
||||
f'UPDATE "{self.table}" SET "models" = array_replace(array_remove("models", $2), $1, $2) '
|
||||
f'WHERE $1 = ANY("models") RETURNING {self.returning}'
|
||||
)
|
||||
|
||||
def append_sql(self) -> str:
|
||||
return (
|
||||
f'UPDATE "{self.table}" SET "models" = array_append("models", $2) '
|
||||
f'WHERE $1 = ANY("models") AND NOT ($2 = ANY("models")) RETURNING {self.returning}'
|
||||
f'{self.kind}_rows AS (UPDATE "{self.table}" SET "models" = {set_clause} WHERE {where_clause} '
|
||||
f"RETURNING '{self.kind}' AS kind, \"{self.id_column}\" AS object_id, {alias} AS team_alias)"
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -63,27 +62,28 @@ def _user_cache_keys(row: _TouchedRow) -> tuple[str, ...]:
|
|||
|
||||
|
||||
_ALLOWLIST_TABLES: Final = (
|
||||
_AllowlistTable("LiteLLM_TeamTable", '"team_id" AS object_id, "team_alias"', _team_cache_keys),
|
||||
_AllowlistTable("LiteLLM_VerificationToken", '"token" AS object_id', _key_cache_keys),
|
||||
_AllowlistTable("LiteLLM_OrganizationTable", '"organization_id" AS object_id', _org_cache_keys),
|
||||
_AllowlistTable("LiteLLM_ProjectTable", '"project_id" AS object_id', _project_cache_keys),
|
||||
_AllowlistTable("LiteLLM_UserTable", '"user_id" AS object_id', _user_cache_keys),
|
||||
_AllowlistTable("team", "LiteLLM_TeamTable", "team_id", _team_cache_keys, alias_column="team_alias"),
|
||||
_AllowlistTable("key", "LiteLLM_VerificationToken", "token", _key_cache_keys),
|
||||
_AllowlistTable("org", "LiteLLM_OrganizationTable", "organization_id", _org_cache_keys),
|
||||
_AllowlistTable("project", "LiteLLM_ProjectTable", "project_id", _project_cache_keys),
|
||||
_AllowlistTable("user", "LiteLLM_UserTable", "user_id", _user_cache_keys),
|
||||
)
|
||||
|
||||
_CACHE_KEYS_BY_KIND: Final = MappingProxyType({table.kind: table.cache_keys for table in _ALLOWLIST_TABLES})
|
||||
|
||||
async def _rewrite_allowlist(
|
||||
executor: RawExecutor,
|
||||
allowlist: _AllowlistTable,
|
||||
sql: str,
|
||||
old_name: str,
|
||||
new_name: str,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
) -> None:
|
||||
touched_rows: Final = await executor.query_raw(sql, old_name, new_name)
|
||||
await evict_and_broadcast(
|
||||
tuple(cache_key for row in touched_rows for cache_key in allowlist.cache_keys(_TouchedRow.model_validate(row))),
|
||||
user_api_key_cache,
|
||||
|
||||
def _rewrite_sql(set_clause: str, where_clause: str) -> str:
|
||||
"""One statement touching every allowlist table, so the rewrite lands everywhere or nowhere."""
|
||||
ctes: Final = ", ".join(table.update_cte(set_clause, where_clause) for table in _ALLOWLIST_TABLES)
|
||||
rows: Final = " UNION ALL ".join(
|
||||
f"SELECT kind, object_id, team_alias FROM {table.kind}_rows" for table in _ALLOWLIST_TABLES
|
||||
)
|
||||
return f"WITH {ctes} {rows}"
|
||||
|
||||
|
||||
_REPLACE_SQL: Final = _rewrite_sql('array_replace(array_remove("models", $2), $1, $2)', '$1 = ANY("models")')
|
||||
|
||||
_APPEND_SQL: Final = _rewrite_sql('array_append("models", $2)', '$1 = ANY("models") AND NOT ($2 = ANY("models"))')
|
||||
|
||||
|
||||
async def sync_model_allowlists_for_renamed_model(
|
||||
|
|
@ -99,12 +99,11 @@ async def sync_model_allowlists_for_renamed_model(
|
|||
return
|
||||
executor: Final = raw_executor(prisma_client)
|
||||
old_name_still_backed: Final = await still_backed(executor, llm_router, old_name, model_id)
|
||||
for allowlist in _ALLOWLIST_TABLES:
|
||||
await _rewrite_allowlist(
|
||||
executor,
|
||||
allowlist,
|
||||
allowlist.append_sql() if old_name_still_backed else allowlist.replace_sql(),
|
||||
old_name,
|
||||
new_name,
|
||||
user_api_key_cache,
|
||||
)
|
||||
touched_rows: Final = await executor.query_raw(
|
||||
_APPEND_SQL if old_name_still_backed else _REPLACE_SQL, old_name, new_name
|
||||
)
|
||||
touched: Final = tuple(_TouchedRow.model_validate(row) for row in touched_rows)
|
||||
await evict_and_broadcast(
|
||||
tuple(cache_key for row in touched for cache_key in _CACHE_KEYS_BY_KIND[row.kind](row)),
|
||||
user_api_key_cache,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -6064,13 +6064,21 @@ class TestAccessGroupModelSync:
|
|||
_MOD = "litellm.proxy.management_endpoints.model_management_endpoints"
|
||||
_INVALIDATE = "litellm.proxy.management_helpers.access_group_model_sync.invalidate_access_group_caches"
|
||||
_EVICT = "litellm.proxy.management_helpers.model_allowlist_rename_sync.evict_and_broadcast"
|
||||
_ALLOWLIST_ROWS = {
|
||||
"LiteLLM_TeamTable": [{"object_id": "team-1", "team_alias": "alias-1"}, {"object_id": "team-2", "team_alias": None}],
|
||||
"LiteLLM_VerificationToken": [{"object_id": "hashed-token-1"}],
|
||||
"LiteLLM_OrganizationTable": [{"object_id": "org-1"}],
|
||||
"LiteLLM_ProjectTable": [{"object_id": "proj-1"}],
|
||||
"LiteLLM_UserTable": [{"object_id": "user-1"}],
|
||||
}
|
||||
_ALLOWLIST_TABLES = (
|
||||
"LiteLLM_TeamTable",
|
||||
"LiteLLM_VerificationToken",
|
||||
"LiteLLM_OrganizationTable",
|
||||
"LiteLLM_ProjectTable",
|
||||
"LiteLLM_UserTable",
|
||||
)
|
||||
_ALLOWLIST_ROWS = [
|
||||
{"kind": "team", "object_id": "team-1", "team_alias": "alias-1"},
|
||||
{"kind": "team", "object_id": "team-2", "team_alias": None},
|
||||
{"kind": "key", "object_id": "hashed-token-1", "team_alias": None},
|
||||
{"kind": "org", "object_id": "org-1", "team_alias": None},
|
||||
{"kind": "project", "object_id": "proj-1", "team_alias": None},
|
||||
{"kind": "user", "object_id": "user-1", "team_alias": None},
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def _admin():
|
||||
|
|
@ -6092,7 +6100,8 @@ class TestAccessGroupModelSync:
|
|||
return [{"deployment_count": deployment_count}]
|
||||
if sql.startswith('UPDATE "LiteLLM_AccessGroupTable"'):
|
||||
return [{"access_group_id": "ag-1"}]
|
||||
return TestAccessGroupModelSync._ALLOWLIST_ROWS[sql.split('"')[1]]
|
||||
assert sql.startswith("WITH ")
|
||||
return TestAccessGroupModelSync._ALLOWLIST_ROWS
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db = MagicMock()
|
||||
|
|
@ -6113,11 +6122,11 @@ class TestAccessGroupModelSync:
|
|||
|
||||
@staticmethod
|
||||
def _allowlist_updates(mock_prisma):
|
||||
return {
|
||||
call.args[0].split('"')[1]: call
|
||||
return [
|
||||
call
|
||||
for call in mock_prisma.db.query_raw.await_args_list
|
||||
if call.args[0].startswith('UPDATE "') and 'SET "models"' in call.args[0]
|
||||
}
|
||||
if call.args[0].startswith("WITH ") and 'SET "models"' in call.args[0]
|
||||
]
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _endpoint_env(self, mock_prisma, router, evict=None):
|
||||
|
|
@ -6130,7 +6139,9 @@ class TestAccessGroupModelSync:
|
|||
patch(f"{self._PS}.proxy_logging_obj", MagicMock()),
|
||||
patch(f"{self._PS}.user_api_key_cache", MagicMock()),
|
||||
patch(self._EVICT, new=evict or AsyncMock()),
|
||||
patch(f"{self._MOD}.ModelManagementAuthChecks.can_user_make_model_call", new=AsyncMock(return_value=None)),
|
||||
patch(
|
||||
f"{self._MOD}.ModelManagementAuthChecks.can_user_make_model_call", new=AsyncMock(return_value=None)
|
||||
),
|
||||
patch(
|
||||
f"{self._MOD}.clear_cache",
|
||||
new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)),
|
||||
|
|
@ -6190,7 +6201,9 @@ class TestAccessGroupModelSync:
|
|||
router.get_model_ids.return_value = ["m-same"]
|
||||
|
||||
with self._endpoint_env(mock_prisma, router) as invalidate:
|
||||
await patch_model(model_id="m-same", patch_data=updateDeployment(blocked=True), user_api_key_dict=self._admin())
|
||||
await patch_model(
|
||||
model_id="m-same", patch_data=updateDeployment(blocked=True), user_api_key_dict=self._admin()
|
||||
)
|
||||
|
||||
mock_prisma.db.query_raw.assert_not_awaited()
|
||||
invalidate.assert_not_awaited()
|
||||
|
|
@ -6278,20 +6291,24 @@ class TestAccessGroupModelSync:
|
|||
user_api_key_dict=self._admin(),
|
||||
)
|
||||
|
||||
updates = self._allowlist_updates(mock_prisma)
|
||||
assert set(updates) == set(self._ALLOWLIST_ROWS)
|
||||
for update_call in updates.values():
|
||||
assert 'SET "models" = array_replace(array_remove("models", $2), $1, $2)' in update_call.args[0]
|
||||
assert 'WHERE $1 = ANY("models")' in update_call.args[0]
|
||||
assert update_call.args[1:] == ("gpt-5.6", "gpt-5.6-eu")
|
||||
evicted = [call.args[0] for call in evict.await_args_list]
|
||||
assert evicted == [
|
||||
("team_id:team-1", "team_alias:alias-1", "team_id:team-2"),
|
||||
("hashed-token-1",),
|
||||
("org_id:org-1", "org_id:org-1:with_budget"),
|
||||
("project_id:proj-1",),
|
||||
("user-1",),
|
||||
]
|
||||
(update_call,) = self._allowlist_updates(mock_prisma)
|
||||
for table in self._ALLOWLIST_TABLES:
|
||||
assert (
|
||||
f'UPDATE "{table}" SET "models" = array_replace(array_remove("models", $2), $1, $2) '
|
||||
'WHERE $1 = ANY("models") RETURNING'
|
||||
) in update_call.args[0]
|
||||
assert update_call.args[1:] == ("gpt-5.6", "gpt-5.6-eu")
|
||||
evict.assert_awaited_once()
|
||||
assert evict.await_args.args[0] == (
|
||||
"team_id:team-1",
|
||||
"team_alias:alias-1",
|
||||
"team_id:team-2",
|
||||
"hashed-token-1",
|
||||
"org_id:org-1",
|
||||
"org_id:org-1:with_budget",
|
||||
"project_id:proj-1",
|
||||
"user-1",
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rename_appends_to_allowlists_when_a_sibling_deployment_keeps_the_old_name(self):
|
||||
|
|
@ -6308,12 +6325,13 @@ class TestAccessGroupModelSync:
|
|||
user_api_key_dict=self._admin(),
|
||||
)
|
||||
|
||||
updates = self._allowlist_updates(mock_prisma)
|
||||
assert set(updates) == set(self._ALLOWLIST_ROWS)
|
||||
for update_call in updates.values():
|
||||
assert 'SET "models" = array_append("models", $2)' in update_call.args[0]
|
||||
assert 'WHERE $1 = ANY("models") AND NOT ($2 = ANY("models"))' in update_call.args[0]
|
||||
assert update_call.args[1:] == ("gpt-5.6", "gpt-5.6-eu")
|
||||
(update_call,) = self._allowlist_updates(mock_prisma)
|
||||
for table in self._ALLOWLIST_TABLES:
|
||||
assert (
|
||||
f'UPDATE "{table}" SET "models" = array_append("models", $2) '
|
||||
'WHERE $1 = ANY("models") AND NOT ($2 = ANY("models")) RETURNING'
|
||||
) in update_call.args[0]
|
||||
assert update_call.args[1:] == ("gpt-5.6", "gpt-5.6-eu")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unchanged_name_never_touches_allowlists(self):
|
||||
|
|
@ -6334,7 +6352,7 @@ class TestAccessGroupModelSync:
|
|||
user_api_key_cache=MagicMock(),
|
||||
)
|
||||
|
||||
assert self._allowlist_updates(mock_prisma) == {}
|
||||
assert self._allowlist_updates(mock_prisma) == []
|
||||
evict.assert_not_awaited()
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue