Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/auth/auth_checks.py: 42%
1963 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1# What is this?
2## Common auth checks between jwt + key based auth
3"""
4Got Valid Token from Cache, DB
5Run checks for:
71. If user can call model
82. If user is in budget
93. If end_user ('user' passed to /chat/completions, /embeddings endpoint) is in budget
10"""
12import asyncio
13import math
14import re
15import time
16from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
17from types import MappingProxyType
18from typing import TYPE_CHECKING, Any, Final, Generic, Literal, Optional, Protocol, TypeAlias
20from fastapi import HTTPException, Request, status
21from pydantic import BaseModel, TypeAdapter
22from typing_extensions import NotRequired, ReadOnly, Required, TypedDict, Unpack
24import litellm
25from litellm._logging import verbose_proxy_logger
26from litellm.caching.dual_cache import LimitedSizeOrderedDict
27from litellm.constants import (
28 CLI_JWT_EXPIRATION_HOURS,
29 CLI_SESSION_KEY_PREFIX,
30 DEFAULT_ACCESS_GROUP_CACHE_TTL,
31 DEFAULT_IN_MEMORY_TTL,
32 DEFAULT_MAX_RECURSE_DEPTH,
33 EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE,
34 END_USER_RESTRICTED_REGISTRY_MAX_SIZE,
35 MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE,
36 REGISTRY_ERROR_NEGATIVE_CACHE_TTL,
37 TAG_REGISTRY_MAX_SIZE,
38)
39from litellm.litellm_core_utils.dd_tracing import tracer
40from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
41from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
42from litellm.models.project import LiteLLM_ProjectTable
43from litellm.proxy._types import (
44 RBAC_ROLES,
45 CallInfo,
46 LiteLLM_AccessGroupTable,
47 LiteLLM_BudgetTable,
48 LiteLLM_EndUserTable,
49 Litellm_EntityType,
50 LiteLLM_JWTAuth,
51 LiteLLM_ManagedVectorStoresTable,
52 LiteLLM_ObjectPermissionTable,
53 LiteLLM_OrganizationMembershipTable,
54 LiteLLM_OrganizationTable,
55 LiteLLM_ProjectTableCachedObj,
56 LiteLLM_TagTable,
57 LiteLLM_TeamMembership,
58 LiteLLM_TeamTable,
59 LiteLLM_TeamTableCachedObj,
60 LiteLLM_UserTable,
61 LiteLLMRoutes,
62 LitellmUserRoles,
63 ModelAccessDeniedProxyException,
64 NewTeamRequest,
65 ProxyErrorTypes,
66 ProxyException,
67 RoleBasedPermissions,
68 SpecialModelNames,
69 UserAPIKeyAuth,
70)
71from litellm.proxy.agent_endpoints.auth.agent_access_groups import (
72 CeilingResolver,
73 resolve_agent_access_group_ceiling,
74)
75from litellm.proxy.agent_endpoints.auth.agent_caller import (
76 agent_caller_auth,
77 load_agent_caller_team,
78 load_agent_caller_user,
79)
80from litellm.proxy.auth.budget_throttle import (
81 budget_throttle_percentage,
82 should_throttle_budget_exceeded,
83)
84from litellm.proxy.auth.model_access_denied import model_access_denied_client_message
85from litellm.proxy.auth.route_checks import RouteChecks
86from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import publish_auth_cache_invalidation
87from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec
88from litellm.proxy.common_utils.http_parsing_utils import (
89 _safe_get_request_headers,
90 _safe_get_request_query_params,
91)
92from litellm.proxy.common_utils.model_listing_utils import alias_map
93from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
94from litellm.proxy.common_utils.user_api_key_cache import (
95 END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL,
96 MODEL_ACCESS_GROUP_REGISTRY_OVERFLOW_SENTINEL,
97 NO_TEAM_MEMBERSHIP_SENTINEL,
98 TAG_REGISTRY_OVERFLOW_SENTINEL,
99 UserApiKeyCache,
100 end_user_cache_key,
101 end_user_restricted_registry_cache_key,
102 get_management_object_ttl,
103 model_access_group_cache_key,
104 model_access_group_registry_cache_key,
105 model_access_group_spend_counter_key,
106 object_permission_cache_key,
107 project_cache_key,
108 project_spend_counter_key,
109 tag_cache_key,
110 tag_registry_cache_key,
111 team_membership_auth_cache_key,
112 team_membership_reservation_cache_key,
113)
114from litellm.proxy.db.db_lookup_gate import bounded_db_lookup, db_lookup_gate
115from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
116from litellm.proxy.guardrails.tool_name_extraction import (
117 TOOL_CAPABLE_CALL_TYPES,
118 extract_request_tool_names,
119)
120from litellm.proxy.route_llm_request import route_request
121from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start
122from litellm.proxy.spend_tracking.carried_budget_state import carry_organization_budget_state
123from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics
124from litellm.repositories.budget_repository import BudgetRepository
125from litellm.repositories.object_permission_repository import ObjectPermissionRepository
126from litellm.repositories.organization_repository import OrganizationRepository
127from litellm.repositories.prisma_protocols import DatabaseClient, RowT_co
128from litellm.repositories.project_repository import ProjectRepository
129from litellm.repositories.table_repositories import (
130 AccessGroupRepository,
131 EndUserRepository,
132 JWTKeyMappingRepository,
133 ManagedVectorStoresRepository,
134 ModelAccessGroupBudgetRepository,
135 TagRepository,
136 TeamMembershipRepository,
137)
138from litellm.repositories.team_repository import TeamRepository
139from litellm.repositories.user_repository import UserRepository
140from litellm.router import Router
141from litellm.types.proxy.auth.auth_checks import UserNotFoundError
142from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget
143from litellm.utils import get_utc_datetime
145from .auth_checks_organization import (
146 add_team_org_context_to_request_body,
147 organization_role_based_access_check,
148)
149from .auth_utils import get_model_from_request, get_request_route_template
151if TYPE_CHECKING: 151 ↛ 152line 151 didn't jump to line 152 because the condition on line 151 was never true
152 from opentelemetry.trace import Span as _Span
154 Span = _Span
155else:
156 Span = Any
159class _PrismaDictableRow(Protocol):
160 def dict(self) -> Mapping[str, object]: ... 160 ↛ exitline 160 didn't return from function 'dict' because
163class _PrismaJWTKeyMappingRow(Protocol):
164 token: str
165 jwt_issuer: str
166 jwt_claim_name: str
167 jwt_claim_value: str
170class _PrismaModelDumpRow(Protocol):
171 def model_dump(self) -> Mapping[str, object]: ... 171 ↛ exitline 171 didn't return from function 'model_dump' because
174class _PrismaTeamRow(Protocol):
175 def dict(self) -> Mapping[str, object]: ... 175 ↛ exitline 175 didn't return from function 'dict' because
177 def model_dump(self) -> Mapping[str, object]: ... 177 ↛ exitline 177 didn't return from function 'model_dump' because
180class _PrismaVectorStoreRow(Protocol):
181 def dict(self) -> Mapping[str, object]: ... 181 ↛ exitline 181 didn't return from function 'dict' because
183 def model_dump(self) -> Mapping[str, object]: ... 183 ↛ exitline 183 didn't return from function 'model_dump' because
185 def __iter__(self) -> Iterator[tuple[str, object]]: ... 185 ↛ exitline 185 didn't return from function '__iter__' because
188class _PrismaUserRow(Protocol):
189 user_id: str
191 @property
192 def organization_memberships(self) -> Sequence[_PrismaModelDumpRow | None] | None: ... 192 ↛ exitline 192 didn't return from function 'organization_memberships' because
194 @organization_memberships.setter
195 def organization_memberships(self, value: Sequence[_PrismaModelDumpRow] | None) -> None: ... 195 ↛ exitline 195 didn't return from function 'organization_memberships' because
197 def __iter__(self) -> Iterator[tuple[str, object]]: ... 197 ↛ exitline 197 didn't return from function '__iter__' because
200class _PrismaAuthTable(Protocol[RowT_co]):
201 async def find_unique( 201 ↛ exitline 201 didn't return from function 'find_unique' because
202 self, *, where: Mapping[str, object], include: Mapping[str, object] | None = None
203 ) -> RowT_co | None: ...
205 async def find_first( 205 ↛ exitline 205 didn't return from function 'find_first' because
206 self, *, where: Mapping[str, object], include: Mapping[str, object] | None = None
207 ) -> RowT_co | None: ...
209 async def find_many( 209 ↛ exitline 209 didn't return from function 'find_many' because
210 self,
211 *,
212 where: Mapping[str, object] | None = None,
213 include: Mapping[str, object] | None = None,
214 take: int | None = None,
215 ) -> Sequence[RowT_co]: ...
217 async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> RowT_co | None: ... 217 ↛ exitline 217 didn't return from function 'update' because
219 async def create(self, *, data: Mapping[str, object], include: Mapping[str, object] | None = None) -> RowT_co: ... 219 ↛ exitline 219 didn't return from function 'create' because
222class _PrismaTableHolder(Protocol[RowT_co]):
223 @property
224 def table(self) -> _PrismaAuthTable[RowT_co]: ... 224 ↛ exitline 224 didn't return from function 'table' because
227class _FindOneKwargs(TypedDict):
228 where: ReadOnly[Required[Mapping[str, object]]]
229 include: ReadOnly[NotRequired[Mapping[str, object] | None]]
232class _FindManyKwargs(TypedDict):
233 where: ReadOnly[NotRequired[Mapping[str, object] | None]]
234 include: ReadOnly[NotRequired[Mapping[str, object] | None]]
235 take: ReadOnly[NotRequired[int | None]]
238class _DeadlineBoundedTable(Generic[RowT_co]):
239 """Every read on the wrapped table fails with ``DBLookupDeadlineExceeded`` once
240 ``PROXY_DB_LOOKUP_DEADLINE_SECONDS`` passes, so a stalled database fails the
241 request fast instead of parking it in the pod until it fills its memory."""
243 __slots__ = ("_lookup", "_table")
245 def __init__(self, table: _PrismaAuthTable[RowT_co], lookup: str) -> None:
246 self._table: Final = table
247 self._lookup: Final = lookup
249 async def find_unique(
250 self,
251 **kwargs: Unpack[_FindOneKwargs], # kwargs-ok: typed pass-through that forwards exactly what the caller passed
252 ) -> RowT_co | None:
253 return await bounded_db_lookup(self._table.find_unique(**kwargs), name=self._lookup)
255 async def find_first(
256 self,
257 **kwargs: Unpack[_FindOneKwargs], # kwargs-ok: typed pass-through that forwards exactly what the caller passed
258 ) -> RowT_co | None:
259 return await bounded_db_lookup(self._table.find_first(**kwargs), name=self._lookup)
261 async def find_many(
262 self,
263 **kwargs: Unpack[_FindManyKwargs], # kwargs-ok: typed pass-through that forwards exactly what the caller passed
264 ) -> Sequence[RowT_co]:
265 return await bounded_db_lookup(self._table.find_many(**kwargs), name=self._lookup)
267 async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> RowT_co | None:
268 return await self._table.update(where=where, data=data)
270 async def create(self, *, data: Mapping[str, object], include: Mapping[str, object] | None = None) -> RowT_co:
271 return await self._table.create(data=data, include=include)
274def _dictable_table(repo: _PrismaTableHolder[_PrismaDictableRow], lookup: str) -> _PrismaAuthTable[_PrismaDictableRow]:
275 return _DeadlineBoundedTable(repo.table, lookup)
278def _jwt_key_mapping_table(
279 repo: _PrismaTableHolder[_PrismaJWTKeyMappingRow],
280) -> _PrismaAuthTable[_PrismaJWTKeyMappingRow]:
281 return _DeadlineBoundedTable(repo.table, "jwt_key_mapping")
284def _model_dump_table(
285 repo: _PrismaTableHolder[_PrismaModelDumpRow], lookup: str
286) -> _PrismaAuthTable[_PrismaModelDumpRow]:
287 return _DeadlineBoundedTable(repo.table, lookup)
290def _team_table(repo: _PrismaTableHolder[_PrismaTeamRow]) -> _PrismaAuthTable[_PrismaTeamRow]:
291 return _DeadlineBoundedTable(repo.table, "team")
294def _vector_store_table(repo: _PrismaTableHolder[_PrismaVectorStoreRow]) -> _PrismaAuthTable[_PrismaVectorStoreRow]:
295 return _DeadlineBoundedTable(repo.table, "vector_store")
298def _user_table(repo: _PrismaTableHolder[_PrismaUserRow]) -> _PrismaAuthTable[_PrismaUserRow]:
299 return _DeadlineBoundedTable(repo.table, "user")
302class _VectorStorePermissionsRow(Protocol):
303 @property
304 def vector_stores(self) -> Sequence[str] | None: ... 304 ↛ exitline 304 didn't return from function 'vector_stores' because
307def _object_permission_table(
308 repo: _PrismaTableHolder[_VectorStorePermissionsRow],
309) -> _PrismaAuthTable[_VectorStorePermissionsRow]:
310 return _DeadlineBoundedTable(repo.table, "object_permission")
313class _PrismaTagRow(Protocol):
314 tag_name: str
316 def dict(self) -> Mapping[str, object]: ... 316 ↛ exitline 316 didn't return from function 'dict' because
319def _tag_table(repo: _PrismaTableHolder[_PrismaTagRow]) -> _PrismaAuthTable[_PrismaTagRow]:
320 return repo.table
323class _PrismaEndUserRow(Protocol):
324 user_id: str
326 def dict(self) -> Mapping[str, object]: ... 326 ↛ exitline 326 didn't return from function 'dict' because
329def _end_user_table(repo: _PrismaTableHolder[_PrismaEndUserRow]) -> _PrismaAuthTable[_PrismaEndUserRow]:
330 return repo.table
333class _PrismaMaxBudgetRow(Protocol):
334 @property
335 def max_budget(self) -> float | None: ... 335 ↛ exitline 335 didn't return from function 'max_budget' because
338class _PrismaModelAccessGroupBudgetRow(Protocol):
339 access_group_name: str
341 @property
342 def spend(self) -> float | None: ... 342 ↛ exitline 342 didn't return from function 'spend' because
344 @property
345 def litellm_budget_table(self) -> _PrismaMaxBudgetRow | None: ... 345 ↛ exitline 345 didn't return from function 'litellm_budget_table' because
348def _model_access_group_budget_table(
349 repo: _PrismaTableHolder[_PrismaModelAccessGroupBudgetRow],
350) -> _PrismaAuthTable[_PrismaModelAccessGroupBudgetRow]:
351 return repo.table
354class _MemberModelScope(Protocol):
355 @property
356 def allowed_models(self) -> Sequence[str] | None: ... 356 ↛ exitline 356 didn't return from function 'allowed_models' because
359class _TeamMembershipModelScope(Protocol):
360 @property
361 def litellm_budget_table(self) -> _MemberModelScope | None: ... 361 ↛ exitline 361 didn't return from function 'litellm_budget_table' because
364def _member_allowed_models(membership: _TeamMembershipModelScope) -> Sequence[str]:
365 """The member's own model scope, read through a narrowed view of the membership row."""
366 budget_table: Final = membership.litellm_budget_table
367 return () if budget_table is None else (budget_table.allowed_models or ())
370class _RawCacheRead(Protocol):
371 async def async_get_cache(self, *, key: str) -> object: ... 371 ↛ exitline 371 didn't return from function 'async_get_cache' because
374def _raw_cache(cache: _RawCacheRead) -> _RawCacheRead:
375 return cache
378def _typed_request_body(request_body: dict) -> Mapping[str, object]:
379 return request_body
382class _JsonLoadsObj(Protocol):
383 def __call__(self, data: str) -> object: ... 383 ↛ exitline 383 didn't return from function '__call__' because
386def _typed_json_loads(fn: _JsonLoadsObj) -> _JsonLoadsObj:
387 return fn
390_safe_json_loads_obj: Final = _typed_json_loads(safe_json_loads)
393last_db_access_time: Final = LimitedSizeOrderedDict(max_size=100)
394db_cache_expiry: Final = DEFAULT_IN_MEMORY_TTL # refresh every 5s
396_TEAM_MEMBERSHIP_INFLIGHT_MAX: Final = 10000
397_team_membership_inflight: Final = LimitedSizeOrderedDict(max_size=_TEAM_MEMBERSHIP_INFLIGHT_MAX)
400class _TeamMembershipCacheMiss:
401 __slots__ = ()
404_TEAM_MEMBERSHIP_CACHE_MISS: Final = _TeamMembershipCacheMiss()
406all_routes: Final = LiteLLMRoutes.openai_routes.value + LiteLLMRoutes.management_routes.value
409def _membership_from_shared_load(result: object) -> LiteLLM_TeamMembership | None:
410 return result if isinstance(result, LiteLLM_TeamMembership) else None
413def _log_budget_lookup_failure(entity: str, error: Exception) -> None:
414 """
415 Log a warning when budget lookup fails; cache will not be populated.
417 Skips logging for expected "user not found" cases (bare Exception from
418 get_user_object when user_id_upsert=False). Adds a schema migration hint
419 when the error appears schema-related.
420 """
421 # Skip logging for expected "user not found" - not caching is correct
422 if str(error) == "" and type(error).__name__ == "Exception":
423 return
424 err_str: Final = str(error).lower()
425 hint = ""
426 if any(x in err_str for x in ("column", "schema", "does not exist", "prisma", "migrate")):
427 hint = " Run `prisma db push` or `prisma migrate deploy` to fix schema mismatches."
428 verbose_proxy_logger.error(
429 "Budget lookup failed for %s; cache will not be populated. Each request will hit the database. Error: %s.%s",
430 entity,
431 error,
432 hint,
433 )
436def _get_router_zero_cost_cache(llm_router: Router) -> dict[str, bool] | None:
437 """
438 Return the router's per-instance zero-cost cache, or ``None`` for objects
439 that don't expose one (e.g. ``MagicMock`` stand-ins in unit tests).
441 The cache lives on the ``Router`` instance so it:
442 * is invalidated by ``Router._invalidate_model_group_info_cache`` on
443 any model add/remove/upsert (including in-place pricing changes via
444 ``/model/update``, which go through ``upsert_deployment``);
445 * dies with the router itself — no risk of CPython reusing the
446 previous router's ``id()`` and serving its cached entries.
447 """
448 cache: Final = getattr(llm_router, "_zero_cost_cache", None)
449 return cache if isinstance(cache, dict) else None
452def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None) -> bool:
453 """
454 Check if a model has zero cost (no configured pricing).
456 Uses the router's get_model_group_info method to get pricing information.
458 Args:
459 model: The model name or list of model names
460 llm_router: The LiteLLM router instance
462 Returns:
463 bool: True if all costs for the model are zero, False otherwise
464 """
465 if model is None or llm_router is None: 465 ↛ 466line 465 didn't jump to line 466 because the condition on line 465 was never true
466 return False
468 # Handle list of models
469 model_list: Final = [model] if isinstance(model, str) else model
471 zero_cost_cache: Final = _get_router_zero_cost_cache(llm_router)
473 for model_name in model_list: 473 ↛ 558line 473 didn't jump to line 558 because the loop on line 473 didn't complete
474 if zero_cost_cache is not None: 474 ↛ 480line 474 didn't jump to line 480 because the condition on line 474 was always true
475 cached = zero_cost_cache.get(model_name)
476 if cached is not None:
477 if cached is False: 477 ↛ 479line 477 didn't jump to line 479 because the condition on line 477 was always true
478 return False
479 continue
480 try:
481 # Use router's get_model_group_info method directly for better reliability
482 model_group_info = llm_router.get_model_group_info(model_group=model_name)
484 if model_group_info is None: 484 ↛ 494line 484 didn't jump to line 494 because the condition on line 484 was always true
485 # Model not found or no pricing info available
486 # Conservative approach: assume it has cost
487 verbose_proxy_logger.debug("No model group info found for %s, assuming it has cost", model_name)
488 if zero_cost_cache is not None: 488 ↛ 490line 488 didn't jump to line 490 because the condition on line 488 was always true
489 zero_cost_cache[model_name] = False
490 return False
492 # Check costs for this model
493 # Only allow bypass if BOTH costs are explicitly set to 0 (not None)
494 input_cost = model_group_info.input_cost_per_token
495 output_cost = model_group_info.output_cost_per_token
497 # If costs are not explicitly configured (None), assume it has cost
498 if input_cost is None or output_cost is None:
499 verbose_proxy_logger.debug(
500 "Model %s has undefined cost (input: %s, output: %s), assuming it has cost",
501 model_name,
502 input_cost,
503 output_cost,
504 )
505 if zero_cost_cache is not None:
506 zero_cost_cache[model_name] = False
507 return False
509 # If either cost is non-zero, return False
510 if input_cost > 0 or output_cost > 0:
511 verbose_proxy_logger.debug(
512 "Model %s has non-zero cost (input: %s, output: %s)", model_name, input_cost, output_cost
513 )
514 if zero_cost_cache is not None:
515 zero_cost_cache[model_name] = False
516 return False
518 # Costs are 0 — verify this is from explicit configuration,
519 # not from defaulted sparse auto-registration entries.
520 # See: https://github.com/BerriAI/litellm/issues/24770
521 safe_name = str(model_name).replace("\n", "").replace("\r", "")
522 if not _is_cost_explicitly_configured(model_name, llm_router):
523 verbose_proxy_logger.debug(
524 "Model %s has zero cost but no explicit cost "
525 "configuration in model_cost entry — treating as unknown "
526 "cost (enforce budget)",
527 safe_name,
528 )
529 if zero_cost_cache is not None:
530 zero_cost_cache[model_name] = False
531 return False
533 if _has_ptu_flat_cost(model_name, llm_router):
534 verbose_proxy_logger.debug(
535 "Model %s prices reserved PTU capacity as a flat cost, so its zero per-token "
536 "rate is not a free model (enforce budget)",
537 safe_name,
538 )
539 if zero_cost_cache is not None:
540 zero_cost_cache[model_name] = False
541 return False
543 verbose_proxy_logger.debug(
544 "Model %s has zero cost explicitly configured (input: %s, output: %s)",
545 safe_name,
546 input_cost,
547 output_cost,
548 )
549 if zero_cost_cache is not None:
550 zero_cost_cache[model_name] = True
552 except Exception as e:
553 # If we can't determine the cost, assume it has cost (conservative approach)
554 verbose_proxy_logger.debug("Error checking cost for model %s: %s, assuming it has cost", model_name, e)
555 return False
557 # All models checked have zero cost
558 return True
561_NO_MODEL_INFO: Final[Mapping[str, object]] = MappingProxyType({})
562_TEAM_GRANT_RELATIONS: Final[Mapping[str, object]] = MappingProxyType({"litellm_model_table": True})
565def _has_ptu_flat_cost(model: str, llm_router: "Router") -> bool:
566 """Whether any deployment in the model group bills reserved PTU capacity as a flat cost.
568 Such a deployment carries an explicit zero per-token price so the flat cost is not charged
569 twice, which otherwise reads here as a free model and waives every budget check for it.
570 """
571 for deployment in llm_router.model_list:
572 if deployment.get("model_name") != model:
573 continue
574 model_info = deployment.get("model_info") or _NO_MODEL_INFO
575 if model_info.get("ptu_count") is not None and model_info.get("cost_per_ptu_per_hour") is not None:
576 return True
577 return False
580def _is_cost_explicitly_configured(model: str, llm_router: "Router") -> bool:
581 """
582 Check if any deployment in the model group has cost fields explicitly
583 set in its litellm.model_cost entry.
585 When Router._create_deployment() registers a model not in the global
586 cost map, it creates a sparse entry like {"id": "<hash>"} with no cost
587 fields. _get_model_info_helper() then defaults missing costs to 0.
588 This function detects that scenario by checking the raw model_cost entry.
589 """
590 for deployment in llm_router.model_list:
591 if deployment.get("model_name") != model:
592 continue
593 model_id = deployment.get("model_info", {}).get("id")
594 if model_id is None:
595 continue
596 raw_entry = litellm.model_cost.get(model_id, {})
597 if "input_cost_per_token" in raw_entry or "output_cost_per_token" in raw_entry:
598 return True
599 return False
602_EMPTY_COST_ENTRY: Final[Mapping[str, object]] = MappingProxyType({})
605def _is_positive_cost(value: object) -> bool:
606 return isinstance(value, (int, float)) and not isinstance(value, bool) and value > 0
609def _entry_has_priced_metric(entry: Mapping[str, object]) -> bool:
610 if entry.get("tiered_pricing") is not None:
611 return True
612 for key, value in entry.items():
613 if "cost_per" not in key:
614 continue
615 if _is_positive_cost(value):
616 return True
617 if isinstance(value, dict) and any(_is_positive_cost(nested) for nested in value.values()):
618 return True
619 return False
622def _entry_declares_price(entry: Mapping[str, object]) -> bool:
623 return any("cost_per" in key or key == "tiered_pricing" for key in entry)
626def _model_group_has_pricing(model: str, llm_router: "Router") -> bool:
627 """
628 A model group counts as priced when a deployment overrides any *cost_per* field or
629 tiered_pricing in its litellm_params, even at zero, or when its resolved model info carries
630 tiered_pricing or a positive price on any billed metric (tokens, characters, seconds, pages,
631 images, queries, ...), so models billed by a non-token metric are not treated as unpriced.
632 """
633 for deployment in llm_router.get_model_list(model_name=model) or ():
634 litellm_params = deployment.get("litellm_params") or _EMPTY_COST_ENTRY
635 if _entry_declares_price(litellm_params):
636 return True
638 model_id = (deployment.get("model_info") or _EMPTY_COST_ENTRY).get("id")
639 if model_id is None:
640 continue
642 model_info = llm_router.get_deployment_model_info(
643 model_id=model_id, model_name=litellm_params.get("model") or ""
644 )
645 if model_info is not None and _entry_has_priced_metric(model_info):
646 return True
648 return False
651def _group_declares_explicit_cost(model: str, llm_router: "Router") -> bool:
652 """
653 Alias-aware counterpart to ``_is_cost_explicitly_configured``, which resolves the model group
654 the same way ``_model_group_has_pricing`` does. A deployment that prices itself through its
655 ``model_info`` block lands in the cost map under its deployment id rather than in its
656 litellm_params, and reaching that entry through the router's own resolution keeps an alias
657 pointing at such a group from being read as unpriced.
658 """
659 for deployment in llm_router.get_model_list(model_name=model) or ():
660 model_id = (deployment.get("model_info") or _EMPTY_COST_ENTRY).get("id")
661 if model_id is None:
662 continue
663 raw_entry = litellm.model_cost.get(model_id, _EMPTY_COST_ENTRY)
664 if "input_cost_per_token" in raw_entry or "output_cost_per_token" in raw_entry:
665 return True
666 return False
669def model_has_no_cost_mapping(model: str | None, llm_router: Router | None) -> bool:
670 if not model or llm_router is None:
671 return False
673 if llm_router.get_model_group_info(model_group=model) is None:
674 return False
676 if _model_group_has_pricing(model=model, llm_router=llm_router):
677 return False
679 return not _group_declares_explicit_cost(model=model, llm_router=llm_router)
682def _unpriced_models_in_request(model: str | list[str] | None, llm_router: Router | None) -> tuple[str, ...]:
683 candidates: Final = (model,) if isinstance(model, str) else tuple(model or ())
684 return tuple(
685 candidate for candidate in candidates if model_has_no_cost_mapping(model=candidate, llm_router=llm_router)
686 )
689def _unpriced_models_block_message(models: tuple[str, ...]) -> str:
690 names: Final = ", ".join(f"'{model}'" for model in models)
691 subject: Final = f"Model {names} has" if len(models) == 1 else f"Models {names} have"
692 return (
693 f"{subject} no pricing in the cost map, so litellm cannot price the request. "
694 "Requests for unpriced models are blocked because 'block_requests_for_models_without_pricing' "
695 "is enabled. Add pricing (input_cost_per_token/output_cost_per_token) to allow the request."
696 )
699async def _run_project_checks(
700 project_object: LiteLLM_ProjectTableCachedObj | None,
701 _model: str | list[str] | None,
702 llm_router: Router | None,
703 skip_budget_checks: bool,
704 valid_token: UserAPIKeyAuth | None,
705 proxy_logging_obj: ProxyLogging,
706) -> None:
707 """
708 Run all project-level checks: blocked, model access, budget, soft budget.
709 Extracted from common_checks() to keep statement count manageable.
710 """
711 if project_object is None: 711 ↛ 715line 711 didn't jump to line 715 because the condition on line 711 was always true
712 return
714 # 1.1. If project is blocked
715 if project_object.blocked is True:
716 raise Exception(
717 f"Project={project_object.project_id} is blocked. Update via `/project/update` if you're an admin."
718 )
720 # 2.2 If project can call model
721 if _model and len(project_object.models) > 0:
722 can_project_access_model(
723 model=_model,
724 project_object=project_object,
725 llm_router=llm_router,
726 key_model_aliases=key_model_aliases_for_auth_check(valid_token),
727 )
729 if not skip_budget_checks:
730 # 3.0.2. If project is in budget
731 await _project_max_budget_check(
732 project_object=project_object,
733 valid_token=valid_token,
734 proxy_logging_obj=proxy_logging_obj,
735 )
737 # 3.0.3. If project is over soft budget (alert only, doesn't block)
738 await _project_soft_budget_check(
739 project_object=project_object,
740 valid_token=valid_token,
741 proxy_logging_obj=proxy_logging_obj,
742 )
745def _enforce_user_param_check(general_settings: dict, request: Request, request_body: dict, route: str) -> None:
746 if not general_settings.get("enforce_user_param", False): 746 ↛ 749line 746 didn't jump to line 749 because the condition on line 746 was always true
747 return
749 http_method: Final = request.method if hasattr(request, "method") else None
750 is_post_method: Final = http_method and http_method.upper() == "POST"
751 is_openai_route: Final = RouteChecks.is_llm_api_route(route=route)
752 is_mcp_route: Final = route in LiteLLMRoutes.mcp_routes.value or RouteChecks.check_route_access(
753 route=route, allowed_routes=LiteLLMRoutes.mcp_routes.value
754 )
756 if is_post_method and is_openai_route and not is_mcp_route and "user" not in request_body:
757 raise Exception(f"'user' param not passed in. 'enforce_user_param'={general_settings['enforce_user_param']}")
760def _reject_clientside_metadata_tags_check(general_settings: dict, request_body: dict, route: str) -> None:
761 if not general_settings.get("reject_clientside_metadata_tags", False): 761 ↛ 764line 761 didn't jump to line 764 because the condition on line 761 was always true
762 return
764 if (
765 RouteChecks.is_llm_api_route(route=route)
766 and "metadata" in request_body
767 and isinstance(request_body["metadata"], dict)
768 and "tags" in request_body["metadata"]
769 ):
770 raise ProxyException(
771 message=f"Client-side 'metadata.tags' not allowed in request. 'reject_clientside_metadata_tags'={general_settings['reject_clientside_metadata_tags']}. Tags can only be set via API key metadata.",
772 type=ProxyErrorTypes.bad_request_error,
773 param="metadata.tags",
774 code=status.HTTP_400_BAD_REQUEST,
775 )
778def _global_proxy_budget_check(global_proxy_spend: float | None, skip_budget_checks: bool, route: str) -> None:
779 if ( 779 ↛ 787line 779 didn't jump to line 787 because the condition on line 779 was never true
780 litellm.max_budget > 0
781 and not skip_budget_checks
782 and global_proxy_spend is not None
783 and RouteChecks.is_llm_api_route(route=route)
784 and route != "/v1/models"
785 and route != "/models"
786 ):
787 if math.isfinite(litellm.max_budget) and global_proxy_spend > litellm.max_budget:
788 raise litellm.BudgetExceededError(
789 current_cost=global_proxy_spend,
790 max_budget=litellm.max_budget,
791 entity_type=Litellm_EntityType.PROXY.value,
792 )
795_GUARDRAIL_MODIFICATION_KEYS: Final[tuple] = (
796 "guardrails",
797 "disable_global_guardrails",
798 "disable_global_guardrail",
799 "opted_out_global_guardrails",
800)
803def _guardrail_modification_check(request_body: Mapping[str, object], team_object: LiteLLM_TeamTable | None) -> None:
804 """
805 Reject user-supplied metadata flags that would modify guardrail behavior
806 unless the team has explicit permission. Checked keys include the plural
807 ``guardrails`` list plus the per-request toggles that influence whether
808 default-on guardrails run (``disable_global_guardrails``,
809 ``disable_global_guardrail`` singular, and ``opted_out_global_guardrails``).
811 User-supplied values for the bypass toggles are also silently ignored by
812 ``_get_admin_metadata`` at read time; this check adds defense in depth by
813 failing loudly at the auth layer so operators see an explicit 403 instead
814 of a confusing silent-ignore.
815 """
816 from litellm.proxy.guardrails.guardrail_helpers import can_modify_guardrails
818 def _coerce_to_dict(container: object) -> dict | None:
819 """Accept dict or JSON-string (from multipart/form-data or extra_body).
821 Without this, an attacker can smuggle guardrail keys past the check by
822 sending ``{"metadata": "{\\"disable_global_guardrails\\": true}"}`` —
823 ``isinstance(dict)`` on the string returns False, the check returns
824 no-modification, and ``add_litellm_data_to_request`` parses the string
825 to a dict downstream.
826 """
827 if isinstance(container, dict):
828 return container
829 if isinstance(container, str):
830 parsed: Final = _safe_json_loads_obj(container)
831 return parsed if isinstance(parsed, dict) else None
832 return None
834 def _user_requested_modification(container: object) -> bool:
835 coerced: Final = _coerce_to_dict(container)
836 if coerced is None:
837 return False
838 return any(key in coerced for key in _GUARDRAIL_MODIFICATION_KEYS)
840 # Check both metadata keys — callers can populate either depending on the
841 # endpoint. Cover the top-level too so root-level injection is rejected.
842 modifies: Final = (
843 _user_requested_modification(request_body.get("metadata"))
844 or _user_requested_modification(request_body.get("litellm_metadata"))
845 or _user_requested_modification(request_body)
846 )
847 if not modifies:
848 return
850 if not can_modify_guardrails(team_object): 850 ↛ 851line 850 didn't jump to line 851 because the condition on line 850 was never true
851 raise HTTPException(
852 status_code=403,
853 detail={"error": "Your team does not have permission to modify guardrails."},
854 )
857async def check_tools_allowlist(
858 request_body: dict,
859 valid_token: UserAPIKeyAuth | None,
860 team_object: LiteLLM_TeamTable | None,
861 route: str,
862) -> None:
863 """
864 Enforce key/team tool allowlist (metadata.allowed_tools). No DB in hot path —
865 effective allowlist is read from valid_token.metadata and valid_token.team_metadata.
866 Raises ProxyException with tool_access_denied if a tool is not allowed.
867 """
868 from litellm.litellm_core_utils.api_route_to_call_types import (
869 get_call_types_for_route,
870 )
872 if valid_token is None: 872 ↛ 873line 872 didn't jump to line 873 because the condition on line 872 was never true
873 return
874 call_types: Final = get_call_types_for_route(route)
875 if not call_types or not any(ct.value in TOOL_CAPABLE_CALL_TYPES for ct in call_types):
876 return
877 tool_names: Final = extract_request_tool_names(route, request_body)
878 if not tool_names: 878 ↛ 880line 878 didn't jump to line 880 because the condition on line 878 was always true
879 return
880 key_meta: Final = (valid_token.metadata or {}) if isinstance(valid_token.metadata, dict) else {}
881 team_meta: Final = (valid_token.team_metadata or {}) if isinstance(valid_token.team_metadata, dict) else {}
882 key_allowed: Final = key_meta.get("allowed_tools")
883 team_allowed: Final = team_meta.get("allowed_tools")
884 effective: Final = key_allowed if (isinstance(key_allowed, list) and len(key_allowed) > 0) else team_allowed
885 if not isinstance(effective, list) or len(effective) == 0:
886 return
887 allowed_set: Final = {str(t) for t in effective}
888 disallowed: Final = [n for n in tool_names if n not in allowed_set]
889 if disallowed:
890 raise ProxyException(
891 message=f"Tool(s) {disallowed} are not in the allowed tools list for this key/team.",
892 type=ProxyErrorTypes.tool_access_denied,
893 param="tools",
894 code=status.HTTP_403_FORBIDDEN,
895 )
898# Read-only discovery routes that incur no spend. Kept narrower than info_routes so an exhausted
899# budget cannot reach side-effectful routes like /health/services (Slack/email/webhook). See #27923.
900MODEL_DISCOVERY_ROUTES: Final = frozenset(
901 {
902 "/v1/models",
903 "/models",
904 "/model/info",
905 "/v1/model/info",
906 "/v2/model/info",
907 "/model_group/info",
908 "/utils/model_info",
909 }
910)
912BUDGET_ENFORCED_SIDE_EFFECT_ROUTES: Final = frozenset(
913 {
914 "/health",
915 "/health/services",
916 "/health/test_connection",
917 "/auto_router/test_routing",
918 }
919)
922def route_skips_budget_checks(route: str) -> bool:
923 return route not in BUDGET_ENFORCED_SIDE_EFFECT_ROUTES and (
924 route in MODEL_DISCOVERY_ROUTES or not RouteChecks.is_llm_api_route(route=route)
925 )
928def request_skips_budget_checks(route: str, model: str | list[str] | None, llm_router: Router | None) -> bool:
929 return route_skips_budget_checks(route=route) or _is_model_cost_zero(model=model, llm_router=llm_router)
932async def common_checks(
933 request_body: dict,
934 team_object: LiteLLM_TeamTable | None,
935 user_object: LiteLLM_UserTable | None,
936 end_user_object: LiteLLM_EndUserTable | None,
937 global_proxy_spend: float | None,
938 general_settings: dict,
939 route: str,
940 llm_router: Router | None,
941 proxy_logging_obj: ProxyLogging,
942 valid_token: UserAPIKeyAuth | None,
943 request: Request,
944 skip_budget_checks: bool = False,
945 project_object: LiteLLM_ProjectTableCachedObj | None = None,
946) -> bool:
947 """
948 Common checks across jwt + key-based auth.
950 1. If team is blocked
951 1.1. If project is blocked
952 2. If team can call model
953 2.2 If project can call model
954 2.3 Which model access groups authorized this request
955 3. If team is in budget
956 3.0.2. If project is in budget
957 3.0.3. If project is over soft budget (alert only)
958 4. If user passed in (JWT or key.user_id) - is in budget
959 5. If end_user (either via JWT or 'user' passed to /chat/completions, /embeddings endpoint) is in budget
960 6. [OPTIONAL] If 'enforce_end_user' enabled - did developer pass in 'user' param for openai endpoints
961 7. [OPTIONAL] If 'litellm.max_budget' is set (>0), is proxy under budget
962 8. [OPTIONAL] If guardrails modified - is request allowed to change this
963 9. Check if request body is safe
964 10. [OPTIONAL] Organization checks - is user_object.organization_id is set, run these checks
965 11. [OPTIONAL] Vector store checks - is the object allowed to access the vector store
966 """
967 from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
969 _model: Final[str | list[str] | None] = get_model_from_request(
970 request_data=request_body,
971 route=route,
972 request_headers=_safe_get_request_headers(request=request),
973 request_query_params=_safe_get_request_query_params(request=request),
974 llm_router=llm_router,
975 request=request,
976 team_id=valid_token.team_id if valid_token is not None else None,
977 )
979 skip_all_budget_checks: Final = skip_budget_checks or route_skips_budget_checks(route=route)
981 membership_user_id: Final = (
982 valid_token.user_id if valid_token is not None and (bool(_model) or not skip_all_budget_checks) else None
983 )
984 team_membership_loaded: Final = team_object is not None and membership_user_id is not None
985 loaded_team_membership: Final = (
986 await get_team_membership(
987 user_id=membership_user_id,
988 team_id=team_object.team_id,
989 prisma_client=prisma_client,
990 user_api_key_cache=user_api_key_cache,
991 proxy_logging_obj=proxy_logging_obj,
992 )
993 if team_object is not None and membership_user_id is not None
994 else None
995 )
997 unpriced_models: Final = (
998 _unpriced_models_in_request(model=_model, llm_router=llm_router)
999 if litellm.block_requests_for_models_without_pricing and RouteChecks.is_llm_api_route(route=route)
1000 else ()
1001 )
1002 if unpriced_models: 1002 ↛ 1003line 1002 didn't jump to line 1003 because the condition on line 1002 was never true
1003 raise ProxyException(
1004 message=_unpriced_models_block_message(unpriced_models),
1005 type=ProxyErrorTypes.model_cost_map_missing,
1006 param="model",
1007 code=status.HTTP_403_FORBIDDEN,
1008 )
1010 # 1. If team is blocked
1011 if team_object is not None and team_object.blocked is True: 1011 ↛ 1012line 1011 didn't jump to line 1012 because the condition on line 1011 was never true
1012 raise Exception(f"Team={team_object.team_id} is blocked. Update via `/team/unblock` if you're an admin.")
1014 # 2. If team can call model (or key's access_group_ids grant it)
1015 if _model and team_object: 1015 ↛ 1016line 1015 didn't jump to line 1016 because the condition on line 1015 was never true
1016 with tracer.trace("litellm.proxy.auth.common_checks.can_team_access_model"):
1017 try:
1018 await can_team_access_model(
1019 model=_model,
1020 team_object=team_object,
1021 llm_router=llm_router,
1022 team_model_aliases=(valid_token.team_model_aliases if valid_token else None),
1023 key_model_aliases=key_model_aliases_for_auth_check(valid_token),
1024 )
1025 except ProxyException as team_denial:
1026 if team_denial.type != ProxyErrorTypes.team_model_access_denied:
1027 raise
1028 if not await _key_access_group_grants_model(
1029 model=_model,
1030 valid_token=valid_token,
1031 team_object=team_object,
1032 llm_router=llm_router,
1033 key_model_aliases=key_model_aliases_for_auth_check(valid_token),
1034 ):
1035 raise
1037 # 2.2. If team member has per-member model scope, enforce it
1038 if _model and team_object and valid_token and valid_token.user_id: 1038 ↛ 1039line 1038 didn't jump to line 1039 because the condition on line 1038 was never true
1039 with tracer.trace("litellm.proxy.auth.common_checks.check_team_member_model_access"):
1040 await _check_team_member_model_access(
1041 model=_model,
1042 team_object=team_object,
1043 valid_token=valid_token,
1044 llm_router=llm_router,
1045 prisma_client=prisma_client,
1046 user_api_key_cache=user_api_key_cache,
1047 proxy_logging_obj=proxy_logging_obj,
1048 team_membership=loaded_team_membership,
1049 team_membership_loaded=team_membership_loaded,
1050 key_model_aliases=key_model_aliases_for_auth_check(valid_token),
1051 )
1053 # Require trace id for agent keys when agent has require_trace_id_on_calls_by_agent
1054 if valid_token is not None and valid_token.agent_id: 1054 ↛ 1055line 1054 didn't jump to line 1055 because the condition on line 1054 was never true
1055 from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
1056 from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers
1058 agent: Final = global_agent_registry.get_agent_by_id(agent_id=valid_token.agent_id)
1059 if agent is not None:
1060 require_trace_id: Final = (agent.litellm_params or {}).get("require_trace_id_on_calls_by_agent")
1061 if require_trace_id:
1062 headers_dict: Final = dict(request.headers)
1063 trace_id: Final = get_chain_id_from_headers(headers_dict)
1064 if not trace_id:
1065 raise ProxyException(
1066 message="Requests made with this agent's key must include the x-litellm-trace-id header.",
1067 type=ProxyErrorTypes.bad_request_error,
1068 param=None,
1069 code=status.HTTP_400_BAD_REQUEST,
1070 )
1072 await _check_agent_access_group_model_access(model=_model, valid_token=valid_token, llm_router=llm_router)
1073 await _check_agent_caller_model_access(
1074 model=_model,
1075 valid_token=valid_token,
1076 llm_router=llm_router,
1077 prisma_client=prisma_client,
1078 user_api_key_cache=user_api_key_cache,
1079 proxy_logging_obj=proxy_logging_obj,
1080 )
1082 ## 2.1 If user can call model (if personal key)
1083 if _model and team_object is None and user_object is not None:
1084 with tracer.trace("litellm.proxy.auth.common_checks.can_user_call_model"):
1085 await can_user_call_model(
1086 model=_model,
1087 llm_router=llm_router,
1088 user_object=user_object,
1089 key_model_aliases=key_model_aliases_for_auth_check(valid_token),
1090 )
1092 # 1.1 - 2.2 - 3.0.2 - 3.0.3: Project checks (blocked, model access, budget)
1093 with tracer.trace("litellm.proxy.auth.common_checks.run_project_checks"):
1094 await _run_project_checks(
1095 project_object=project_object,
1096 _model=_model,
1097 llm_router=llm_router,
1098 skip_budget_checks=skip_all_budget_checks,
1099 valid_token=valid_token,
1100 proxy_logging_obj=proxy_logging_obj,
1101 )
1103 # 2.3 Which model access groups authorized this request
1104 matched_model_access_groups: Final = await stamp_matched_model_access_groups(
1105 model=_model,
1106 valid_token=valid_token,
1107 team_object=team_object,
1108 project_object=project_object,
1109 llm_router=llm_router,
1110 prisma_client=prisma_client,
1111 user_api_key_cache=user_api_key_cache,
1112 proxy_logging_obj=proxy_logging_obj,
1113 team_membership=loaded_team_membership,
1114 team_membership_loaded=team_membership_loaded,
1115 )
1117 # Run before apply_key_tags_pre_auth injects key metadata.tags into request_body.
1118 _reject_clientside_metadata_tags_check(general_settings, request_body, route)
1120 # If this is a free model, skip all budget checks
1121 if not skip_all_budget_checks:
1122 # Key metadata.tags are injected into request_body here so the tag budget
1123 # check can read them; this mutation must run before the gathered checks.
1124 if valid_token is not None: 1124 ↛ 1137line 1124 didn't jump to line 1137 because the condition on line 1124 was always true
1125 from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
1127 LiteLLMProxyRequestSetup.pre_seed_litellm_metadata_for_route(
1128 request_data=request_body,
1129 route=route,
1130 )
1132 LiteLLMProxyRequestSetup.apply_key_tags_pre_auth(
1133 request_data=request_body,
1134 user_api_key_dict=valid_token,
1135 )
1137 async def _user_max_budget_check() -> None:
1138 # 4.1 personal budget
1139 if user_object is None or user_object.max_budget is None: 1139 ↛ 1141line 1139 didn't jump to line 1141 because the condition on line 1139 was always true
1140 return
1141 is_team_key: Final = team_object is not None and team_object.team_id is not None
1142 if is_team_key and general_settings.get("apply_user_budget_to_team_keys") is not True:
1143 return
1145 from litellm.proxy.proxy_server import get_current_spend
1147 user_budget: Final = user_object.max_budget
1148 user_spend: Final = await get_current_spend(
1149 counter_key=f"spend:user:{user_object.user_id}",
1150 fallback_spend=user_object.spend or 0.0,
1151 max_budget=user_budget,
1152 )
1153 call_info: Final = CallInfo(
1154 spend=user_spend,
1155 max_budget=user_budget,
1156 user_id=user_object.user_id,
1157 user_email=user_object.user_email,
1158 event_group=Litellm_EntityType.USER,
1159 )
1160 asyncio.create_task(
1161 proxy_logging_obj.budget_alerts(
1162 type="user_budget",
1163 user_info=call_info,
1164 )
1165 )
1166 if math.isfinite(user_budget) and user_spend >= user_budget:
1167 raise litellm.BudgetExceededError(
1168 current_cost=user_spend,
1169 max_budget=user_budget,
1170 message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}",
1171 entity_type=Litellm_EntityType.USER.value,
1172 entity_id=user_object.user_id,
1173 )
1175 # Each scope reads a distinct counter key with no cross-scope ordering
1176 # dependency, so the per-scope Redis-first reads run concurrently instead
1177 # of one sequential await per scope. return_exceptions lets every scope
1178 # settle, then the first error in scope-priority order propagates exactly
1179 # as the sequential path raised.
1180 budget_check_coros: Final = tuple(
1181 coro
1182 for coro in (
1183 _team_max_budget_check(
1184 team_object=team_object,
1185 proxy_logging_obj=proxy_logging_obj,
1186 valid_token=valid_token,
1187 ),
1188 _team_multi_budget_check(team_object=team_object),
1189 _virtual_key_multi_budget_check(valid_token=valid_token) if valid_token is not None else None,
1190 _team_soft_budget_check(
1191 team_object=team_object,
1192 proxy_logging_obj=proxy_logging_obj,
1193 valid_token=valid_token,
1194 ),
1195 _organization_max_budget_check(
1196 valid_token=valid_token,
1197 team_object=team_object,
1198 prisma_client=prisma_client,
1199 user_api_key_cache=user_api_key_cache,
1200 proxy_logging_obj=proxy_logging_obj,
1201 ),
1202 _tag_max_budget_check(
1203 request_body=request_body,
1204 prisma_client=prisma_client,
1205 user_api_key_cache=user_api_key_cache,
1206 proxy_logging_obj=proxy_logging_obj,
1207 valid_token=valid_token,
1208 ),
1209 _model_access_group_max_budget_check(
1210 matched_model_access_groups=matched_model_access_groups,
1211 prisma_client=prisma_client,
1212 user_api_key_cache=user_api_key_cache,
1213 )
1214 if matched_model_access_groups
1215 else None,
1216 _user_max_budget_check(),
1217 _check_team_member_budget(
1218 team_object=team_object,
1219 user_object=user_object,
1220 valid_token=valid_token,
1221 prisma_client=prisma_client,
1222 user_api_key_cache=user_api_key_cache,
1223 proxy_logging_obj=proxy_logging_obj,
1224 team_membership=loaded_team_membership,
1225 team_membership_loaded=team_membership_loaded,
1226 ),
1227 _check_end_user_budget(end_user_obj=end_user_object, route=route)
1228 if end_user_object is not None and end_user_object.litellm_budget_table is not None
1229 else None,
1230 )
1231 if coro is not None
1232 )
1234 with tracer.trace("litellm.proxy.auth.common_checks.budget_checks"):
1235 budget_results: Final = await asyncio.gather(*budget_check_coros, return_exceptions=True)
1236 budget_error: Final = next((r for r in budget_results if isinstance(r, BaseException)), None)
1237 if budget_error is not None: 1237 ↛ 1238line 1237 didn't jump to line 1238 because the condition on line 1237 was never true
1238 raise budget_error
1240 _enforce_user_param_check(general_settings, request, request_body, route)
1241 _global_proxy_budget_check(global_proxy_spend, skip_all_budget_checks, route)
1242 _guardrail_modification_check(_typed_request_body(request_body), team_object)
1244 # 10 [OPTIONAL] Organization RBAC checks
1245 organization_role_based_access_check(user_object=user_object, route=route, request_body=request_body)
1247 async def _fetch_team_org_id(team_id: str) -> str | None:
1248 try:
1249 team: Final = await get_team_object(
1250 team_id=team_id,
1251 prisma_client=prisma_client,
1252 user_api_key_cache=user_api_key_cache,
1253 proxy_logging_obj=proxy_logging_obj,
1254 )
1255 except HTTPException:
1256 return None
1257 return team.organization_id
1259 request_body_for_route_check: Final = await add_team_org_context_to_request_body(
1260 route=route,
1261 request_body=request_body,
1262 fetch_team_org_id=_fetch_team_org_id,
1263 route_template=get_request_route_template(request),
1264 )
1266 _is_route_allowed: Final = _is_api_route_allowed(
1267 route=route,
1268 request=request,
1269 request_data=request_body_for_route_check,
1270 valid_token=valid_token,
1271 user_obj=user_object,
1272 )
1274 # 11. [OPTIONAL] Vector store checks - is the object allowed to access the vector store
1275 with tracer.trace("litellm.proxy.auth.common_checks.vector_store_access_check"):
1276 await vector_store_access_check(
1277 request_body=request_body,
1278 team_object=team_object,
1279 valid_token=valid_token,
1280 )
1282 # 12. [OPTIONAL] Tool allowlist - key/team allowed_tools (no DB in hot path)
1283 with tracer.trace("litellm.proxy.auth.common_checks.check_tools_allowlist"):
1284 await check_tools_allowlist(
1285 request_body=request_body,
1286 valid_token=valid_token,
1287 team_object=team_object,
1288 route=route,
1289 )
1291 return True
1294def effective_user_role(user_role: str | None) -> LitellmUserRoles:
1295 try:
1296 return LitellmUserRoles(user_role)
1297 except ValueError:
1298 return LitellmUserRoles.INTERNAL_USER
1301def _get_user_role(
1302 user_obj: LiteLLM_UserTable | None,
1303) -> LitellmUserRoles | None:
1304 if user_obj is None:
1305 return None
1306 return effective_user_role(user_obj.user_role)
1309def _is_api_route_allowed(
1310 route: str,
1311 request: Request,
1312 request_data: dict,
1313 valid_token: UserAPIKeyAuth | None,
1314 user_obj: LiteLLM_UserTable | None = None,
1315) -> bool:
1316 """
1317 - Route b/w api token check and normal token check
1318 """
1319 _user_role: Final = _get_user_role(user_obj=user_obj)
1321 if valid_token is None: 1321 ↛ 1322line 1321 didn't jump to line 1322 because the condition on line 1321 was never true
1322 raise Exception("Invalid proxy server token passed. valid_token=None.")
1324 if not _is_user_proxy_admin(user_obj=user_obj): # if non-admin
1325 RouteChecks.non_proxy_admin_allowed_routes_check(
1326 user_obj=user_obj,
1327 _user_role=_user_role,
1328 route=route,
1329 request=request,
1330 request_data=request_data,
1331 valid_token=valid_token,
1332 )
1333 return True
1336def _is_user_proxy_admin(user_obj: LiteLLM_UserTable | None):
1337 if user_obj is None:
1338 return False
1340 if user_obj.user_role is not None and user_obj.user_role == LitellmUserRoles.PROXY_ADMIN.value:
1341 return True
1343 return False
1346def _allowed_routes_check(user_route: str, allowed_routes: list) -> bool:
1347 """
1348 Return if a user is allowed to access route. Helper function for `allowed_routes_check`.
1350 Parameters:
1351 - user_route: str - the route the user is trying to call
1352 - allowed_routes: List[str|LiteLLMRoutes] - the list of allowed routes for the user. Entries are a route group name
1353 (e.g. "openai_routes"), an exact route, or a trailing-wildcard prefix (e.g. "/internal-models/*").
1354 """
1355 from starlette.routing import compile_path
1357 for allowed_route in allowed_routes:
1358 if allowed_route in LiteLLMRoutes.__members__:
1359 for template in LiteLLMRoutes[allowed_route].value:
1360 regex, _, _ = compile_path(template)
1361 if regex.match(user_route):
1362 return True
1363 elif RouteChecks.route_matches_wildcard_pattern(route=user_route, pattern=allowed_route):
1364 return True
1365 return False
1368def allowed_routes_check(
1369 user_role: LitellmUserRoles,
1370 user_route: str,
1371 litellm_proxy_roles: LiteLLM_JWTAuth,
1372) -> bool:
1373 """
1374 Check if user -> not admin - allowed to access these routes
1375 """
1377 if user_role == LitellmUserRoles.PROXY_ADMIN:
1378 is_allowed = _allowed_routes_check(
1379 user_route=user_route,
1380 allowed_routes=litellm_proxy_roles.admin_allowed_routes,
1381 )
1382 return is_allowed
1384 elif user_role == LitellmUserRoles.TEAM:
1385 if litellm_proxy_roles.team_allowed_routes is None:
1386 """
1387 By default allow a team to call openai + info routes
1388 """
1389 is_allowed = _allowed_routes_check(user_route=user_route, allowed_routes=["openai_routes", "info_routes"])
1390 return is_allowed
1391 elif litellm_proxy_roles.team_allowed_routes is not None:
1392 is_allowed = _allowed_routes_check(
1393 user_route=user_route,
1394 allowed_routes=litellm_proxy_roles.team_allowed_routes,
1395 )
1396 return is_allowed
1397 return False
1400def allowed_route_check_inside_route(
1401 user_api_key_dict: UserAPIKeyAuth,
1402 requested_user_id: str | None,
1403) -> bool:
1404 ret_val = True
1405 if (
1406 user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN
1407 and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY
1408 ):
1409 ret_val = False
1410 if requested_user_id is not None and user_api_key_dict.user_id is not None:
1411 if user_api_key_dict.user_id == requested_user_id:
1412 ret_val = True
1413 return ret_val
1416def get_actual_routes(allowed_routes: list) -> list:
1417 actual_routes: Final[list] = []
1418 for route_name in allowed_routes:
1419 try:
1420 route_value = LiteLLMRoutes[route_name].value
1421 if isinstance(route_value, set):
1422 actual_routes.extend(list(route_value))
1423 else:
1424 actual_routes.extend(route_value)
1426 except KeyError:
1427 actual_routes.append(route_name)
1428 return actual_routes
1431KEY_END_USER_BUDGET_ID_METADATA_FIELD: Final = "end_user_budget_id"
1434def get_key_end_user_budget_id(key_metadata: Mapping[str, object] | None) -> str | None:
1435 """The default budget a key assigns to end users that carry no budget of their own."""
1436 if key_metadata is None: 1436 ↛ 1437line 1436 didn't jump to line 1437 because the condition on line 1436 was never true
1437 return None
1438 budget_id: Final = key_metadata.get(KEY_END_USER_BUDGET_ID_METADATA_FIELD)
1439 return budget_id if isinstance(budget_id, str) and budget_id != "" else None
1442async def get_default_end_user_budget(
1443 prisma_client: PrismaClient | None,
1444 user_api_key_cache: UserApiKeyCache,
1445 parent_otel_span: Span | None = None,
1446 budget_id: str | None = None,
1447) -> LiteLLM_BudgetTable | None:
1448 """
1449 Fetches the default end user budget from the database.
1451 ``budget_id`` selects the budget row; when omitted the proxy-wide
1452 ``litellm.max_end_user_budget_id`` is used. This budget is applied to end
1453 users who don't have an explicit budget_id set. Results are cached for performance.
1455 Args:
1456 prisma_client: Database client instance
1457 user_api_key_cache: Cache for storing/retrieving budget data
1458 parent_otel_span: Optional OpenTelemetry span for tracing
1459 budget_id: Budget row to load instead of the proxy-wide default
1461 Returns:
1462 LiteLLM_BudgetTable if configured and found, None otherwise
1463 """
1464 default_budget_id: Final = budget_id if budget_id is not None else litellm.max_end_user_budget_id
1465 if prisma_client is None or default_budget_id is None:
1466 return None
1468 cache_key: Final = f"default_end_user_budget:{default_budget_id}"
1470 # Check cache first
1471 cached_budget: Final = await user_api_key_cache.async_get_cache(
1472 key=cache_key,
1473 model_type=LiteLLM_BudgetTable,
1474 )
1475 if cached_budget is not None:
1476 return cached_budget
1478 # Fetch from database
1479 try:
1480 budget_record: Final = await _dictable_table(BudgetRepository(prisma_client), "budget").find_unique(
1481 where={"budget_id": default_budget_id} # mutable-ok: prisma where clause
1482 )
1484 if budget_record is None:
1485 verbose_proxy_logger.warning(
1486 "Default end user budget not found in database: %s",
1487 default_budget_id.replace("\r", "").replace("\n", ""),
1488 )
1489 return None
1491 _budget_obj: Final = LiteLLM_BudgetTable.model_validate(budget_record.dict())
1492 # Cache the budget for 60 seconds
1493 await user_api_key_cache.async_set_cache(
1494 key=cache_key,
1495 value=_budget_obj,
1496 model_type=LiteLLM_BudgetTable,
1497 ttl=get_management_object_ttl(user_api_key_cache),
1498 )
1500 return _budget_obj
1502 except Exception as e:
1503 verbose_proxy_logger.error("Error fetching default end user budget: %s", e)
1504 return None
1507@log_db_metrics
1508async def get_team_member_default_budget(
1509 budget_id: str,
1510 prisma_client: PrismaClient | None,
1511 user_api_key_cache: UserApiKeyCache,
1512) -> LiteLLM_BudgetTable | None:
1513 """
1514 Fetches the team-level default per-member budget referenced by team.metadata["team_member_budget_id"].
1516 This budget is applied to team members whose TeamMembership row has no
1517 linked budget, or whose linked budget has max_budget=NULL. Results are
1518 cached for performance.
1520 Args:
1521 budget_id: The budget_id pulled from team.metadata["team_member_budget_id"]
1522 prisma_client: Database client instance
1523 user_api_key_cache: Cache for storing/retrieving budget data
1525 Returns:
1526 LiteLLM_BudgetTable if found, None otherwise
1527 """
1528 if prisma_client is None:
1529 return None
1531 cache_key: Final = f"team_member_default_budget:{budget_id}"
1533 cached_budget: Final = await user_api_key_cache.async_get_cache(
1534 key=cache_key,
1535 model_type=LiteLLM_BudgetTable,
1536 )
1537 if cached_budget is not None:
1538 return cached_budget
1540 try:
1541 budget_record: Final = await _dictable_table(BudgetRepository(prisma_client), "budget").find_unique(
1542 where={"budget_id": budget_id}
1543 )
1544 except Exception:
1545 verbose_proxy_logger.exception("Error fetching team-default member budget %s", budget_id)
1546 return None
1548 if budget_record is None:
1549 verbose_proxy_logger.warning("Team-default member budget not found in database: %s", budget_id)
1550 return None
1552 budget: Final = LiteLLM_BudgetTable.model_validate(budget_record.dict())
1553 await user_api_key_cache.async_set_cache(
1554 key=cache_key,
1555 value=budget,
1556 model_type=LiteLLM_BudgetTable,
1557 ttl=get_management_object_ttl(user_api_key_cache),
1558 )
1560 return budget
1563async def resolve_default_end_user_budget(
1564 prisma_client: PrismaClient,
1565 user_api_key_cache: UserApiKeyCache,
1566 key_end_user_budget_id: str | None,
1567 parent_otel_span: Span | None = None,
1568) -> LiteLLM_BudgetTable | None:
1569 """
1570 The default budget for an end user with no budget of its own.
1572 The key's ``end_user_budget_id`` takes precedence over the proxy-wide
1573 ``litellm.max_end_user_budget_id``; the proxy-wide default is the fallback when the key
1574 names no budget or its budget row is missing.
1575 """
1576 if key_end_user_budget_id is not None:
1577 key_budget: Final = await get_default_end_user_budget(
1578 prisma_client=prisma_client,
1579 user_api_key_cache=user_api_key_cache,
1580 parent_otel_span=parent_otel_span,
1581 budget_id=key_end_user_budget_id,
1582 )
1583 if key_budget is not None:
1584 return key_budget
1586 if litellm.max_end_user_budget_id is None:
1587 return None
1589 return await get_default_end_user_budget(
1590 prisma_client=prisma_client,
1591 user_api_key_cache=user_api_key_cache,
1592 parent_otel_span=parent_otel_span,
1593 )
1596async def _apply_default_budget_to_end_user(
1597 end_user_obj: LiteLLM_EndUserTable,
1598 prisma_client: PrismaClient,
1599 user_api_key_cache: UserApiKeyCache,
1600 parent_otel_span: Span | None = None,
1601 key_end_user_budget_id: str | None = None,
1602) -> LiteLLM_EndUserTable:
1603 """
1604 Returns the end user with the resolved default budget when it has no budget of its own.
1606 A row whose own ``budget_id`` resolved to a budget is returned unchanged. Otherwise the
1607 default is resolved on every call and set on a copy: the cached row carries at most the
1608 proxy-wide default (readers such as the Prometheus customer gauges rely on that), never a
1609 key's, so requests through keys with different defaults never observe each other's budget.
1611 Args:
1612 end_user_obj: The end user object to potentially apply default budget to
1613 prisma_client: Database client instance
1614 user_api_key_cache: Cache for storing/retrieving data
1615 parent_otel_span: Optional OpenTelemetry span for tracing
1616 key_end_user_budget_id: The requesting key's ``end_user_budget_id``, if any
1617 """
1618 if end_user_obj.budget_id is not None and end_user_obj.litellm_budget_table is not None:
1619 return end_user_obj
1621 if key_end_user_budget_id is None and litellm.max_end_user_budget_id is None:
1622 return end_user_obj
1624 default_budget: Final = await resolve_default_end_user_budget(
1625 prisma_client=prisma_client,
1626 user_api_key_cache=user_api_key_cache,
1627 key_end_user_budget_id=key_end_user_budget_id,
1628 parent_otel_span=parent_otel_span,
1629 )
1631 if default_budget is None:
1632 return end_user_obj
1634 verbose_proxy_logger.debug(
1635 "Applied default budget %s to end user %s", default_budget.budget_id, end_user_obj.user_id
1636 )
1637 return end_user_obj.model_copy(update=MappingProxyType({"litellm_budget_table": default_budget}))
1640async def _check_end_user_budget(
1641 end_user_obj: LiteLLM_EndUserTable,
1642 route: str,
1643) -> None:
1644 """
1645 Check if end user is within their budget limit.
1647 Args:
1648 end_user_obj: The end user object to check
1649 route: The request route
1651 Raises:
1652 litellm.BudgetExceededError: If end user has exceeded their budget
1653 """
1654 if RouteChecks.is_info_route(route):
1655 return
1657 if end_user_obj.litellm_budget_table is None:
1658 return
1660 end_user_budget: Final = end_user_obj.litellm_budget_table.max_budget
1661 if end_user_budget is None:
1662 return
1664 from litellm.proxy.proxy_server import get_current_spend
1666 end_user_spend: Final = await get_current_spend(
1667 counter_key=f"spend:end_user:{end_user_obj.user_id}",
1668 fallback_spend=end_user_obj.spend or 0.0,
1669 max_budget=end_user_budget,
1670 fallback_authoritative=True,
1671 )
1672 if end_user_spend > end_user_budget:
1673 raise litellm.BudgetExceededError(
1674 current_cost=end_user_spend,
1675 max_budget=end_user_budget,
1676 message=f"ExceededBudget: End User={end_user_obj.user_id} over budget. Spend={end_user_spend}, Budget={end_user_budget}",
1677 entity_type=Litellm_EntityType.END_USER.value,
1678 entity_id=end_user_obj.user_id,
1679 )
1682#: Columns whose non-null value makes an end-user row restrict something auth enforces. ``blocked``
1683#: is separate: it restricts when true rather than when merely set.
1684_RESTRICTED_COLUMNS: Final = ("budget_id", "allowed_model_region", "default_model", "object_permission_id")
1687def _column_is_set(column: str) -> Mapping[str, object]:
1688 """``column IS NOT NULL`` as a plain dict, which is the only shape prisma's builder accepts."""
1689 return {column: {"not": None}} # mutable-ok: prisma's query builder isinstance-checks for dict
1692def _restricted_end_user_where() -> Mapping[str, object]:
1693 """Prisma filter selecting every end-user row that carries a restriction auth enforces."""
1694 return {"OR": [{"blocked": True}, *map(_column_is_set, _RESTRICTED_COLUMNS)]} # mutable-ok: prisma needs dict/list
1697class _RegistryNotCached:
1698 """No cached registry answer, as distinct from the cached answer ``None`` (registry unusable)."""
1701_REGISTRY_NOT_CACHED: Final = _RegistryNotCached()
1703#: One lock per registry; module-level because the stampede to collapse is worker-wide.
1704_TAG_REGISTRY_LOAD_LOCK: Final = asyncio.Lock()
1705_END_USER_REGISTRY_LOAD_LOCK: Final = asyncio.Lock()
1706_MODEL_ACCESS_GROUP_REGISTRY_LOAD_LOCK: Final = asyncio.Lock()
1709async def _cached_registry(
1710 cache_key: str,
1711 overflow_sentinel: str,
1712 user_api_key_cache: UserApiKeyCache,
1713) -> frozenset[str] | None | _RegistryNotCached:
1714 """The cached registry answer, or ``_REGISTRY_NOT_CACHED`` when the caller has to query."""
1715 cached: Final = await _raw_cache(user_api_key_cache).async_get_cache(key=cache_key)
1716 if cached == overflow_sentinel: 1716 ↛ 1717line 1716 didn't jump to line 1717 because the condition on line 1716 was never true
1717 return None
1718 # Memory hands back the tuple that was written; Redis round-trips it through JSON as a list.
1719 if isinstance(cached, (list, tuple)):
1720 return frozenset(entry for entry in cached if isinstance(entry, str))
1721 return _REGISTRY_NOT_CACHED
1724async def _cache_registry_answer(
1725 cache_key: str,
1726 value: tuple[str, ...] | str,
1727 ttl: float,
1728 user_api_key_cache: UserApiKeyCache,
1729) -> None:
1730 """Best-effort: a cache backend failure must not turn a registry load into a failed request."""
1731 try:
1732 await user_api_key_cache.async_set_cache(key=cache_key, value=value, ttl=ttl)
1733 except Exception as e: # noqa: BLE001 # best-effort cache write: auth must survive a cache backend error
1734 verbose_proxy_logger.warning("Failed to cache registry %s: %s", cache_key, e)
1737async def _fetch_and_cache_registry(
1738 cache_key: str,
1739 overflow_sentinel: str,
1740 max_size: int,
1741 fetch_ids: Callable[[], Awaitable[tuple[str, ...]]],
1742 user_api_key_cache: UserApiKeyCache,
1743) -> frozenset[str] | None:
1744 """The registry as the database has it, cached whole, or ``None`` when it is unusable."""
1745 try:
1746 registry_ids: Final = await fetch_ids()
1747 except Exception as e: # noqa: BLE001 # fail-safe: any registry load error must degrade to per-id lookups, never break auth
1748 verbose_proxy_logger.warning(
1749 "Registry %s could not be loaded from the database, so per-id lookups will run and the "
1750 "registry query is suppressed for %ss: %s",
1751 cache_key,
1752 REGISTRY_ERROR_NEGATIVE_CACHE_TTL,
1753 e,
1754 )
1755 await _cache_registry_answer(
1756 cache_key=cache_key,
1757 value=overflow_sentinel,
1758 ttl=REGISTRY_ERROR_NEGATIVE_CACHE_TTL,
1759 user_api_key_cache=user_api_key_cache,
1760 )
1761 return None
1763 if len(registry_ids) > max_size: 1763 ↛ 1764line 1763 didn't jump to line 1764 because the condition on line 1763 was never true
1764 await _cache_registry_answer(
1765 cache_key=cache_key,
1766 value=overflow_sentinel,
1767 ttl=get_management_object_ttl(user_api_key_cache),
1768 user_api_key_cache=user_api_key_cache,
1769 )
1770 return None
1772 await _cache_registry_answer(
1773 cache_key=cache_key,
1774 value=registry_ids,
1775 ttl=get_management_object_ttl(user_api_key_cache),
1776 user_api_key_cache=user_api_key_cache,
1777 )
1778 return frozenset(registry_ids)
1781async def _load_bounded_registry(
1782 cache_key: str,
1783 overflow_sentinel: str,
1784 max_size: int,
1785 load_lock: asyncio.Lock,
1786 fetch_ids: Callable[[], Awaitable[tuple[str, ...]]],
1787 user_api_key_cache: UserApiKeyCache,
1788) -> frozenset[str] | None:
1789 """
1790 A bounded id set under one cache key, so an id outside it costs no DB read.
1792 ``None`` = unusable (overflow or recent DB error): fall back to per-id lookups. An empty
1793 frozenset is a real, cacheable answer. Loads are single-flighted to stop TTL-expiry stampedes.
1794 """
1795 cached: Final = await _cached_registry(cache_key, overflow_sentinel, user_api_key_cache)
1796 if not isinstance(cached, _RegistryNotCached):
1797 return cached
1799 async with load_lock:
1800 # The request that held the lock has since cached an answer for everyone waiting on it.
1801 cached_after_wait: Final = await _cached_registry(cache_key, overflow_sentinel, user_api_key_cache)
1802 if not isinstance(cached_after_wait, _RegistryNotCached): 1802 ↛ 1803line 1802 didn't jump to line 1803 because the condition on line 1802 was never true
1803 return cached_after_wait
1805 return await _fetch_and_cache_registry(
1806 cache_key=cache_key,
1807 overflow_sentinel=overflow_sentinel,
1808 max_size=max_size,
1809 fetch_ids=fetch_ids,
1810 user_api_key_cache=user_api_key_cache,
1811 )
1814async def _load_end_user_restricted_registry(
1815 prisma_client: PrismaClient,
1816 user_api_key_cache: UserApiKeyCache,
1817) -> frozenset[str] | None:
1818 """The set of end-user ids whose ``LiteLLM_EndUserTable`` row carries a restriction."""
1820 async def fetch_ids() -> tuple[str, ...]:
1821 restricted_rows: Final = await _end_user_table(EndUserRepository(prisma_client)).find_many(
1822 where=_restricted_end_user_where(),
1823 take=END_USER_RESTRICTED_REGISTRY_MAX_SIZE + 1,
1824 )
1825 return tuple(row.user_id for row in restricted_rows)
1827 return await _load_bounded_registry(
1828 cache_key=end_user_restricted_registry_cache_key(),
1829 overflow_sentinel=END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL,
1830 max_size=END_USER_RESTRICTED_REGISTRY_MAX_SIZE,
1831 load_lock=_END_USER_REGISTRY_LOAD_LOCK,
1832 fetch_ids=fetch_ids,
1833 user_api_key_cache=user_api_key_cache,
1834 )
1837async def _end_user_is_known_unrestricted(
1838 end_user_id: str,
1839 prisma_client: PrismaClient,
1840 user_api_key_cache: UserApiKeyCache,
1841 token_end_user_max_budget: float | None,
1842 key_end_user_budget_id: str | None = None,
1843) -> bool:
1844 """
1845 True when the cached registry proves the id restricts nothing, so its row need not be read.
1847 Every field ``get_end_user_object`` callers consume (budget, spend under that budget, region,
1848 default model, object permission, blocked) is part of the registry predicate, so an id outside
1849 it is indistinguishable from one with no row at all. The skip is off whenever mere existence of
1850 the row is meaningful: ``max_end_user_budget_id`` or the key's ``end_user_budget_id`` grafts a
1851 default budget onto any row that exists, ``validate_end_user_id_in_db`` rejects ids that resolve
1852 to no row, and a token-supplied ``end_user_max_budget`` (a ``user_custom_auth`` callable can set
1853 one against an otherwise unrestricted row) is enforced against the row's recorded spend.
1854 """
1855 if ( 1855 ↛ 1861line 1855 didn't jump to line 1861 because the condition on line 1855 was never true
1856 litellm.max_end_user_budget_id is not None
1857 or key_end_user_budget_id is not None
1858 or litellm.validate_end_user_id_in_db
1859 or token_end_user_max_budget is not None
1860 ):
1861 return False
1863 registry: Final = await _load_end_user_restricted_registry(
1864 prisma_client=prisma_client,
1865 user_api_key_cache=user_api_key_cache,
1866 )
1867 return registry is not None and end_user_id not in registry
1870@log_db_metrics
1871async def get_end_user_object(
1872 end_user_id: str | None,
1873 prisma_client: PrismaClient | None,
1874 user_api_key_cache: UserApiKeyCache,
1875 route: str | None = "",
1876 parent_otel_span: Span | None = None,
1877 proxy_logging_obj: ProxyLogging | None = None,
1878 token_end_user_max_budget: float | None = None,
1879 key_end_user_budget_id: str | None = None,
1880) -> LiteLLM_EndUserTable | None:
1881 """
1882 Returns end user object from database or cache.
1884 If end user exists but has no budget_id, applies the default budget: the key's
1885 ``end_user_budget_id`` when set, otherwise ``litellm.max_end_user_budget_id``.
1887 Args:
1888 end_user_id: The ID of the end user
1889 prisma_client: Database client instance
1890 user_api_key_cache: Cache for storing/retrieving data
1891 route: The request route
1892 parent_otel_span: Optional OpenTelemetry span for tracing
1893 proxy_logging_obj: Optional proxy logging object
1894 token_end_user_max_budget: ``valid_token.end_user_max_budget``, when the caller holds a
1895 token. Budget enforcement reads the row's spend, so a row that restricts nothing on
1896 its own must still be loaded when the token carries a budget for it.
1897 key_end_user_budget_id: The requesting key's default end-user budget, if any
1899 Returns:
1900 LiteLLM_EndUserTable if found, None otherwise
1901 """
1902 if prisma_client is None: 1902 ↛ 1903line 1902 didn't jump to line 1903 because the condition on line 1902 was never true
1903 raise Exception("No db connected")
1905 if end_user_id is None: 1905 ↛ 1906line 1905 didn't jump to line 1906 because the condition on line 1905 was never true
1906 return None
1908 _key: Final = end_user_cache_key(end_user_id)
1910 # Check cache first
1911 cached_user_obj: Final = await user_api_key_cache.async_get_cache(
1912 key=_key,
1913 model_type=LiteLLM_EndUserTable,
1914 )
1915 if cached_user_obj is not None: 1915 ↛ 1916line 1915 didn't jump to line 1916 because the condition on line 1915 was never true
1916 return await _apply_default_budget_to_end_user(
1917 end_user_obj=cached_user_obj,
1918 prisma_client=prisma_client,
1919 user_api_key_cache=user_api_key_cache,
1920 parent_otel_span=parent_otel_span,
1921 key_end_user_budget_id=key_end_user_budget_id,
1922 )
1924 if await _end_user_is_known_unrestricted( 1924 ↛ 1934line 1924 didn't jump to line 1934 because the condition on line 1924 was always true
1925 end_user_id=end_user_id,
1926 prisma_client=prisma_client,
1927 user_api_key_cache=user_api_key_cache,
1928 token_end_user_max_budget=token_end_user_max_budget,
1929 key_end_user_budget_id=key_end_user_budget_id,
1930 ):
1931 return None
1933 # Fetch from database
1934 try:
1935 response: Final = await _dictable_table(EndUserRepository(prisma_client), "end_user").find_unique(
1936 where={"user_id": end_user_id},
1937 include={"litellm_budget_table": True, "object_permission": True},
1938 )
1940 if response is None:
1941 raise Exception
1943 end_user_row: Final = await _apply_default_budget_to_end_user(
1944 end_user_obj=LiteLLM_EndUserTable.model_validate(response.dict()),
1945 prisma_client=prisma_client,
1946 user_api_key_cache=user_api_key_cache,
1947 parent_otel_span=parent_otel_span,
1948 )
1950 await user_api_key_cache.async_set_cache(
1951 key=_key,
1952 value=end_user_row,
1953 model_type=LiteLLM_EndUserTable,
1954 ttl=get_management_object_ttl(user_api_key_cache),
1955 )
1957 if key_end_user_budget_id is None:
1958 return end_user_row
1960 return await _apply_default_budget_to_end_user(
1961 end_user_obj=end_user_row,
1962 prisma_client=prisma_client,
1963 user_api_key_cache=user_api_key_cache,
1964 parent_otel_span=parent_otel_span,
1965 key_end_user_budget_id=key_end_user_budget_id,
1966 )
1968 except Exception:
1969 return None
1972_END_USER_VALIDATION_NEGATIVE_TTL: Final = 60
1973_END_USER_VALIDATION_POSITIVE_TTL: Final = 300
1976async def resolve_and_validate_end_user_id(
1977 raw_end_user_id: str | None,
1978 prisma_client: PrismaClient | None,
1979 user_api_key_cache: UserApiKeyCache,
1980 parent_otel_span: Span | None = None,
1981 proxy_logging_obj: ProxyLogging | None = None,
1982 route: str = "",
1983 key_end_user_budget_id: str | None = None,
1984) -> str | None:
1985 """Optionally drop end-user ids that don't resolve to a known DB row.
1987 Default: pass-through. LiteLLM's documented pattern is that the `user`
1988 field is an arbitrary caller-supplied identifier, so validation is
1989 opt-in behind ``litellm.validate_end_user_id_in_db`` to preserve
1990 backwards compatibility.
1992 When the flag is set: accept the id when it matches any of
1993 - LiteLLM_EndUserTable.user_id
1994 - LiteLLM_UserTable.user_id
1995 - LiteLLM_UserTable.user_email (case-insensitive)
1997 If the id doesn't match but a default end-user budget is configured
1998 (``litellm.max_end_user_budget_id`` or the key's ``end_user_budget_id``),
1999 we still preserve the id so that budget is applied downstream; otherwise
2000 we return None.
2002 DB lookups reuse ``get_end_user_object`` / ``get_user_object`` so they
2003 share the same cache as the rest of the auth path instead of adding new
2004 raw Prisma queries.
2005 """
2006 if raw_end_user_id is None:
2007 return None
2008 if not litellm.validate_end_user_id_in_db: 2008 ↛ 2010line 2008 didn't jump to line 2010 because the condition on line 2008 was always true
2009 return raw_end_user_id
2010 if prisma_client is None:
2011 return raw_end_user_id
2013 has_default_budget: Final = bool(litellm.max_end_user_budget_id) or key_end_user_budget_id is not None
2014 cache_key: Final = f"end_user_validation:{raw_end_user_id}"
2015 cached: Final = await _raw_cache(user_api_key_cache).async_get_cache(key=cache_key)
2016 if cached == "valid":
2017 return raw_end_user_id
2018 if cached == "invalid":
2019 return raw_end_user_id if has_default_budget else None
2021 is_valid: Final = await _end_user_id_exists_in_db(
2022 end_user_id=raw_end_user_id,
2023 prisma_client=prisma_client,
2024 user_api_key_cache=user_api_key_cache,
2025 parent_otel_span=parent_otel_span,
2026 proxy_logging_obj=proxy_logging_obj,
2027 route=route,
2028 )
2030 await user_api_key_cache.async_set_cache(
2031 key=cache_key,
2032 value="valid" if is_valid else "invalid",
2033 ttl=(_END_USER_VALIDATION_POSITIVE_TTL if is_valid else _END_USER_VALIDATION_NEGATIVE_TTL),
2034 )
2036 return raw_end_user_id if is_valid or has_default_budget else None
2039async def _end_user_id_exists_in_db(
2040 end_user_id: str,
2041 prisma_client: PrismaClient,
2042 user_api_key_cache: UserApiKeyCache,
2043 parent_otel_span: Span | None = None,
2044 proxy_logging_obj: ProxyLogging | None = None,
2045 route: str = "",
2046) -> bool:
2047 """True when the id matches an EndUser, User, or user_email row."""
2048 try:
2049 end_user_obj: Final = await get_end_user_object(
2050 end_user_id=end_user_id,
2051 prisma_client=prisma_client,
2052 user_api_key_cache=user_api_key_cache,
2053 parent_otel_span=parent_otel_span,
2054 proxy_logging_obj=proxy_logging_obj,
2055 route=route,
2056 )
2057 if end_user_obj is not None:
2058 return True
2059 except Exception as e:
2060 verbose_proxy_logger.debug("end_user validation: get_end_user_object lookup failed: %s", e)
2062 try:
2063 user_obj: Final = await get_user_object(
2064 user_id=end_user_id,
2065 prisma_client=prisma_client,
2066 user_api_key_cache=user_api_key_cache,
2067 user_id_upsert=False,
2068 parent_otel_span=parent_otel_span,
2069 proxy_logging_obj=proxy_logging_obj,
2070 check_db_only=False,
2071 user_email=end_user_id if "@" in end_user_id else None,
2072 )
2073 if user_obj is not None:
2074 return True
2075 except Exception as e:
2076 verbose_proxy_logger.debug("end_user validation: get_user_object lookup failed: %s", e)
2078 return False
2081async def _load_tag_registry(
2082 prisma_client: PrismaClient,
2083 user_api_key_cache: UserApiKeyCache,
2084) -> frozenset[str] | None:
2085 """The set of tag names that have a row in ``LiteLLM_TagTable``."""
2087 async def fetch_ids() -> tuple[str, ...]:
2088 registry_rows: Final = await _tag_table(TagRepository(prisma_client)).find_many(
2089 take=TAG_REGISTRY_MAX_SIZE + 1,
2090 )
2091 return tuple(row.tag_name for row in registry_rows)
2093 return await _load_bounded_registry(
2094 cache_key=tag_registry_cache_key(),
2095 overflow_sentinel=TAG_REGISTRY_OVERFLOW_SENTINEL,
2096 max_size=TAG_REGISTRY_MAX_SIZE,
2097 load_lock=_TAG_REGISTRY_LOAD_LOCK,
2098 fetch_ids=fetch_ids,
2099 user_api_key_cache=user_api_key_cache,
2100 )
2103async def _load_model_access_group_registry(
2104 prisma_client: PrismaClient,
2105 user_api_key_cache: UserApiKeyCache,
2106) -> frozenset[str] | None:
2107 """The set of model access group names that have a row in ``LiteLLM_ModelAccessGroupBudgetTable``."""
2109 async def fetch_ids() -> tuple[str, ...]:
2110 registry_rows: Final = await _model_access_group_budget_table(
2111 ModelAccessGroupBudgetRepository(prisma_client)
2112 ).find_many(take=MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE + 1)
2113 return tuple(row.access_group_name for row in registry_rows)
2115 return await _load_bounded_registry(
2116 cache_key=model_access_group_registry_cache_key(),
2117 overflow_sentinel=MODEL_ACCESS_GROUP_REGISTRY_OVERFLOW_SENTINEL,
2118 max_size=MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE,
2119 load_lock=_MODEL_ACCESS_GROUP_REGISTRY_LOAD_LOCK,
2120 fetch_ids=fetch_ids,
2121 user_api_key_cache=user_api_key_cache,
2122 )
2125async def _fetch_uncached_model_access_group_budgets(
2126 uncached_groups: Sequence[str],
2127 prisma_client: PrismaClient,
2128 user_api_key_cache: UserApiKeyCache,
2129) -> tuple[tuple[str, ModelAccessGroupBudget], ...]:
2130 """Budget rows for the groups a cache probe missed.
2132 No registry gate here, unlike the tag path: the names only ever come from
2133 ``matched_model_access_groups``, which :func:`collect_matched_model_access_groups` already
2134 intersected with the registry, so a name that has no row cannot reach this.
2135 """
2136 if not uncached_groups:
2137 return ()
2139 try:
2140 db_rows: Final = await _model_access_group_budget_table(
2141 ModelAccessGroupBudgetRepository(prisma_client)
2142 ).find_many(
2143 where={"access_group_name": {"in": list(uncached_groups)}},
2144 include={"litellm_budget_table": True},
2145 )
2146 fetched: Final = tuple((row.access_group_name, _model_access_group_budget(row)) for row in db_rows)
2147 for fetched_name, fetched_obj in fetched:
2148 await user_api_key_cache.async_set_cache(
2149 key=model_access_group_cache_key(fetched_name),
2150 value=fetched_obj,
2151 model_type=ModelAccessGroupBudget,
2152 ttl=get_management_object_ttl(user_api_key_cache),
2153 )
2154 except Exception as e: # noqa: BLE001 # fail-safe: a budget fetch error must yield "no budget rows", never break auth
2155 verbose_proxy_logger.debug("Error batch fetching model access group budgets from database: %s", e)
2156 return ()
2157 else:
2158 return fetched
2161def _model_access_group_budget(row: _PrismaModelAccessGroupBudgetRow) -> ModelAccessGroupBudget:
2162 budget_table: Final = row.litellm_budget_table
2163 return ModelAccessGroupBudget(
2164 access_group_name=row.access_group_name,
2165 spend=row.spend or 0.0,
2166 max_budget=None if budget_table is None else budget_table.max_budget,
2167 )
2170@log_db_metrics
2171async def get_model_access_group_budgets_batch(
2172 access_group_names: Sequence[str],
2173 prisma_client: PrismaClient | None,
2174 user_api_key_cache: UserApiKeyCache,
2175) -> dict[str, ModelAccessGroupBudget]:
2176 """Budget rows for the given model access groups, served from cache where possible.
2178 Shared by the two enforcement paths so they read one row per group per request: the
2179 reservation counters when reservations are on, and :func:`_model_access_group_max_budget_check`
2180 when ``disable_budget_reservation`` turns them off.
2181 """
2182 if prisma_client is None or not access_group_names:
2183 return {}
2185 probed: Final = [
2186 (
2187 group,
2188 await user_api_key_cache.async_get_cache(
2189 key=model_access_group_cache_key(group), model_type=ModelAccessGroupBudget
2190 ),
2191 )
2192 for group in access_group_names
2193 ]
2194 fetched: Final = await _fetch_uncached_model_access_group_budgets(
2195 uncached_groups=tuple(group for group, budget in probed if budget is None),
2196 prisma_client=prisma_client,
2197 user_api_key_cache=user_api_key_cache,
2198 )
2199 return {group: budget for group, budget in (*probed, *fetched) if budget is not None}
2202async def _fetch_uncached_tags(
2203 uncached_tags: Sequence[str],
2204 prisma_client: PrismaClient,
2205 user_api_key_cache: UserApiKeyCache,
2206) -> tuple[tuple[str, LiteLLM_TagTable], ...]:
2207 """Rows for the tags a cache probe missed; names absent from the registry never reach the DB."""
2208 if not uncached_tags:
2209 return ()
2211 registry: Final = await _load_tag_registry(
2212 prisma_client=prisma_client,
2213 user_api_key_cache=user_api_key_cache,
2214 )
2215 tags_to_fetch: Final = (
2216 tuple(uncached_tags) if registry is None else tuple(tag for tag in uncached_tags if tag in registry)
2217 )
2218 if not tags_to_fetch:
2219 return ()
2221 try:
2222 db_tags: Final = await _tag_table(TagRepository(prisma_client)).find_many(
2223 where={"tag_name": {"in": list(tags_to_fetch)}},
2224 include={"litellm_budget_table": True},
2225 )
2226 fetched: Final = tuple((db_tag.tag_name, LiteLLM_TagTable.model_validate(db_tag.dict())) for db_tag in db_tags)
2227 for fetched_name, fetched_obj in fetched:
2228 await user_api_key_cache.async_set_cache(
2229 key=tag_cache_key(fetched_name),
2230 value=fetched_obj,
2231 model_type=LiteLLM_TagTable,
2232 ttl=get_management_object_ttl(user_api_key_cache),
2233 )
2234 except Exception as e: # noqa: BLE001 # fail-safe: a tag fetch error must yield "no budget objects", never break auth
2235 verbose_proxy_logger.debug("Error batch fetching tags from database: %s", e)
2236 return ()
2237 else:
2238 return fetched
2241@log_db_metrics
2242async def get_tag_objects_batch(
2243 tag_names: Sequence[str],
2244 prisma_client: PrismaClient | None,
2245 user_api_key_cache: UserApiKeyCache,
2246 parent_otel_span: Span | None = None,
2247 proxy_logging_obj: ProxyLogging | None = None,
2248) -> dict[str, LiteLLM_TagTable]:
2249 """
2250 Batch fetch multiple tag objects from cache and db.
2252 Optimizes for latency by:
2253 1. Serving already-cached tags without touching the DB
2254 2. Skipping tags that no ``LiteLLM_TagTable`` row exists for, via the cached name registry
2255 3. Batch fetching the remaining uncached tags in one DB query
2257 Args:
2258 tag_names: List of tag names to fetch
2259 prisma_client: Prisma database client
2260 user_api_key_cache: Cache for storing tag objects
2261 parent_otel_span: Optional OpenTelemetry span for tracing
2262 proxy_logging_obj: Optional proxy logging object
2264 Returns:
2265 Dictionary mapping tag_name to LiteLLM_TagTable object
2266 """
2267 if prisma_client is None or not tag_names: 2267 ↛ 2268line 2267 didn't jump to line 2268 because the condition on line 2267 was never true
2268 return {}
2270 probed: Final = [
2271 (
2272 tag_name,
2273 await user_api_key_cache.async_get_cache(key=tag_cache_key(tag_name), model_type=LiteLLM_TagTable),
2274 )
2275 for tag_name in tag_names
2276 ]
2277 fetched: Final = await _fetch_uncached_tags(
2278 uncached_tags=tuple(tag_name for tag_name, tag_obj in probed if tag_obj is None),
2279 prisma_client=prisma_client,
2280 user_api_key_cache=user_api_key_cache,
2281 )
2282 return {tag_name: tag_obj for tag_name, tag_obj in (*probed, *fetched) if tag_obj is not None}
2285@log_db_metrics
2286async def get_tag_object(
2287 tag_name: str | None,
2288 prisma_client: PrismaClient | None,
2289 user_api_key_cache: UserApiKeyCache,
2290 parent_otel_span: Span | None = None,
2291 proxy_logging_obj: ProxyLogging | None = None,
2292) -> LiteLLM_TagTable | None:
2293 """
2294 Returns tag object from cache or db.
2296 Uses default cache TTL (same as end_user objects) to avoid drift.
2298 Args:
2299 tag_name: Name of the tag to fetch
2300 prisma_client: Prisma database client
2301 user_api_key_cache: Cache for storing tag objects
2302 parent_otel_span: Optional OpenTelemetry span for tracing
2303 proxy_logging_obj: Optional proxy logging object
2305 Returns:
2306 LiteLLM_TagTable object if found, None otherwise
2307 """
2308 if prisma_client is None or tag_name is None:
2309 return None
2311 # Use batch helper for consistency
2312 tag_objects: Final = await get_tag_objects_batch(
2313 tag_names=[tag_name],
2314 prisma_client=prisma_client,
2315 user_api_key_cache=user_api_key_cache,
2316 parent_otel_span=parent_otel_span,
2317 proxy_logging_obj=proxy_logging_obj,
2318 )
2320 return tag_objects.get(tag_name)
2323def _membership_from_cached_payload(
2324 cached: object,
2325) -> LiteLLM_TeamMembership | None | _TeamMembershipCacheMiss:
2326 if cached is None:
2327 return _TEAM_MEMBERSHIP_CACHE_MISS
2328 if cached == NO_TEAM_MEMBERSHIP_SENTINEL: 2328 ↛ 2329line 2328 didn't jump to line 2329 because the condition on line 2328 was never true
2329 return None
2330 cached_membership: Final = CacheCodec.deserialize(cached, model_type=LiteLLM_TeamMembership)
2331 return cached_membership if cached_membership is not None else _TEAM_MEMBERSHIP_CACHE_MISS
2334@log_db_metrics
2335async def _fetch_team_membership_from_db(
2336 user_id: str,
2337 team_id: str,
2338 prisma_client: PrismaClient,
2339 user_api_key_cache: UserApiKeyCache,
2340 parent_otel_span: Span | None = None,
2341 proxy_logging_obj: ProxyLogging | None = None,
2342) -> LiteLLM_TeamMembership | None:
2343 _ = parent_otel_span, proxy_logging_obj
2344 response: Final = await _dictable_table(TeamMembershipRepository(prisma_client), "team_membership").find_unique(
2345 where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}},
2346 include={"litellm_budget_table": True},
2347 )
2348 membership: Final = None if response is None else LiteLLM_TeamMembership.model_validate(response.dict())
2349 _key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id)
2350 if membership is None: 2350 ↛ 2351line 2350 didn't jump to line 2351 because the condition on line 2350 was never true
2351 await user_api_key_cache.async_set_cache(
2352 key=_key,
2353 value=NO_TEAM_MEMBERSHIP_SENTINEL,
2354 ttl=get_management_object_ttl(user_api_key_cache),
2355 )
2356 else:
2357 await user_api_key_cache.async_set_cache(
2358 key=_key,
2359 value=membership,
2360 model_type=LiteLLM_TeamMembership,
2361 )
2362 return membership
2365async def _load_team_membership_on_cache_miss(
2366 user_id: str,
2367 team_id: str,
2368 cache_key: str,
2369 prisma_client: PrismaClient,
2370 user_api_key_cache: UserApiKeyCache,
2371 parent_otel_span: Span | None,
2372 proxy_logging_obj: ProxyLogging | None,
2373) -> LiteLLM_TeamMembership | None:
2374 redis_cached: Final[object] = await user_api_key_cache.async_get_cache(key=cache_key)
2375 redis_membership: Final = _membership_from_cached_payload(redis_cached)
2376 if not isinstance(redis_membership, _TeamMembershipCacheMiss): 2376 ↛ 2377line 2376 didn't jump to line 2377 because the condition on line 2376 was never true
2377 return redis_membership
2379 return await _fetch_team_membership_from_db(
2380 user_id=user_id,
2381 team_id=team_id,
2382 prisma_client=prisma_client,
2383 user_api_key_cache=user_api_key_cache,
2384 parent_otel_span=parent_otel_span,
2385 proxy_logging_obj=proxy_logging_obj,
2386 )
2389async def get_team_membership(
2390 user_id: str,
2391 team_id: str,
2392 prisma_client: PrismaClient | None,
2393 user_api_key_cache: UserApiKeyCache,
2394 parent_otel_span: Span | None = None,
2395 proxy_logging_obj: ProxyLogging | None = None,
2396) -> Optional["LiteLLM_TeamMembership"]:
2397 """
2398 Returns team membership object if user is member of team.
2400 Do a isolated check for team membership vs. doing a combined key + team + user + team-membership check, as key might come in frequently for different users/teams. Larger call will slowdown query time. This way we get to cache the constant (key/team/user info) and only update based on the changing value (team membership).
2401 """
2402 if user_id is None or team_id is None: 2402 ↛ 2403line 2402 didn't jump to line 2403 because the condition on line 2402 was never true
2403 return None
2405 _key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id)
2407 l1_cached: Final[object] = await user_api_key_cache.async_get_cache(key=_key, local_only=True)
2408 l1_membership: Final = _membership_from_cached_payload(l1_cached)
2409 if not isinstance(l1_membership, _TeamMembershipCacheMiss):
2410 return l1_membership
2412 inflight: Final[object] = _team_membership_inflight.get(_key)
2413 if isinstance(inflight, asyncio.Task): 2413 ↛ 2414line 2413 didn't jump to line 2414 because the condition on line 2413 was never true
2414 return _membership_from_shared_load(await asyncio.shield(inflight))
2416 if prisma_client is None: 2416 ↛ 2417line 2416 didn't jump to line 2417 because the condition on line 2416 was never true
2417 raise Exception("No db connected")
2419 task: Final = asyncio.ensure_future(
2420 _load_team_membership_on_cache_miss(
2421 user_id=user_id,
2422 team_id=team_id,
2423 cache_key=_key,
2424 prisma_client=prisma_client,
2425 user_api_key_cache=user_api_key_cache,
2426 parent_otel_span=parent_otel_span,
2427 proxy_logging_obj=proxy_logging_obj,
2428 )
2429 )
2430 _team_membership_inflight[_key] = task
2432 def _clear_inflight(_done: object) -> None:
2433 if _team_membership_inflight.get(_key) is task: 2433 ↛ exitline 2433 didn't return from function '_clear_inflight' because the condition on line 2433 was always true
2434 _team_membership_inflight.pop(_key, None)
2436 task.add_done_callback(_clear_inflight)
2437 return _membership_from_shared_load(await asyncio.shield(task))
2440def model_in_access_group(model: str, team_models: list[str] | None, llm_router: Router | None) -> bool:
2441 from collections import defaultdict
2443 if team_models is None:
2444 return True
2445 if model in team_models:
2446 return True
2448 access_groups: dict[str, list[str]] = defaultdict(list)
2449 if llm_router:
2450 access_groups = llm_router.get_model_access_groups(model_name=model)
2452 if len(access_groups) > 0: # check if token contains any model access groups
2453 for idx, m in enumerate(
2454 team_models
2455 ): # loop token models, if any of them are an access group add the access group
2456 if m in access_groups:
2457 return True
2459 # Filter out models that are access_groups
2460 filtered_models: Final = [m for m in team_models if m not in access_groups]
2462 if model in filtered_models:
2463 return True
2465 return False
2468def _should_check_db(key: str, last_db_access_time: LimitedSizeOrderedDict, db_cache_expiry: int) -> bool:
2469 """
2470 Prevent calling db repeatedly for items that don't exist in the db.
2471 """
2472 current_time: Final = time.time()
2473 # if key doesn't exist in last_db_access_time -> check db
2474 if key not in last_db_access_time or last_db_access_time[key][0] is not None:
2475 return True
2476 elif last_db_access_time[key][0] is None: 2476 ↛ 2479line 2476 didn't jump to line 2479 because the condition on line 2476 was always true
2477 if current_time - last_db_access_time[key][1] >= db_cache_expiry:
2478 return True
2479 return False
2482def _update_last_db_access_time(key: str, value: object | None, last_db_access_time: LimitedSizeOrderedDict):
2483 last_db_access_time[key] = (value, time.time())
2486ROLE_BASED_PERMISSIONS_ADAPTER: Final[TypeAdapter[list[RoleBasedPermissions]]] = TypeAdapter(list[RoleBasedPermissions])
2489def _get_role_based_permissions(
2490 rbac_role: RBAC_ROLES,
2491 general_settings: Mapping[str, object],
2492 key: Literal["models", "routes"],
2493) -> list[str] | None:
2494 """
2495 Get the role based permissions from the general settings.
2496 """
2497 configured: Final = general_settings.get("role_permissions")
2498 if configured is None:
2499 return None
2501 for role_based_permission in ROLE_BASED_PERMISSIONS_ADAPTER.validate_python(configured):
2502 if role_based_permission.role == rbac_role:
2503 return role_based_permission.models if key == "models" else role_based_permission.routes
2505 return None
2508def get_role_based_models(
2509 rbac_role: RBAC_ROLES,
2510 general_settings: Mapping[str, object],
2511) -> list[str] | None:
2512 """
2513 Get the models allowed for a user role.
2515 Used by JWT Auth.
2516 """
2518 return _get_role_based_permissions(
2519 rbac_role=rbac_role,
2520 general_settings=general_settings,
2521 key="models",
2522 )
2525def get_role_based_routes(
2526 rbac_role: RBAC_ROLES,
2527 general_settings: Mapping[str, object],
2528) -> list[str] | None:
2529 """
2530 Get the routes allowed for a user role.
2531 """
2533 return _get_role_based_permissions(
2534 rbac_role=rbac_role,
2535 general_settings=general_settings,
2536 key="routes",
2537 )
2540async def _get_fuzzy_user_object(
2541 prisma_client: PrismaClient,
2542 sso_user_id: str | None = None,
2543 user_email: str | None = None,
2544) -> "_PrismaUserRow | None":
2545 """
2546 Checks if sso user is in db.
2548 Called when user id match is not found in db.
2550 - Check if sso_user_id is user_id in db
2551 - Check if sso_user_id is sso_user_id in db
2552 - Check if user_email is user_email in db
2553 - If not, create new user with user_email and sso_user_id and user_id = sso_user_id
2554 """
2556 response = None
2557 if sso_user_id is not None: 2557 ↛ 2558line 2557 didn't jump to line 2558 because the condition on line 2557 was never true
2558 response = await _user_table(UserRepository(prisma_client)).find_unique(
2559 where={"sso_user_id": sso_user_id},
2560 include={"organization_memberships": True},
2561 )
2563 if response is None and user_email is not None: 2563 ↛ 2566line 2563 didn't jump to line 2566 because the condition on line 2563 was never true
2564 # Use case-insensitive query to handle emails with different casing
2565 # This matches the pattern used in _check_duplicate_user_email
2566 response = await _user_table(UserRepository(prisma_client)).find_first(
2567 where={"user_email": {"equals": user_email, "mode": "insensitive"}},
2568 include={"organization_memberships": True},
2569 )
2571 if response is not None and sso_user_id is not None: # update sso_user_id
2572 asyncio.create_task( # background task to update user with sso id
2573 _user_table(UserRepository(prisma_client)).update(
2574 where={"user_id": response.user_id},
2575 data={"sso_user_id": sso_user_id},
2576 )
2577 )
2579 return response
2582async def _backfill_null_user_email(
2583 prisma_client: PrismaClient | None,
2584 user_api_key_cache: UserApiKeyCache,
2585 user_row: LiteLLM_UserTable,
2586 user_email: str | None,
2587) -> LiteLLM_UserTable:
2588 if user_email is None or user_row.user_email is not None or prisma_client is None: 2588 ↛ 2591line 2588 didn't jump to line 2591 because the condition on line 2588 was always true
2589 return user_row
2591 user_repo: Final = UserRepository(prisma_client)
2592 await user_repo.backfill_null_user_email(
2593 user_id=user_row.user_id,
2594 user_email=user_email,
2595 )
2596 db_row: Final = await user_repo.find_by_id(user_row.user_id)
2597 if db_row is None:
2598 return user_row
2599 email_update: Final = {"user_email": db_row.user_email} # mutable-ok: model_copy update payload is dict-shaped
2600 updated_row: Final = user_row.model_copy(update=email_update)
2601 await user_api_key_cache.async_set_cache(
2602 key=user_row.user_id,
2603 value=updated_row,
2604 model_type=LiteLLM_UserTable,
2605 ttl=get_management_object_ttl(user_api_key_cache),
2606 )
2607 return updated_row
2610@log_db_metrics
2611async def get_user_object(
2612 user_id: str | None,
2613 prisma_client: PrismaClient | None,
2614 user_api_key_cache: UserApiKeyCache,
2615 user_id_upsert: bool,
2616 parent_otel_span: Span | None = None,
2617 proxy_logging_obj: ProxyLogging | None = None,
2618 sso_user_id: str | None = None,
2619 user_email: str | None = None,
2620 check_db_only: bool | None = None,
2621) -> LiteLLM_UserTable | None:
2622 """
2623 - Check if user id in proxy User Table
2624 - if valid, return LiteLLM_UserTable object with defined limits
2625 - if not, then raise an error
2626 """
2628 if user_id is None: 2628 ↛ 2629line 2628 didn't jump to line 2629 because the condition on line 2628 was never true
2629 return None
2631 # check if in cache
2632 if not check_db_only: 2632 ↛ 2645line 2632 didn't jump to line 2645 because the condition on line 2632 was always true
2633 cached_user_obj: Final = await user_api_key_cache.async_get_cache(
2634 key=user_id,
2635 model_type=LiteLLM_UserTable,
2636 )
2637 if cached_user_obj is not None:
2638 return await _backfill_null_user_email(
2639 prisma_client=prisma_client,
2640 user_api_key_cache=user_api_key_cache,
2641 user_row=cached_user_obj,
2642 user_email=user_email,
2643 )
2644 # else, check db
2645 if prisma_client is None: 2645 ↛ 2646line 2645 didn't jump to line 2646 because the condition on line 2645 was never true
2646 raise Exception("No db connected")
2647 try:
2648 db_access_time_key: Final = f"user_id:{user_id}"
2649 should_check_db: Final = bool(check_db_only) or _should_check_db(
2650 key=db_access_time_key,
2651 last_db_access_time=last_db_access_time,
2652 db_cache_expiry=db_cache_expiry,
2653 )
2655 if should_check_db:
2656 response = await _user_table(UserRepository(prisma_client)).find_unique(
2657 where={"user_id": user_id}, include={"organization_memberships": True}
2658 )
2660 if response is None:
2661 response = await _get_fuzzy_user_object(
2662 prisma_client=prisma_client,
2663 sso_user_id=sso_user_id,
2664 user_email=user_email,
2665 )
2667 else:
2668 response = None
2670 if response is None:
2671 if user_id_upsert: 2671 ↛ 2672line 2671 didn't jump to line 2672 because the condition on line 2671 was never true
2672 from litellm.proxy.management_endpoints.internal_user_endpoints import (
2673 add_new_user_to_default_team,
2674 check_if_default_team_set,
2675 )
2677 default_params: Final = litellm.default_internal_user_params or {}
2678 scalar_default_params: Final = {
2679 key: value for key, value in default_params.items() if key not in ("teams", "available_teams")
2680 }
2681 new_user_params: Final[dict[str, Any]] = {
2682 "user_id": user_id,
2683 **({"user_email": user_email} if user_email is not None else {}),
2684 **scalar_default_params,
2685 }
2686 if (
2687 new_user_params.get("budget_duration") is not None
2688 and new_user_params.get("budget_reset_at") is None
2689 ):
2690 new_user_params["budget_reset_at"] = get_budget_reset_time(
2691 budget_duration=new_user_params["budget_duration"]
2692 )
2694 response = await _user_table(UserRepository(prisma_client)).create(
2695 data=new_user_params,
2696 include={"organization_memberships": True},
2697 )
2699 default_teams: Final = check_if_default_team_set()
2700 if default_teams:
2701 await add_new_user_to_default_team(
2702 user_id=user_id,
2703 user_email=user_email,
2704 user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
2705 teams=default_teams,
2706 prisma_client=prisma_client,
2707 )
2708 else:
2709 if should_check_db:
2710 _update_last_db_access_time(
2711 key=db_access_time_key,
2712 value=None,
2713 last_db_access_time=last_db_access_time,
2714 )
2715 raise UserNotFoundError(user_id=user_id)
2717 if response.organization_memberships is not None and len(response.organization_memberships) > 0: 2717 ↛ 2719line 2717 didn't jump to line 2719 because the condition on line 2717 was never true
2718 # dump each organization membership to type LiteLLM_OrganizationMembershipTable
2719 _dumped_memberships: Final = [
2720 LiteLLM_OrganizationMembershipTable.model_validate(membership.model_dump())
2721 for membership in response.organization_memberships
2722 if membership is not None
2723 ]
2724 response.organization_memberships = _dumped_memberships
2726 _response = LiteLLM_UserTable.model_validate(dict(response))
2727 _response = await _backfill_null_user_email(
2728 prisma_client=prisma_client,
2729 user_api_key_cache=user_api_key_cache,
2730 user_row=_response,
2731 user_email=user_email,
2732 )
2733 response_dict: Final = _response.model_dump()
2735 # save the user object to cache
2736 await user_api_key_cache.async_set_cache(
2737 key=user_id,
2738 value=_response,
2739 model_type=LiteLLM_UserTable,
2740 ttl=get_management_object_ttl(user_api_key_cache),
2741 )
2743 # save to db access time
2744 _update_last_db_access_time(
2745 key=db_access_time_key,
2746 value=response_dict,
2747 last_db_access_time=last_db_access_time,
2748 )
2750 return _response
2751 except UserNotFoundError:
2752 raise
2753 except Exception as e:
2754 _log_budget_lookup_failure("user", e)
2755 raise _user_read_failure(user_id=user_id, error=e)
2758def _user_read_failure(user_id: str, error: Exception) -> Exception:
2759 if PrismaDBExceptionHandler.is_database_service_unavailable_error(error):
2760 return error
2761 return ValueError(
2762 f"User doesn't exist in db. 'user_id'={user_id}. Create user via `/user/new` call. Got error - {error}"
2763 )
2766async def _cache_management_object(
2767 key: str,
2768 value: BaseModel | Mapping[str, object],
2769 user_api_key_cache: UserApiKeyCache,
2770 proxy_logging_obj: ProxyLogging | None,
2771 *,
2772 model_type: type[BaseModel],
2773):
2774 """
2775 Persist management objects via ``UserApiKeyCache`` (in-memory + optional Redis).
2777 ``UserApiKeyCache`` serializes with ``model_type`` so Redis and in-memory stay aligned.
2778 """
2779 await user_api_key_cache.async_set_cache(
2780 key=key,
2781 value=value,
2782 model_type=model_type,
2783 ttl=get_management_object_ttl(user_api_key_cache),
2784 )
2787async def _cache_team_object(
2788 team_id: str,
2789 team_table: LiteLLM_TeamTableCachedObj,
2790 user_api_key_cache: UserApiKeyCache,
2791 proxy_logging_obj: ProxyLogging | None,
2792):
2793 ## CACHE REFRESH TIME!
2794 team_table.last_refreshed_at = time.time()
2796 key: Final = f"team_id:{team_id}"
2798 if proxy_logging_obj is not None:
2799 try:
2800 await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache(key=key)
2801 except Exception as e: # noqa: BLE001 # best-effort invalidation: any cache backend error must not fail the write
2802 verbose_proxy_logger.warning(
2803 "Failed to invalidate internal usage cache entry %s; "
2804 "a stale team object may be served until its TTL expires: %s",
2805 key,
2806 e,
2807 )
2809 # team_id is the table primary key — guaranteed unique, safe to write.
2810 await _cache_management_object(
2811 key=key,
2812 value=team_table,
2813 user_api_key_cache=user_api_key_cache,
2814 proxy_logging_obj=proxy_logging_obj,
2815 model_type=LiteLLM_TeamTableCachedObj,
2816 )
2818 # Invalidate the alias-keyed cache so the JWT auth path with
2819 # `team_alias_jwt_field` (which reads via `get_team_object_by_alias`)
2820 # doesn't keep serving the pre-mutation team after every team-write
2821 # endpoint (team_model_add, team_model_delete, update_team, etc.).
2822 #
2823 # Why DELETE and not WRITE: `team_alias` has no UNIQUE constraint in
2824 # schema.prisma. Writing this cache from the generic refresh path
2825 # would let a team admin who renamed their team to collide with
2826 # another team's alias silently overwrite the cached team for
2827 # JWT-by-alias auth (veria-ai review on #28739). Deleting forces the
2828 # next reader through `get_team_object_by_alias`, which DOES enforce
2829 # uniqueness (len(teams) > 1 raises HTTPException) before populating
2830 # the cache from a verified single row.
2831 if team_table.team_alias:
2832 alias_key: Final = f"team_alias:{team_table.team_alias}"
2833 try:
2834 user_api_key_cache.delete_cache(key=alias_key)
2835 if proxy_logging_obj is not None:
2836 await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache(key=alias_key)
2837 except Exception as e: # noqa: BLE001 # best-effort invalidation: any cache backend error must not fail the mutation
2838 verbose_proxy_logger.warning(
2839 "Failed to invalidate cached team alias entry %s; "
2840 "a stale team object may be served until its TTL expires: %s",
2841 alias_key,
2842 e,
2843 )
2846async def invalidate_team_member_spend_state(
2847 user_id: str,
2848 team_id: str,
2849 user_api_key_cache: UserApiKeyCache,
2850 new_spend: float | None = None,
2851) -> None:
2852 """
2853 Clear every cached read path for one team member's budget so a spend
2854 reset or a raised cap takes effect on the next request instead of
2855 waiting on the membership cache's TTL.
2857 Two independently-keyed cache entries hold the same LiteLLM_TeamMembership
2858 row: user_api_key_auth.py's admission check writes ``{team_id}_{user_id}``,
2859 while budget_reservation.py's pre-call reservation and auth_checks.py's own
2860 get_team_membership() (used by _check_team_member_budget) both write
2861 ``team_membership:{user_id}:{team_id}``. Both formats must be invalidated
2862 explicitly; writing one does not refresh the other. All keys are also
2863 broadcast (LIT-3803): each worker's own in-memory copy (membership object,
2864 spend counter, or the counter's own short-TTL DB-floor marker) survives
2865 eviction elsewhere until its TTL, so the handling worker alone clearing its
2866 copy leaves every other worker still enforcing the pre-reset budget.
2868 ``new_spend`` is only passed by reset_team_member_spend_fn, which knows the
2869 exact post-reset value: it is SET everywhere (matching /key/{key}/reset_spend's
2870 own precedent) rather than deleted, so a worker's next read reflects it
2871 directly instead of re-deriving it through a DB reseed. team_member_update
2872 only changes the budget cap, not the tracked spend, so it passes no
2873 new_spend; the live spend counter is untouched in that case (deleting it
2874 would force a reseed from the DB's own spend column, which lags the live
2875 counter via periodic batch writes, briefly under-enforcing the raised cap
2876 against a spend value lower than what was actually tracked) and only the
2877 membership caches carrying the new cap are invalidated.
2879 The floor marker (``spend_db_floor:``, proxy_server.py's
2880 _authoritative_floor_spend) caches the pre-reset DB spend for
2881 SPEND_DB_FLOOR_CACHE_TTL_SECONDS; left stale after a real reset, a request
2882 landing on the pod that cached it can read that higher floor and raise the
2883 counter right back above the just-reset spend. It is overwritten here with
2884 the post-reset floor (not merely deleted) and _authoritative_floor_spend
2885 re-checks the marker after its DB read, so a floor read already in flight
2886 on this pod when the reset commits cannot clobber it with the pre-reset
2887 value. Both keys are broadcast as SETs carrying new_spend, not deletes:
2888 every subscriber (remote pods AND this pod's own, which receives its own
2889 message) writes the post-reset value, so the self-delivered message cannot
2890 erase the guard just written here.
2892 Raises HTTPException(503) if Redis still holds the stale pre-reset counter
2893 after both the SET and the fallback DELETE fail: budget checks read Redis
2894 first, so returning success would leave the old value authoritative for
2895 every worker despite the DB write having committed.
2896 """
2897 from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
2898 evict_and_broadcast,
2899 publish_auth_cache_invalidation,
2900 )
2902 inflight: Final[object] = _team_membership_inflight.pop(
2903 team_membership_reservation_cache_key(user_id=user_id, team_id=team_id), None
2904 )
2905 if isinstance(inflight, asyncio.Task) and inflight is not asyncio.current_task(): 2905 ↛ 2906line 2905 didn't jump to line 2906 because the condition on line 2905 was never true
2906 await asyncio.wait((inflight,))
2908 if new_spend is not None:
2909 from litellm.proxy.proxy_server import SPEND_DB_FLOOR_CACHE_TTL_SECONDS, spend_counter_cache
2911 spend_counter_key: Final = f"spend:team_member:{user_id}:{team_id}"
2912 spend_db_floor_key: Final = f"spend_db_floor:{spend_counter_key}"
2914 spend_counter_cache.in_memory_cache.set_cache(key=spend_counter_key, value=new_spend, ttl=60)
2915 if spend_counter_cache.redis_cache is not None: 2915 ↛ 2916line 2915 didn't jump to line 2916 because the condition on line 2915 was never true
2916 try:
2917 await spend_counter_cache.redis_cache.async_set_cache(key=spend_counter_key, value=new_spend, ttl=60)
2918 except Exception as e: # noqa: BLE001 # fall back to deleting the stale entry before giving up
2919 verbose_proxy_logger.warning(
2920 "Failed to set spend counter %s in Redis after reset: %s; deleting it instead so the next "
2921 "read reseeds from the DB rather than keeping the stale pre-reset value authoritative",
2922 spend_counter_key,
2923 e,
2924 )
2925 try:
2926 await spend_counter_cache.redis_cache.async_delete_cache(key=spend_counter_key)
2927 except Exception: # noqa: BLE001 # stale value now authoritative in Redis; surface instead of reporting success
2928 verbose_proxy_logger.warning(
2929 "Failed to delete stale spend counter %s in Redis after a failed reset write",
2930 spend_counter_key,
2931 exc_info=True,
2932 )
2933 raise HTTPException(
2934 status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
2935 detail={ # mutable-ok: HTTPException.detail takes a dict
2936 "error": "Spend was reset in the database, but Redis is unreachable and still "
2937 "holds the pre-reset counter. Retry once Redis is reachable."
2938 },
2939 ) from e
2941 spend_counter_cache.in_memory_cache.set_cache(
2942 key=spend_db_floor_key,
2943 value=new_spend,
2944 ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS,
2945 )
2946 await publish_auth_cache_invalidation(cache_key=spend_counter_key, new_value=new_spend, ttl=60)
2947 await publish_auth_cache_invalidation(
2948 cache_key=spend_db_floor_key,
2949 new_value=new_spend,
2950 ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS,
2951 )
2953 await evict_and_broadcast(
2954 cache_keys=(
2955 team_membership_auth_cache_key(team_id=team_id, user_id=user_id),
2956 team_membership_reservation_cache_key(user_id=user_id, team_id=team_id),
2957 ),
2958 user_api_key_cache=user_api_key_cache,
2959 )
2962async def delete_cache_team_object(
2963 team_id: str,
2964 team_alias: str | None,
2965 user_api_key_cache: UserApiKeyCache,
2966 proxy_logging_obj: ProxyLogging | None,
2967) -> None:
2968 """
2969 Evict both keys `_cache_team_object` writes.
2971 `get_team_object` reads the id key and the JWT `team_alias_jwt_field` path reads the alias key,
2972 so leaving either behind keeps a deleted team resolvable for auth until its TTL expires.
2974 Mirrors `delete_cached_project_object`: evicting locally only reaches the worker handling the
2975 delete, so every key is also broadcast to drop the other workers' in-memory copies.
2977 Eviction is best-effort, matching `_cache_team_object`. `delete_team` calls this after the team
2978 rows are already gone, so letting an unreachable cache backend raise here would fail a request
2979 whose delete has committed.
2980 """
2981 keys: Final = (f"team_id:{team_id}", *((f"team_alias:{team_alias}",) if team_alias else ()))
2983 for key in keys:
2984 try:
2985 user_api_key_cache.delete_cache(key=key)
2987 ## UPDATE REDIS CACHE ##
2988 if proxy_logging_obj is not None:
2989 await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache(key=key)
2990 except Exception as e: # noqa: BLE001 # best-effort invalidation: any cache backend error must not abort the delete
2991 verbose_proxy_logger.warning(
2992 "Failed to invalidate cached team entry %s on delete; "
2993 "a deleted team may be served until its TTL expires: %s",
2994 key,
2995 e,
2996 )
2997 await publish_auth_cache_invalidation(cache_key=key)
3000async def _cache_key_object(
3001 hashed_token: str,
3002 user_api_key_obj: UserAPIKeyAuth,
3003 user_api_key_cache: UserApiKeyCache,
3004 proxy_logging_obj: ProxyLogging | None,
3005):
3006 key: Final = hashed_token
3008 ## CACHE REFRESH TIME
3009 user_api_key_obj.last_refreshed_at = time.time()
3011 cached_key_obj: Final = _copy_user_api_key_auth_for_cache(user_api_key_obj=user_api_key_obj)
3012 await _cache_management_object(
3013 key=key,
3014 value=cached_key_obj,
3015 user_api_key_cache=user_api_key_cache,
3016 proxy_logging_obj=proxy_logging_obj,
3017 model_type=UserAPIKeyAuth,
3018 )
3021async def _delete_cache_key_object(
3022 hashed_token: str,
3023 user_api_key_cache: UserApiKeyCache,
3024 proxy_logging_obj: ProxyLogging | None,
3025):
3026 """
3027 Evict one key object, best-effort, matching `delete_cache_team_object` and
3028 `delete_cache_key_objects`.
3030 Every caller runs this after its own write has already committed, and the in-memory entry is
3031 dropped before the Redis round trip. Letting a cache-backend error raise here therefore reports
3032 failure for work that succeeded without making the cache any less stale; the leftover Redis
3033 entry expires at its TTL either way.
3035 Also broadcasts the eviction to every other worker (LIT-3803): auth serves this object
3036 cache-first with no freshness check, so a worker that never receives the broadcast keeps
3037 admitting requests against the pre-mutation object (e.g. a just-reset spend) until its own
3038 copy's TTL expires.
3039 """
3040 key: Final = hashed_token
3042 try:
3043 user_api_key_cache.delete_cache(key=key)
3045 ## UPDATE REDIS CACHE ##
3046 if proxy_logging_obj is not None: 3046 ↛ 3055line 3046 didn't jump to line 3055 because the condition on line 3046 was always true
3047 await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache(key=key)
3048 except Exception as e: # noqa: BLE001 # best-effort: a cache error must not fail a committed write
3049 verbose_proxy_logger.warning(
3050 "Failed to invalidate cached key entry %s; a stale key object may be served until its TTL expires: %s",
3051 key,
3052 e,
3053 )
3055 await publish_auth_cache_invalidation(cache_key=key)
3058async def delete_cache_key_objects(
3059 hashed_tokens: Sequence[str],
3060 user_api_key_cache: UserApiKeyCache,
3061 proxy_logging_obj: ProxyLogging | None,
3062) -> None:
3063 """
3064 Evict a batch of key objects, for callers that delete keys in bulk rather than through
3065 `/key/delete`. Auth resolves a cached key object without re-reading its team, so a key left
3066 cached after its row is gone keeps buying access until its TTL expires.
3068 Evicting locally only reaches this worker; `_delete_cache_key_object` itself broadcasts each
3069 token, so a deleted key left in a peer worker's in-memory cache still authenticates there until
3070 its TTL expires.
3072 Best-effort per key: the rows are already deleted by the time this runs, so an unreachable
3073 cache backend must not abort the caller partway through its own cascade.
3074 """
3075 results: Final = await asyncio.gather(
3076 *(
3077 _delete_cache_key_object(
3078 hashed_token=hashed_token,
3079 user_api_key_cache=user_api_key_cache,
3080 proxy_logging_obj=proxy_logging_obj,
3081 )
3082 for hashed_token in hashed_tokens
3083 ),
3084 return_exceptions=True,
3085 )
3087 for hashed_token, result in zip(hashed_tokens, results): 3087 ↛ 3088line 3087 didn't jump to line 3088 because the loop on line 3087 never started
3088 if isinstance(result, BaseException):
3089 verbose_proxy_logger.warning(
3090 "Failed to evict cached key entry for %s; a deleted key may authenticate until its TTL expires: %s",
3091 hashed_token,
3092 result,
3093 )
3096class _TeamNotFoundDetail(TypedDict):
3097 error: ReadOnly[str]
3100class TeamNotFoundError(HTTPException):
3101 """The team row is provably absent, as opposed to merely unreadable.
3103 ``get_team_object`` reports every failure as a 404, so a deleted team and a
3104 database that would not answer are indistinguishable to its callers. Callers
3105 that must not treat a degraded read as a definitive answer, such as the
3106 authorization fallback in ``user_api_key_auth``, key on this subclass. It
3107 stays a 404 carrying the same detail, so every other caller is unaffected.
3108 """
3110 def __init__(self, team_id: str) -> None:
3111 detail: Final[_TeamNotFoundDetail] = {
3112 "error": f"Team doesn't exist in db. Team={team_id}. Create team via `/team/new` call."
3113 }
3114 super().__init__(status_code=404, detail=detail)
3117@log_db_metrics
3118async def _get_team_db_check(
3119 team_id: str, prisma_client: PrismaClient, team_id_upsert: bool | None = None
3120) -> "_PrismaTeamRow | None":
3121 response = await _team_table(TeamRepository(prisma_client)).find_unique(
3122 where={"team_id": team_id}, include=_TEAM_GRANT_RELATIONS
3123 )
3125 if response is None and team_id_upsert: 3125 ↛ 3126line 3125 didn't jump to line 3126 because the condition on line 3125 was never true
3126 from litellm.proxy.management_endpoints.team_endpoints import new_team
3128 new_team_data: Final = NewTeamRequest(team_id=team_id)
3130 mock_request: Final = Request(scope={"type": "http"})
3131 system_admin_user: Final = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
3133 created_team_dict: Final = await new_team(
3134 data=new_team_data,
3135 http_request=mock_request,
3136 user_api_key_dict=system_admin_user,
3137 )
3138 response = LiteLLM_TeamTable.model_validate(created_team_dict)
3139 return response
3142async def _get_team_object_from_db(team_id: str, prisma_client: PrismaClient) -> "_PrismaTeamRow | None":
3143 return await _team_table(TeamRepository(prisma_client)).find_unique(where={"team_id": team_id})
3146async def _get_team_object_from_user_api_key_cache(
3147 team_id: str,
3148 prisma_client: PrismaClient,
3149 user_api_key_cache: UserApiKeyCache,
3150 last_db_access_time: LimitedSizeOrderedDict,
3151 db_cache_expiry: int,
3152 proxy_logging_obj: ProxyLogging | None,
3153 key: str,
3154 team_id_upsert: bool | None = None,
3155) -> LiteLLM_TeamTableCachedObj:
3156 db_access_time_key: Final = key
3157 should_check_db: Final = _should_check_db(
3158 key=db_access_time_key,
3159 last_db_access_time=last_db_access_time,
3160 db_cache_expiry=db_cache_expiry,
3161 )
3162 if should_check_db: 3162 ↛ 3169line 3162 didn't jump to line 3169 because the condition on line 3162 was always true
3163 response = await _get_team_db_check(team_id=team_id, prisma_client=prisma_client, team_id_upsert=team_id_upsert)
3164 # The database answered and the row is not there. Distinct from every
3165 # other failure here, which leaves the team's grant unknown.
3166 if response is None:
3167 raise TeamNotFoundError(team_id=team_id)
3168 else:
3169 response = None
3171 if response is None: 3171 ↛ 3172line 3171 didn't jump to line 3172 because the condition on line 3171 was never true
3172 raise Exception
3174 _response: Final = LiteLLM_TeamTableCachedObj.model_validate(response.dict())
3176 # Load object_permission if object_permission_id exists but object_permission is not loaded
3177 if _response.object_permission_id and not _response.object_permission:
3178 try:
3179 _response.object_permission = await get_object_permission(
3180 object_permission_id=_response.object_permission_id,
3181 prisma_client=prisma_client,
3182 user_api_key_cache=user_api_key_cache,
3183 parent_otel_span=None,
3184 proxy_logging_obj=proxy_logging_obj,
3185 )
3186 except Exception as e:
3187 verbose_proxy_logger.debug(
3188 "Failed to load object_permission for team %s with object_permission_id=%s: %s",
3189 team_id,
3190 _response.object_permission_id,
3191 e,
3192 )
3194 # save the team object to cache
3195 await _cache_team_object(
3196 team_id=team_id,
3197 team_table=_response,
3198 user_api_key_cache=user_api_key_cache,
3199 proxy_logging_obj=proxy_logging_obj,
3200 )
3202 # save to db access time
3203 _update_last_db_access_time(
3204 key=db_access_time_key,
3205 value=_response,
3206 last_db_access_time=last_db_access_time,
3207 )
3209 return _response
3212async def _get_team_object_from_cache(
3213 key: str,
3214 user_api_key_cache: UserApiKeyCache,
3215 parent_otel_span: Span | None,
3216) -> LiteLLM_TeamTableCachedObj | None:
3217 decoded: Final = await user_api_key_cache.async_get_cache(
3218 key=key,
3219 parent_otel_span=parent_otel_span,
3220 model_type=LiteLLM_TeamTableCachedObj,
3221 )
3222 return decoded
3225async def get_team_object(
3226 team_id: str,
3227 prisma_client: PrismaClient | None,
3228 user_api_key_cache: UserApiKeyCache,
3229 parent_otel_span: Span | None = None,
3230 proxy_logging_obj: ProxyLogging | None = None,
3231 check_cache_only: bool | None = None,
3232 check_db_only: bool | None = None,
3233 team_id_upsert: bool | None = None,
3234) -> LiteLLM_TeamTableCachedObj:
3235 """
3236 - Check if team id in proxy Team Table
3237 - if valid, return LiteLLM_TeamTable object with defined limits
3238 - if not, then raise an error
3240 Raises:
3241 - HTTPException: If team doesn't exist in db or cache (status_code=404)
3242 """
3243 if prisma_client is None: 3243 ↛ 3244line 3243 didn't jump to line 3244 because the condition on line 3243 was never true
3244 raise Exception("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys")
3246 # check if in cache
3247 key: Final = f"team_id:{team_id}"
3249 if not check_db_only:
3250 cached_team_obj: Final = await _get_team_object_from_cache(
3251 key=key,
3252 user_api_key_cache=user_api_key_cache,
3253 parent_otel_span=parent_otel_span,
3254 )
3256 if cached_team_obj is not None:
3257 return cached_team_obj
3259 if check_cache_only: 3259 ↛ 3260line 3259 didn't jump to line 3260 because the condition on line 3259 was never true
3260 raise HTTPException(
3261 status_code=404,
3262 detail={"error": f"Team doesn't exist in cache + check_cache_only=True. Team={team_id}."},
3263 )
3265 # else, check db
3266 try:
3267 return await _get_team_object_from_user_api_key_cache(
3268 team_id=team_id,
3269 prisma_client=prisma_client,
3270 user_api_key_cache=user_api_key_cache,
3271 proxy_logging_obj=proxy_logging_obj,
3272 last_db_access_time=last_db_access_time,
3273 db_cache_expiry=db_cache_expiry,
3274 key=key,
3275 team_id_upsert=team_id_upsert,
3276 )
3277 except TeamNotFoundError:
3278 raise
3279 except Exception:
3280 raise HTTPException(
3281 status_code=404,
3282 detail={"error": f"Team doesn't exist in db. Team={team_id}. Create team via `/team/new` call."},
3283 )
3286async def _cache_access_object(
3287 access_group_id: str,
3288 access_group_table: LiteLLM_AccessGroupTable,
3289 user_api_key_cache: UserApiKeyCache,
3290 proxy_logging_obj: ProxyLogging | None = None,
3291):
3292 key: Final = f"access_group_id:{access_group_id}"
3293 await user_api_key_cache.async_set_cache(
3294 key=key,
3295 value=access_group_table,
3296 model_type=LiteLLM_AccessGroupTable,
3297 ttl=DEFAULT_ACCESS_GROUP_CACHE_TTL,
3298 )
3301async def _delete_cache_access_object(
3302 access_group_id: str,
3303 user_api_key_cache: UserApiKeyCache,
3304 proxy_logging_obj: ProxyLogging | None = None,
3305):
3306 key: Final = f"access_group_id:{access_group_id}"
3308 user_api_key_cache.delete_cache(key=key)
3310 ## UPDATE REDIS CACHE ##
3311 if proxy_logging_obj is not None: 3311 ↛ exitline 3311 didn't return from function '_delete_cache_access_object' because the condition on line 3311 was always true
3312 await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache(key=key)
3315@log_db_metrics
3316async def get_access_object(
3317 access_group_id: str,
3318 prisma_client: DatabaseClient | None,
3319 user_api_key_cache: UserApiKeyCache,
3320 proxy_logging_obj: ProxyLogging | None = None,
3321) -> LiteLLM_AccessGroupTable:
3322 """
3323 - Check if access_group_id in proxy AccessGroupTable
3324 - Always checks cache first, then DB only when not found in cache
3325 - if valid, return LiteLLM_AccessGroupTable object
3326 - if not, then raise an error
3328 Unlike get_team_object, this has no check_cache_only or check_db_only flags;
3329 it always follows cache-first-then-db semantics.
3331 Raises:
3332 - HTTPException: If access group doesn't exist in db or cache (status_code=404)
3333 """
3334 if prisma_client is None: 3334 ↛ 3335line 3334 didn't jump to line 3335 because the condition on line 3334 was never true
3335 raise Exception("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys")
3337 key: Final = f"access_group_id:{access_group_id}"
3339 cached_access_obj: Final = await user_api_key_cache.async_get_cache(
3340 key=key,
3341 model_type=LiteLLM_AccessGroupTable,
3342 )
3343 if cached_access_obj is not None: 3343 ↛ 3344line 3343 didn't jump to line 3344 because the condition on line 3343 was never true
3344 return cached_access_obj
3346 # Not in cache - fetch from DB
3347 try:
3348 response: Final = await _dictable_table(AccessGroupRepository(prisma_client), "access_group").find_unique(
3349 where={"access_group_id": access_group_id}
3350 )
3352 if response is None: 3352 ↛ 3358line 3352 didn't jump to line 3358 because the condition on line 3352 was always true
3353 raise HTTPException(
3354 status_code=404,
3355 detail={"error": f"Access group doesn't exist in db. Access group={access_group_id}."},
3356 )
3358 _response: Final = LiteLLM_AccessGroupTable.model_validate(response.dict())
3360 # Save to cache
3361 await _cache_access_object(
3362 access_group_id=access_group_id,
3363 access_group_table=_response,
3364 user_api_key_cache=user_api_key_cache,
3365 proxy_logging_obj=proxy_logging_obj,
3366 )
3368 return _response
3369 except HTTPException:
3370 raise
3371 except Exception as e:
3372 verbose_proxy_logger.exception(
3373 "Error getting access group for access_group_id: %s",
3374 access_group_id,
3375 )
3376 raise HTTPException(
3377 status_code=404,
3378 detail={"error": f"Access group doesn't exist in db. Access group={access_group_id}. Error: {e}"},
3379 )
3382@log_db_metrics
3383async def get_team_object_by_alias(
3384 team_alias: str,
3385 prisma_client: PrismaClient | None,
3386 user_api_key_cache: UserApiKeyCache,
3387 parent_otel_span: Optional["Span"] = None,
3388 proxy_logging_obj: ProxyLogging | None = None,
3389) -> LiteLLM_TeamTableCachedObj:
3390 """
3391 Look up a team by its team_alias (name) in the database.
3393 Args:
3394 team_alias: The team name/alias to look up
3395 prisma_client: Database client
3396 user_api_key_cache: Cache for storing results
3397 parent_otel_span: Optional OpenTelemetry span
3398 proxy_logging_obj: Optional proxy logging object
3400 Returns:
3401 LiteLLM_TeamTableCachedObj: The team object if found
3403 Raises:
3404 HTTPException: If team doesn't exist or multiple teams have the same alias
3405 """
3406 if prisma_client is None:
3407 raise Exception("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys")
3409 # Check cache first (keyed by alias)
3410 cache_key: Final = f"team_alias:{team_alias}"
3412 cached_team_obj: Final = await _get_team_object_from_cache(
3413 key=cache_key,
3414 user_api_key_cache=user_api_key_cache,
3415 parent_otel_span=parent_otel_span,
3416 )
3418 if cached_team_obj is not None:
3419 return cached_team_obj
3421 # Query database by team_alias
3422 try:
3423 teams: Final = await _team_table(TeamRepository(prisma_client)).find_many(
3424 where={"team_alias": team_alias}, include=_TEAM_GRANT_RELATIONS
3425 )
3427 if not teams:
3428 raise HTTPException(
3429 status_code=404,
3430 detail={
3431 "error": f"Team with alias '{team_alias}' doesn't exist in db. Create team via `/team/new` call."
3432 },
3433 )
3435 if len(teams) > 1:
3436 raise HTTPException(
3437 status_code=400,
3438 detail={
3439 "error": f"Multiple teams found with alias '{team_alias}'. Please use team_id_jwt_field instead or ensure team aliases are unique."
3440 },
3441 )
3443 team: Final = teams[0]
3444 team_obj: Final = LiteLLM_TeamTableCachedObj.model_validate(team.model_dump())
3446 # Load object_permission if object_permission_id exists but object_permission is not loaded
3447 if team_obj.object_permission_id and not team_obj.object_permission:
3448 try:
3449 team_obj.object_permission = await get_object_permission(
3450 object_permission_id=team_obj.object_permission_id,
3451 prisma_client=prisma_client,
3452 user_api_key_cache=user_api_key_cache,
3453 parent_otel_span=parent_otel_span,
3454 proxy_logging_obj=proxy_logging_obj,
3455 )
3456 except Exception as e:
3457 verbose_proxy_logger.debug(
3458 "Failed to load object_permission for team %s with object_permission_id=%s: %s",
3459 team_obj.team_id,
3460 team_obj.object_permission_id,
3461 e,
3462 )
3464 # Cache the result by both alias and team_id
3465 await user_api_key_cache.async_set_cache(
3466 key=cache_key,
3467 value=team_obj,
3468 model_type=LiteLLM_TeamTableCachedObj,
3469 ttl=DEFAULT_IN_MEMORY_TTL,
3470 )
3471 # Also cache by team_id for consistency
3472 team_id_cache_key: Final = f"team_id:{team_obj.team_id}"
3473 await user_api_key_cache.async_set_cache(
3474 key=team_id_cache_key,
3475 value=team_obj,
3476 model_type=LiteLLM_TeamTableCachedObj,
3477 ttl=DEFAULT_IN_MEMORY_TTL,
3478 )
3480 return team_obj
3482 except HTTPException:
3483 raise
3484 except Exception as e:
3485 verbose_proxy_logger.exception("Error looking up team by alias: %s", team_alias)
3486 raise HTTPException(
3487 status_code=500,
3488 detail={"error": f"Error looking up team by alias '{team_alias}': {e}"},
3489 )
3492@log_db_metrics
3493async def get_org_object_by_alias(
3494 org_alias: str,
3495 prisma_client: PrismaClient | None,
3496 user_api_key_cache: UserApiKeyCache,
3497 parent_otel_span: Optional["Span"] = None,
3498 proxy_logging_obj: ProxyLogging | None = None,
3499) -> LiteLLM_OrganizationTable | None:
3500 """
3501 Look up an organization by its organization_alias in the database.
3503 Args:
3504 org_alias: The organization name/alias to look up
3505 prisma_client: Database client
3506 user_api_key_cache: Cache for storing results
3507 parent_otel_span: Optional OpenTelemetry span
3508 proxy_logging_obj: Optional proxy logging object
3510 Returns:
3511 LiteLLM_OrganizationTable if found, None otherwise
3513 Raises:
3514 HTTPException: If organization not found or multiple orgs have the same alias
3515 """
3516 if prisma_client is None:
3517 raise Exception("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys")
3519 # Check cache first (keyed by alias)
3520 cache_key: Final = f"org_alias:{org_alias}"
3521 cached_org_obj: Final = await user_api_key_cache.async_get_cache(
3522 key=cache_key,
3523 model_type=LiteLLM_OrganizationTable,
3524 )
3525 if cached_org_obj is not None:
3526 return cached_org_obj
3528 # Query database by organization_alias
3529 try:
3530 orgs = await _model_dump_table(OrganizationRepository(prisma_client), "organization").find_many(
3531 where={"organization_alias": org_alias}
3532 )
3534 if not orgs:
3535 raise HTTPException(
3536 status_code=404,
3537 detail={
3538 "error": f"Organization with alias '{org_alias}' doesn't exist in db. Create organization via `/organization/new` call."
3539 },
3540 )
3542 if len(orgs) > 1:
3543 raise HTTPException(
3544 status_code=400,
3545 detail={
3546 "error": f"Multiple organizations found with alias '{org_alias}'. Please use org_id_jwt_field instead or ensure organization aliases are unique."
3547 },
3548 )
3550 org: Final = orgs[0]
3551 org_obj: Final = LiteLLM_OrganizationTable.model_validate(org.model_dump())
3553 # Cache the result
3554 await user_api_key_cache.async_set_cache(
3555 key=cache_key,
3556 value=org_obj,
3557 model_type=LiteLLM_OrganizationTable,
3558 ttl=DEFAULT_IN_MEMORY_TTL,
3559 )
3560 # Also cache by org_id for consistency
3561 await user_api_key_cache.async_set_cache(
3562 key=f"org_id:{org_obj.organization_id}",
3563 value=org_obj,
3564 model_type=LiteLLM_OrganizationTable,
3565 ttl=DEFAULT_IN_MEMORY_TTL,
3566 )
3568 return org_obj
3570 except HTTPException:
3571 raise
3572 except Exception as e:
3573 verbose_proxy_logger.exception("Error looking up organization by alias: %s", org_alias)
3574 raise HTTPException(
3575 status_code=500,
3576 detail={"error": f"Error looking up organization by alias '{org_alias}': {e}"},
3577 )
3580LITELLM_SESSION_TOKEN_PREFIX: Final = "litellm_login_"
3583class ExperimentalUIJWTToken:
3584 @staticmethod
3585 def get_experimental_ui_login_jwt_auth_token(user_info: LiteLLM_UserTable) -> str:
3586 from datetime import timedelta
3588 from litellm.proxy.common_utils.encrypt_decrypt_utils import (
3589 encrypt_bearer_token,
3590 )
3592 if user_info.user_role is None:
3593 raise Exception("User role is required for experimental UI login")
3595 # Experimental UI flow uses fixed 10-min expiry for security (does not use LITELLM_UI_SESSION_DURATION)
3596 expiration_time: Final = get_utc_datetime() + timedelta(minutes=10)
3598 # Format the expiration time as ISO 8601 string
3599 expires: Final = expiration_time.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "+00:00"
3601 valid_token: Final = UserAPIKeyAuth(
3602 token="ui-token",
3603 key_name="ui-token",
3604 key_alias="ui-token",
3605 max_budget=litellm.max_ui_session_budget,
3606 rpm_limit=100, # allow user to have a conversation on test key pane of UI
3607 expires=expires,
3608 user_id=user_info.user_id,
3609 team_id="litellm-dashboard",
3610 models=user_info.models,
3611 max_parallel_requests=None,
3612 user_role=LitellmUserRoles(user_info.user_role),
3613 )
3615 return encrypt_bearer_token(valid_token.model_dump_json(exclude_none=True), prefix=LITELLM_SESSION_TOKEN_PREFIX)
3617 @staticmethod
3618 def get_cli_jwt_auth_token(
3619 user_info: LiteLLM_UserTable,
3620 team_id: str | None = None,
3621 team_alias: str | None = None,
3622 team_models: Sequence[str] | None = None,
3623 team_model_aliases: Mapping[str, str] | None = None,
3624 max_budget: float | None = None,
3625 ) -> str:
3626 """
3627 Generate a JWT token for CLI authentication with configurable expiration.
3629 The expiration time can be controlled via the LITELLM_CLI_JWT_EXPIRATION_HOURS
3630 environment variable (defaults to 24 hours).
3632 Args:
3633 user_info: User information from the database
3634 team_id: Team ID for the user (optional, uses user's team if available)
3635 team_alias: Team alias for the selected team, if available
3636 team_models: Model allowlist granted by the selected team
3637 team_model_aliases: Team model aliases for the selected team
3639 Returns:
3640 Encrypted JWT token string
3641 """
3642 import secrets
3643 from datetime import timedelta
3645 from litellm.proxy.common_utils.encrypt_decrypt_utils import (
3646 encrypt_bearer_token,
3647 )
3649 if user_info.user_role is None:
3650 raise Exception("User role is required for CLI JWT login")
3652 # Calculate expiration time (configurable via LITELLM_CLI_JWT_EXPIRATION_HOURS env var)
3653 expiration_time: Final = get_utc_datetime() + timedelta(hours=CLI_JWT_EXPIRATION_HOURS)
3655 # Format the expiration time as ISO 8601 string
3656 expires: Final = expiration_time.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "+00:00"
3658 # Use provided team_id, or fall back to user's teams if available
3659 _team_id = team_id
3660 if _team_id is None and hasattr(user_info, "teams") and user_info.teams:
3661 # Use first team if user has teams
3662 _team_id = user_info.teams[0] if len(user_info.teams) > 0 else None
3664 session_token: Final = f"{CLI_SESSION_KEY_PREFIX}-{secrets.token_urlsafe(16)}"
3665 session_alias: Final = f"{CLI_SESSION_KEY_PREFIX}-{user_info.user_id}"
3667 valid_token: Final = UserAPIKeyAuth(
3668 token=session_token,
3669 key_name=session_alias,
3670 key_alias=session_alias,
3671 expires=expires,
3672 max_budget=max_budget,
3673 user_id=user_info.user_id,
3674 team_id=_team_id,
3675 team_alias=team_alias,
3676 team_models=list(team_models) if team_models is not None else [],
3677 team_model_aliases=dict(team_model_aliases) if team_model_aliases is not None else None,
3678 models=[] if _team_id is not None else user_info.models,
3679 max_parallel_requests=None,
3680 user_role=LitellmUserRoles(user_info.user_role),
3681 is_session_token=True,
3682 )
3684 return encrypt_bearer_token(valid_token.model_dump_json(exclude_none=True), prefix=LITELLM_SESSION_TOKEN_PREFIX)
3686 @staticmethod
3687 def get_key_object_from_ui_hash_key(
3688 hashed_token: str,
3689 ) -> UserAPIKeyAuth | None:
3690 import json
3692 from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
3693 from litellm.proxy.common_utils.encrypt_decrypt_utils import (
3694 decrypt_bearer_token,
3695 )
3697 decrypted_token: Final = decrypt_bearer_token(hashed_token, prefix=LITELLM_SESSION_TOKEN_PREFIX)
3698 if decrypted_token is None: 3698 ↛ 3700line 3698 didn't jump to line 3700 because the condition on line 3698 was always true
3699 return None
3700 try:
3701 return UserAPIKeyAuth.model_validate(json.loads(decrypted_token))
3702 except Exception as e:
3703 raise Exception(f"Invalid hash key. Hash key={hashed_token}. Decrypted token={decrypted_token}. Error: {e}")
3706async def _fetch_key_object_from_db_with_reconnect(
3707 hashed_token: str,
3708 prisma_client: PrismaClient,
3709 parent_otel_span: Span | None,
3710 proxy_logging_obj: ProxyLogging | None,
3711 deadline_seconds: float | None = None,
3712) -> BaseModel | None:
3713 """
3714 Fetch key object from DB and retry once if a DB connection error can be healed.
3715 The gate wait, the query, the reconnect, and the retry share one deadline, so a
3716 stalled database fails the request with ``DBLookupDeadlineExceeded`` instead of
3717 parking it.
3718 """
3719 return await bounded_db_lookup(
3720 _fetch_key_object_from_db_unbounded(
3721 hashed_token=hashed_token,
3722 prisma_client=prisma_client,
3723 parent_otel_span=parent_otel_span,
3724 proxy_logging_obj=proxy_logging_obj,
3725 ),
3726 name="key",
3727 deadline_seconds=deadline_seconds,
3728 )
3731async def _fetch_key_object_from_db_unbounded(
3732 hashed_token: str,
3733 prisma_client: PrismaClient,
3734 parent_otel_span: Span | None,
3735 proxy_logging_obj: ProxyLogging | None,
3736) -> BaseModel | None:
3737 async with db_lookup_gate.current():
3738 try:
3739 return await prisma_client.get_data(
3740 token=hashed_token,
3741 table_name="combined_view",
3742 parent_otel_span=parent_otel_span,
3743 proxy_logging_obj=proxy_logging_obj,
3744 )
3745 except Exception as e:
3746 if PrismaDBExceptionHandler.is_database_transport_error(e):
3747 did_reconnect = False
3748 if hasattr(prisma_client, "attempt_db_reconnect"):
3749 auth_reconnect_timeout = getattr(prisma_client, "_db_auth_reconnect_timeout_seconds", 2.0)
3750 if not isinstance(auth_reconnect_timeout, (int, float)):
3751 auth_reconnect_timeout = 2.0
3752 auth_reconnect_lock_timeout = getattr(prisma_client, "_db_auth_reconnect_lock_timeout_seconds", 0.1)
3753 if not isinstance(auth_reconnect_lock_timeout, (int, float)):
3754 auth_reconnect_lock_timeout = 0.1
3755 did_reconnect = await prisma_client.attempt_db_reconnect(
3756 reason="auth_get_key_object_lookup_failure",
3757 timeout_seconds=auth_reconnect_timeout,
3758 lock_timeout_seconds=auth_reconnect_lock_timeout,
3759 )
3760 if did_reconnect:
3761 return await prisma_client.get_data(
3762 token=hashed_token,
3763 table_name="combined_view",
3764 parent_otel_span=parent_otel_span,
3765 proxy_logging_obj=proxy_logging_obj,
3766 )
3767 raise
3770def jwt_key_mapping_cache_key(jwt_claim_name: str, jwt_claim_value: str, jwt_issuer: str | None = None) -> str:
3771 """Cache key under which a JWT-claim-to-key mapping is stored, scoped to one
3772 issuer (or the issuer-agnostic/global scope when ``jwt_issuer`` is falsy).
3774 Scoped by issuer (when one is configured) so a cached hit or ``__NO_MAPPING__`` miss
3775 for one issuer's claim value can never be served to a different issuer whose claim
3776 value happens to collide. Unchanged for the global scope, keeping the single-issuer
3777 (no ``litellm_jwtauth.issuers`` configured) cache key format stable across this fix.
3778 """
3779 if not jwt_issuer: 3779 ↛ 3781line 3779 didn't jump to line 3781 because the condition on line 3779 was always true
3780 return f"jwt_key_mapping:{jwt_claim_name}:{jwt_claim_value}"
3781 return f"jwt_key_mapping:{jwt_issuer}:{jwt_claim_name}:{jwt_claim_value}"
3784@log_db_metrics
3785async def get_jwt_key_mapping_cache_keys_for_token(
3786 hashed_token: str,
3787 prisma_client: PrismaClient,
3788) -> tuple[str, ...]:
3789 """Cache keys of every JWT claim mapped to the given virtual key."""
3790 mappings: Final = await _jwt_key_mapping_table(JWTKeyMappingRepository(prisma_client)).find_many(
3791 where={"token": hashed_token}
3792 )
3793 return tuple(jwt_key_mapping_cache_key(m.jwt_claim_name, m.jwt_claim_value, m.jwt_issuer) for m in mappings)
3796class _TokenInFilter(TypedDict):
3797 token: ReadOnly[Mapping[str, Sequence[str]]]
3800async def get_jwt_key_mapping_cache_keys_for_tokens(
3801 hashed_tokens: Sequence[str],
3802 prisma_client: PrismaClient,
3803) -> tuple[str, ...]:
3804 """Cache keys of every JWT claim mapped to any of the given virtual keys."""
3805 if not hashed_tokens: 3805 ↛ 3807line 3805 didn't jump to line 3807 because the condition on line 3805 was always true
3806 return ()
3807 token_filter: Final[_TokenInFilter] = {"token": {"in": tuple(hashed_tokens)}}
3808 mappings: Final = await _jwt_key_mapping_table(JWTKeyMappingRepository(prisma_client)).find_many(where=token_filter)
3809 return tuple(jwt_key_mapping_cache_key(m.jwt_claim_name, m.jwt_claim_value, m.jwt_issuer) for m in mappings)
3812@log_db_metrics
3813async def get_jwt_key_mapping_object(
3814 jwt_claim_name: str,
3815 jwt_claim_value: str,
3816 prisma_client: PrismaClient,
3817 jwt_issuer: str | None = None,
3818) -> str | None:
3819 """
3820 Lookup a JWT-to-virtual-key mapping from the database for one exact scope:
3821 ``jwt_issuer`` (or the global/issuer-agnostic scope when falsy). Does not fall
3822 back to the global scope itself -- a caller that wants "issuer-scoped mapping,
3823 else the global one" queries both scopes itself, so each result can be cached
3824 under its own scope's key (see ``_resolve_jwt_to_virtual_key``).
3826 Returns the hashed token (str) if a matching active mapping is found, else None.
3827 """
3828 mapping: Final = await _jwt_key_mapping_table(JWTKeyMappingRepository(prisma_client)).find_first(
3829 where={
3830 "jwt_claim_name": jwt_claim_name,
3831 "jwt_claim_value": jwt_claim_value,
3832 "jwt_issuer": jwt_issuer or "",
3833 "is_active": True,
3834 }
3835 )
3836 if mapping is not None:
3837 return mapping.token
3838 return None
3841@log_db_metrics
3842async def get_key_object(
3843 hashed_token: str,
3844 prisma_client: PrismaClient | None,
3845 user_api_key_cache: UserApiKeyCache,
3846 parent_otel_span: Span | None = None,
3847 proxy_logging_obj: ProxyLogging | None = None,
3848 check_cache_only: bool | None = None,
3849) -> UserAPIKeyAuth:
3850 """
3851 - Check if team id in proxy Team Table
3852 - if valid, return LiteLLM_TeamTable object with defined limits
3853 - if not, then raise an error
3854 """
3855 if prisma_client is None: 3855 ↛ 3856line 3855 didn't jump to line 3856 because the condition on line 3855 was never true
3856 raise Exception("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys")
3858 # check if in cache
3859 key: Final = hashed_token
3861 # Same flow as before: use cache only when we have a hit we can turn into UserAPIKeyAuth
3862 # (dict from Redis / model_dump, or UserAPIKeyAuth from in-memory). Otherwise fall through to DB.
3863 user_api_key_auth: Final = await user_api_key_cache.async_get_cache(
3864 key=key,
3865 model_type=UserAPIKeyAuth,
3866 )
3867 if user_api_key_auth is not None:
3868 return _copy_user_api_key_auth_for_cache(user_api_key_obj=user_api_key_auth)
3870 if check_cache_only: 3870 ↛ 3871line 3870 didn't jump to line 3871 because the condition on line 3870 was never true
3871 raise Exception(f"Key doesn't exist in cache + check_cache_only=True. key={key}.")
3873 # else, check db
3874 _valid_token: Final[BaseModel | None] = await _fetch_key_object_from_db_with_reconnect(
3875 hashed_token=hashed_token,
3876 prisma_client=prisma_client,
3877 parent_otel_span=parent_otel_span,
3878 proxy_logging_obj=proxy_logging_obj,
3879 )
3881 if _valid_token is None: 3881 ↛ 3889line 3881 didn't jump to line 3889 because the condition on line 3881 was always true
3882 raise ProxyException(
3883 message=f"Authentication Error, Invalid proxy server token passed. key={hashed_token}, not found in db. Create key via `/key/generate` call.",
3884 type=ProxyErrorTypes.token_not_found_in_db,
3885 param="key",
3886 code=status.HTTP_401_UNAUTHORIZED,
3887 )
3889 _response: Final = UserAPIKeyAuth.model_validate(_valid_token.model_dump(exclude_none=True))
3891 # Load object_permission if object_permission_id exists but object_permission is not loaded
3892 if _response.object_permission_id and not _response.object_permission:
3893 try:
3894 _response.object_permission = await get_object_permission(
3895 object_permission_id=_response.object_permission_id,
3896 prisma_client=prisma_client,
3897 user_api_key_cache=user_api_key_cache,
3898 parent_otel_span=parent_otel_span,
3899 proxy_logging_obj=proxy_logging_obj,
3900 )
3901 except Exception as e:
3902 verbose_proxy_logger.debug(
3903 "Failed to load object_permission for key with object_permission_id=%s: %s",
3904 _response.object_permission_id,
3905 e,
3906 )
3908 # save the key object to cache
3909 await _cache_key_object(
3910 hashed_token=hashed_token,
3911 user_api_key_obj=_response,
3912 user_api_key_cache=user_api_key_cache,
3913 proxy_logging_obj=proxy_logging_obj,
3914 )
3916 return _response
3919def _copy_user_api_key_auth_for_cache(
3920 user_api_key_obj: UserAPIKeyAuth,
3921) -> UserAPIKeyAuth:
3922 copied_key_obj: Final = user_api_key_obj.model_copy()
3923 copied_key_obj.budget_reservation = None
3924 copied_key_obj.budget_throttle_pct = None
3925 copied_key_obj.parent_otel_span = None
3926 copied_key_obj.request_route = None
3927 return copied_key_obj
3930@log_db_metrics
3931async def get_object_permission(
3932 object_permission_id: str,
3933 prisma_client: PrismaClient | None,
3934 user_api_key_cache: UserApiKeyCache,
3935 parent_otel_span: Span | None = None,
3936 proxy_logging_obj: ProxyLogging | None = None,
3937) -> LiteLLM_ObjectPermissionTable | None:
3938 """
3939 - Check if object permission id in proxy ObjectPermissionTable
3940 - if valid, return LiteLLM_ObjectPermissionTable object
3941 - if not, then raise an error
3942 """
3943 if prisma_client is None: 3943 ↛ 3944line 3943 didn't jump to line 3944 because the condition on line 3943 was never true
3944 raise Exception("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys")
3946 # check if in cache
3947 key: Final = object_permission_cache_key(object_permission_id)
3948 deserialized_perm: Final = await user_api_key_cache.async_get_cache(
3949 key=key,
3950 model_type=LiteLLM_ObjectPermissionTable,
3951 )
3952 if deserialized_perm is not None: 3952 ↛ 3953line 3952 didn't jump to line 3953 because the condition on line 3952 was never true
3953 return deserialized_perm
3955 # else, check db
3956 try:
3957 response: Final = await _dictable_table(
3958 ObjectPermissionRepository(prisma_client), "object_permission"
3959 ).find_unique(where={"object_permission_id": object_permission_id})
3961 if response is None: 3961 ↛ 3962line 3961 didn't jump to line 3962 because the condition on line 3961 was never true
3962 return None
3964 _perm_obj: Final = LiteLLM_ObjectPermissionTable.model_validate(response.dict())
3965 await user_api_key_cache.async_set_cache(
3966 key=key,
3967 value=_perm_obj,
3968 model_type=LiteLLM_ObjectPermissionTable,
3969 ttl=get_management_object_ttl(user_api_key_cache),
3970 )
3972 return _perm_obj
3973 except Exception:
3974 return None
3977@log_db_metrics
3978async def get_managed_vector_store_rows_by_uuids(
3979 uuids: list[str],
3980 prisma_client: PrismaClient | None,
3981 user_api_key_cache: UserApiKeyCache,
3982 parent_otel_span: Span | None = None,
3983 proxy_logging_obj: ProxyLogging | None = None,
3984) -> list[LiteLLM_ManagedVectorStoresTable]:
3985 """
3986 Fetch managed vector store rows by their internal UUIDs.
3988 Follows the get_team_object / get_key_object / get_object_permission pattern:
3989 cache-first lookup (in-memory / Redis), DB fallback only on cache miss.
3990 Critical-path DB access must go through this helper to avoid raw Prisma
3991 calls on the hot request path.
3992 """
3993 if not uuids or prisma_client is None: 3993 ↛ 3994line 3993 didn't jump to line 3994 because the condition on line 3993 was never true
3994 return []
3996 result: Final[list[LiteLLM_ManagedVectorStoresTable]] = []
3997 cache_misses: Final[list[str]] = []
3999 for uuid in uuids:
4000 key = f"managed_vector_store_id:{uuid}"
4001 deserialized_vs = await user_api_key_cache.async_get_cache(
4002 key=key,
4003 model_type=LiteLLM_ManagedVectorStoresTable,
4004 )
4005 if deserialized_vs is not None:
4006 result.append(deserialized_vs)
4007 else:
4008 cache_misses.append(uuid)
4010 if not cache_misses:
4011 return result
4013 rows: Final = await _vector_store_table(ManagedVectorStoresRepository(prisma_client)).find_many(
4014 where={"vector_store_id": {"in": cache_misses}},
4015 take=len(cache_misses),
4016 )
4018 for row in rows:
4019 row_dict = row.model_dump() if hasattr(row, "model_dump") else (row.dict() if hasattr(row, "dict") else None)
4020 if not isinstance(row_dict, dict) or not row_dict: 4020 ↛ 4021line 4020 didn't jump to line 4021 because the condition on line 4020 was never true
4021 row_dict = dict(row) if hasattr(row, "__dict__") else {}
4022 if not row_dict: 4022 ↛ 4023line 4022 didn't jump to line 4023 because the condition on line 4022 was never true
4023 continue
4024 cached_obj = LiteLLM_ManagedVectorStoresTable.model_validate(row_dict)
4025 key = f"managed_vector_store_id:{cached_obj.vector_store_id}"
4026 await user_api_key_cache.async_set_cache(
4027 key=key,
4028 value=cached_obj,
4029 model_type=LiteLLM_ManagedVectorStoresTable,
4030 ttl=get_management_object_ttl(user_api_key_cache),
4031 )
4032 result.append(cached_obj)
4034 return result
4037class OrganizationNotFoundError(Exception):
4038 """The organization row is CONFIRMED absent, as opposed to a lookup that failed.
4040 Subclasses Exception so every existing except Exception caller keeps its current
4041 behavior; it exists so a caller that wants to treat "no such org" as "no restriction" can do
4042 that WITHOUT also swallowing an outage and silently dropping a real org ceiling.
4043 """
4046@log_db_metrics
4047async def get_org_object(
4048 org_id: str,
4049 prisma_client: PrismaClient | None,
4050 user_api_key_cache: UserApiKeyCache,
4051 parent_otel_span: Span | None = None,
4052 proxy_logging_obj: ProxyLogging | None = None,
4053 include_budget_table: bool = False,
4054) -> LiteLLM_OrganizationTable | None:
4055 """
4056 - Check if org id in proxy Org Table
4057 - if valid, return LiteLLM_OrganizationTable object
4058 - if not, then raise an error
4060 Args:
4061 org_id: Organization ID to look up
4062 prisma_client: Database client
4063 user_api_key_cache: Cache for storing results
4064 parent_otel_span: Optional OpenTelemetry span
4065 proxy_logging_obj: Optional proxy logging object
4066 include_budget_table: If True, includes litellm_budget_table in the query
4067 """
4068 if prisma_client is None:
4069 raise Exception("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys")
4070 if not isinstance(org_id, str):
4071 return None
4073 # Use different cache key if budget table is included
4074 cache_key = f"org_id:{org_id}"
4075 if include_budget_table:
4076 cache_key = f"org_id:{org_id}:with_budget"
4078 # check if in cache
4079 deserialized_org: Final = await user_api_key_cache.async_get_cache(
4080 key=cache_key,
4081 model_type=LiteLLM_OrganizationTable,
4082 )
4083 if deserialized_org is not None:
4084 return deserialized_org
4085 # else, check db
4086 try:
4087 query_kwargs: Final[dict[str, Mapping[str, object]]] = {"where": {"organization_id": org_id}}
4088 if include_budget_table:
4089 query_kwargs["include"] = {"litellm_budget_table": True}
4091 response: Final = await _model_dump_table(OrganizationRepository(prisma_client), "organization").find_unique(
4092 **query_kwargs
4093 )
4094 except Exception:
4095 # An operational failure (DB down, timeout, cache fault) is NOT the same fact as a confirmed
4096 # missing row, and relabelling it as "doesn't exist" made every caller unable to tell them
4097 # apart — a caller that treats absence as "this org places no restriction" then drops a real
4098 # org ceiling during an outage. Propagate the real error; callers that already catch
4099 # Exception are unaffected.
4100 raise
4102 if response is None:
4103 raise OrganizationNotFoundError(
4104 f"Organization doesn't exist in db. Organization={org_id}. Create organization via `/organization/new` call."
4105 )
4107 _org_obj: Final = LiteLLM_OrganizationTable.model_validate(response.model_dump())
4108 # Cache the result
4109 await user_api_key_cache.async_set_cache(
4110 key=cache_key,
4111 value=_org_obj,
4112 model_type=LiteLLM_OrganizationTable,
4113 ttl=DEFAULT_IN_MEMORY_TTL,
4114 )
4116 return _org_obj
4119def _last_known_org_cache_key(org_id: str) -> str:
4120 return f"org_id:{org_id}:with_budget:last_known"
4123async def _keep_last_known_org(
4124 org: LiteLLM_OrganizationTable, org_id: str, user_api_key_cache: UserApiKeyCache
4125) -> None:
4126 cache_key: Final = _last_known_org_cache_key(org_id)
4127 held_locally: Final = await user_api_key_cache.async_get_cache(
4128 key=cache_key, local_only=True, model_type=LiteLLM_OrganizationTable
4129 )
4130 if held_locally is not None:
4131 return
4132 await user_api_key_cache.async_set_cache(
4133 key=cache_key,
4134 value=org,
4135 model_type=LiteLLM_OrganizationTable,
4136 ttl=get_management_object_ttl(user_api_key_cache),
4137 )
4140async def get_org_object_for_request(
4141 org_id: str,
4142 prisma_client: PrismaClient,
4143 user_api_key_cache: UserApiKeyCache,
4144 parent_otel_span: Span | None,
4145 proxy_logging_obj: ProxyLogging | None,
4146) -> LiteLLM_OrganizationTable | None:
4147 try:
4148 org: Final = await get_org_object(
4149 org_id=org_id,
4150 prisma_client=prisma_client,
4151 user_api_key_cache=user_api_key_cache,
4152 parent_otel_span=parent_otel_span,
4153 proxy_logging_obj=proxy_logging_obj,
4154 include_budget_table=True,
4155 )
4156 except OrganizationNotFoundError:
4157 return None
4158 except Exception as e:
4159 if not PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(e):
4160 verbose_proxy_logger.debug("org lookup failed, continuing without org limits", exc_info=True)
4161 return None
4162 last_known_org: Final = await user_api_key_cache.async_get_cache(
4163 key=_last_known_org_cache_key(org_id),
4164 model_type=LiteLLM_OrganizationTable,
4165 )
4166 if last_known_org is not None:
4167 return last_known_org
4168 if PrismaDBExceptionHandler.should_allow_request_on_db_unavailable():
4169 return None
4170 raise
4171 if org is None:
4172 return None
4173 await _keep_last_known_org(org, org_id, user_api_key_cache)
4174 return org
4177async def _get_resources_from_access_groups(
4178 access_group_ids: Sequence[str],
4179 resource_field: Literal["access_model_names", "access_mcp_server_ids", "access_agent_ids"],
4180 prisma_client: DatabaseClient | None = None,
4181 user_api_key_cache: UserApiKeyCache | None = None,
4182 proxy_logging_obj: ProxyLogging | None = None,
4183) -> list[str]:
4184 """
4185 Fetch access groups by their IDs (from cache or DB) and collect
4186 the specified resource field across all of them.
4188 Args:
4189 access_group_ids: List of access group IDs to fetch
4190 resource_field: Which resource list to extract from each access group
4191 - "access_model_names": model names (for model access checks)
4192 - "access_mcp_server_ids": MCP server IDs (for MCP access checks)
4193 - "access_agent_ids": agent IDs (for agent access checks)
4194 prisma_client: Optional PrismaClient (lazy-imported from proxy_server if None)
4195 user_api_key_cache: Optional DualCache (lazy-imported from proxy_server if None)
4196 proxy_logging_obj: Optional ProxyLogging (lazy-imported from proxy_server if None)
4198 Returns:
4199 Deduplicated list of resource identifiers from all resolved access groups.
4200 """
4201 if not access_group_ids:
4202 return []
4204 # Lazy import to avoid circular imports
4205 if prisma_client is None or user_api_key_cache is None:
4206 from litellm.proxy.proxy_server import prisma_client as _prisma_client
4207 from litellm.proxy.proxy_server import proxy_logging_obj as _proxy_logging_obj
4208 from litellm.proxy.proxy_server import user_api_key_cache as _user_api_key_cache
4210 prisma_client = prisma_client or _prisma_client
4211 user_api_key_cache = user_api_key_cache or _user_api_key_cache
4212 proxy_logging_obj = proxy_logging_obj or _proxy_logging_obj
4214 if user_api_key_cache is None: 4214 ↛ 4215line 4214 didn't jump to line 4215 because the condition on line 4214 was never true
4215 return []
4217 resources: Final[list[str]] = []
4218 for ag_id in access_group_ids:
4219 try:
4220 ag = await get_access_object(
4221 access_group_id=ag_id,
4222 prisma_client=prisma_client,
4223 user_api_key_cache=user_api_key_cache,
4224 proxy_logging_obj=proxy_logging_obj,
4225 )
4226 resources.extend(getattr(ag, resource_field, []))
4227 except Exception:
4228 verbose_proxy_logger.debug(
4229 "Could not fetch access group %s for resource field %s",
4230 ag_id,
4231 resource_field,
4232 )
4233 return list(set(resources))
4236async def _get_models_from_access_groups(
4237 access_group_ids: Sequence[str],
4238 prisma_client: DatabaseClient | None = None,
4239 user_api_key_cache: UserApiKeyCache | None = None,
4240 proxy_logging_obj: ProxyLogging | None = None,
4241) -> list[str]:
4242 """
4243 Collect model names from unified access groups.
4244 Models are matched by model name for backwards compatibility.
4245 """
4246 return await _get_resources_from_access_groups(
4247 access_group_ids=access_group_ids,
4248 resource_field="access_model_names",
4249 prisma_client=prisma_client,
4250 user_api_key_cache=user_api_key_cache,
4251 proxy_logging_obj=proxy_logging_obj,
4252 )
4255async def _get_mcp_server_ids_from_access_groups(
4256 access_group_ids: list[str],
4257 prisma_client: PrismaClient | None = None,
4258 user_api_key_cache: UserApiKeyCache | None = None,
4259 proxy_logging_obj: ProxyLogging | None = None,
4260) -> list[str]:
4261 """
4262 Collect MCP server IDs from unified access groups.
4263 MCPs are matched by server ID.
4264 """
4265 return await _get_resources_from_access_groups(
4266 access_group_ids=access_group_ids,
4267 resource_field="access_mcp_server_ids",
4268 prisma_client=prisma_client,
4269 user_api_key_cache=user_api_key_cache,
4270 proxy_logging_obj=proxy_logging_obj,
4271 )
4274async def _get_agent_ids_from_access_groups(
4275 access_group_ids: list[str],
4276 prisma_client: PrismaClient | None = None,
4277 user_api_key_cache: UserApiKeyCache | None = None,
4278 proxy_logging_obj: ProxyLogging | None = None,
4279) -> list[str]:
4280 """
4281 Collect agent IDs from unified access groups.
4282 Agents are matched by agent ID.
4283 """
4284 return await _get_resources_from_access_groups(
4285 access_group_ids=access_group_ids,
4286 resource_field="access_agent_ids",
4287 prisma_client=prisma_client,
4288 user_api_key_cache=user_api_key_cache,
4289 proxy_logging_obj=proxy_logging_obj,
4290 )
4293def _resolve_all_team_model_sentinel_for_auth_check(
4294 models: list[str],
4295 llm_router: Router | None,
4296 team_id: str | None,
4297) -> list[str]:
4298 if SpecialModelNames.all_team_models.value not in models or team_id is None or llm_router is None: 4298 ↛ 4300line 4298 didn't jump to line 4300 because the condition on line 4298 was always true
4299 return models
4300 proxy_models: Final = llm_router.get_model_names()
4301 non_sentinel_models: Final = [model for model in models if model != SpecialModelNames.all_team_models.value]
4302 if not proxy_models:
4303 return non_sentinel_models or models
4304 return list(dict.fromkeys(non_sentinel_models + proxy_models))
4307def _check_model_access_helper(
4308 model: str,
4309 llm_router: Router | None,
4310 models: list[str],
4311 team_model_aliases: dict[str, str] | None = None,
4312 team_id: str | None = None,
4313) -> bool:
4314 ## check if model in allowed model names
4315 from collections import defaultdict
4317 access_groups: dict[str, list[str]] = defaultdict(list)
4319 if llm_router: 4319 ↛ 4322line 4319 didn't jump to line 4322 because the condition on line 4319 was always true
4320 access_groups = llm_router.get_model_access_groups(model_name=model, team_id=team_id)
4322 models = _resolve_all_team_model_sentinel_for_auth_check(
4323 models=models,
4324 llm_router=llm_router,
4325 team_id=team_id,
4326 )
4328 if len(access_groups) > 0 and llm_router is not None: # check if token contains any model access groups 4328 ↛ 4329line 4328 didn't jump to line 4329 because the condition on line 4328 was never true
4329 for idx, m in enumerate(models): # loop token models, if any of them are an access group add the access group
4330 if m in access_groups:
4331 return True
4333 # Filter out models that are access_groups
4334 filtered_models: Final = [m for m in models if m not in access_groups]
4336 if _model_in_team_aliases(model=model, team_model_aliases=team_model_aliases): 4336 ↛ 4337line 4336 didn't jump to line 4337 because the condition on line 4336 was never true
4337 return True
4339 if _model_matches_any_wildcard_pattern_in_list(model=model, allowed_model_list=filtered_models): 4339 ↛ 4340line 4339 didn't jump to line 4340 because the condition on line 4339 was never true
4340 return True
4342 all_model_access: bool = False
4344 if (len(filtered_models) == 0 and len(models) == 0) or "*" in filtered_models: 4344 ↛ 4347line 4344 didn't jump to line 4347 because the condition on line 4344 was always true
4345 all_model_access = True
4347 if SpecialModelNames.all_proxy_models.value in filtered_models: 4347 ↛ 4348line 4347 didn't jump to line 4348 because the condition on line 4347 was never true
4348 all_model_access = True
4350 if model is not None and model not in filtered_models and all_model_access is False: 4350 ↛ 4351line 4350 didn't jump to line 4351 because the condition on line 4350 was never true
4351 return False
4352 return True
4355def _can_object_call_model(
4356 model: str | list[str],
4357 llm_router: Router | None,
4358 models: list[str],
4359 team_model_aliases: dict[str, str] | None = None,
4360 team_id: str | None = None,
4361 key_model_aliases: Mapping[str, str] | None = None,
4362 object_type: Literal["user", "team", "key", "org", "project", "agent"] = "user",
4363 fallback_depth: int = 0,
4364) -> Literal[True]:
4365 """
4366 Checks if token can call a given model
4368 Args:
4369 - model: str
4370 - llm_router: Optional[Router]
4371 - models: List[str]
4372 - team_model_aliases: Optional[Dict[str, str]]
4373 - object_type: Literal["user", "team", "key", "org"]. We use the object type to raise the correct exception type
4375 Returns:
4376 - True: if token allowed to call model
4378 Raises:
4379 - Exception: If token not allowed to call model
4380 """
4381 if fallback_depth >= DEFAULT_MAX_RECURSE_DEPTH: 4381 ↛ 4382line 4381 didn't jump to line 4382 because the condition on line 4381 was never true
4382 raise Exception(f"Unable to parse model, max fallback depth exceeded - received model: {model}")
4383 if isinstance(model, list):
4384 for m in model:
4385 _can_object_call_model(
4386 model=m,
4387 llm_router=llm_router,
4388 models=models,
4389 team_model_aliases=team_model_aliases,
4390 team_id=team_id,
4391 key_model_aliases=key_model_aliases,
4392 object_type=object_type,
4393 fallback_depth=fallback_depth + 1,
4394 )
4395 return True
4397 from litellm.router_strategy.complexity_router.context_compaction import native_compaction_parent
4399 compaction_parent: Final = native_compaction_parent(model)
4400 global_or_router_alias_target: Final = (
4401 litellm.model_alias_map[model]
4402 if model in litellm.model_alias_map
4403 else (
4404 llm_router._get_model_from_alias(model)
4405 if llm_router is not None and model in llm_router.model_group_alias
4406 else None
4407 )
4408 )
4409 after_team_alias: Final = team_model_aliases.get(model, model) if team_model_aliases else model
4410 after_key_alias: Final = (
4411 key_model_aliases.get(after_team_alias, after_team_alias) if key_model_aliases else after_team_alias
4412 )
4413 after_global_alias: Final = litellm.model_alias_map.get(after_key_alias, after_key_alias)
4414 dispatched_model: Final = (
4415 key_model_aliases.get(after_global_alias, after_global_alias) if key_model_aliases else after_global_alias
4416 )
4417 key_alias_applied: Final = after_key_alias != after_team_alias or dispatched_model != after_global_alias
4418 potential_models: Final = (
4419 (dispatched_model,)
4420 if key_alias_applied
4421 else (
4422 *((model, compaction_parent) if compaction_parent is not None else (model,)),
4423 *((global_or_router_alias_target,) if global_or_router_alias_target else ()),
4424 )
4425 )
4427 ## check model access for alias + underlying model - allow if either is in allowed models
4428 for m in potential_models: 4428 ↛ 4438line 4428 didn't jump to line 4438 because the loop on line 4428 didn't complete
4429 if _check_model_access_helper( 4429 ↛ 4428line 4429 didn't jump to line 4428 because the condition on line 4429 was always true
4430 model=m,
4431 llm_router=llm_router,
4432 models=models,
4433 team_model_aliases=team_model_aliases,
4434 team_id=team_id,
4435 ):
4436 return True
4438 internal_message: Final = (
4439 f"{object_type} not allowed to access model. This {object_type} can only access models={models}. "
4440 f"Tried to access {model}"
4441 )
4442 raise ModelAccessDeniedProxyException(
4443 message=model_access_denied_client_message(model=model),
4444 internal_message=internal_message,
4445 type=ProxyErrorTypes.get_model_access_error_type_for_object(object_type=object_type),
4446 param="model",
4447 code=status.HTTP_403_FORBIDDEN,
4448 )
4451def _resolve_team_alias(
4452 model: str | list[str],
4453 team_model_aliases: dict[str, str] | None,
4454 team_id: str | None,
4455 llm_router: Router | None,
4456) -> str | list[str]:
4457 if not team_model_aliases:
4458 return model
4459 if isinstance(model, str):
4460 return _live_team_alias_target(model, team_model_aliases, team_id, llm_router)
4461 return [ # mutable-ok: _can_object_call_model takes list[str]
4462 _live_team_alias_target(name, team_model_aliases, team_id, llm_router) for name in model
4463 ]
4466def _live_team_alias_target(
4467 model: str, team_model_aliases: dict[str, str], team_id: str | None, llm_router: Router | None
4468) -> str:
4469 target: Final = team_model_aliases.get(model)
4470 if target is None:
4471 return model
4472 deleted_team_deployment: Final = (
4473 llm_router is not None
4474 and target.startswith(f"model_name_{team_id}_")
4475 and target not in llm_router.model_name_to_deployment_indices
4476 )
4477 return model if deleted_team_deployment else target
4480async def _check_agent_access_group_model_access(
4481 model: str | list[str] | None, # mutable-ok: _can_object_call_model and the client message helper take list[str]
4482 valid_token: UserAPIKeyAuth | None,
4483 llm_router: Router | None,
4484 resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling,
4485) -> Literal[True]:
4486 """Attached groups naming no model deny every model; the empty allowlist in ``_can_object_call_model`` allows."""
4487 if not model or valid_token is None or not valid_token.agent_id: 4487 ↛ 4489line 4487 didn't jump to line 4489 because the condition on line 4487 was always true
4488 return True
4489 ceiling: Final = await resolve_ceiling(valid_token.agent_id)
4490 if ceiling is None:
4491 return True
4492 if not ceiling.models:
4493 raise ModelAccessDeniedProxyException(
4494 message=model_access_denied_client_message(model=model),
4495 internal_message=f"agent {valid_token.agent_id} access groups {ceiling.access_group_ids} grant no models",
4496 type=ProxyErrorTypes.agent_model_access_denied,
4497 param="model",
4498 code=status.HTTP_403_FORBIDDEN,
4499 )
4500 dispatched: Final = _resolve_team_alias(model, valid_token.team_model_aliases, valid_token.team_id, llm_router)
4501 return _can_object_call_model(
4502 model=dispatched,
4503 llm_router=llm_router,
4504 models=sorted(ceiling.models),
4505 team_id=valid_token.team_id,
4506 object_type="agent",
4507 key_model_aliases=key_model_aliases_for_auth_check(valid_token),
4508 )
4511LoadedCallerTeam: TypeAlias = LiteLLM_TeamTable | None
4512LoadedCallerUser: TypeAlias = LiteLLM_UserTable | None
4513CallerTeamLoader: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[LoadedCallerTeam]] # mutable-ok: Callable params
4514CallerUserLoader: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[LoadedCallerUser]] # mutable-ok: Callable params
4517async def _check_agent_caller_model_access(
4518 model: str | list[str] | None, # mutable-ok: the model checks it delegates to take list[str]
4519 valid_token: UserAPIKeyAuth | None,
4520 llm_router: Router | None,
4521 prisma_client: Optional["PrismaClient"],
4522 user_api_key_cache: UserApiKeyCache,
4523 proxy_logging_obj: ProxyLogging,
4524 load_team: CallerTeamLoader = load_agent_caller_team,
4525 load_user: CallerUserLoader = load_agent_caller_user,
4526) -> None:
4527 """An agent key acting for an invoking user may call only what that user's own key could: the
4528 invoking team's models (and per-member scope) when a team was echoed, else the user's models."""
4529 if not model or valid_token is None:
4530 return
4531 caller_auth: Final = agent_caller_auth(valid_token)
4532 if caller_auth is None: 4532 ↛ 4534line 4532 didn't jump to line 4534 because the condition on line 4532 was always true
4533 return
4534 caller_team: Final = await load_team(valid_token)
4535 caller_key_model_aliases: Final = key_model_aliases_for_auth_check(valid_token)
4536 if caller_team is not None:
4537 await can_team_access_model(
4538 model=model,
4539 team_object=caller_team,
4540 llm_router=llm_router,
4541 prisma_client=prisma_client,
4542 key_model_aliases=caller_key_model_aliases,
4543 )
4544 await _check_team_member_model_access(
4545 model=model,
4546 team_object=caller_team,
4547 valid_token=caller_auth,
4548 llm_router=llm_router,
4549 prisma_client=prisma_client,
4550 user_api_key_cache=user_api_key_cache,
4551 proxy_logging_obj=proxy_logging_obj,
4552 key_model_aliases=caller_key_model_aliases,
4553 )
4554 return
4555 caller_user: Final = await load_user(valid_token)
4556 if caller_user is None:
4557 return
4558 await can_user_call_model(
4559 model=model,
4560 llm_router=llm_router,
4561 user_object=caller_user,
4562 key_model_aliases=caller_key_model_aliases,
4563 )
4566def _model_in_team_aliases(model: str, team_model_aliases: dict[str, str] | None = None) -> bool:
4567 """
4568 Returns True if `model` being accessed is an alias of a team model
4570 - `model=gpt-4o`
4571 - `team_model_aliases={"gpt-4o": "gpt-4o-team-1"}`
4572 - returns True
4574 - `model=gp-4o`
4575 - `team_model_aliases={"o-3": "o3-preview"}`
4576 - returns False
4577 """
4578 if team_model_aliases: 4578 ↛ 4579line 4578 didn't jump to line 4579 because the condition on line 4578 was never true
4579 if model in team_model_aliases:
4580 return True
4581 return False
4584def key_model_aliases_for_auth_check(valid_token: UserAPIKeyAuth | None) -> Mapping[str, str] | None:
4585 return alias_map(valid_token.aliases) if valid_token is not None and valid_token.aliases else None
4588def _resolve_key_models_for_auth_check(valid_token: UserAPIKeyAuth) -> list[str]:
4589 """
4590 Expand key model sentinels before auth checks.
4592 ``all-team-models`` means inherit the parent team's allowlist -- same
4593 semantics as ``get_key_models`` in ``model_checks.py``.
4595 If the key has no team_id, it inherits the full proxy model list
4596 (equivalent to an empty models field, i.e. unrestricted access).
4597 """
4598 models: Final = list(valid_token.models or [])
4599 if SpecialModelNames.all_team_models.value in models: 4599 ↛ 4600line 4599 didn't jump to line 4600 because the condition on line 4599 was never true
4600 if valid_token.team_id is None:
4601 return []
4602 return list(valid_token.team_models or [])
4603 return models
4606def _model_access_groups_serving_model(
4607 model: str | Sequence[str],
4608 llm_router: Router,
4609 team_id: str | None,
4610) -> frozenset[str]:
4611 """Every model access group whose deployments serve the requested model(s)."""
4612 requested: Final = (model,) if isinstance(model, str) else tuple(model)
4613 return frozenset(
4614 group
4615 for requested_model in requested
4616 for group in llm_router.get_model_access_groups(model_name=requested_model, team_id=team_id)
4617 )
4620async def _team_member_granted_models(
4621 valid_token: UserAPIKeyAuth,
4622 team_object: LiteLLM_TeamTable | None,
4623 prisma_client: PrismaClient,
4624 user_api_key_cache: UserApiKeyCache,
4625 proxy_logging_obj: ProxyLogging,
4626 team_membership: LiteLLM_TeamMembership | None = None,
4627 team_membership_loaded: bool = False,
4628) -> Sequence[str]:
4629 """The member's own ``allowed_models`` scope; empty when the member is not narrowed below the team."""
4630 if team_object is None or valid_token.user_id is None:
4631 return ()
4633 if not team_membership_loaded:
4634 team_membership = await get_team_membership(
4635 user_id=valid_token.user_id,
4636 team_id=team_object.team_id,
4637 prisma_client=prisma_client,
4638 user_api_key_cache=user_api_key_cache,
4639 proxy_logging_obj=proxy_logging_obj,
4640 )
4641 return () if team_membership is None else _member_allowed_models(team_membership)
4644async def _org_granted_models(
4645 valid_token: UserAPIKeyAuth,
4646 team_object: LiteLLM_TeamTable | None,
4647 prisma_client: PrismaClient,
4648 user_api_key_cache: UserApiKeyCache,
4649 proxy_logging_obj: ProxyLogging,
4650) -> Sequence[str]:
4651 """The org allowlist reached through the key, or through its team when the key names no org."""
4652 org_id: Final = valid_token.org_id or (team_object.organization_id if team_object is not None else None)
4653 if org_id is None:
4654 return ()
4656 try:
4657 org_object: Final = await get_org_object(
4658 org_id=org_id,
4659 prisma_client=prisma_client,
4660 user_api_key_cache=user_api_key_cache,
4661 proxy_logging_obj=proxy_logging_obj,
4662 )
4663 except Exception as e: # noqa: BLE001 # fail-safe: attribution degrades to "no org grant", it must never break auth
4664 verbose_proxy_logger.debug("access group attribution: org lookup failed: %s", e)
4665 return ()
4666 return org_object.models if org_object is not None else ()
4669async def _granted_model_lists(
4670 valid_token: UserAPIKeyAuth,
4671 team_object: LiteLLM_TeamTable | None,
4672 project_object: LiteLLM_ProjectTableCachedObj | None,
4673 prisma_client: PrismaClient,
4674 user_api_key_cache: UserApiKeyCache,
4675 proxy_logging_obj: ProxyLogging,
4676 team_membership: LiteLLM_TeamMembership | None = None,
4677 team_membership_loaded: bool = False,
4678) -> tuple[Sequence[str], ...]:
4679 """One model allowlist per level that participates in authorizing the request."""
4680 return (
4681 _resolve_key_models_for_auth_check(valid_token=valid_token),
4682 team_object.models if team_object is not None else (),
4683 await _team_member_granted_models(
4684 valid_token=valid_token,
4685 team_object=team_object,
4686 prisma_client=prisma_client,
4687 user_api_key_cache=user_api_key_cache,
4688 proxy_logging_obj=proxy_logging_obj,
4689 team_membership=team_membership,
4690 team_membership_loaded=team_membership_loaded,
4691 ),
4692 project_object.models if project_object is not None else (),
4693 await _org_granted_models(
4694 valid_token=valid_token,
4695 team_object=team_object,
4696 prisma_client=prisma_client,
4697 user_api_key_cache=user_api_key_cache,
4698 proxy_logging_obj=proxy_logging_obj,
4699 ),
4700 )
4703async def _user_object_or_none(
4704 valid_token: UserAPIKeyAuth,
4705 prisma_client: PrismaClient,
4706 user_api_key_cache: UserApiKeyCache,
4707 proxy_logging_obj: ProxyLogging,
4708) -> LiteLLM_UserTable | None:
4709 try:
4710 return await get_user_object(
4711 user_id=valid_token.user_id,
4712 prisma_client=prisma_client,
4713 user_api_key_cache=user_api_key_cache,
4714 user_id_upsert=False,
4715 proxy_logging_obj=proxy_logging_obj,
4716 )
4717 except UserNotFoundError:
4718 return None
4721async def enforced_model_allowlists(
4722 valid_token: UserAPIKeyAuth,
4723 prisma_client: PrismaClient | None,
4724 user_api_key_cache: UserApiKeyCache,
4725 proxy_logging_obj: ProxyLogging,
4726) -> tuple[Sequence[str], ...]:
4727 """One model allowlist per level that ``common_checks`` enforces on a request from this identity."""
4728 key_models: Final = _resolve_key_models_for_auth_check(valid_token=valid_token)
4729 if prisma_client is None:
4730 return (key_models, tuple(valid_token.team_models or ()))
4731 team_object: Final = (
4732 None
4733 if valid_token.team_id is None
4734 else await get_team_object(
4735 team_id=valid_token.team_id,
4736 prisma_client=prisma_client,
4737 user_api_key_cache=user_api_key_cache,
4738 proxy_logging_obj=proxy_logging_obj,
4739 )
4740 )
4741 user_object: Final = (
4742 None
4743 if team_object is not None
4744 else await _user_object_or_none(
4745 valid_token=valid_token,
4746 prisma_client=prisma_client,
4747 user_api_key_cache=user_api_key_cache,
4748 proxy_logging_obj=proxy_logging_obj,
4749 )
4750 )
4751 project_object: Final = (
4752 None
4753 if valid_token.project_id is None
4754 else await get_project_object(
4755 project_id=valid_token.project_id,
4756 prisma_client=prisma_client,
4757 user_api_key_cache=user_api_key_cache,
4758 proxy_logging_obj=proxy_logging_obj,
4759 )
4760 )
4761 return (
4762 key_models,
4763 team_object.models if team_object is not None else (),
4764 await _team_member_granted_models(
4765 valid_token=valid_token,
4766 team_object=team_object,
4767 prisma_client=prisma_client,
4768 user_api_key_cache=user_api_key_cache,
4769 proxy_logging_obj=proxy_logging_obj,
4770 ),
4771 user_object.models if user_object is not None else (),
4772 project_object.models if project_object is not None else (),
4773 )
4776async def collect_matched_model_access_groups(
4777 model: str | Sequence[str] | None,
4778 valid_token: UserAPIKeyAuth | None,
4779 team_object: LiteLLM_TeamTable | None,
4780 project_object: LiteLLM_ProjectTableCachedObj | None,
4781 llm_router: Router | None,
4782 prisma_client: PrismaClient | None,
4783 user_api_key_cache: UserApiKeyCache,
4784 proxy_logging_obj: ProxyLogging,
4785 team_membership: LiteLLM_TeamMembership | None = None,
4786 team_membership_loaded: bool = False,
4787) -> tuple[str, ...]:
4788 """
4789 The budgeted model access groups that authorized this request, sorted and deduplicated.
4791 A group is charged only when its name appears on an allowlist the caller was granted -- key,
4792 team, team-member scope, project or org -- *and* that group serves the requested model. Asking
4793 for a model that merely belongs to a group attributes nothing, because nothing about the caller
4794 named the group.
4796 Levels are unioned, never ranked: a team granted ``*`` whose member is scoped to one group is
4797 still a caller gated by that group. An unrestricted allowlist (empty, ``*``) names no group and
4798 so contributes nothing.
4800 The whole walk is gated on the budget registry, because collecting every match costs a full scan
4801 of each allowlist where the plain access check stops at the first hit. An empty registry means no
4802 group carries a budget, so there is nothing to attribute and no work worth doing.
4803 """
4804 if model is None or valid_token is None or llm_router is None or prisma_client is None:
4805 return ()
4807 registry: Final = await _load_model_access_group_registry(
4808 prisma_client=prisma_client,
4809 user_api_key_cache=user_api_key_cache,
4810 )
4811 if registry is not None and not registry: 4811 ↛ 4814line 4811 didn't jump to line 4814 because the condition on line 4811 was always true
4812 return ()
4814 covering_groups: Final = _model_access_groups_serving_model(
4815 model=model,
4816 llm_router=llm_router,
4817 team_id=valid_token.team_id,
4818 )
4819 budgeted_groups: Final = covering_groups if registry is None else covering_groups & registry
4820 if not budgeted_groups:
4821 return ()
4823 granted: Final = frozenset(
4824 granted_model
4825 for granted_models in await _granted_model_lists(
4826 valid_token=valid_token,
4827 team_object=team_object,
4828 project_object=project_object,
4829 prisma_client=prisma_client,
4830 user_api_key_cache=user_api_key_cache,
4831 proxy_logging_obj=proxy_logging_obj,
4832 team_membership=team_membership,
4833 team_membership_loaded=team_membership_loaded,
4834 )
4835 for granted_model in granted_models
4836 )
4837 return tuple(sorted(budgeted_groups & granted))
4840async def stamp_matched_model_access_groups(
4841 model: str | Sequence[str] | None,
4842 valid_token: UserAPIKeyAuth | None,
4843 team_object: LiteLLM_TeamTable | None,
4844 project_object: LiteLLM_ProjectTableCachedObj | None,
4845 llm_router: Router | None,
4846 prisma_client: PrismaClient | None,
4847 user_api_key_cache: UserApiKeyCache,
4848 proxy_logging_obj: ProxyLogging,
4849 team_membership: LiteLLM_TeamMembership | None = None,
4850 team_membership_loaded: bool = False,
4851) -> tuple[str, ...]:
4852 """Record the groups that authorized this request on its auth object, for the post-call spend
4853 writer and the reservation counters, and hand them back for the budget check."""
4854 if valid_token is None: 4854 ↛ 4855line 4854 didn't jump to line 4855 because the condition on line 4854 was never true
4855 return ()
4857 try:
4858 matched: Final = await collect_matched_model_access_groups(
4859 model=model,
4860 valid_token=valid_token,
4861 team_object=team_object,
4862 project_object=project_object,
4863 llm_router=llm_router,
4864 prisma_client=prisma_client,
4865 user_api_key_cache=user_api_key_cache,
4866 proxy_logging_obj=proxy_logging_obj,
4867 team_membership=team_membership,
4868 team_membership_loaded=team_membership_loaded,
4869 )
4870 except Exception as e: # noqa: BLE001 # fail-safe: attribution is spend telemetry, it must never break auth
4871 verbose_proxy_logger.debug("model access group attribution failed: %s", e)
4872 return ()
4873 if not matched: 4873 ↛ 4875line 4873 didn't jump to line 4875 because the condition on line 4873 was always true
4874 return ()
4875 matched_groups: Final = list(matched) # mutable-ok: the auth field is typed list[str] | None
4876 valid_token.matched_model_access_groups = matched_groups # rebind-ok: request-scoped carrier for the writer
4877 return matched
4880async def can_key_call_model(
4881 model: str | list[str],
4882 llm_model_list: Sequence[object] | None,
4883 valid_token: UserAPIKeyAuth,
4884 llm_router: litellm.Router | None,
4885 prisma_client: DatabaseClient | None = None,
4886) -> Literal[True]:
4887 """
4888 Checks if token can call a given model
4890 1. First checks native key-level model permissions (current implementation)
4891 2. If not allowed natively, falls back to access_group_ids on the key
4893 Returns:
4894 - True: if token allowed to call model
4896 Raises:
4897 - Exception: If token not allowed to call model
4898 """
4899 key_models: Final = _resolve_key_models_for_auth_check(valid_token=valid_token)
4900 try:
4901 return _can_object_call_model(
4902 model=model,
4903 llm_router=llm_router,
4904 models=key_models,
4905 team_model_aliases=valid_token.team_model_aliases,
4906 team_id=valid_token.team_id,
4907 key_model_aliases=key_model_aliases_for_auth_check(valid_token),
4908 object_type="key",
4909 )
4910 except ProxyException:
4911 # Fallback: check key's access_group_ids
4912 key_access_group_ids: Final = valid_token.access_group_ids or []
4913 if key_access_group_ids:
4914 models_from_groups: Final = await _get_models_from_access_groups(
4915 access_group_ids=key_access_group_ids,
4916 prisma_client=prisma_client,
4917 )
4918 if models_from_groups:
4919 return _can_object_call_model(
4920 model=model,
4921 llm_router=llm_router,
4922 models=models_from_groups,
4923 team_model_aliases=valid_token.team_model_aliases,
4924 team_id=valid_token.team_id,
4925 key_model_aliases=key_model_aliases_for_auth_check(valid_token),
4926 object_type="key",
4927 )
4928 raise
4931async def can_key_call_resolved_model(
4932 model: str,
4933 llm_model_list: Sequence[object] | None,
4934 valid_token: UserAPIKeyAuth,
4935 llm_router: litellm.Router | None,
4936) -> None:
4937 from litellm.proxy.proxy_server import (
4938 prisma_client,
4939 proxy_logging_obj,
4940 user_api_key_cache,
4941 )
4943 skip_key_model_check: Final = valid_token.config or (
4944 isinstance(valid_token.models, list) and SpecialModelNames.all_team_models.value in valid_token.models
4945 )
4946 if not skip_key_model_check: 4946 ↛ 4954line 4946 didn't jump to line 4954 because the condition on line 4946 was always true
4947 await can_key_call_model(
4948 model=model,
4949 llm_model_list=llm_model_list,
4950 valid_token=valid_token,
4951 llm_router=llm_router,
4952 )
4954 team_object: LiteLLM_TeamTableCachedObj | None = None
4955 team_object_from_lookup = False
4956 if valid_token.team_id is not None: 4956 ↛ 4957line 4956 didn't jump to line 4957 because the condition on line 4956 was never true
4957 try:
4958 team_object = await get_team_object(
4959 team_id=valid_token.team_id,
4960 prisma_client=prisma_client,
4961 user_api_key_cache=user_api_key_cache,
4962 parent_otel_span=valid_token.parent_otel_span,
4963 proxy_logging_obj=proxy_logging_obj,
4964 )
4965 team_object_from_lookup = True
4966 except Exception:
4967 team_object = LiteLLM_TeamTableCachedObj(
4968 team_id=valid_token.team_id,
4969 models=valid_token.team_models,
4970 blocked=valid_token.team_blocked,
4971 team_alias=valid_token.team_alias,
4972 metadata=valid_token.team_metadata,
4973 object_permission_id=valid_token.team_object_permission_id,
4974 object_permission=valid_token.team_object_permission,
4975 )
4977 if team_object is not None: 4977 ↛ 4978line 4977 didn't jump to line 4978 because the condition on line 4977 was never true
4978 try:
4979 await can_team_access_model(
4980 model=model,
4981 team_object=team_object,
4982 llm_router=llm_router,
4983 team_model_aliases=valid_token.team_model_aliases,
4984 key_model_aliases=key_model_aliases_for_auth_check(valid_token),
4985 )
4986 except ProxyException as team_denial:
4987 if team_denial.type != ProxyErrorTypes.team_model_access_denied:
4988 raise
4989 if not await _key_access_group_grants_model(
4990 model=model,
4991 valid_token=valid_token,
4992 team_object=team_object,
4993 llm_router=llm_router,
4994 key_model_aliases=key_model_aliases_for_auth_check(valid_token),
4995 ):
4996 raise
4998 if valid_token.user_id is not None and team_object_from_lookup:
4999 await _check_team_member_model_access(
5000 model=model,
5001 team_object=team_object,
5002 valid_token=valid_token,
5003 llm_router=llm_router,
5004 prisma_client=prisma_client,
5005 user_api_key_cache=user_api_key_cache,
5006 proxy_logging_obj=proxy_logging_obj,
5007 key_model_aliases=key_model_aliases_for_auth_check(valid_token),
5008 )
5010 if valid_token.project_id is not None: 5010 ↛ 5011line 5010 didn't jump to line 5011 because the condition on line 5010 was never true
5011 project_object: Final = await get_project_object(
5012 project_id=valid_token.project_id,
5013 prisma_client=prisma_client,
5014 user_api_key_cache=user_api_key_cache,
5015 proxy_logging_obj=proxy_logging_obj,
5016 )
5017 if project_object is not None and len(project_object.models) > 0:
5018 can_project_access_model(
5019 model=model,
5020 project_object=project_object,
5021 llm_router=llm_router,
5022 key_model_aliases=key_model_aliases_for_auth_check(valid_token),
5023 )
5026def can_org_access_model(
5027 model: str,
5028 org_object: LiteLLM_OrganizationTable | None,
5029 llm_router: Router | None,
5030 team_model_aliases: dict[str, str] | None = None,
5031) -> Literal[True]:
5032 """
5033 Returns True if the team can access a specific model.
5035 """
5036 return _can_object_call_model(
5037 model=model,
5038 llm_router=llm_router,
5039 models=org_object.models if org_object else [],
5040 team_model_aliases=team_model_aliases,
5041 object_type="org",
5042 )
5045async def can_team_access_model(
5046 model: str | list[str],
5047 team_object: LiteLLM_TeamTable | None,
5048 llm_router: Router | None,
5049 team_model_aliases: dict[str, str] | None = None,
5050 key_model_aliases: Mapping[str, str] | None = None,
5051 prisma_client: DatabaseClient | None = None,
5052) -> Literal[True]:
5053 """
5054 Returns True if the team can access a specific model.
5056 1. First checks native team-level model permissions (current implementation)
5057 2. If not allowed natively, falls back to access_group_ids on the team
5058 """
5059 try:
5060 return _can_object_call_model(
5061 model=model,
5062 llm_router=llm_router,
5063 models=team_object.models if team_object else [],
5064 team_model_aliases=team_model_aliases,
5065 team_id=team_object.team_id if team_object else None,
5066 key_model_aliases=key_model_aliases,
5067 object_type="team",
5068 )
5069 except ProxyException:
5070 # Fallback: check team's access_group_ids
5071 team_access_group_ids: Final = (team_object.access_group_ids or []) if team_object else []
5072 if team_access_group_ids:
5073 models_from_groups: Final = await _get_models_from_access_groups(
5074 access_group_ids=team_access_group_ids,
5075 prisma_client=prisma_client,
5076 )
5077 if models_from_groups:
5078 return _can_object_call_model(
5079 model=model,
5080 llm_router=llm_router,
5081 models=list(dict.fromkeys([*(team_object.models if team_object else []), *models_from_groups])),
5082 team_model_aliases=team_model_aliases,
5083 team_id=team_object.team_id if team_object else None,
5084 key_model_aliases=key_model_aliases,
5085 object_type="team",
5086 )
5087 raise
5090async def get_authorized_resources_from_key_access_groups(
5091 valid_token: UserAPIKeyAuth | None,
5092 team_object: LiteLLM_TeamTable | None,
5093 resource_field: Literal["access_model_names", "access_mcp_server_ids", "access_agent_ids"],
5094) -> list[str]:
5095 """
5096 For each access_group_id on the key, fetch the LiteLLM_AccessGroupTable row
5097 and contribute its `resource_field` only if the group authorizes the caller
5098 as an owner — that is, the group's `assigned_team_ids` includes the key's
5099 `team_id`, or the group's `assigned_key_ids` includes the key's token. This
5100 preserves the team-as-owner boundary while still letting a group reach the
5101 key without first being added to the team's `access_group_ids` list.
5102 """
5103 if valid_token is None: 5103 ↛ 5104line 5103 didn't jump to line 5104 because the condition on line 5103 was never true
5104 return []
5105 key_access_group_ids: Final = list(valid_token.access_group_ids or [])
5106 if not key_access_group_ids: 5106 ↛ 5109line 5106 didn't jump to line 5109 because the condition on line 5106 was always true
5107 return []
5109 from litellm.proxy.proxy_server import prisma_client as _prisma_client
5110 from litellm.proxy.proxy_server import proxy_logging_obj as _proxy_logging_obj
5111 from litellm.proxy.proxy_server import user_api_key_cache as _user_api_key_cache
5113 if _prisma_client is None or _user_api_key_cache is None:
5114 return []
5116 key_team_id: Final = valid_token.team_id or (team_object.team_id if team_object is not None else None)
5117 key_token: Final = valid_token.token
5119 authorized_resources: Final[list[str]] = []
5120 for ag_id in key_access_group_ids:
5121 try:
5122 ag = await get_access_object(
5123 access_group_id=ag_id,
5124 prisma_client=_prisma_client,
5125 user_api_key_cache=_user_api_key_cache,
5126 proxy_logging_obj=_proxy_logging_obj,
5127 )
5128 except Exception:
5129 continue
5130 team_authorized = bool(key_team_id and key_team_id in (ag.assigned_team_ids or []))
5131 key_authorized = bool(key_token and key_token in (ag.assigned_key_ids or []))
5132 if team_authorized or key_authorized:
5133 authorized_resources.extend(getattr(ag, resource_field, []) or [])
5135 return list(set(authorized_resources))
5138async def _key_access_group_grants_model(
5139 model: str | list[str],
5140 valid_token: UserAPIKeyAuth | None,
5141 team_object: LiteLLM_TeamTable | None,
5142 llm_router: Router | None,
5143 key_model_aliases: Mapping[str, str] | None = None,
5144) -> bool:
5145 """
5146 Returns True if the key's `access_group_ids` expand to models that grant
5147 access to `model`. Used to let a key's access group override a team's
5148 model restriction in `common_checks`.
5149 """
5150 authorized_models: Final = await get_authorized_resources_from_key_access_groups(
5151 valid_token=valid_token,
5152 team_object=team_object,
5153 resource_field="access_model_names",
5154 )
5155 if not authorized_models:
5156 return False
5157 try:
5158 _can_object_call_model(
5159 model=model,
5160 llm_router=llm_router,
5161 models=authorized_models,
5162 team_model_aliases=valid_token.team_model_aliases if valid_token else None,
5163 team_id=valid_token.team_id if valid_token else None,
5164 key_model_aliases=key_model_aliases,
5165 object_type="key",
5166 )
5167 return True
5168 except ProxyException:
5169 return False
5172def can_project_access_model(
5173 model: str | list[str],
5174 project_object: LiteLLM_ProjectTable,
5175 llm_router: Router | None,
5176 key_model_aliases: Mapping[str, str] | None = None,
5177) -> Literal[True]:
5178 """
5179 Returns True if the project can access a specific model.
5181 Raises ProxyException if access is denied.
5182 """
5183 return _can_object_call_model(
5184 model=model,
5185 llm_router=llm_router,
5186 models=project_object.models if project_object else [],
5187 key_model_aliases=key_model_aliases,
5188 object_type="project",
5189 )
5192async def can_user_call_model(
5193 model: str | list[str],
5194 llm_router: Router | None,
5195 user_object: LiteLLM_UserTable | None,
5196 key_model_aliases: Mapping[str, str] | None = None,
5197) -> Literal[True]:
5198 if user_object is None: 5198 ↛ 5199line 5198 didn't jump to line 5199 because the condition on line 5198 was never true
5199 return True
5201 if SpecialModelNames.no_default_models.value in user_object.models: 5201 ↛ 5202line 5201 didn't jump to line 5202 because the condition on line 5201 was never true
5202 internal_message: Final = (
5203 f"User not allowed to access model. No default model access, only team models allowed. "
5204 f"Tried to access {model}"
5205 )
5206 raise ModelAccessDeniedProxyException(
5207 message=model_access_denied_client_message(model=model),
5208 internal_message=internal_message,
5209 type=ProxyErrorTypes.key_model_access_denied,
5210 param="model",
5211 code=status.HTTP_403_FORBIDDEN,
5212 )
5214 return _can_object_call_model(
5215 model=model,
5216 llm_router=llm_router,
5217 models=user_object.models,
5218 key_model_aliases=key_model_aliases,
5219 object_type="user",
5220 )
5223def _search_tool_names_from_object_permission(
5224 object_permission: LiteLLM_ObjectPermissionTable | None,
5225) -> list[str]:
5226 """Return allowlisted search tool names from object_permission (empty = unrestricted)."""
5227 if object_permission is None:
5228 return []
5229 raw: Final = object_permission.search_tools
5230 if not raw:
5231 return []
5232 return list(raw)
5235def _can_object_call_search_tool(
5236 search_tool_name: str,
5237 allowed_search_tools: list[str],
5238 object_type: Literal["key", "team", "project"],
5239) -> Literal[True]:
5240 """
5241 Check if an object (key/team/project) can access a specific search tool.
5243 Similar to _can_object_call_model but for search tools.
5245 Args:
5246 search_tool_name: The search tool being requested
5247 allowed_search_tools: List of allowed search tool names for this object
5248 object_type: Type of object for error messaging
5250 Returns:
5251 True if access is allowed
5253 Raises:
5254 ProxyException if access is denied
5255 """
5256 # Empty list means all search tools are allowed
5257 if not allowed_search_tools:
5258 return True
5260 # Check if the search tool is in the allowlist
5261 if search_tool_name in allowed_search_tools:
5262 return True
5264 # Access denied
5265 raise ProxyException(
5266 message=f"{object_type.capitalize()} not allowed to access search tool: {search_tool_name}. "
5267 f"Allowed search tools: {allowed_search_tools}",
5268 type=ProxyErrorTypes.key_model_access_denied,
5269 param="search_tool_name",
5270 code=status.HTTP_403_FORBIDDEN,
5271 )
5274async def can_key_call_search_tool(
5275 search_tool_name: str,
5276 valid_token: UserAPIKeyAuth,
5277) -> Literal[True]:
5278 """
5279 Check if a key can access a specific search tool.
5281 Similar to can_key_call_model but for search tools.
5283 Args:
5284 search_tool_name: The search tool being requested
5285 valid_token: The authenticated key
5287 Returns:
5288 True if access is allowed
5290 Raises:
5291 ProxyException if access is denied
5292 """
5293 return _can_object_call_search_tool(
5294 search_tool_name=search_tool_name,
5295 allowed_search_tools=_search_tool_names_from_object_permission(valid_token.object_permission),
5296 object_type="key",
5297 )
5300async def can_team_call_search_tool(
5301 search_tool_name: str,
5302 team_object: LiteLLM_TeamTable | None,
5303) -> Literal[True]:
5304 """
5305 Check if a team can access a specific search tool.
5307 Similar to can_team_access_model but for search tools.
5309 Args:
5310 search_tool_name: The search tool being requested
5311 team_object: The team object
5313 Returns:
5314 True if access is allowed
5316 Raises:
5317 ProxyException if access is denied
5318 """
5319 if team_object is None:
5320 return True
5322 return _can_object_call_search_tool(
5323 search_tool_name=search_tool_name,
5324 allowed_search_tools=_search_tool_names_from_object_permission(team_object.object_permission),
5325 object_type="team",
5326 )
5329async def can_user_view_search_tool(
5330 search_tool_name: str,
5331 valid_token: UserAPIKeyAuth,
5332 team_object: LiteLLM_TeamTable | None,
5333) -> bool:
5334 """
5335 Boolean variant of the key + team authorization enforced on /search, used to
5336 scope /search_tools/list so a non-admin caller only sees tools it may invoke.
5337 """
5338 try:
5339 await can_key_call_search_tool(
5340 search_tool_name=search_tool_name,
5341 valid_token=valid_token,
5342 )
5343 await can_team_call_search_tool(
5344 search_tool_name=search_tool_name,
5345 team_object=team_object,
5346 )
5347 except ProxyException:
5348 return False
5349 return True
5352async def is_valid_fallback_model(
5353 model: str,
5354 llm_router: Router | None,
5355 user_model: str | None,
5356) -> Literal[True]:
5357 """
5358 Try to route the fallback model request.
5360 Validate if it can't be routed.
5362 Help catch invalid fallback models.
5363 """
5364 await route_request(
5365 data={
5366 "model": model,
5367 "messages": [{"role": "user", "content": "Who was Alexander?"}],
5368 },
5369 llm_router=llm_router,
5370 user_model=user_model,
5371 route_type="acompletion", # route type shouldn't affect the fallback model check
5372 )
5374 return True
5377# The shape abbreviate_api_key writes into LiteLLM_VerificationToken.key_name. The
5378# last four characters are only barred from being whitespace or a control code,
5379# because a custom key's can be anything else, punctuation and non-ASCII included;
5380# a real key is at least MINIMUM_CUSTOM_KEY_LENGTH long, so it never fullmatches.
5381_MASKED_KEY_NAME_RE: Final = re.compile(r"sk-\.\.\.(?:[^\s\x00-\x1f\x7f-\x9f]{4})?")
5384def _apply_budget_exceeded_throttle(valid_token: UserAPIKeyAuth) -> bool:
5385 """
5386 Throttle an over-budget key instead of blocking it, when the key opted in
5387 via `throttle_on_budget_exceeded` and a global percentage is configured.
5389 Records the percentage on the request-scoped `budget_throttle_pct` so the
5390 rate limiter scales the key's TPM/RPM down to it; the persistent limits are
5391 left untouched so the throttle never compounds across requests. Returns True
5392 when the key was throttled (caller skips raising), False when it should still
5393 be hard-blocked.
5394 """
5395 pct: Final = budget_throttle_percentage()
5396 if pct is None or not should_throttle_budget_exceeded(valid_token):
5397 return False
5398 valid_token.budget_throttle_pct = pct
5399 return True
5402async def _virtual_key_max_budget_check(
5403 valid_token: UserAPIKeyAuth,
5404 proxy_logging_obj: ProxyLogging,
5405 user_obj: LiteLLM_UserTable | None = None,
5406):
5407 """
5408 Raises:
5409 BudgetExceededError if the token is over it's max budget.
5410 Triggers a budget alert if the token is over it's max budget.
5412 """
5413 if valid_token.max_budget is not None: 5413 ↛ 5414line 5413 didn't jump to line 5414 because the condition on line 5413 was never true
5414 from litellm.proxy.proxy_server import get_current_spend
5416 fallback_spend: Final = valid_token.spend or 0.0
5417 counter_key: Final = f"spend:key:{valid_token.token}"
5419 # Read spend from cross-pod counter (Redis-first) or cached object (fallback)
5420 spend: Final = await get_current_spend(
5421 counter_key=counter_key,
5422 fallback_spend=fallback_spend,
5423 max_budget=valid_token.max_budget,
5424 )
5426 ####################################
5427 # collect information for alerting #
5428 ####################################
5430 user_email = None
5431 # Check if the token has any user id information
5432 if user_obj is not None:
5433 user_email = user_obj.user_email
5435 call_info: Final = CallInfo(
5436 token=valid_token.token,
5437 spend=spend,
5438 max_budget=valid_token.max_budget,
5439 soft_budget=valid_token.soft_budget,
5440 user_id=valid_token.user_id,
5441 team_id=valid_token.team_id,
5442 organization_id=valid_token.org_id,
5443 user_email=user_email,
5444 key_alias=valid_token.key_alias,
5445 event_group=Litellm_EntityType.KEY,
5446 )
5447 asyncio.create_task(
5448 proxy_logging_obj.budget_alerts(
5449 type="token_budget",
5450 user_info=call_info,
5451 )
5452 )
5454 ####################################
5455 # collect information for alerting #
5456 ####################################
5458 # Defense-in-depth (GHSA-2rv4-xv66-fpjg): spend >= NaN is always False,
5459 # so a NaN max_budget would silently disable enforcement. Treat a
5460 # non-finite max_budget as "no configured limit" rather than as a bypass.
5461 if math.isfinite(valid_token.max_budget) and spend >= valid_token.max_budget:
5462 if _apply_budget_exceeded_throttle(valid_token):
5463 return
5464 # This message is returned to the caller, and key_name has no enforced
5465 # shape (a direct DB write bypasses abbreviate_api_key), so echo it only
5466 # when it still looks masked and fall back to the alias otherwise.
5467 key_label: Final = valid_token.key_alias or "key"
5468 key_descriptor: Final = (
5469 f"{key_label} ({valid_token.key_name})"
5470 if valid_token.key_name and _MASKED_KEY_NAME_RE.fullmatch(valid_token.key_name)
5471 else key_label
5472 )
5473 raise litellm.BudgetExceededError(
5474 current_cost=spend,
5475 max_budget=valid_token.max_budget,
5476 message=f"Budget has been exceeded! Key={key_descriptor} Current cost: {spend}, Max budget: {valid_token.max_budget}",
5477 entity_type=Litellm_EntityType.KEY.value,
5478 entity_id=valid_token.token,
5479 )
5482async def _virtual_key_multi_budget_check(
5483 valid_token: UserAPIKeyAuth,
5484):
5485 """
5486 Raises BudgetExceededError if any budget window in valid_token.budget_limits is exceeded.
5488 Each window has its own Redis counter keyed by spend:key:{token}:window:{budget_duration}.
5489 Using budget_duration (not list index) keeps counters stable when windows are reordered
5490 or removed during a key update.
5492 Note: counters are not seeded from DB on Redis cold-start. After a Redis flush,
5493 per-window spend resets to zero within the current window period. This is an acceptable
5494 trade-off: the DB stores reset_at timestamps but not per-window accumulated spend.
5495 """
5496 if not valid_token.budget_limits: 5496 ↛ 5499line 5496 didn't jump to line 5499 because the condition on line 5496 was always true
5497 return
5499 from litellm.proxy.proxy_server import get_current_spend
5501 for window in valid_token.budget_limits:
5502 w: dict = window if isinstance(window, dict) else window.model_dump()
5503 counter_key = f"spend:key:{valid_token.token}:window:{w['budget_duration']}"
5504 window_spend = await get_current_spend(
5505 counter_key=counter_key,
5506 fallback_spend=0.0,
5507 max_budget=w["max_budget"],
5508 window_entity_type="Key",
5509 window_entity_id=valid_token.token,
5510 window_duration=str(w["budget_duration"]),
5511 window_start=get_budget_window_start(w),
5512 )
5513 if math.isfinite(w["max_budget"]) and window_spend >= w["max_budget"]:
5514 raise litellm.BudgetExceededError(
5515 current_cost=window_spend,
5516 max_budget=w["max_budget"],
5517 message=(
5518 f"ExceededBudget: Key over {w['budget_duration']} budget. "
5519 f"Spend=${window_spend:.4f}, Limit=${w['max_budget']:.2f}"
5520 ),
5521 entity_type=Litellm_EntityType.KEY.value,
5522 entity_id=valid_token.token,
5523 )
5526async def _virtual_key_soft_budget_check(
5527 valid_token: UserAPIKeyAuth,
5528 proxy_logging_obj: ProxyLogging,
5529 user_obj: LiteLLM_UserTable | None = None,
5530):
5531 """
5532 Triggers a budget alert if the token is over it's soft budget.
5534 """
5536 if valid_token.soft_budget and valid_token.spend >= valid_token.soft_budget: 5536 ↛ 5537line 5536 didn't jump to line 5537 because the condition on line 5536 was never true
5537 verbose_proxy_logger.debug(
5538 "Crossed Soft Budget for token %s, spend %s, soft_budget %s",
5539 valid_token.token,
5540 valid_token.spend,
5541 valid_token.soft_budget,
5542 )
5543 call_info: Final = CallInfo(
5544 token=valid_token.token,
5545 spend=valid_token.spend,
5546 max_budget=valid_token.max_budget,
5547 soft_budget=valid_token.soft_budget,
5548 user_id=valid_token.user_id,
5549 team_id=valid_token.team_id,
5550 team_alias=valid_token.team_alias,
5551 organization_id=valid_token.org_id,
5552 user_email=user_obj.user_email if user_obj else None,
5553 key_alias=valid_token.key_alias,
5554 event_group=Litellm_EntityType.KEY,
5555 )
5557 asyncio.create_task(
5558 proxy_logging_obj.budget_alerts(
5559 type="soft_budget",
5560 user_info=call_info,
5561 )
5562 )
5565def _parse_email_list(raw: str | Sequence[object] | None) -> list[str]:
5566 """Parse emails from a list or comma-separated string."""
5567 if isinstance(raw, list):
5568 return [e.strip() for e in raw if isinstance(e, str) and e.strip()]
5569 elif isinstance(raw, str):
5570 return [e.strip() for e in raw.split(",") if e.strip()]
5571 return []
5574def _normalize_alert_emails(
5575 cfg: Mapping[str, str | Sequence[object] | None] | None,
5576) -> dict[str, list[str]]:
5577 """Coerce user-supplied threshold→recipients mapping to Dict[str, List[str]].
5579 Values may legitimately arrive as list, comma-separated string, or None
5580 from YAML/metadata; _parse_email_list tolerates all three.
5581 """
5582 if not cfg:
5583 return {}
5584 return {k: _parse_email_list(v) for k, v in cfg.items()}
5587def _merge_budget_alert_email_configs(
5588 global_cfg: Mapping[str, str | Sequence[object] | None] | None,
5589 per_key_cfg: Mapping[str, str | Sequence[object] | None] | None,
5590) -> dict[str, list[str]] | None:
5591 """
5592 Per-threshold additive merge: each threshold's recipient list is the union
5593 of global + per-key entries (deduped, global-first ordering). Missing
5594 thresholds on one side are inherited from the other.
5595 """
5596 global_cfg_normalized: Final = _normalize_alert_emails(global_cfg)
5597 per_key_cfg_normalized: Final = _normalize_alert_emails(per_key_cfg)
5598 if not global_cfg_normalized and not per_key_cfg_normalized:
5599 return None
5600 thresholds: Final = set(global_cfg_normalized) | set(per_key_cfg_normalized)
5601 return {
5602 t: list(dict.fromkeys(global_cfg_normalized.get(t, []) + per_key_cfg_normalized.get(t, []))) for t in thresholds
5603 }
5606async def _virtual_key_max_budget_alert_check(
5607 valid_token: UserAPIKeyAuth,
5608 proxy_logging_obj: ProxyLogging,
5609 user_obj: LiteLLM_UserTable | None = None,
5610):
5611 """
5612 Triggers a budget alert if the token has reached EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE
5613 (default 80%) of its max budget.
5614 This is a warning alert before the token actually exceeds the max budget.
5616 """
5618 if valid_token.max_budget is not None and valid_token.spend is not None and valid_token.spend > 0: 5618 ↛ 5619line 5618 didn't jump to line 5619 because the condition on line 5618 was never true
5619 owner_email: Final = user_obj.user_email if user_obj else None
5620 alert_email_config: Final[dict[str, list[str]] | None] = _merge_budget_alert_email_configs(
5621 global_cfg=litellm.default_key_max_budget_alert_emails,
5622 per_key_cfg=(valid_token.metadata or {}).get("max_budget_alert_emails"),
5623 )
5625 if isinstance(alert_email_config, dict) and alert_email_config:
5626 # New path: only create task if spend has crossed the lowest threshold
5627 min_pct: Final = min(
5628 (int(k) for k in alert_email_config if k.isdigit()),
5629 default=None,
5630 )
5631 if min_pct is None or valid_token.spend < valid_token.max_budget * (min_pct / 100.0):
5632 return
5634 call_info = CallInfo(
5635 token=valid_token.token,
5636 spend=valid_token.spend,
5637 max_budget=valid_token.max_budget,
5638 soft_budget=valid_token.soft_budget,
5639 user_id=valid_token.user_id,
5640 team_id=valid_token.team_id,
5641 team_alias=valid_token.team_alias,
5642 organization_id=valid_token.org_id,
5643 user_email=owner_email,
5644 key_alias=valid_token.key_alias,
5645 event_group=Litellm_EntityType.KEY,
5646 max_budget_alert_emails=alert_email_config,
5647 )
5648 asyncio.create_task(
5649 proxy_logging_obj.budget_alerts(
5650 type="max_budget_alert",
5651 user_info=call_info,
5652 )
5653 )
5654 else:
5655 # Old path: existing single 80% threshold — completely unchanged
5656 alert_threshold: Final = valid_token.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE
5658 if valid_token.spend >= alert_threshold and valid_token.spend < valid_token.max_budget:
5659 verbose_proxy_logger.debug(
5660 "Reached Max Budget Alert Threshold for token %s, spend %s, max_budget %s, alert_threshold %s",
5661 valid_token.token,
5662 valid_token.spend,
5663 valid_token.max_budget,
5664 alert_threshold,
5665 )
5666 call_info = CallInfo(
5667 token=valid_token.token,
5668 spend=valid_token.spend,
5669 max_budget=valid_token.max_budget,
5670 soft_budget=valid_token.soft_budget,
5671 user_id=valid_token.user_id,
5672 team_id=valid_token.team_id,
5673 team_alias=valid_token.team_alias,
5674 organization_id=valid_token.org_id,
5675 user_email=owner_email,
5676 key_alias=valid_token.key_alias,
5677 event_group=Litellm_EntityType.KEY,
5678 )
5680 asyncio.create_task(
5681 proxy_logging_obj.budget_alerts(
5682 type="max_budget_alert",
5683 user_info=call_info,
5684 )
5685 )
5688async def _check_team_member_budget(
5689 team_object: LiteLLM_TeamTable | None,
5690 user_object: LiteLLM_UserTable | None,
5691 valid_token: UserAPIKeyAuth | None,
5692 prisma_client: PrismaClient | None,
5693 user_api_key_cache: UserApiKeyCache,
5694 proxy_logging_obj: ProxyLogging,
5695 team_membership: LiteLLM_TeamMembership | None = None,
5696 team_membership_loaded: bool = False,
5697):
5698 """Check if team member is over their max budget within the team."""
5699 if ( 5699 ↛ 5705line 5699 didn't jump to line 5705 because the condition on line 5699 was never true
5700 team_object is not None
5701 and team_object.team_id is not None
5702 and valid_token is not None
5703 and valid_token.user_id is not None
5704 ):
5705 if not team_membership_loaded:
5706 team_membership = await get_team_membership(
5707 user_id=valid_token.user_id,
5708 team_id=team_object.team_id,
5709 prisma_client=prisma_client,
5710 user_api_key_cache=user_api_key_cache,
5711 proxy_logging_obj=proxy_logging_obj,
5712 )
5713 loaded_membership = team_membership
5715 # Per-member override wins; otherwise fall back to the team-level
5716 # default configured via team.metadata["team_member_budget_id"].
5717 team_member_budget: float | None = None
5718 member_budget_row: Final = loaded_membership.litellm_budget_table if loaded_membership is not None else None
5719 now: Final = get_utc_datetime()
5720 if member_budget_row is not None and member_budget_row.max_budget is not None:
5721 team_member_budget = member_budget_row.effective_max_budget(now=now)
5722 else:
5723 default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id")
5724 if isinstance(default_budget_id, str):
5725 default_budget: Final = await get_team_member_default_budget(
5726 budget_id=default_budget_id,
5727 prisma_client=prisma_client,
5728 user_api_key_cache=user_api_key_cache,
5729 )
5730 # Treat 0 on the team default as "no cap".
5731 # Per-member rows still respect 0 as an explicit admin disable.
5732 if (
5733 default_budget is not None
5734 and default_budget.max_budget is not None
5735 and default_budget.max_budget > 0
5736 ):
5737 team_member_budget = default_budget.max_budget + (
5738 member_budget_row.active_temp_budget_increase(now=now) if member_budget_row is not None else 0.0
5739 )
5741 if team_member_budget is not None:
5742 team_member_spend = (loaded_membership.spend if loaded_membership is not None else 0.0) or 0.0
5744 # Read from cross-pod counter (Redis-first) if available
5745 from litellm.proxy.proxy_server import get_current_spend
5747 team_member_spend = await get_current_spend(
5748 counter_key=f"spend:team_member:{valid_token.user_id}:{team_object.team_id}",
5749 fallback_spend=team_member_spend,
5750 max_budget=team_member_budget,
5751 )
5753 if math.isfinite(team_member_budget) and team_member_spend >= team_member_budget:
5754 raise litellm.BudgetExceededError(
5755 current_cost=team_member_spend,
5756 max_budget=team_member_budget,
5757 message=f"Budget has been exceeded! User={valid_token.user_id} in Team={team_object.team_id} Current cost: {team_member_spend}, Max budget: {team_member_budget}",
5758 entity_type=Litellm_EntityType.TEAM_MEMBER.value,
5759 entity_id=f"{valid_token.user_id}:{team_object.team_id}",
5760 )
5763async def _check_team_member_model_access(
5764 model: str | list[str],
5765 team_object: LiteLLM_TeamTable,
5766 valid_token: UserAPIKeyAuth,
5767 llm_router: Router | None,
5768 prisma_client: Optional["PrismaClient"],
5769 user_api_key_cache: UserApiKeyCache,
5770 proxy_logging_obj: ProxyLogging,
5771 team_membership: LiteLLM_TeamMembership | None = None,
5772 team_membership_loaded: bool = False,
5773 key_model_aliases: Mapping[str, str] | None = None,
5774) -> None:
5775 """
5776 Check if a team member's per-member model scope allows access to the requested model.
5778 Only enforced when the member's budget table has a non-empty allowed_models list.
5779 If allowed_models is empty or absent, the team-level models list applies (no extra restriction).
5780 """
5781 if valid_token.user_id is None or team_object.team_id is None:
5782 return
5784 if not team_membership_loaded:
5785 team_membership = await get_team_membership(
5786 user_id=valid_token.user_id,
5787 team_id=team_object.team_id,
5788 prisma_client=prisma_client,
5789 user_api_key_cache=user_api_key_cache,
5790 proxy_logging_obj=proxy_logging_obj,
5791 )
5792 loaded_membership = team_membership
5794 if (
5795 loaded_membership is None
5796 or loaded_membership.litellm_budget_table is None
5797 or not loaded_membership.litellm_budget_table.allowed_models
5798 ):
5799 return # no per-member restriction — inherit team-level check
5801 member_allowed_models: Final[list[str]] = loaded_membership.litellm_budget_table.allowed_models
5802 try:
5803 _can_object_call_model(
5804 model=model,
5805 llm_router=llm_router,
5806 models=member_allowed_models,
5807 object_type="team",
5808 team_id=team_object.team_id,
5809 key_model_aliases=key_model_aliases,
5810 )
5811 except ProxyException:
5812 internal_message: Final = (
5813 f"Team member not allowed to access model. User={valid_token.user_id}, Team={team_object.team_id}, "
5814 f"Model={model}. Allowed member models = {member_allowed_models}"
5815 )
5816 raise ModelAccessDeniedProxyException(
5817 message=model_access_denied_client_message(model=model),
5818 internal_message=internal_message,
5819 type=ProxyErrorTypes.team_model_access_denied,
5820 param="model",
5821 code=status.HTTP_403_FORBIDDEN,
5822 )
5825async def _team_max_budget_check(
5826 team_object: LiteLLM_TeamTable | None,
5827 valid_token: UserAPIKeyAuth | None,
5828 proxy_logging_obj: ProxyLogging,
5829):
5830 """
5831 Check if the team is over it's max budget.
5833 Raises:
5834 BudgetExceededError if the team is over it's max budget.
5835 Triggers a budget alert if the team is over it's max budget.
5836 """
5837 if team_object is not None and team_object.max_budget is not None: 5837 ↛ 5838line 5837 didn't jump to line 5838 because the condition on line 5837 was never true
5838 from litellm.proxy.proxy_server import get_current_spend
5840 # Read spend from cross-pod counter (Redis-first) or cached object (fallback)
5841 spend: Final = await get_current_spend(
5842 counter_key=f"spend:team:{team_object.team_id}",
5843 fallback_spend=team_object.spend or 0.0,
5844 max_budget=team_object.max_budget,
5845 )
5847 if math.isfinite(team_object.max_budget) and spend > team_object.max_budget:
5848 if valid_token:
5849 call_info: Final = CallInfo(
5850 token=valid_token.token,
5851 spend=spend,
5852 max_budget=team_object.max_budget,
5853 user_id=valid_token.user_id,
5854 team_id=valid_token.team_id,
5855 team_alias=valid_token.team_alias,
5856 organization_id=valid_token.org_id,
5857 event_group=Litellm_EntityType.TEAM,
5858 )
5859 asyncio.create_task(
5860 proxy_logging_obj.budget_alerts(
5861 type="team_budget",
5862 user_info=call_info,
5863 )
5864 )
5866 raise litellm.BudgetExceededError(
5867 current_cost=spend,
5868 max_budget=team_object.max_budget,
5869 message=f"Budget has been exceeded! Team={team_object.team_id} Current cost: {spend}, Max budget: {team_object.max_budget}",
5870 entity_type=Litellm_EntityType.TEAM.value,
5871 entity_id=team_object.team_id,
5872 )
5875async def _team_multi_budget_check(
5876 team_object: LiteLLM_TeamTable | None,
5877):
5878 """
5879 Raises BudgetExceededError if any budget window in team_object.budget_limits is exceeded.
5881 Each window has its own Redis counter keyed by spend:team:{team_id}:window:{budget_duration}.
5882 Using budget_duration (not list index) keeps counters stable when windows are reordered
5883 or removed during a team update.
5884 """
5885 if team_object is None or not team_object.budget_limits: 5885 ↛ 5888line 5885 didn't jump to line 5888 because the condition on line 5885 was always true
5886 return
5888 from litellm.proxy.proxy_server import get_current_spend
5890 for window in team_object.budget_limits:
5891 w: dict = window if isinstance(window, dict) else window.model_dump()
5892 counter_key = f"spend:team:{team_object.team_id}:window:{w['budget_duration']}"
5893 window_spend = await get_current_spend(
5894 counter_key=counter_key,
5895 fallback_spend=0.0,
5896 max_budget=w["max_budget"],
5897 window_entity_type="Team",
5898 window_entity_id=team_object.team_id,
5899 window_duration=str(w["budget_duration"]),
5900 window_start=get_budget_window_start(w),
5901 )
5902 if math.isfinite(w["max_budget"]) and window_spend >= w["max_budget"]:
5903 raise litellm.BudgetExceededError(
5904 current_cost=window_spend,
5905 max_budget=w["max_budget"],
5906 message=(
5907 f"ExceededBudget: Team={team_object.team_id} over {w['budget_duration']} budget. "
5908 f"Spend=${window_spend:.4f}, Limit=${w['max_budget']:.2f}"
5909 ),
5910 entity_type=Litellm_EntityType.TEAM.value,
5911 entity_id=team_object.team_id,
5912 )
5915async def _team_soft_budget_check(
5916 team_object: LiteLLM_TeamTable | None,
5917 valid_token: UserAPIKeyAuth | None,
5918 proxy_logging_obj: ProxyLogging,
5919):
5920 """
5921 Triggers a budget alert if the team is over it's soft budget.
5922 """
5923 if ( 5923 ↛ 5929line 5923 didn't jump to line 5929 because the condition on line 5923 was never true
5924 team_object is not None
5925 and team_object.soft_budget is not None
5926 and team_object.spend is not None
5927 and team_object.spend >= team_object.soft_budget
5928 ):
5929 verbose_proxy_logger.debug(
5930 "Crossed Soft Budget for team %s, spend %s, soft_budget %s",
5931 team_object.team_id,
5932 team_object.spend,
5933 team_object.soft_budget,
5934 )
5935 if valid_token:
5936 # Extract alert emails from team metadata
5937 alert_emails: list[str] | None = None
5938 if team_object.metadata is not None and isinstance(team_object.metadata, dict):
5939 soft_budget_alert_emails: Final = team_object.metadata.get("soft_budget_alerting_emails")
5940 if soft_budget_alert_emails is not None:
5941 if isinstance(soft_budget_alert_emails, list):
5942 alert_emails = [
5943 email for email in soft_budget_alert_emails if isinstance(email, str) and email.strip()
5944 ]
5945 elif isinstance(soft_budget_alert_emails, str):
5946 # Handle comma-separated string
5947 alert_emails = [email.strip() for email in soft_budget_alert_emails.split(",") if email.strip()]
5948 # Filter out empty strings
5949 if alert_emails:
5950 alert_emails = [email for email in alert_emails if email]
5951 else:
5952 alert_emails = None
5954 # Only send team soft budget alerts if alert_emails are configured
5955 # Team soft budget alerts are sent via metadata.soft_budget_alerting_emails, not global alerting
5956 if alert_emails is None or len(alert_emails) == 0:
5957 verbose_proxy_logger.debug(
5958 "Skipping team soft budget alert for team %s: no alert_emails configured in metadata.soft_budget_alerting_emails",
5959 team_object.team_id,
5960 )
5961 return
5963 call_info: Final = CallInfo(
5964 token=valid_token.token,
5965 spend=team_object.spend,
5966 max_budget=team_object.max_budget,
5967 soft_budget=team_object.soft_budget,
5968 user_id=valid_token.user_id,
5969 team_id=valid_token.team_id,
5970 team_alias=valid_token.team_alias,
5971 organization_id=valid_token.org_id,
5972 user_email=None, # Team-level alert, no specific user email
5973 key_alias=valid_token.key_alias,
5974 event_group=Litellm_EntityType.TEAM,
5975 alert_emails=alert_emails,
5976 )
5978 asyncio.create_task(
5979 proxy_logging_obj.budget_alerts(
5980 type="soft_budget",
5981 user_info=call_info,
5982 )
5983 )
5986async def _project_max_budget_check(
5987 project_object: LiteLLM_ProjectTableCachedObj | None,
5988 valid_token: UserAPIKeyAuth | None,
5989 proxy_logging_obj: ProxyLogging,
5990):
5991 """
5992 Check if the project is over its max budget.
5994 Raises:
5995 BudgetExceededError if the project is over its max budget.
5996 Triggers a budget alert if the project is over its max budget.
5997 """
5998 if project_object is None:
5999 return
6001 max_budget = None
6002 if project_object.litellm_budget_table is not None:
6003 max_budget = project_object.litellm_budget_table.max_budget
6005 if max_budget is None or not math.isfinite(max_budget):
6006 return
6008 from litellm.proxy.proxy_server import get_current_spend
6010 project_spend: Final = await get_current_spend(
6011 counter_key=project_spend_counter_key(project_object.project_id),
6012 fallback_spend=project_object.spend or 0.0,
6013 max_budget=max_budget,
6014 )
6016 if project_spend >= max_budget:
6017 if valid_token:
6018 call_info: Final = CallInfo(
6019 token=valid_token.token,
6020 spend=project_spend,
6021 max_budget=max_budget,
6022 user_id=valid_token.user_id,
6023 team_id=valid_token.team_id,
6024 team_alias=valid_token.team_alias,
6025 organization_id=valid_token.org_id,
6026 event_group=Litellm_EntityType.PROJECT,
6027 )
6028 asyncio.create_task(
6029 proxy_logging_obj.budget_alerts(
6030 type="project_budget",
6031 user_info=call_info,
6032 )
6033 )
6035 raise litellm.BudgetExceededError(
6036 current_cost=project_spend,
6037 max_budget=max_budget,
6038 message=f"Budget has been exceeded! Project={project_object.project_id} Current cost: {project_spend}, Max budget: {max_budget}",
6039 entity_type=Litellm_EntityType.PROJECT.value,
6040 entity_id=project_object.project_id,
6041 )
6044async def _project_soft_budget_check(
6045 project_object: LiteLLM_ProjectTableCachedObj | None,
6046 valid_token: UserAPIKeyAuth | None,
6047 proxy_logging_obj: ProxyLogging,
6048):
6049 """
6050 Triggers a budget alert if the project is over its soft budget.
6052 Mirrors _team_soft_budget_check() pattern.
6053 """
6054 if project_object is None:
6055 return
6057 soft_budget = None
6058 if project_object.litellm_budget_table is not None:
6059 soft_budget = project_object.litellm_budget_table.soft_budget
6061 if soft_budget is not None and project_object.spend is not None and project_object.spend >= soft_budget:
6062 verbose_proxy_logger.debug(
6063 "Crossed Soft Budget for project %s, spend %s, soft_budget %s",
6064 project_object.project_id,
6065 project_object.spend,
6066 soft_budget,
6067 )
6068 if valid_token:
6069 call_info: Final = CallInfo(
6070 token=valid_token.token,
6071 spend=project_object.spend,
6072 max_budget=None,
6073 soft_budget=soft_budget,
6074 user_id=valid_token.user_id,
6075 team_id=valid_token.team_id,
6076 team_alias=valid_token.team_alias,
6077 organization_id=valid_token.org_id,
6078 event_group=Litellm_EntityType.PROJECT,
6079 )
6080 asyncio.create_task(
6081 proxy_logging_obj.budget_alerts(
6082 type="soft_budget",
6083 user_info=call_info,
6084 )
6085 )
6088async def get_project_object(
6089 project_id: str,
6090 prisma_client: PrismaClient | None,
6091 user_api_key_cache: UserApiKeyCache,
6092 proxy_logging_obj: ProxyLogging | None = None,
6093) -> LiteLLM_ProjectTableCachedObj | None:
6094 """
6095 Fetch project object from cache or DB.
6097 Follows get_team_object() caching pattern with TTL and last_refreshed_at.
6099 Returns LiteLLM_ProjectTableCachedObj or None if not found.
6100 """
6101 if prisma_client is None: 6101 ↛ 6102line 6101 didn't jump to line 6102 because the condition on line 6101 was never true
6102 return None
6104 # Check cache first
6105 cache_key: Final = project_cache_key(project_id)
6106 deserialized_project: Final = await user_api_key_cache.async_get_cache(
6107 key=cache_key,
6108 model_type=LiteLLM_ProjectTableCachedObj,
6109 )
6110 if deserialized_project is not None: 6110 ↛ 6111line 6110 didn't jump to line 6111 because the condition on line 6110 was never true
6111 return deserialized_project
6113 # Fetch from DB
6114 project_row: Final = await _model_dump_table(ProjectRepository(prisma_client), "project").find_unique(
6115 where={"project_id": project_id},
6116 include={"litellm_budget_table": True},
6117 )
6118 if project_row is None: 6118 ↛ 6121line 6118 didn't jump to line 6121 because the condition on line 6118 was always true
6119 return None
6121 project_obj: Final = LiteLLM_ProjectTableCachedObj.model_validate(project_row.model_dump())
6123 # Cache with TTL following _cache_management_object pattern
6124 project_obj.last_refreshed_at = time.time()
6125 await _cache_management_object(
6126 key=cache_key,
6127 value=project_obj,
6128 user_api_key_cache=user_api_key_cache,
6129 proxy_logging_obj=proxy_logging_obj,
6130 model_type=LiteLLM_ProjectTableCachedObj,
6131 )
6133 return project_obj
6136async def delete_cached_project_object(
6137 project_id: str,
6138 user_api_key_cache: UserApiKeyCache,
6139) -> None:
6140 """
6141 Every endpoint that mutates litellm_projecttable must call this, or a stale project (e.g. a
6142 pre-update empty model allowlist) keeps being enforced until the TTL expires (LIT-3803).
6143 """
6144 from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast
6146 await evict_and_broadcast(
6147 cache_keys=(project_cache_key(project_id),),
6148 user_api_key_cache=user_api_key_cache,
6149 )
6152async def _organization_max_budget_check(
6153 valid_token: UserAPIKeyAuth | None,
6154 team_object: LiteLLM_TeamTable | None,
6155 prisma_client: PrismaClient | None,
6156 user_api_key_cache: UserApiKeyCache,
6157 proxy_logging_obj: ProxyLogging,
6158):
6159 """
6160 Check if the organization is over its max budget.
6162 This function checks the organization budget using:
6163 1. First, tries to use valid_token.org_id (if key has organization_id set)
6164 2. Falls back to team_object.organization_id (if key doesn't have org_id but team does)
6166 This ensures organization budget checks work even when keys don't have organization_id
6167 set directly, as long as their team belongs to an organization.
6169 Raises:
6170 BudgetExceededError if the organization is over its max budget.
6171 Triggers a budget alert if the organization is over its max budget.
6172 """
6173 if valid_token is None or prisma_client is None: 6173 ↛ 6174line 6173 didn't jump to line 6174 because the condition on line 6173 was never true
6174 return
6176 # Determine organization_id: first try from token, then fallback to team
6177 org_id: str | None = None
6178 if valid_token.org_id is not None: 6178 ↛ 6179line 6178 didn't jump to line 6179 because the condition on line 6178 was never true
6179 org_id = valid_token.org_id
6180 elif team_object is not None and team_object.organization_id is not None: 6180 ↛ 6181line 6180 didn't jump to line 6181 because the condition on line 6180 was never true
6181 org_id = team_object.organization_id
6183 # If no organization_id found, skip the check
6184 if org_id is None: 6184 ↛ 6188line 6184 didn't jump to line 6188 because the condition on line 6184 was always true
6185 return
6187 # Get organization object with budget table - use get_org_object so it can be mocked in tests
6188 try:
6189 org_table: Final = await get_org_object(
6190 org_id=org_id,
6191 prisma_client=prisma_client,
6192 user_api_key_cache=user_api_key_cache,
6193 proxy_logging_obj=proxy_logging_obj,
6194 include_budget_table=True,
6195 )
6196 except Exception:
6197 # If organization lookup fails, skip the check
6198 return
6200 if org_table is None:
6201 return
6203 carry_organization_budget_state(valid_token=valid_token, org_table=org_table)
6205 # Get max_budget from organization's budget table
6206 org_max_budget: float | None = None
6207 if org_table.litellm_budget_table is not None:
6208 org_max_budget = org_table.litellm_budget_table.max_budget
6210 if org_max_budget is None:
6211 return
6213 # Read spend from cross-pod counter (Redis-first) or cached object (fallback)
6214 from litellm.proxy.proxy_server import get_current_spend
6216 org_spend: Final = await get_current_spend(
6217 counter_key=f"spend:org:{org_id}",
6218 fallback_spend=org_table.spend or 0.0,
6219 max_budget=org_max_budget,
6220 )
6222 # Check if organization spend exceeds max budget
6223 if math.isfinite(org_max_budget) and org_spend >= org_max_budget:
6224 # Trigger budget alert
6225 call_info: Final = CallInfo(
6226 token=valid_token.token,
6227 spend=org_spend,
6228 max_budget=org_max_budget,
6229 user_id=valid_token.user_id,
6230 team_id=valid_token.team_id,
6231 team_alias=valid_token.team_alias,
6232 organization_id=org_id,
6233 event_group=Litellm_EntityType.ORGANIZATION,
6234 )
6235 asyncio.create_task(
6236 proxy_logging_obj.budget_alerts(
6237 type="organization_budget",
6238 user_info=call_info,
6239 )
6240 )
6242 raise litellm.BudgetExceededError(
6243 current_cost=org_spend,
6244 max_budget=org_max_budget,
6245 message=f"Budget has been exceeded! Organization={org_id} Current cost: {org_spend}, Max budget: {org_max_budget}",
6246 entity_type=Litellm_EntityType.ORGANIZATION.value,
6247 entity_id=org_id,
6248 )
6251async def _tag_max_budget_check(
6252 request_body: dict,
6253 prisma_client: PrismaClient | None,
6254 user_api_key_cache: UserApiKeyCache,
6255 proxy_logging_obj: ProxyLogging,
6256 valid_token: UserAPIKeyAuth | None,
6257):
6258 """
6259 Check if any tags in the request are over their max budget.
6261 Raises:
6262 BudgetExceededError if any tag is over its max budget.
6263 Triggers a budget alert if any tag is over its max budget.
6264 """
6265 from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body
6267 await tag_max_budget_check_for_tags(
6268 tags=get_tags_from_request_body(request_body=request_body),
6269 prisma_client=prisma_client,
6270 user_api_key_cache=user_api_key_cache,
6271 proxy_logging_obj=proxy_logging_obj,
6272 valid_token=valid_token,
6273 )
6276async def tag_max_budget_check_for_tags(
6277 tags: Sequence[str],
6278 prisma_client: PrismaClient | None,
6279 user_api_key_cache: UserApiKeyCache,
6280 proxy_logging_obj: ProxyLogging,
6281 valid_token: UserAPIKeyAuth | None,
6282) -> None:
6283 if prisma_client is None or not tags:
6284 return
6286 tag_objects: Final = await get_tag_objects_batch(
6287 tag_names=tags,
6288 prisma_client=prisma_client,
6289 user_api_key_cache=user_api_key_cache,
6290 proxy_logging_obj=proxy_logging_obj,
6291 )
6293 # Check budget for each tag
6294 for tag_name in tags:
6295 tag_object = tag_objects.get(tag_name)
6296 if tag_object is None:
6297 continue
6299 # Check if tag has budget limits
6300 if tag_object.litellm_budget_table is not None and tag_object.litellm_budget_table.max_budget is not None: 6300 ↛ 6301line 6300 didn't jump to line 6301 because the condition on line 6300 was never true
6301 from litellm.proxy.proxy_server import get_current_spend
6303 tag_spend = await get_current_spend(
6304 counter_key=f"spend:tag:{tag_name}",
6305 fallback_spend=tag_object.spend or 0.0,
6306 max_budget=tag_object.litellm_budget_table.max_budget,
6307 fallback_authoritative=True,
6308 )
6309 if tag_spend <= tag_object.litellm_budget_table.max_budget:
6310 continue
6311 raise litellm.BudgetExceededError(
6312 current_cost=tag_spend,
6313 max_budget=tag_object.litellm_budget_table.max_budget,
6314 message=f"Budget has been exceeded! Tag={tag_name} Current cost: {tag_spend}, Max budget: {tag_object.litellm_budget_table.max_budget}",
6315 entity_type=Litellm_EntityType.TAG.value,
6316 entity_id=tag_name,
6317 )
6320async def _model_access_group_max_budget_check(
6321 matched_model_access_groups: Sequence[str],
6322 prisma_client: PrismaClient | None,
6323 user_api_key_cache: UserApiKeyCache,
6324) -> None:
6325 """Block the request when a model access group that authorized it is over its max budget.
6327 Only the groups auth already matched are charged and therefore only they are checked, so a
6328 request that no budgeted group authorized costs nothing here.
6330 Like the tag check this is a plain read with no reservation, so concurrent requests can
6331 overshoot the ceiling slightly. The reservation counters are the precise path; this one covers
6332 the ``disable_budget_reservation`` case.
6334 The ceiling is exclusive, unlike the tag check it otherwise mirrors: a pool whose recorded
6335 spend has reached ``max_budget`` has nothing left to give, so the next request is refused.
6336 Keys and organizations already draw the line there. A non-positive budget means no budget,
6337 matching what the reservation path treats as unbudgeted.
6339 Raises:
6340 BudgetExceededError if a matched group is over its max budget.
6341 """
6342 if prisma_client is None or not matched_model_access_groups:
6343 return
6345 budgets: Final = await get_model_access_group_budgets_batch(
6346 access_group_names=matched_model_access_groups,
6347 prisma_client=prisma_client,
6348 user_api_key_cache=user_api_key_cache,
6349 )
6351 from litellm.proxy.proxy_server import get_current_spend
6353 for group in matched_model_access_groups:
6354 budget = budgets.get(group)
6355 if budget is None or budget.max_budget is None or budget.max_budget <= 0:
6356 continue
6358 group_spend = await get_current_spend(
6359 counter_key=model_access_group_spend_counter_key(group),
6360 fallback_spend=budget.spend,
6361 max_budget=budget.max_budget,
6362 fallback_authoritative=True,
6363 )
6364 if group_spend < budget.max_budget:
6365 continue
6366 raise litellm.BudgetExceededError(
6367 current_cost=group_spend,
6368 max_budget=budget.max_budget,
6369 message=f"Budget has been exceeded! Model access group={group} Current cost: {group_spend}, Max budget: {budget.max_budget}",
6370 entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP.value,
6371 entity_id=group,
6372 )
6375def is_model_allowed_by_pattern(model: str, allowed_model_pattern: str) -> bool:
6376 """
6377 Check if a model matches an allowed pattern.
6378 Handles exact matches and wildcard patterns.
6380 Args:
6381 model (str): The model to check (e.g., "bedrock/anthropic.claude-3-5-sonnet-20240620")
6382 allowed_model_pattern (str): The allowed pattern (e.g., "bedrock/*", "*", "openai/*")
6384 Returns:
6385 bool: True if model matches the pattern, False otherwise
6386 """
6387 if "*" in allowed_model_pattern:
6388 pattern: Final = f"^{allowed_model_pattern.replace('*', '.*')}$"
6389 return bool(re.match(pattern, model))
6391 return False
6394def _model_matches_any_wildcard_pattern_in_list(model: str, allowed_model_list: list) -> bool:
6395 """
6396 Returns True if a model matches any wildcard pattern in a list.
6398 eg.
6399 - model=`bedrock/us.amazon.nova-micro-v1:0`, allowed_models=`bedrock/*` returns True
6400 - model=`bedrock/us.amazon.nova-micro-v1:0`, allowed_models=`bedrock/us.*` returns True
6401 - model=`bedrockzzzz/us.amazon.nova-micro-v1:0`, allowed_models=`bedrock/*` returns False
6402 """
6404 if any( 6404 ↛ 6409line 6404 didn't jump to line 6409 because the condition on line 6404 was never true
6405 _is_wildcard_pattern(allowed_model_pattern)
6406 and is_model_allowed_by_pattern(model=model, allowed_model_pattern=allowed_model_pattern)
6407 for allowed_model_pattern in allowed_model_list
6408 ):
6409 return True
6411 if any( 6411 ↛ 6418line 6411 didn't jump to line 6418 because the condition on line 6411 was never true
6412 _is_wildcard_pattern(allowed_model_pattern)
6413 and _model_custom_llm_provider_matches_wildcard_pattern(
6414 model=model, allowed_model_pattern=allowed_model_pattern
6415 )
6416 for allowed_model_pattern in allowed_model_list
6417 ):
6418 return True
6420 return False
6423def _model_custom_llm_provider_matches_wildcard_pattern(model: str, allowed_model_pattern: str) -> bool:
6424 """
6425 Returns True for this scenario:
6426 - `model=gpt-4o`
6427 - `allowed_model_pattern=openai/*`
6429 or
6430 - `model=claude-3-5-sonnet-20240620`
6431 - `allowed_model_pattern=anthropic/*`
6433 A model that already carries a namespace get_llm_provider did not consume
6434 (e.g. `bedrockz/anthropic.claude-...`) is never granted here: its provider was
6435 inferred from a fragment of the full string, so rebuilding
6436 `{provider}/{model}` would produce `bedrock/bedrockz/...` and slip an
6437 unrecognized namespace through a `bedrock/*` key.
6438 """
6439 try:
6440 stripped_model, custom_llm_provider, _, _ = get_llm_provider(model=model)
6441 except Exception:
6442 return False
6444 if stripped_model == model and "/" in model:
6445 return False
6447 return is_model_allowed_by_pattern(
6448 model=f"{custom_llm_provider}/{stripped_model}",
6449 allowed_model_pattern=allowed_model_pattern,
6450 )
6453def _is_wildcard_pattern(allowed_model_pattern: str) -> bool:
6454 """
6455 Returns True if the pattern is a wildcard pattern.
6457 Checks if `*` is in the pattern.
6458 """
6459 return "*" in allowed_model_pattern
6462async def vector_store_access_check(
6463 request_body: dict,
6464 team_object: LiteLLM_TeamTable | None,
6465 valid_token: UserAPIKeyAuth | None,
6466):
6467 """
6468 Checks if the object (key, team, org) has access to the vector store.
6470 Raises ProxyException if the object (key, team, org) cannot access the specific vector store.
6471 """
6472 from litellm.proxy.proxy_server import prisma_client
6474 #########################################################
6475 # Get the vector store the user is trying to access
6476 #########################################################
6477 if prisma_client is None: 6477 ↛ 6478line 6477 didn't jump to line 6478 because the condition on line 6477 was never true
6478 verbose_proxy_logger.debug("Prisma client not found, skipping vector store access check")
6479 return True
6481 if litellm.vector_store_registry is None:
6482 verbose_proxy_logger.debug("Vector store registry not found, skipping vector store access check")
6483 return True
6485 vector_store_ids_to_run: Final = litellm.vector_store_registry.get_vector_store_ids_to_run(
6486 non_default_params=request_body, tools=request_body.get("tools", None)
6487 )
6488 if not vector_store_ids_to_run:
6489 verbose_proxy_logger.debug("Vector store to run not found, skipping vector store access check")
6490 return True
6492 #########################################################
6493 # Check if the object (key, team, org) has access to the vector store
6494 #########################################################
6495 # Check if the key can access the vector store
6496 if valid_token is not None and valid_token.object_permission_id is not None: 6496 ↛ 6497line 6496 didn't jump to line 6497 because the condition on line 6496 was never true
6497 key_object_permission: Final = await _object_permission_table(
6498 ObjectPermissionRepository(prisma_client)
6499 ).find_unique(
6500 where={"object_permission_id": valid_token.object_permission_id},
6501 )
6502 if key_object_permission is not None:
6503 _can_object_call_vector_stores(
6504 object_type="key",
6505 vector_store_ids_to_run=vector_store_ids_to_run,
6506 object_permissions=key_object_permission,
6507 )
6509 # Check if the team can access the vector store
6510 if team_object is not None and team_object.object_permission_id is not None: 6510 ↛ 6511line 6510 didn't jump to line 6511 because the condition on line 6510 was never true
6511 team_object_permission: Final = await _object_permission_table(
6512 ObjectPermissionRepository(prisma_client)
6513 ).find_unique(
6514 where={"object_permission_id": team_object.object_permission_id},
6515 )
6516 if team_object_permission is not None:
6517 _can_object_call_vector_stores(
6518 object_type="team",
6519 vector_store_ids_to_run=vector_store_ids_to_run,
6520 object_permissions=team_object_permission,
6521 )
6522 return True
6525def _can_object_call_vector_stores(
6526 object_type: Literal["key", "team", "org"],
6527 vector_store_ids_to_run: list[str],
6528 object_permissions: _VectorStorePermissionsRow | None,
6529):
6530 """
6531 Raises ProxyException if the object (key, team, org) cannot access the specific vector store.
6532 """
6533 if object_permissions is None:
6534 return True
6536 if object_permissions.vector_stores is None:
6537 return True
6539 # If length is 0, then the object has access to all vector stores.
6540 if len(object_permissions.vector_stores) == 0:
6541 return True
6543 for vector_store_id in vector_store_ids_to_run:
6544 if vector_store_id not in object_permissions.vector_stores:
6545 raise ProxyException(
6546 message=f"User not allowed to access vector store. Tried to access {vector_store_id}. Only allowed to access {object_permissions.vector_stores}",
6547 type=ProxyErrorTypes.get_vector_store_access_error_type_for_object(object_type),
6548 param="vector_store",
6549 code=status.HTTP_401_UNAUTHORIZED,
6550 )
6552 return True