mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
fix(router): serialize heuristic tuning quota filters for Prisma
This commit is contained in:
parent
a2b7868a5b
commit
f40c3842dc
2 changed files with 22 additions and 3 deletions
|
|
@ -4,7 +4,6 @@ Model repository for database operations on LiteLLM_ProxyModelTable.
|
|||
|
||||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol
|
||||
|
||||
from litellm.models.model import LiteLLM_ProxyModelTable
|
||||
|
|
@ -109,7 +108,7 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]):
|
|||
async def find_all_except(self, model_id: str) -> Sequence[LiteLLM_ProxyModelTable]:
|
||||
"""Find every model except the row currently being updated."""
|
||||
records: Final = await self.table.find_many(
|
||||
where=MappingProxyType({"model_id": MappingProxyType({"not": model_id})})
|
||||
where={"model_id": {"not": model_id}} # mutable-ok: Prisma requires plain dicts for query serialization
|
||||
)
|
||||
return tuple(self._to_model_list(records))
|
||||
|
||||
|
|
|
|||
|
|
@ -4,10 +4,13 @@ Tests for gateway repository layer.
|
|||
|
||||
import json
|
||||
from datetime import datetime
|
||||
from typing import Any, Dict, List, Optional
|
||||
from types import SimpleNamespace
|
||||
from typing import Any, Dict, Final, List, Optional
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from prisma import models as prisma_models
|
||||
from prisma.builder import QueryBuilder
|
||||
|
||||
from litellm.models.base import DomainModel
|
||||
from litellm.models.budget import LiteLLM_BudgetTable
|
||||
|
|
@ -307,6 +310,23 @@ class TestModelRepository:
|
|||
client = MockPrismaClient()
|
||||
return ModelRepository(client)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_find_all_except_serializes_exclusion_for_prisma(self) -> None:
|
||||
find_many: Final = AsyncMock(return_value=[])
|
||||
client: Final = SimpleNamespace(
|
||||
db=SimpleNamespace(litellm_proxymodeltable=SimpleNamespace(find_many=find_many))
|
||||
)
|
||||
|
||||
await ModelRepository(client).find_all_except("current-model")
|
||||
|
||||
find_many.assert_awaited_once()
|
||||
query: Final = QueryBuilder(
|
||||
method="find_many",
|
||||
model=prisma_models.LiteLLM_ProxyModelTable,
|
||||
arguments=find_many.call_args.kwargs,
|
||||
).build_query()
|
||||
assert 'where: { model_id: { not: "current-model" } }' in " ".join(query.split())
|
||||
|
||||
def test_table_is_wrapped_for_config_sync(self, repo):
|
||||
from litellm.proxy.common_utils.config_sync_pubsub import (
|
||||
_PublishOnWriteActions,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue