mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
fix(proxy): order all model-listing fetches deterministically
Greptile flagged that ModelRepository.find_all() (and the paginated search fetch) still ran find_many() without an ORDER BY, so the reshuffle could resurface through those paths. The search fetch is worse: it applies take (LIMIT) with no order, so which rows land on a filtered page is itself non-deterministic and changes on every refresh. Centralize the order spec as MODEL_LIST_ORDER (created_at asc, model_id tie-break) and apply it to find_all and the paginated search fetch as well.
This commit is contained in:
parent
74a57e51ce
commit
ab132cd45d
4 changed files with 128 additions and 16 deletions
|
|
@ -6017,7 +6017,7 @@ class ProxyConfig:
|
|||
"""
|
||||
try:
|
||||
new_models = await ModelRepository(prisma_client).table.find_many(
|
||||
order=[{"created_at": "asc"}, {"model_id": "asc"}]
|
||||
order=MODEL_LIST_ORDER
|
||||
)
|
||||
return new_models
|
||||
except Exception as e:
|
||||
|
|
@ -10946,7 +10946,7 @@ async def run_thread(
|
|||
# async def get_available_routes(user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth)):
|
||||
from litellm.llms.base_llm.base_utils import BaseTokenCounter
|
||||
from litellm.repositories.config_repository import ConfigRepository
|
||||
from litellm.repositories.model_repository import ModelRepository
|
||||
from litellm.repositories.model_repository import MODEL_LIST_ORDER, ModelRepository
|
||||
from litellm.repositories.table_repositories import (
|
||||
AccessGroupRepository,
|
||||
ConfigOverridesRepository,
|
||||
|
|
@ -11786,6 +11786,7 @@ async def _fetch_db_models_for_search(
|
|||
db_models_raw = await ModelRepository(prisma_client).table.find_many(
|
||||
where=db_where_condition,
|
||||
take=take_limit,
|
||||
order=MODEL_LIST_ORDER,
|
||||
)
|
||||
|
||||
# Scope BYOK rows to the caller's allowed teams so non-admin callers
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ Model repository for database operations on LiteLLM_ProxyModelTable.
|
|||
"""
|
||||
|
||||
import json
|
||||
from typing import Any, Dict, List, Optional, Type
|
||||
from typing import Any, Dict, Final, List, Optional, Type
|
||||
|
||||
from litellm.models.model import LiteLLM_ProxyModelTable
|
||||
from litellm.repositories.base_repository import BaseRepository
|
||||
|
|
@ -12,6 +12,11 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
|||
encrypt_value_helper,
|
||||
)
|
||||
|
||||
MODEL_LIST_ORDER: Final[List[Dict[str, str]]] = [
|
||||
{"created_at": "asc"},
|
||||
{"model_id": "asc"},
|
||||
]
|
||||
|
||||
|
||||
class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]):
|
||||
"""Repository for proxy model database operations with encryption support."""
|
||||
|
|
@ -83,7 +88,7 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]):
|
|||
|
||||
async def find_all(self) -> List[LiteLLM_ProxyModelTable]:
|
||||
"""Find all models."""
|
||||
records = await self.table.find_many()
|
||||
records = await self.table.find_many(order=MODEL_LIST_ORDER)
|
||||
return self._to_model_list(records)
|
||||
|
||||
async def find_unblocked(self) -> List[LiteLLM_ProxyModelTable]:
|
||||
|
|
|
|||
|
|
@ -8462,19 +8462,23 @@ class _FakeModelTable:
|
|||
def __init__(self, rows):
|
||||
self._rows = list(rows)
|
||||
|
||||
async def find_many(self, *args, order=None, where=None, **kwargs):
|
||||
async def count(self, where=None):
|
||||
return len(self._rows)
|
||||
|
||||
async def find_many(self, *args, order=None, where=None, take=None, **kwargs):
|
||||
if not order:
|
||||
return list(self._rows)
|
||||
|
||||
def sort_key(row):
|
||||
return tuple(getattr(row, list(o)[0]) for o in order)
|
||||
|
||||
directions = [list(o.values())[0] for o in order]
|
||||
if any(d not in ("asc", "desc") for d in directions):
|
||||
raise ValueError(f"unexpected sort direction in {order}")
|
||||
if any(d == "desc" for d in directions):
|
||||
raise AssertionError("model list is expected ascending, not descending")
|
||||
return sorted(self._rows, key=sort_key)
|
||||
rows = list(self._rows)
|
||||
else:
|
||||
directions = [list(o.values())[0] for o in order]
|
||||
if any(d not in ("asc", "desc") for d in directions):
|
||||
raise ValueError(f"unexpected sort direction in {order}")
|
||||
if any(d == "desc" for d in directions):
|
||||
raise AssertionError("model list is expected ascending, not descending")
|
||||
rows = sorted(
|
||||
self._rows,
|
||||
key=lambda row: tuple(getattr(row, list(o)[0]) for o in order),
|
||||
)
|
||||
return rows[:take] if take is not None else rows
|
||||
|
||||
|
||||
class TestGetModelsFromDbDeterministicOrder:
|
||||
|
|
@ -8507,3 +8511,45 @@ class TestGetModelsFromDbDeterministicOrder:
|
|||
result = await ProxyConfig()._get_models_from_db(prisma_client=prisma_client)
|
||||
|
||||
assert [r.model_id for r in result] == ["a-oldest", "b-oldest", "z-newest"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_search_page_is_stable_under_take_limit(self):
|
||||
"""The paginated search fetch applies `take` (LIMIT); without an ORDER BY
|
||||
which rows land on the page is itself non-deterministic. Ordering must be
|
||||
pushed to the DB so the same slice comes back on every refresh."""
|
||||
from litellm.proxy.proxy_server import _fetch_db_models_for_search
|
||||
|
||||
def _row(model_id, created_at):
|
||||
r = MagicMock()
|
||||
r.model_id = model_id
|
||||
r.created_at = created_at
|
||||
r.model_info = {}
|
||||
return r
|
||||
|
||||
early = datetime(2026, 1, 1, tzinfo=timezone.utc)
|
||||
mid = datetime(2026, 3, 1, tzinfo=timezone.utc)
|
||||
late = datetime(2026, 6, 1, tzinfo=timezone.utc)
|
||||
scrambled = [_row("z", late), _row("m", mid), _row("a", early)]
|
||||
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_proxymodeltable = _FakeModelTable(scrambled)
|
||||
|
||||
proxy_config = MagicMock()
|
||||
proxy_config.decrypt_model_list_from_db = lambda models: [
|
||||
{"model_id": models[0].model_id}
|
||||
]
|
||||
|
||||
result, total = await _fetch_db_models_for_search(
|
||||
prisma_client=prisma_client,
|
||||
proxy_config=proxy_config,
|
||||
search_lower="gpt",
|
||||
db_model_ids_in_router=set(),
|
||||
router_models_count=0,
|
||||
page=1,
|
||||
size=2,
|
||||
sort_by=None,
|
||||
is_byok_outside_caller_teams=lambda info: False,
|
||||
)
|
||||
|
||||
assert total == 3
|
||||
assert [m["model_id"] for m in result] == ["a", "m"]
|
||||
|
|
|
|||
|
|
@ -122,6 +122,28 @@ class MockTable:
|
|||
return MockRecord(self._records[key_value])
|
||||
|
||||
|
||||
class _OrderAwareModelTable:
|
||||
"""Model table mock that honors the `order` argument and returns rows in
|
||||
scrambled insertion order when none is given, so an ordering assertion
|
||||
fails if the caller stops requesting an order."""
|
||||
|
||||
def __init__(self, rows: List[Dict[str, Any]]):
|
||||
self._rows = list(rows)
|
||||
|
||||
async def find_many(
|
||||
self,
|
||||
order: Optional[List[Dict[str, str]]] = None,
|
||||
where: Optional[Dict[str, Any]] = None,
|
||||
skip: Optional[int] = None,
|
||||
take: Optional[int] = None,
|
||||
) -> List[MockRecord]:
|
||||
rows = list(self._rows)
|
||||
for spec in reversed(order or []):
|
||||
((field, direction),) = spec.items()
|
||||
rows.sort(key=lambda r: r[field], reverse=direction == "desc")
|
||||
return [MockRecord(r) for r in rows]
|
||||
|
||||
|
||||
class MockPrismaClient:
|
||||
"""Mock Prisma client for testing."""
|
||||
|
||||
|
|
@ -389,6 +411,44 @@ class TestModelRepository:
|
|||
models = await repo.find_all()
|
||||
assert len(models) == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch(
|
||||
"litellm.repositories.model_repository.decrypt_value_helper",
|
||||
side_effect=lambda v, **kw: v,
|
||||
)
|
||||
async def test_find_all_orders_by_created_at_then_model_id(
|
||||
self, mock_decrypt, repo
|
||||
):
|
||||
early = datetime(2026, 1, 1)
|
||||
late = datetime(2026, 6, 1)
|
||||
scrambled = [
|
||||
{
|
||||
"model_id": "z",
|
||||
"model_name": "n",
|
||||
"litellm_params": "{}",
|
||||
"created_at": late,
|
||||
},
|
||||
{
|
||||
"model_id": "b",
|
||||
"model_name": "n",
|
||||
"litellm_params": "{}",
|
||||
"created_at": early,
|
||||
},
|
||||
{
|
||||
"model_id": "a",
|
||||
"model_name": "n",
|
||||
"litellm_params": "{}",
|
||||
"created_at": early,
|
||||
},
|
||||
]
|
||||
repo._prisma_client.db.litellm_proxymodeltable = _OrderAwareModelTable(
|
||||
scrambled
|
||||
)
|
||||
|
||||
models = await repo.find_all()
|
||||
|
||||
assert [m.model_id for m in models] == ["a", "b", "z"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch(
|
||||
"litellm.repositories.model_repository.decrypt_value_helper",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue