Merge pull request #38205 from BerriAI/litellm_decrease_anys_opus5_round2

refactor(repositories): type prisma table access with one generic protocol

Author: mateo-berri

basedpyright-code-budget.json

@@ -1,18 +1,18 @@
 {
   "reportAny": {
-    "limit": 19949
+    "limit": 18505
   },
   "reportArgumentType": {
-    "limit": 2566
+    "limit": 2564
   },
   "reportAssignmentType": {
     "limit": 320
   },
   "reportAttributeAccessIssue": {
-    "limit": 488
+    "limit": 483
   },
   "reportCallIssue": {
-    "limit": 114
+    "limit": 113
   },
   "reportConstantRedefinition": {
     "limit": 40
@@ -24,7 +24,7 @@
     "limit": 19
   },
   "reportExplicitAny": {
-    "limit": 6049
+    "limit": 5976
   },
   "reportFunctionMemberAccess": {
     "limit": 7
@@ -45,7 +45,7 @@
     "limit": 35
   },
   "reportInvalidTypeForm": {
-    "limit": 35
+    "limit": 34
   },
   "reportInvalidTypeVarUse": {
     "limit": 2
@@ -54,10 +54,10 @@
     "limit": 0
   },
   "reportMissingParameterType": {
-    "limit": 5661
+    "limit": 5659
   },
   "reportMissingTypeArgument": {
-    "limit": 15555
+    "limit": 15504
   },
   "reportMissingTypeStubs": {
     "limit": 40
@@ -72,7 +72,7 @@
     "limit": 0
   },
   "reportOptionalMemberAccess": {
-    "limit": 1061
+    "limit": 1058
   },
   "reportOptionalOperand": {
     "limit": 0
@@ -99,31 +99,31 @@
     "limit": 0
   },
   "reportUnknownArgumentType": {
-    "limit": 44655
+    "limit": 44530
   },
   "reportUnknownLambdaType": {
     "limit": 109
   },
   "reportUnknownMemberType": {
-    "limit": 39009
+    "limit": 38828
   },
   "reportUnknownParameterType": {
-    "limit": 19883
+    "limit": 19847
   },
   "reportUnknownVariableType": {
-    "limit": 30569
+    "limit": 30386
   },
   "reportUnnecessaryCast": {
     "limit": 117
   },
   "reportUnnecessaryComparison": {
-    "limit": 699
+    "limit": 697
   },
   "reportUnnecessaryContains": {
     "limit": 5
   },
   "reportUnnecessaryIsInstance": {
-    "limit": 836
+    "limit": 833
   },
   "reportUntypedBaseClass": {
     "limit": 0

enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py

@@ -14,6 +14,8 @@
 )
 
 if TYPE_CHECKING:
+    from prisma import models as prisma_models
+
     from litellm.integrations.prometheus import PrometheusLogger
     from litellm.proxy._types import LiteLLM_ManagedObjectTable
     from litellm.proxy.utils import PrismaClient, ProxyLogging
@@ -351,7 +353,7 @@ def _is_output_file_gone_at_provider(error: Exception, output_file_id: Optional[
         return isinstance(error, (NotFoundError, openai.NotFoundError)) and output_file_id in str(error)
 
     async def _finalize_unbilled_terminal_job(
-        self, job: "LiteLLM_ManagedObjectTable", response: "LiteLLMBatch"
+        self, job: "prisma_models.LiteLLM_ManagedObjectTable", response: "LiteLLMBatch"
     ) -> None:
         """Persist a terminal batch that has nothing billable, converting any raw
         provider file ids to managed ids, and take it out of the poll page."""

litellm/integrations/prometheus.py

@@ -96,7 +96,10 @@ async def count(self) -> int: ...
 
 def _paginated_table(repository: BaseRepository[_TableRowT]) -> _PaginatedPrismaTable[_TableRowT]:
     """View a repository's prisma table through the pagination surface budget metrics need."""
-    return repository.table
+    return cast(
+        _PaginatedPrismaTable[_TableRowT],
+        repository.table,  # cast-ok: prisma rows carry the budget columns the domain model declares
+    )
 
 
 class _OrgBudgetRow(Protocol):

litellm/models/team.py

@@ -64,8 +64,8 @@ class TeamBase(LiteLLMPydanticObjectBase):
     team_alias: str | None = None
     team_id: str | None = None
     organization_id: str | None = None
-    admins: list = []
-    members: list = []
+    admins: list[str] = []
+    members: list[str] = []
     members_with_roles: list[Member] = []
     team_member_permissions: list[str] | None = None
     metadata: dict | None = None
@@ -75,7 +75,7 @@ class TeamBase(LiteLLMPydanticObjectBase):
     soft_budget: float | None = None
     budget_duration: str | None = None
     budget_limits: list[BudgetLimitEntry] | None = None
-    models: list = []
+    models: list[str] = []
     blocked: bool = False
     router_settings: dict | None = None
     access_group_ids: list[str] | None = None