From ce13ffa4691bab3b32e2a5817bf3d89e3d067e7a Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 10 Aug 2026 14:04:56 +0000 Subject: [PATCH] fix(typing): preserve runtime row shapes at jwt mapping and team fetch seams Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- basedpyright-code-budget.json | 14 +++++++------- .../key_management_endpoints.py | 2 ++ litellm/repositories/prisma_protocols.py | 4 ++++ litellm/repositories/table_repositories.py | 15 ++++++++++----- ruff-strict-budget.json | 10 +++++----- type-discipline-budget.json | 10 +++++----- 6 files changed, 33 insertions(+), 22 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 759b8d6d20b..c245bd4624b 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,6 +1,6 @@ { "reportAny": { - "limit": 27487 + "limit": 26879 }, "reportArgumentType": { "limit": 2626 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 8791 + "limit": 8733 }, "reportFunctionMemberAccess": { "limit": 7 @@ -57,7 +57,7 @@ "limit": 5835 }, "reportMissingTypeArgument": { - "limit": 15779 + "limit": 15757 }, "reportMissingTypeStubs": { "limit": 40 @@ -99,19 +99,19 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 45056 + "limit": 45036 }, "reportUnknownLambdaType": { "limit": 113 }, "reportUnknownMemberType": { - "limit": 39757 + "limit": 39721 }, "reportUnknownParameterType": { - "limit": 20206 + "limit": 20204 }, "reportUnknownVariableType": { - "limit": 31257 + "limit": 31199 }, "reportUnnecessaryCast": { "limit": 122 diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 76fff42b7d6..4644f39964e 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -5238,6 +5238,8 @@ async def _fetch_user_team_objects( teams: Final = await _prisma_table(TeamRepository(prisma_client)).find_many( where={"team_id": {"in": complete_user_info.teams}} ) + if not teams: + return [] return [LiteLLM_TeamTable.model_validate(team.model_dump()) for team in teams] diff --git a/litellm/repositories/prisma_protocols.py b/litellm/repositories/prisma_protocols.py index 9778b6214cd..e2b7e456f62 100644 --- a/litellm/repositories/prisma_protocols.py +++ b/litellm/repositories/prisma_protocols.py @@ -103,6 +103,10 @@ class PrismaTableActions(Protocol): async def count(self, *, where: Mapping[str, object] | None = None) -> int: ... +class JWTKeyMappingRecord(Protocol): + token: str + + class ReadOnlyTable(Protocol): async def find_many(self, *, where: Mapping[str, object]) -> Sequence[PrismaRecord]: ... diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index f5b3b3c2318..c0171c004f1 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -10,7 +10,7 @@ methods; richer repositories live in their own modules. from collections.abc import Mapping, Sequence from typing import Any, Final, Generic, TypeVar -from pydantic import BaseModel, TypeAdapter +from pydantic import BaseModel from litellm.models.access_group import LiteLLM_AccessGroupTable from litellm.models.end_user import LiteLLM_EndUserTable @@ -19,17 +19,17 @@ from litellm.models.team_membership import LiteLLM_TeamMembership from litellm.proxy.common_utils.config_sync_pubsub import wrap_table_actions_for_config_sync from litellm.repositories.prisma_protocols import ( DailyToolSpendTable, + JWTKeyMappingRecord, MemoryTable, PrismaTableActions, SpendLogUsageRecord, SpendLogUsageTable, + TableActions, ToolIndexTable, ) ModelT = TypeVar("ModelT", bound=BaseModel) -_TOKEN_ADAPTER: Final = TypeAdapter(str) - class PrismaTableRepository: """Base for repositories that expose a single Prisma table.""" @@ -167,16 +167,21 @@ class InvitationLinkRepository(PrismaTableRepository): class JWTKeyMappingRepository(PrismaTableRepository): table_name = "litellm_jwtkeymapping" + @property + def rows(self) -> TableActions[JWTKeyMappingRecord]: + rows: Final[TableActions[JWTKeyMappingRecord]] = self.table + return rows + async def find_active_token(self, *, jwt_claim_name: str, jwt_claim_value: str) -> str | None: """Return the key token mapped to an active JWT claim pair.""" - record: Final = await self.typed_table.find_first( + record: Final = await self.rows.find_first( where={ "jwt_claim_name": jwt_claim_name, "jwt_claim_value": jwt_claim_value, "is_active": True, } ) - return None if record is None else _TOKEN_ADAPTER.validate_python(record.dict()["token"]) + return None if record is None else record.token class ManagedFileRepository(PrismaTableRepository): diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 911cc794072..224ed6e85f2 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -12,7 +12,7 @@ "limit": 2031 }, "ANN202": { - "limit": 864 + "limit": 862 }, "ANN204": { "limit": 713 @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 1544 + "limit": 1514 }, "ASYNC230": { "limit": 11 @@ -57,7 +57,7 @@ "limit": 3 }, "BLE001": { - "limit": 2924 + "limit": 2922 }, "C401": { "limit": 8 @@ -192,7 +192,7 @@ "limit": 0 }, "S110": { - "limit": 218 + "limit": 216 }, "S112": { "limit": 22 @@ -234,7 +234,7 @@ "limit": 5 }, "TID251": { - "limit": 1235 + "limit": 1231 }, "TRY002": { "limit": 528 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 93bc8b3e836..693a800b1f5 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 23142 + "limit": 23128 }, "LIT002": { - "limit": 27156 + "limit": 27111 }, "LIT003": { "limit": 269 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1086 + "limit": 1084 }, "LIT007": { "limit": 0 @@ -27,9 +27,9 @@ "limit": 0 }, "LIT010": { - "limit": 16738 + "limit": 16692 }, "LIT011": { - "limit": 5598 + "limit": 5594 } }