# What is this? ## Common auth checks between jwt + key based auth """ Got Valid Token from Cache, DB Run checks for: 1. If user can call model 2. If user is in budget 3. If end_user ('user' passed to /chat/completions, /embeddings endpoint) is in budget """ import asyncio import math import re import time from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Type, Union, cast from fastapi import HTTPException, Request, status from pydantic import BaseModel import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.dual_cache import LimitedSizeOrderedDict from litellm.constants import ( CLI_JWT_EXPIRATION_HOURS, CLI_JWT_TOKEN_NAME, DEFAULT_ACCESS_GROUP_CACHE_TTL, DEFAULT_IN_MEMORY_TTL, DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, DEFAULT_MAX_RECURSE_DEPTH, EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE, ) from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.proxy._types import ( RBAC_ROLES, CallInfo, LiteLLM_AccessGroupTable, LiteLLM_BudgetTable, LiteLLM_EndUserTable, Litellm_EntityType, LiteLLM_JWTAuth, LiteLLM_ManagedVectorStoresTable, LiteLLM_ObjectPermissionTable, LiteLLM_OrganizationMembershipTable, LiteLLM_OrganizationTable, LiteLLM_ProjectTableCachedObj, LiteLLM_TagTable, LiteLLM_TeamMembership, LiteLLM_TeamTable, LiteLLM_TeamTableCachedObj, LiteLLM_UserTable, LiteLLMRoutes, LitellmUserRoles, NewTeamRequest, ProxyErrorTypes, ProxyException, RoleBasedPermissions, SpecialModelNames, UserAPIKeyAuth, ) from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec from litellm.proxy.common_utils.http_parsing_utils import ( _safe_get_request_headers, _safe_get_request_query_params, ) from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.guardrails.tool_name_extraction import ( TOOL_CAPABLE_CALL_TYPES, extract_request_tool_names, ) from litellm.proxy.route_llm_request import route_request from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.project_repository import ProjectRepository from litellm.repositories.table_repositories import ( AccessGroupRepository, EndUserRepository, JWTKeyMappingRepository, ManagedVectorStoresRepository, TagRepository, TeamMembershipRepository, ) from litellm.repositories.team_repository import TeamRepository from litellm.repositories.user_repository import UserRepository from litellm.router import Router from litellm.utils import get_utc_datetime from .auth_checks_organization import organization_role_based_access_check from .auth_utils import get_model_from_request if TYPE_CHECKING: from opentelemetry.trace import Span as _Span Span = Union[_Span, Any] else: Span = Any last_db_access_time = LimitedSizeOrderedDict(max_size=100) db_cache_expiry = DEFAULT_IN_MEMORY_TTL # refresh every 5s all_routes = LiteLLMRoutes.openai_routes.value + LiteLLMRoutes.management_routes.value def _log_budget_lookup_failure(entity: str, error: Exception) -> None: """ Log a warning when budget lookup fails; cache will not be populated. Skips logging for expected "user not found" cases (bare Exception from get_user_object when user_id_upsert=False). Adds a schema migration hint when the error appears schema-related. """ # Skip logging for expected "user not found" - not caching is correct if str(error) == "" and type(error).__name__ == "Exception": return err_str = str(error).lower() hint = "" if any( x in err_str for x in ("column", "schema", "does not exist", "prisma", "migrate") ): hint = ( " Run `prisma db push` or `prisma migrate deploy` to fix schema mismatches." ) verbose_proxy_logger.error( f"Budget lookup failed for {entity}; cache will not be populated. " f"Each request will hit the database. Error: {error}.{hint}" ) def _get_router_zero_cost_cache(llm_router: Router) -> Optional[Dict[str, bool]]: """ Return the router's per-instance zero-cost cache, or ``None`` for objects that don't expose one (e.g. ``MagicMock`` stand-ins in unit tests). The cache lives on the ``Router`` instance so it: * is invalidated by ``Router._invalidate_model_group_info_cache`` on any model add/remove/upsert (including in-place pricing changes via ``/model/update``, which go through ``upsert_deployment``); * dies with the router itself — no risk of CPython reusing the previous router's ``id()`` and serving its cached entries. """ cache = getattr(llm_router, "_zero_cost_cache", None) return cache if isinstance(cache, dict) else None def _is_model_cost_zero( model: Optional[Union[str, List[str]]], llm_router: Optional[Router] ) -> bool: """ Check if a model has zero cost (no configured pricing). Uses the router's get_model_group_info method to get pricing information. Args: model: The model name or list of model names llm_router: The LiteLLM router instance Returns: bool: True if all costs for the model are zero, False otherwise """ if model is None or llm_router is None: return False # Handle list of models model_list = [model] if isinstance(model, str) else model zero_cost_cache = _get_router_zero_cost_cache(llm_router) for model_name in model_list: if zero_cost_cache is not None: cached = zero_cost_cache.get(model_name) if cached is not None: if cached is False: return False continue try: # Use router's get_model_group_info method directly for better reliability model_group_info = llm_router.get_model_group_info(model_group=model_name) if model_group_info is None: # Model not found or no pricing info available # Conservative approach: assume it has cost verbose_proxy_logger.debug( f"No model group info found for {model_name}, assuming it has cost" ) if zero_cost_cache is not None: zero_cost_cache[model_name] = False return False # Check costs for this model # Only allow bypass if BOTH costs are explicitly set to 0 (not None) input_cost = model_group_info.input_cost_per_token output_cost = model_group_info.output_cost_per_token # If costs are not explicitly configured (None), assume it has cost if input_cost is None or output_cost is None: verbose_proxy_logger.debug( f"Model {model_name} has undefined cost (input: {input_cost}, output: {output_cost}), assuming it has cost" ) if zero_cost_cache is not None: zero_cost_cache[model_name] = False return False # If either cost is non-zero, return False if input_cost > 0 or output_cost > 0: verbose_proxy_logger.debug( f"Model {model_name} has non-zero cost (input: {input_cost}, output: {output_cost})" ) if zero_cost_cache is not None: zero_cost_cache[model_name] = False return False # Costs are 0 — verify this is from explicit configuration, # not from defaulted sparse auto-registration entries. # See: https://github.com/BerriAI/litellm/issues/24770 safe_name = str(model_name).replace("\n", "").replace("\r", "") if not _is_cost_explicitly_configured(model_name, llm_router): verbose_proxy_logger.debug( "Model %s has zero cost but no explicit cost " "configuration in model_cost entry — treating as unknown " "cost (enforce budget)", safe_name, ) if zero_cost_cache is not None: zero_cost_cache[model_name] = False return False verbose_proxy_logger.debug( "Model %s has zero cost explicitly configured (input: %s, output: %s)", safe_name, input_cost, output_cost, ) if zero_cost_cache is not None: zero_cost_cache[model_name] = True except Exception as e: # If we can't determine the cost, assume it has cost (conservative approach) verbose_proxy_logger.debug( f"Error checking cost for model {model_name}: {str(e)}, assuming it has cost" ) return False # All models checked have zero cost return True def _is_cost_explicitly_configured(model: str, llm_router: "Router") -> bool: """ Check if any deployment in the model group has cost fields explicitly set in its litellm.model_cost entry. When Router._create_deployment() registers a model not in the global cost map, it creates a sparse entry like {"id": ""} with no cost fields. _get_model_info_helper() then defaults missing costs to 0. This function detects that scenario by checking the raw model_cost entry. """ for deployment in llm_router.model_list: if deployment.get("model_name") != model: continue model_id = deployment.get("model_info", {}).get("id") if model_id is None: continue raw_entry = litellm.model_cost.get(model_id, {}) if "input_cost_per_token" in raw_entry or "output_cost_per_token" in raw_entry: return True return False async def _run_project_checks( project_object: Optional[LiteLLM_ProjectTableCachedObj], _model: Optional[Union[str, List[str]]], llm_router: Optional[Router], skip_budget_checks: bool, valid_token: Optional[UserAPIKeyAuth], proxy_logging_obj: ProxyLogging, ) -> None: """ Run all project-level checks: blocked, model access, budget, soft budget. Extracted from common_checks() to keep statement count manageable. """ if project_object is None: return # 1.1. If project is blocked if project_object.blocked is True: raise Exception( f"Project={project_object.project_id} is blocked. Update via `/project/update` if you're an admin." ) # 2.2 If project can call model if _model and len(project_object.models) > 0: can_project_access_model( model=_model, project_object=project_object, llm_router=llm_router, ) if not skip_budget_checks: # 3.0.2. If project is in budget await _project_max_budget_check( project_object=project_object, valid_token=valid_token, proxy_logging_obj=proxy_logging_obj, ) # 3.0.3. If project is over soft budget (alert only, doesn't block) await _project_soft_budget_check( project_object=project_object, valid_token=valid_token, proxy_logging_obj=proxy_logging_obj, ) def _enforce_user_param_check( general_settings: dict, request: Request, request_body: dict, route: str ) -> None: if not general_settings.get("enforce_user_param", False): return http_method = request.method if hasattr(request, "method") else None is_post_method = http_method and http_method.upper() == "POST" is_openai_route = RouteChecks.is_llm_api_route(route=route) is_mcp_route = ( route in LiteLLMRoutes.mcp_routes.value or RouteChecks.check_route_access( route=route, allowed_routes=LiteLLMRoutes.mcp_routes.value ) ) if ( is_post_method and is_openai_route and not is_mcp_route and "user" not in request_body ): raise Exception( f"'user' param not passed in. 'enforce_user_param'={general_settings['enforce_user_param']}" ) def _reject_clientside_metadata_tags_check( general_settings: dict, request_body: dict, route: str ) -> None: if not general_settings.get("reject_clientside_metadata_tags", False): return if ( RouteChecks.is_llm_api_route(route=route) and "metadata" in request_body and isinstance(request_body["metadata"], dict) and "tags" in request_body["metadata"] ): raise ProxyException( 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.", type=ProxyErrorTypes.bad_request_error, param="metadata.tags", code=status.HTTP_400_BAD_REQUEST, ) def _global_proxy_budget_check( global_proxy_spend: Optional[float], skip_budget_checks: bool, route: str ) -> None: if ( litellm.max_budget > 0 and not skip_budget_checks and global_proxy_spend is not None and RouteChecks.is_llm_api_route(route=route) and route != "/v1/models" and route != "/models" ): if ( math.isfinite(litellm.max_budget) and global_proxy_spend > litellm.max_budget ): raise litellm.BudgetExceededError( current_cost=global_proxy_spend, max_budget=litellm.max_budget ) _GUARDRAIL_MODIFICATION_KEYS: tuple = ( "guardrails", "disable_global_guardrails", "disable_global_guardrail", "opted_out_global_guardrails", ) def _guardrail_modification_check( request_body: dict, team_object: Optional[LiteLLM_TeamTable] ) -> None: """ Reject user-supplied metadata flags that would modify guardrail behavior unless the team has explicit permission. Checked keys include the plural ``guardrails`` list plus the per-request toggles that influence whether default-on guardrails run (``disable_global_guardrails``, ``disable_global_guardrail`` singular, and ``opted_out_global_guardrails``). User-supplied values for the bypass toggles are also silently ignored by ``_get_admin_metadata`` at read time; this check adds defense in depth by failing loudly at the auth layer so operators see an explicit 403 instead of a confusing silent-ignore. """ from litellm.proxy.guardrails.guardrail_helpers import can_modify_guardrails def _coerce_to_dict(container: Any) -> Optional[dict]: """Accept dict or JSON-string (from multipart/form-data or extra_body). Without this, an attacker can smuggle guardrail keys past the check by sending ``{"metadata": "{\\"disable_global_guardrails\\": true}"}`` — ``isinstance(dict)`` on the string returns False, the check returns no-modification, and ``add_litellm_data_to_request`` parses the string to a dict downstream. """ if isinstance(container, dict): return container if isinstance(container, str): parsed = safe_json_loads(container) return parsed if isinstance(parsed, dict) else None return None def _user_requested_modification(container: Any) -> bool: coerced = _coerce_to_dict(container) if coerced is None: return False return any(key in coerced for key in _GUARDRAIL_MODIFICATION_KEYS) # Check both metadata keys — callers can populate either depending on the # endpoint. Cover the top-level too so root-level injection is rejected. modifies = ( _user_requested_modification(request_body.get("metadata")) or _user_requested_modification(request_body.get("litellm_metadata")) or _user_requested_modification(request_body) ) if not modifies: return if not can_modify_guardrails(team_object): raise HTTPException( status_code=403, detail={ "error": "Your team does not have permission to modify guardrails." }, ) async def check_tools_allowlist( request_body: dict, valid_token: Optional[UserAPIKeyAuth], team_object: Optional[LiteLLM_TeamTable], route: str, ) -> None: """ Enforce key/team tool allowlist (metadata.allowed_tools). No DB in hot path — effective allowlist is read from valid_token.metadata and valid_token.team_metadata. Raises ProxyException with tool_access_denied if a tool is not allowed. """ from litellm.litellm_core_utils.api_route_to_call_types import ( get_call_types_for_route, ) if valid_token is None: return call_types = get_call_types_for_route(route) if not call_types or not any( ct.value in TOOL_CAPABLE_CALL_TYPES for ct in call_types ): return tool_names = extract_request_tool_names(route, request_body) if not tool_names: return key_meta = ( (valid_token.metadata or {}) if isinstance(valid_token.metadata, dict) else {} ) team_meta = ( (valid_token.team_metadata or {}) if isinstance(valid_token.team_metadata, dict) else {} ) key_allowed = key_meta.get("allowed_tools") team_allowed = team_meta.get("allowed_tools") effective = ( key_allowed if (isinstance(key_allowed, list) and len(key_allowed) > 0) else team_allowed ) if not isinstance(effective, list) or len(effective) == 0: return allowed_set = {str(t) for t in effective} disallowed = [n for n in tool_names if n not in allowed_set] if disallowed: raise ProxyException( message=f"Tool(s) {disallowed} are not in the allowed tools list for this key/team.", type=ProxyErrorTypes.tool_access_denied, param="tools", code=status.HTTP_403_FORBIDDEN, ) # Read-only discovery routes that incur no spend. Kept narrower than info_routes so an exhausted # budget cannot reach side-effectful routes like /health/services (Slack/email/webhook). See #27923. MODEL_DISCOVERY_ROUTES = frozenset( { "/v1/models", "/models", "/model/info", "/v1/model/info", "/v2/model/info", "/model_group/info", } ) async def common_checks( request_body: dict, team_object: Optional[LiteLLM_TeamTable], user_object: Optional[LiteLLM_UserTable], end_user_object: Optional[LiteLLM_EndUserTable], global_proxy_spend: Optional[float], general_settings: dict, route: str, llm_router: Optional[Router], proxy_logging_obj: ProxyLogging, valid_token: Optional[UserAPIKeyAuth], request: Request, skip_budget_checks: bool = False, project_object: Optional[LiteLLM_ProjectTableCachedObj] = None, ) -> bool: """ Common checks across jwt + key-based auth. 1. If team is blocked 1.1. If project is blocked 2. If team can call model 2.2 If project can call model 3. If team is in budget 3.0.2. If project is in budget 3.0.3. If project is over soft budget (alert only) 4. If user passed in (JWT or key.user_id) - is in budget 5. If end_user (either via JWT or 'user' passed to /chat/completions, /embeddings endpoint) is in budget 6. [OPTIONAL] If 'enforce_end_user' enabled - did developer pass in 'user' param for openai endpoints 7. [OPTIONAL] If 'litellm.max_budget' is set (>0), is proxy under budget 8. [OPTIONAL] If guardrails modified - is request allowed to change this 9. Check if request body is safe 10. [OPTIONAL] Organization checks - is user_object.organization_id is set, run these checks 11. [OPTIONAL] Vector store checks - is the object allowed to access the vector store """ from litellm.proxy.proxy_server import prisma_client, user_api_key_cache _model: Optional[Union[str, List[str]]] = get_model_from_request( request_data=request_body, route=route, request_headers=_safe_get_request_headers(request=request), request_query_params=_safe_get_request_query_params(request=request), llm_router=llm_router, ) if route in MODEL_DISCOVERY_ROUTES: skip_budget_checks = True # 1. If team is blocked if team_object is not None and team_object.blocked is True: raise Exception( f"Team={team_object.team_id} is blocked. Update via `/team/unblock` if you're an admin." ) # 2. If team can call model (or key's access_group_ids grant it) if _model and team_object: with tracer.trace("litellm.proxy.auth.common_checks.can_team_access_model"): try: await can_team_access_model( model=_model, team_object=team_object, llm_router=llm_router, team_model_aliases=( valid_token.team_model_aliases if valid_token else None ), ) except ProxyException as team_denial: if team_denial.type != ProxyErrorTypes.team_model_access_denied: raise if not await _key_access_group_grants_model( model=_model, valid_token=valid_token, team_object=team_object, llm_router=llm_router, ): raise # 2.2. If team member has per-member model scope, enforce it if _model and team_object and valid_token and valid_token.user_id: with tracer.trace( "litellm.proxy.auth.common_checks.check_team_member_model_access" ): await _check_team_member_model_access( model=_model, team_object=team_object, valid_token=valid_token, llm_router=llm_router, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) # Require trace id for agent keys when agent has require_trace_id_on_calls_by_agent if valid_token is not None and valid_token.agent_id: from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers agent = global_agent_registry.get_agent_by_id(agent_id=valid_token.agent_id) if agent is not None: require_trace_id = (agent.litellm_params or {}).get( "require_trace_id_on_calls_by_agent" ) if require_trace_id: headers_dict = dict(request.headers) trace_id = get_chain_id_from_headers(headers_dict) if not trace_id: raise ProxyException( message="Requests made with this agent's key must include the x-litellm-trace-id header.", type=ProxyErrorTypes.bad_request_error, param=None, code=status.HTTP_400_BAD_REQUEST, ) ## 2.1 If user can call model (if personal key) if _model and team_object is None and user_object is not None: with tracer.trace("litellm.proxy.auth.common_checks.can_user_call_model"): await can_user_call_model( model=_model, llm_router=llm_router, user_object=user_object, ) # 1.1 - 2.2 - 3.0.2 - 3.0.3: Project checks (blocked, model access, budget) with tracer.trace("litellm.proxy.auth.common_checks.run_project_checks"): await _run_project_checks( project_object=project_object, _model=_model, llm_router=llm_router, skip_budget_checks=skip_budget_checks, valid_token=valid_token, proxy_logging_obj=proxy_logging_obj, ) # Run before apply_key_tags_pre_auth injects key metadata.tags into request_body. _reject_clientside_metadata_tags_check(general_settings, request_body, route) # If this is a free model, skip all budget checks if not skip_budget_checks: # 3. If team is in budget with tracer.trace("litellm.proxy.auth.common_checks.team_max_budget_check"): await _team_max_budget_check( team_object=team_object, proxy_logging_obj=proxy_logging_obj, valid_token=valid_token, ) # 3.1. Multi-window budget check for team with tracer.trace("litellm.proxy.auth.common_checks.team_multi_budget_check"): await _team_multi_budget_check(team_object=team_object) # 3.2. Multi-window budget check for key with tracer.trace( "litellm.proxy.auth.common_checks.virtual_key_multi_budget_check" ): if valid_token is not None: await _virtual_key_multi_budget_check(valid_token=valid_token) # 3.0.5. If team is over soft budget (alert only, doesn't block) with tracer.trace("litellm.proxy.auth.common_checks.team_soft_budget_check"): await _team_soft_budget_check( team_object=team_object, proxy_logging_obj=proxy_logging_obj, valid_token=valid_token, ) # 3.1. If organization is in budget with tracer.trace( "litellm.proxy.auth.common_checks.organization_max_budget_check" ): await _organization_max_budget_check( valid_token=valid_token, team_object=team_object, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) if valid_token is not None: from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup LiteLLMProxyRequestSetup.apply_key_tags_pre_auth( request_data=request_body, user_api_key_dict=valid_token, ) with tracer.trace("litellm.proxy.auth.common_checks.tag_max_budget_check"): await _tag_max_budget_check( request_body=request_body, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, valid_token=valid_token, ) # 4. If user is in budget ## 4.1 check personal budget, if personal key if ( (team_object is None or team_object.team_id is None) and user_object is not None and user_object.max_budget is not None ): user_budget = user_object.max_budget from litellm.proxy.proxy_server import get_current_spend user_spend = await get_current_spend( counter_key=f"spend:user:{user_object.user_id}", fallback_spend=user_object.spend or 0.0, max_budget=user_budget, ) if math.isfinite(user_budget) and user_spend >= user_budget: raise litellm.BudgetExceededError( current_cost=user_spend, max_budget=user_budget, message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}", ) ## 4.2 check team member budget, if team key with tracer.trace("litellm.proxy.auth.common_checks.check_team_member_budget"): await _check_team_member_budget( team_object=team_object, user_object=user_object, valid_token=valid_token, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) # 5. If end_user ('user' passed to /chat/completions, /embeddings endpoint) is in budget if ( end_user_object is not None and end_user_object.litellm_budget_table is not None ): await _check_end_user_budget(end_user_obj=end_user_object, route=route) _enforce_user_param_check(general_settings, request, request_body, route) _global_proxy_budget_check(global_proxy_spend, skip_budget_checks, route) _guardrail_modification_check(request_body, team_object) # 10 [OPTIONAL] Organization RBAC checks organization_role_based_access_check( user_object=user_object, route=route, request_body=request_body ) _is_route_allowed = _is_api_route_allowed( route=route, request=request, request_data=request_body, valid_token=valid_token, user_obj=user_object, ) # 11. [OPTIONAL] Vector store checks - is the object allowed to access the vector store with tracer.trace("litellm.proxy.auth.common_checks.vector_store_access_check"): await vector_store_access_check( request_body=request_body, team_object=team_object, valid_token=valid_token, ) # 12. [OPTIONAL] Tool allowlist - key/team allowed_tools (no DB in hot path) with tracer.trace("litellm.proxy.auth.common_checks.check_tools_allowlist"): await check_tools_allowlist( request_body=request_body, valid_token=valid_token, team_object=team_object, route=route, ) return True def _get_user_role( user_obj: Optional[LiteLLM_UserTable], ) -> Optional[LitellmUserRoles]: if user_obj is None: return None _user = user_obj _user_role = _user.user_role try: role = LitellmUserRoles(_user_role) except ValueError: return LitellmUserRoles.INTERNAL_USER return role def _is_api_route_allowed( route: str, request: Request, request_data: dict, valid_token: Optional[UserAPIKeyAuth], user_obj: Optional[LiteLLM_UserTable] = None, ) -> bool: """ - Route b/w api token check and normal token check """ _user_role = _get_user_role(user_obj=user_obj) if valid_token is None: raise Exception("Invalid proxy server token passed. valid_token=None.") if not _is_user_proxy_admin(user_obj=user_obj): # if non-admin RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=user_obj, _user_role=_user_role, route=route, request=request, request_data=request_data, valid_token=valid_token, ) return True def _is_user_proxy_admin(user_obj: Optional[LiteLLM_UserTable]): if user_obj is None: return False if ( user_obj.user_role is not None and user_obj.user_role == LitellmUserRoles.PROXY_ADMIN.value ): return True if ( user_obj.user_role is not None and user_obj.user_role == LitellmUserRoles.PROXY_ADMIN.value ): return True return False def _allowed_routes_check(user_route: str, allowed_routes: list) -> bool: """ Return if a user is allowed to access route. Helper function for `allowed_routes_check`. Parameters: - user_route: str - the route the user is trying to call - allowed_routes: List[str|LiteLLMRoutes] - the list of allowed routes for the user. """ from starlette.routing import compile_path for allowed_route in allowed_routes: if allowed_route in LiteLLMRoutes.__members__: for template in LiteLLMRoutes[allowed_route].value: regex, _, _ = compile_path(template) if regex.match(user_route): return True elif allowed_route == user_route: return True return False def allowed_routes_check( user_role: LitellmUserRoles, user_route: str, litellm_proxy_roles: LiteLLM_JWTAuth, ) -> bool: """ Check if user -> not admin - allowed to access these routes """ if user_role == LitellmUserRoles.PROXY_ADMIN: is_allowed = _allowed_routes_check( user_route=user_route, allowed_routes=litellm_proxy_roles.admin_allowed_routes, ) return is_allowed elif user_role == LitellmUserRoles.TEAM: if litellm_proxy_roles.team_allowed_routes is None: """ By default allow a team to call openai + info routes """ is_allowed = _allowed_routes_check( user_route=user_route, allowed_routes=["openai_routes", "info_routes"] ) return is_allowed elif litellm_proxy_roles.team_allowed_routes is not None: is_allowed = _allowed_routes_check( user_route=user_route, allowed_routes=litellm_proxy_roles.team_allowed_routes, ) return is_allowed return False def allowed_route_check_inside_route( user_api_key_dict: UserAPIKeyAuth, requested_user_id: Optional[str], ) -> bool: ret_val = True if ( user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY ): ret_val = False if requested_user_id is not None and user_api_key_dict.user_id is not None: if user_api_key_dict.user_id == requested_user_id: ret_val = True return ret_val def get_actual_routes(allowed_routes: list) -> list: actual_routes: list = [] for route_name in allowed_routes: try: route_value = LiteLLMRoutes[route_name].value if isinstance(route_value, set): actual_routes.extend(list(route_value)) else: actual_routes.extend(route_value) except KeyError: actual_routes.append(route_name) return actual_routes async def get_default_end_user_budget( prisma_client: Optional[PrismaClient], user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span] = None, ) -> Optional[LiteLLM_BudgetTable]: """ Fetches the default end user budget from the database if litellm.max_end_user_budget_id is configured. This budget is applied to end users who don't have an explicit budget_id set. Results are cached for performance. Args: prisma_client: Database client instance user_api_key_cache: Cache for storing/retrieving budget data parent_otel_span: Optional OpenTelemetry span for tracing Returns: LiteLLM_BudgetTable if configured and found, None otherwise """ if prisma_client is None or litellm.max_end_user_budget_id is None: return None cache_key = f"default_end_user_budget:{litellm.max_end_user_budget_id}" # Check cache first cached_budget = await user_api_key_cache.async_get_cache( key=cache_key, model_type=LiteLLM_BudgetTable, ) if cached_budget is not None: return cached_budget # Fetch from database try: budget_record = await BudgetRepository(prisma_client).table.find_unique( where={"budget_id": litellm.max_end_user_budget_id} ) if budget_record is None: verbose_proxy_logger.warning( f"Default end user budget not found in database: {litellm.max_end_user_budget_id}" ) return None _budget_obj = LiteLLM_BudgetTable(**budget_record.dict()) # Cache the budget for 60 seconds await user_api_key_cache.async_set_cache( key=cache_key, value=_budget_obj, model_type=LiteLLM_BudgetTable, ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) return _budget_obj except Exception as e: verbose_proxy_logger.error(f"Error fetching default end user budget: {str(e)}") return None @log_db_metrics async def get_team_member_default_budget( budget_id: str, prisma_client: Optional[PrismaClient], user_api_key_cache: UserApiKeyCache, ) -> Optional[LiteLLM_BudgetTable]: """ Fetches the team-level default per-member budget referenced by team.metadata["team_member_budget_id"]. This budget is applied to team members whose TeamMembership row has no linked budget, or whose linked budget has max_budget=NULL. Results are cached for performance. Args: budget_id: The budget_id pulled from team.metadata["team_member_budget_id"] prisma_client: Database client instance user_api_key_cache: Cache for storing/retrieving budget data Returns: LiteLLM_BudgetTable if found, None otherwise """ if prisma_client is None: return None cache_key = f"team_member_default_budget:{budget_id}" cached_budget = await user_api_key_cache.async_get_cache(key=cache_key) if isinstance(cached_budget, LiteLLM_BudgetTable): return cached_budget if isinstance(cached_budget, dict): return LiteLLM_BudgetTable(**cached_budget) try: budget_record = await BudgetRepository(prisma_client).table.find_unique( where={"budget_id": budget_id} ) if budget_record is None: verbose_proxy_logger.warning( f"Team-default member budget not found in database: {budget_id}" ) return None await user_api_key_cache.async_set_cache( key=cache_key, value=budget_record.dict(), ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) return LiteLLM_BudgetTable(**budget_record.dict()) except Exception: verbose_proxy_logger.exception( f"Error fetching team-default member budget {budget_id}" ) return None async def _apply_default_budget_to_end_user( end_user_obj: LiteLLM_EndUserTable, prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span] = None, ) -> LiteLLM_EndUserTable: """ Helper function to apply default budget to end user if they don't have a budget assigned. Args: end_user_obj: The end user object to potentially apply default budget to prisma_client: Database client instance user_api_key_cache: Cache for storing/retrieving data parent_otel_span: Optional OpenTelemetry span for tracing Returns: Updated end user object with default budget applied if applicable """ # If end user already has a budget assigned, no need to apply default if end_user_obj.litellm_budget_table is not None: return end_user_obj # If no default budget configured, return as-is if litellm.max_end_user_budget_id is None: return end_user_obj # Fetch and apply default budget default_budget = await get_default_end_user_budget( prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, parent_otel_span=parent_otel_span, ) if default_budget is not None: # Apply default budget to end user object end_user_obj.litellm_budget_table = default_budget verbose_proxy_logger.debug( f"Applied default budget {litellm.max_end_user_budget_id} to end user {end_user_obj.user_id}" ) return end_user_obj async def _check_end_user_budget( end_user_obj: LiteLLM_EndUserTable, route: str, ) -> None: """ Check if end user is within their budget limit. Args: end_user_obj: The end user object to check route: The request route Raises: litellm.BudgetExceededError: If end user has exceeded their budget """ if RouteChecks.is_info_route(route): return if end_user_obj.litellm_budget_table is None: return end_user_budget = end_user_obj.litellm_budget_table.max_budget if end_user_budget is None: return from litellm.proxy.proxy_server import get_current_spend end_user_spend = await get_current_spend( counter_key=f"spend:end_user:{end_user_obj.user_id}", fallback_spend=end_user_obj.spend or 0.0, max_budget=end_user_budget, fallback_authoritative=True, ) if end_user_spend > end_user_budget: raise litellm.BudgetExceededError( current_cost=end_user_spend, max_budget=end_user_budget, message=f"ExceededBudget: End User={end_user_obj.user_id} over budget. Spend={end_user_spend}, Budget={end_user_budget}", ) @log_db_metrics async def get_end_user_object( end_user_id: Optional[str], prisma_client: Optional[PrismaClient], user_api_key_cache: UserApiKeyCache, route: Optional[str] = "", parent_otel_span: Optional[Span] = None, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> Optional[LiteLLM_EndUserTable]: """ Returns end user object from database or cache. If end user exists but has no budget_id, applies the default budget (if configured via litellm.max_end_user_budget_id). Args: end_user_id: The ID of the end user prisma_client: Database client instance user_api_key_cache: Cache for storing/retrieving data route: The request route parent_otel_span: Optional OpenTelemetry span for tracing proxy_logging_obj: Optional proxy logging object Returns: LiteLLM_EndUserTable if found, None otherwise """ if prisma_client is None: raise Exception("No db connected") if end_user_id is None: return None _key = "end_user_id:{}".format(end_user_id) # Check cache first cached_user_obj = await user_api_key_cache.async_get_cache( key=_key, model_type=LiteLLM_EndUserTable, ) if cached_user_obj is not None: return_obj = cached_user_obj # Apply default budget if needed return_obj = await _apply_default_budget_to_end_user( end_user_obj=return_obj, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, parent_otel_span=parent_otel_span, ) return return_obj # Fetch from database try: response = await EndUserRepository(prisma_client).table.find_unique( where={"user_id": end_user_id}, include={"litellm_budget_table": True, "object_permission": True}, ) if response is None: raise Exception # Convert to LiteLLM_EndUserTable object _response = LiteLLM_EndUserTable(**response.dict()) # Apply default budget if needed _response = await _apply_default_budget_to_end_user( end_user_obj=_response, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, parent_otel_span=parent_otel_span, ) # Save to cache await user_api_key_cache.async_set_cache( key="end_user_id:{}".format(end_user_id), value=_response, model_type=LiteLLM_EndUserTable, ) return _response except Exception: return None _END_USER_VALIDATION_NEGATIVE_TTL = 60 _END_USER_VALIDATION_POSITIVE_TTL = 300 async def resolve_and_validate_end_user_id( raw_end_user_id: Optional[str], prisma_client: Optional[PrismaClient], user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span] = None, proxy_logging_obj: Optional[ProxyLogging] = None, route: str = "", ) -> Optional[str]: """Optionally drop end-user ids that don't resolve to a known DB row. Default: pass-through. LiteLLM's documented pattern is that the `user` field is an arbitrary caller-supplied identifier, so validation is opt-in behind ``litellm.validate_end_user_id_in_db`` to preserve backwards compatibility. When the flag is set: accept the id when it matches any of - LiteLLM_EndUserTable.user_id - LiteLLM_UserTable.user_id - LiteLLM_UserTable.user_email (case-insensitive) If the id doesn't match but ``litellm.max_end_user_budget_id`` is set, we still preserve the id so the default end-user budget is applied downstream; otherwise we return None. DB lookups reuse ``get_end_user_object`` / ``get_user_object`` so they share the same cache as the rest of the auth path instead of adding new raw Prisma queries. """ if raw_end_user_id is None: return None if not litellm.validate_end_user_id_in_db: return raw_end_user_id if prisma_client is None: return raw_end_user_id cache_key = f"end_user_validation:{raw_end_user_id}" cached = await user_api_key_cache.async_get_cache(key=cache_key) if cached == "valid": return raw_end_user_id if cached == "invalid": return raw_end_user_id if litellm.max_end_user_budget_id else None is_valid = await _end_user_id_exists_in_db( end_user_id=raw_end_user_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, route=route, ) await user_api_key_cache.async_set_cache( key=cache_key, value="valid" if is_valid else "invalid", ttl=( _END_USER_VALIDATION_POSITIVE_TTL if is_valid else _END_USER_VALIDATION_NEGATIVE_TTL ), ) if is_valid: return raw_end_user_id # Preserve id so the caller can still apply litellm.max_end_user_budget_id. if litellm.max_end_user_budget_id: return raw_end_user_id return None async def _end_user_id_exists_in_db( end_user_id: str, prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span] = None, proxy_logging_obj: Optional[ProxyLogging] = None, route: str = "", ) -> bool: """True when the id matches an EndUser, User, or user_email row.""" try: end_user_obj = await get_end_user_object( end_user_id=end_user_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, route=route, ) if end_user_obj is not None: return True except Exception as e: verbose_proxy_logger.debug( f"end_user validation: get_end_user_object lookup failed: {e}" ) try: user_obj = await get_user_object( user_id=end_user_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, user_id_upsert=False, parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, check_db_only=False, user_email=end_user_id if "@" in end_user_id else None, ) if user_obj is not None: return True except Exception as e: verbose_proxy_logger.debug( f"end_user validation: get_user_object lookup failed: {e}" ) return False @log_db_metrics async def get_tag_objects_batch( tag_names: List[str], prisma_client: Optional[PrismaClient], user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span] = None, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> Dict[str, LiteLLM_TagTable]: """ Batch fetch multiple tag objects from cache and db. Optimizes for latency by: 1. Fetching all cached tags in parallel 2. Batch fetching uncached tags in one DB query Args: tag_names: List of tag names to fetch prisma_client: Prisma database client user_api_key_cache: Cache for storing tag objects parent_otel_span: Optional OpenTelemetry span for tracing proxy_logging_obj: Optional proxy logging object Returns: Dictionary mapping tag_name to LiteLLM_TagTable object """ if prisma_client is None: return {} if not tag_names: return {} tag_objects = {} uncached_tags = [] # Try to get all tags from cache first for tag_name in tag_names: cache_key = f"tag:{tag_name}" cached_tag = await user_api_key_cache.async_get_cache( key=cache_key, model_type=LiteLLM_TagTable, ) if cached_tag is not None: tag_objects[tag_name] = cached_tag else: uncached_tags.append(tag_name) # Batch fetch uncached tags from DB in one query if uncached_tags: try: db_tags = await TagRepository(prisma_client).table.find_many( where={"tag_name": {"in": uncached_tags}}, include={"litellm_budget_table": True}, ) # Cache and add to tag_objects for db_tag in db_tags: tag_name = db_tag.tag_name cache_key = f"tag:{tag_name}" _tag_obj = LiteLLM_TagTable(**db_tag.dict()) await user_api_key_cache.async_set_cache( key=cache_key, value=_tag_obj, model_type=LiteLLM_TagTable, ) tag_objects[tag_name] = _tag_obj except Exception as e: verbose_proxy_logger.debug(f"Error batch fetching tags from database: {e}") return tag_objects @log_db_metrics async def get_tag_object( tag_name: Optional[str], prisma_client: Optional[PrismaClient], user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span] = None, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> Optional[LiteLLM_TagTable]: """ Returns tag object from cache or db. Uses default cache TTL (same as end_user objects) to avoid drift. Args: tag_name: Name of the tag to fetch prisma_client: Prisma database client user_api_key_cache: Cache for storing tag objects parent_otel_span: Optional OpenTelemetry span for tracing proxy_logging_obj: Optional proxy logging object Returns: LiteLLM_TagTable object if found, None otherwise """ if prisma_client is None or tag_name is None: return None # Use batch helper for consistency tag_objects = await get_tag_objects_batch( tag_names=[tag_name], prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, ) return tag_objects.get(tag_name) @log_db_metrics async def get_team_membership( user_id: str, team_id: str, prisma_client: Optional[PrismaClient], user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span] = None, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> Optional["LiteLLM_TeamMembership"]: """ Returns team membership object if user is member of team. 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). """ from litellm.proxy._types import LiteLLM_TeamMembership if prisma_client is None: raise Exception("No db connected") if user_id is None or team_id is None: return None _key = "team_membership:{}:{}".format(user_id, team_id) # check if in cache cached_membership_obj = await user_api_key_cache.async_get_cache( key=_key, model_type=LiteLLM_TeamMembership, ) if cached_membership_obj is not None: return cached_membership_obj # else, check db try: response = await TeamMembershipRepository(prisma_client).table.find_unique( where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, include={"litellm_budget_table": True}, ) if response is None: return None _response = LiteLLM_TeamMembership(**response.dict()) await user_api_key_cache.async_set_cache( key=_key, value=_response, model_type=LiteLLM_TeamMembership, ) return _response except Exception: verbose_proxy_logger.exception( "Error getting team membership for user_id: %s, team_id: %s", user_id, team_id, ) return None def model_in_access_group( model: str, team_models: Optional[List[str]], llm_router: Optional[Router] ) -> bool: from collections import defaultdict if team_models is None: return True if model in team_models: return True access_groups: dict[str, list[str]] = defaultdict(list) if llm_router: access_groups = llm_router.get_model_access_groups(model_name=model) if len(access_groups) > 0: # check if token contains any model access groups for idx, m in enumerate( team_models ): # loop token models, if any of them are an access group add the access group if m in access_groups: return True # Filter out models that are access_groups filtered_models = [m for m in team_models if m not in access_groups] if model in filtered_models: return True return False def _should_check_db( key: str, last_db_access_time: LimitedSizeOrderedDict, db_cache_expiry: int ) -> bool: """ Prevent calling db repeatedly for items that don't exist in the db. """ current_time = time.time() # if key doesn't exist in last_db_access_time -> check db if key not in last_db_access_time: return True elif ( last_db_access_time[key][0] is not None ): # check db for non-null values (for refresh operations) return True elif last_db_access_time[key][0] is None: if current_time - last_db_access_time[key] >= db_cache_expiry: return True return False def _update_last_db_access_time( key: str, value: Optional[Any], last_db_access_time: LimitedSizeOrderedDict ): last_db_access_time[key] = (value, time.time()) def _get_role_based_permissions( rbac_role: RBAC_ROLES, general_settings: dict, key: Literal["models", "routes"], ) -> Optional[List[str]]: """ Get the role based permissions from the general settings. """ role_based_permissions = cast( Optional[List[RoleBasedPermissions]], general_settings.get("role_permissions", []), ) if role_based_permissions is None: return None for role_based_permission in role_based_permissions: if role_based_permission.role == rbac_role: return getattr(role_based_permission, key) return None def get_role_based_models( rbac_role: RBAC_ROLES, general_settings: dict, ) -> Optional[List[str]]: """ Get the models allowed for a user role. Used by JWT Auth. """ return _get_role_based_permissions( rbac_role=rbac_role, general_settings=general_settings, key="models", ) def get_role_based_routes( rbac_role: RBAC_ROLES, general_settings: dict, ) -> Optional[List[str]]: """ Get the routes allowed for a user role. """ return _get_role_based_permissions( rbac_role=rbac_role, general_settings=general_settings, key="routes", ) async def _get_fuzzy_user_object( prisma_client: PrismaClient, sso_user_id: Optional[str] = None, user_email: Optional[str] = None, ) -> Optional[LiteLLM_UserTable]: """ Checks if sso user is in db. Called when user id match is not found in db. - Check if sso_user_id is user_id in db - Check if sso_user_id is sso_user_id in db - Check if user_email is user_email in db - If not, create new user with user_email and sso_user_id and user_id = sso_user_id """ response = None if sso_user_id is not None: response = await UserRepository(prisma_client).table.find_unique( where={"sso_user_id": sso_user_id}, include={"organization_memberships": True}, ) if response is None and user_email is not None: # Use case-insensitive query to handle emails with different casing # This matches the pattern used in _check_duplicate_user_email response = await UserRepository(prisma_client).table.find_first( where={"user_email": {"equals": user_email, "mode": "insensitive"}}, include={"organization_memberships": True}, ) if response is not None and sso_user_id is not None: # update sso_user_id asyncio.create_task( # background task to update user with sso id UserRepository(prisma_client).table.update( where={"user_id": response.user_id}, data={"sso_user_id": sso_user_id}, ) ) return response @log_db_metrics async def get_user_object( user_id: Optional[str], prisma_client: Optional[PrismaClient], user_api_key_cache: UserApiKeyCache, user_id_upsert: bool, parent_otel_span: Optional[Span] = None, proxy_logging_obj: Optional[ProxyLogging] = None, sso_user_id: Optional[str] = None, user_email: Optional[str] = None, check_db_only: Optional[bool] = None, ) -> Optional[LiteLLM_UserTable]: """ - Check if user id in proxy User Table - if valid, return LiteLLM_UserTable object with defined limits - if not, then raise an error """ if user_id is None: return None # check if in cache if not check_db_only: cached_user_obj = await user_api_key_cache.async_get_cache( key=user_id, model_type=LiteLLM_UserTable, ) if cached_user_obj is not None: return cached_user_obj # else, check db if prisma_client is None: raise Exception("No db connected") try: db_access_time_key = "user_id:{}".format(user_id) should_check_db = _should_check_db( key=db_access_time_key, last_db_access_time=last_db_access_time, db_cache_expiry=db_cache_expiry, ) if should_check_db: response = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_id}, include={"organization_memberships": True} ) if response is None: response = await _get_fuzzy_user_object( prisma_client=prisma_client, sso_user_id=sso_user_id, user_email=user_email, ) else: response = None if response is None: if user_id_upsert: new_user_params: Dict[str, Any] = { "user_id": user_id, } if user_email is not None: new_user_params["user_email"] = user_email if litellm.default_internal_user_params is not None: new_user_params.update(litellm.default_internal_user_params) response = await UserRepository(prisma_client).table.create( data=new_user_params, include={"organization_memberships": True}, ) else: raise Exception if ( response.organization_memberships is not None and len(response.organization_memberships) > 0 ): # dump each organization membership to type LiteLLM_OrganizationMembershipTable _dumped_memberships = [ LiteLLM_OrganizationMembershipTable(**membership.model_dump()) for membership in response.organization_memberships if membership is not None ] response.organization_memberships = _dumped_memberships _response = LiteLLM_UserTable(**dict(response)) response_dict = _response.model_dump() # save the user object to cache await user_api_key_cache.async_set_cache( key=user_id, value=_response, model_type=LiteLLM_UserTable, ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) # save to db access time _update_last_db_access_time( key=db_access_time_key, value=response_dict, last_db_access_time=last_db_access_time, ) return _response except Exception as e: # if user not in db _log_budget_lookup_failure("user", e) raise ValueError( f"User doesn't exist in db. 'user_id'={user_id}. Create user via `/user/new` call. Got error - {e}" ) async def _cache_management_object( key: str, value: Union[BaseModel, Dict[str, Any]], user_api_key_cache: UserApiKeyCache, proxy_logging_obj: Optional[ProxyLogging], *, model_type: Type[BaseModel], ): """ Persist management objects via ``UserApiKeyCache`` (in-memory + optional Redis). ``UserApiKeyCache`` serializes with ``model_type`` so Redis and in-memory stay aligned. """ await user_api_key_cache.async_set_cache( key=key, value=value, model_type=model_type, ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) async def _cache_team_object( team_id: str, team_table: LiteLLM_TeamTableCachedObj, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: Optional[ProxyLogging], ): ## CACHE REFRESH TIME! team_table.last_refreshed_at = time.time() # team_id is the table primary key — guaranteed unique, safe to write. await _cache_management_object( key="team_id:{}".format(team_id), value=team_table, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, model_type=LiteLLM_TeamTableCachedObj, ) # Invalidate the alias-keyed cache so the JWT auth path with # `team_alias_jwt_field` (which reads via `get_team_object_by_alias`) # doesn't keep serving the pre-mutation team after every team-write # endpoint (team_model_add, team_model_delete, update_team, etc.). # # Why DELETE and not WRITE: `team_alias` has no UNIQUE constraint in # schema.prisma. Writing this cache from the generic refresh path # would let a team admin who renamed their team to collide with # another team's alias silently overwrite the cached team for # JWT-by-alias auth (veria-ai review on #28739). Deleting forces the # next reader through `get_team_object_by_alias`, which DOES enforce # uniqueness (len(teams) > 1 raises HTTPException) before populating # the cache from a verified single row. if team_table.team_alias: alias_key = "team_alias:{}".format(team_table.team_alias) user_api_key_cache.delete_cache(key=alias_key) if proxy_logging_obj is not None: await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache( key=alias_key ) async def _cache_key_object( hashed_token: str, user_api_key_obj: UserAPIKeyAuth, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: Optional[ProxyLogging], ): key = hashed_token ## CACHE REFRESH TIME user_api_key_obj.last_refreshed_at = time.time() cached_key_obj = _copy_user_api_key_auth_for_cache( user_api_key_obj=user_api_key_obj ) await _cache_management_object( key=key, value=cached_key_obj, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, model_type=UserAPIKeyAuth, ) async def _delete_cache_key_object( hashed_token: str, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: Optional[ProxyLogging], ): key = hashed_token user_api_key_cache.delete_cache(key=key) ## UPDATE REDIS CACHE ## if proxy_logging_obj is not None: await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache( key=key ) @log_db_metrics async def _get_team_db_check( team_id: str, prisma_client: PrismaClient, team_id_upsert: Optional[bool] = None ): response = await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id} ) if response is None and team_id_upsert: from litellm.proxy.management_endpoints.team_endpoints import new_team new_team_data = NewTeamRequest(team_id=team_id) mock_request = Request(scope={"type": "http"}) system_admin_user = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) created_team_dict = await new_team( data=new_team_data, http_request=mock_request, user_api_key_dict=system_admin_user, ) response = LiteLLM_TeamTable(**created_team_dict) return response async def _get_team_object_from_db(team_id: str, prisma_client: PrismaClient): return await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id} ) async def _get_team_object_from_user_api_key_cache( team_id: str, prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache, last_db_access_time: LimitedSizeOrderedDict, db_cache_expiry: int, proxy_logging_obj: Optional[ProxyLogging], key: str, team_id_upsert: Optional[bool] = None, ) -> LiteLLM_TeamTableCachedObj: db_access_time_key = key should_check_db = _should_check_db( key=db_access_time_key, last_db_access_time=last_db_access_time, db_cache_expiry=db_cache_expiry, ) if should_check_db: response = await _get_team_db_check( team_id=team_id, prisma_client=prisma_client, team_id_upsert=team_id_upsert ) else: response = None if response is None: raise Exception _response = LiteLLM_TeamTableCachedObj(**response.dict()) # Load object_permission if object_permission_id exists but object_permission is not loaded if _response.object_permission_id and not _response.object_permission: try: _response.object_permission = await get_object_permission( object_permission_id=_response.object_permission_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, parent_otel_span=None, proxy_logging_obj=proxy_logging_obj, ) except Exception as e: verbose_proxy_logger.debug( f"Failed to load object_permission for team {team_id} with object_permission_id={_response.object_permission_id}: {e}" ) # save the team object to cache await _cache_team_object( team_id=team_id, team_table=_response, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) # save to db access time _update_last_db_access_time( key=db_access_time_key, value=_response, last_db_access_time=last_db_access_time, ) return _response async def _get_team_object_from_cache( key: str, proxy_logging_obj: Optional[ProxyLogging], user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span], ) -> Optional[LiteLLM_TeamTableCachedObj]: ## INTERNAL USAGE CACHE (plain DualCache) — checked before UserApiKeyCache stores ## if ( proxy_logging_obj is not None and proxy_logging_obj.internal_usage_cache.dual_cache ): cached_raw = ( await proxy_logging_obj.internal_usage_cache.dual_cache.async_get_cache( key=key, parent_otel_span=parent_otel_span ) ) if cached_raw is not None: from_internal = CacheCodec.deserialize( cached_raw, LiteLLM_TeamTableCachedObj ) if from_internal is not None: return from_internal decoded = await user_api_key_cache.async_get_cache( key=key, parent_otel_span=parent_otel_span, model_type=LiteLLM_TeamTableCachedObj, ) return decoded async def get_team_object( team_id: str, prisma_client: Optional[PrismaClient], user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span] = None, proxy_logging_obj: Optional[ProxyLogging] = None, check_cache_only: Optional[bool] = None, check_db_only: Optional[bool] = None, team_id_upsert: Optional[bool] = None, ) -> LiteLLM_TeamTableCachedObj: """ - Check if team id in proxy Team Table - if valid, return LiteLLM_TeamTable object with defined limits - if not, then raise an error Raises: - HTTPException: If team doesn't exist in db or cache (status_code=404) """ if prisma_client is None: raise Exception( "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys" ) # check if in cache key = "team_id:{}".format(team_id) if not check_db_only: cached_team_obj = await _get_team_object_from_cache( key=key, proxy_logging_obj=proxy_logging_obj, user_api_key_cache=user_api_key_cache, parent_otel_span=parent_otel_span, ) if cached_team_obj is not None: return cached_team_obj if check_cache_only: raise HTTPException( status_code=404, detail={ "error": f"Team doesn't exist in cache + check_cache_only=True. Team={team_id}." }, ) # else, check db try: return await _get_team_object_from_user_api_key_cache( team_id=team_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, last_db_access_time=last_db_access_time, db_cache_expiry=db_cache_expiry, key=key, team_id_upsert=team_id_upsert, ) except Exception: raise HTTPException( status_code=404, detail={ "error": f"Team doesn't exist in db. Team={team_id}. Create team via `/team/new` call." }, ) async def _cache_access_object( access_group_id: str, access_group_table: LiteLLM_AccessGroupTable, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: Optional[ProxyLogging] = None, ): key = "access_group_id:{}".format(access_group_id) await user_api_key_cache.async_set_cache( key=key, value=access_group_table, model_type=LiteLLM_AccessGroupTable, ttl=DEFAULT_ACCESS_GROUP_CACHE_TTL, ) async def _delete_cache_access_object( access_group_id: str, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: Optional[ProxyLogging] = None, ): key = "access_group_id:{}".format(access_group_id) user_api_key_cache.delete_cache(key=key) ## UPDATE REDIS CACHE ## if proxy_logging_obj is not None: await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache( key=key ) @log_db_metrics async def get_access_object( access_group_id: str, prisma_client: Optional[PrismaClient], user_api_key_cache: UserApiKeyCache, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> LiteLLM_AccessGroupTable: """ - Check if access_group_id in proxy AccessGroupTable - Always checks cache first, then DB only when not found in cache - if valid, return LiteLLM_AccessGroupTable object - if not, then raise an error Unlike get_team_object, this has no check_cache_only or check_db_only flags; it always follows cache-first-then-db semantics. Raises: - HTTPException: If access group doesn't exist in db or cache (status_code=404) """ if prisma_client is None: raise Exception( "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys" ) key = "access_group_id:{}".format(access_group_id) cached_access_obj = await user_api_key_cache.async_get_cache( key=key, model_type=LiteLLM_AccessGroupTable, ) if cached_access_obj is not None: return cached_access_obj # Not in cache - fetch from DB try: response = await AccessGroupRepository(prisma_client).table.find_unique( where={"access_group_id": access_group_id} ) if response is None: raise HTTPException( status_code=404, detail={ "error": f"Access group doesn't exist in db. Access group={access_group_id}." }, ) _response = LiteLLM_AccessGroupTable(**response.dict()) # Save to cache await _cache_access_object( access_group_id=access_group_id, access_group_table=_response, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) return _response except HTTPException: raise except Exception as e: verbose_proxy_logger.exception( "Error getting access group for access_group_id: %s", access_group_id, ) raise HTTPException( status_code=404, detail={ "error": f"Access group doesn't exist in db. Access group={access_group_id}. Error: {e}" }, ) @log_db_metrics async def get_team_object_by_alias( team_alias: str, prisma_client: Optional[PrismaClient], user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional["Span"] = None, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> LiteLLM_TeamTableCachedObj: """ Look up a team by its team_alias (name) in the database. Args: team_alias: The team name/alias to look up prisma_client: Database client user_api_key_cache: Cache for storing results parent_otel_span: Optional OpenTelemetry span proxy_logging_obj: Optional proxy logging object Returns: LiteLLM_TeamTableCachedObj: The team object if found Raises: HTTPException: If team doesn't exist or multiple teams have the same alias """ if prisma_client is None: raise Exception( "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys" ) # Check cache first (keyed by alias) cache_key = "team_alias:{}".format(team_alias) cached_team_obj = await _get_team_object_from_cache( key=cache_key, proxy_logging_obj=proxy_logging_obj, user_api_key_cache=user_api_key_cache, parent_otel_span=parent_otel_span, ) if cached_team_obj is not None: return cached_team_obj # Query database by team_alias try: teams = await TeamRepository(prisma_client).table.find_many( where={"team_alias": team_alias} ) if not teams: raise HTTPException( status_code=404, detail={ "error": f"Team with alias '{team_alias}' doesn't exist in db. Create team via `/team/new` call." }, ) if len(teams) > 1: raise HTTPException( status_code=400, detail={ "error": f"Multiple teams found with alias '{team_alias}'. Please use team_id_jwt_field instead or ensure team aliases are unique." }, ) team = teams[0] team_obj = LiteLLM_TeamTableCachedObj(**team.model_dump()) # Load object_permission if object_permission_id exists but object_permission is not loaded if team_obj.object_permission_id and not team_obj.object_permission: try: team_obj.object_permission = await get_object_permission( object_permission_id=team_obj.object_permission_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, ) except Exception as e: verbose_proxy_logger.debug( f"Failed to load object_permission for team {team_obj.team_id} with object_permission_id={team_obj.object_permission_id}: {e}" ) # Cache the result by both alias and team_id await user_api_key_cache.async_set_cache( key=cache_key, value=team_obj, model_type=LiteLLM_TeamTableCachedObj, ttl=DEFAULT_IN_MEMORY_TTL, ) # Also cache by team_id for consistency team_id_cache_key = "team_id:{}".format(team_obj.team_id) await user_api_key_cache.async_set_cache( key=team_id_cache_key, value=team_obj, model_type=LiteLLM_TeamTableCachedObj, ttl=DEFAULT_IN_MEMORY_TTL, ) return team_obj except HTTPException: raise except Exception as e: verbose_proxy_logger.exception("Error looking up team by alias: %s", team_alias) raise HTTPException( status_code=500, detail={ "error": f"Error looking up team by alias '{team_alias}': {str(e)}" }, ) @log_db_metrics async def get_org_object_by_alias( org_alias: str, prisma_client: Optional[PrismaClient], user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional["Span"] = None, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> Optional[LiteLLM_OrganizationTable]: """ Look up an organization by its organization_alias in the database. Args: org_alias: The organization name/alias to look up prisma_client: Database client user_api_key_cache: Cache for storing results parent_otel_span: Optional OpenTelemetry span proxy_logging_obj: Optional proxy logging object Returns: LiteLLM_OrganizationTable if found, None otherwise Raises: HTTPException: If organization not found or multiple orgs have the same alias """ if prisma_client is None: raise Exception( "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys" ) # Check cache first (keyed by alias) cache_key = "org_alias:{}".format(org_alias) cached_org_obj = await user_api_key_cache.async_get_cache( key=cache_key, model_type=LiteLLM_OrganizationTable, ) if cached_org_obj is not None: return cached_org_obj # Query database by organization_alias try: orgs = await OrganizationRepository(prisma_client).table.find_many( where={"organization_alias": org_alias} ) if not orgs: raise HTTPException( status_code=404, detail={ "error": f"Organization with alias '{org_alias}' doesn't exist in db. Create organization via `/organization/new` call." }, ) if len(orgs) > 1: raise HTTPException( status_code=400, detail={ "error": f"Multiple organizations found with alias '{org_alias}'. Please use org_id_jwt_field instead or ensure organization aliases are unique." }, ) org = orgs[0] org_obj = LiteLLM_OrganizationTable(**org.model_dump()) # Cache the result await user_api_key_cache.async_set_cache( key=cache_key, value=org_obj, model_type=LiteLLM_OrganizationTable, ttl=DEFAULT_IN_MEMORY_TTL, ) # Also cache by org_id for consistency await user_api_key_cache.async_set_cache( key="org_id:{}".format(org_obj.organization_id), value=org_obj, model_type=LiteLLM_OrganizationTable, ttl=DEFAULT_IN_MEMORY_TTL, ) return org_obj except HTTPException: raise except Exception as e: verbose_proxy_logger.exception( "Error looking up organization by alias: %s", org_alias ) raise HTTPException( status_code=500, detail={ "error": f"Error looking up organization by alias '{org_alias}': {str(e)}" }, ) class ExperimentalUIJWTToken: @staticmethod def get_experimental_ui_login_jwt_auth_token(user_info: LiteLLM_UserTable) -> str: from datetime import timedelta from litellm.proxy.common_utils.encrypt_decrypt_utils import ( encrypt_value_helper, ) if user_info.user_role is None: raise Exception("User role is required for experimental UI login") # Experimental UI flow uses fixed 10-min expiry for security (does not use LITELLM_UI_SESSION_DURATION) expiration_time = get_utc_datetime() + timedelta(minutes=10) # Format the expiration time as ISO 8601 string expires = expiration_time.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "+00:00" valid_token = UserAPIKeyAuth( token="ui-token", key_name="ui-token", key_alias="ui-token", max_budget=litellm.max_ui_session_budget, rpm_limit=100, # allow user to have a conversation on test key pane of UI expires=expires, user_id=user_info.user_id, team_id="litellm-dashboard", models=user_info.models, max_parallel_requests=None, user_role=LitellmUserRoles(user_info.user_role), ) return encrypt_value_helper(valid_token.model_dump_json(exclude_none=True)) @staticmethod def get_cli_jwt_auth_token( user_info: LiteLLM_UserTable, team_id: Optional[str] = None, team_alias: Optional[str] = None, ) -> str: """ Generate a JWT token for CLI authentication with configurable expiration. The expiration time can be controlled via the LITELLM_CLI_JWT_EXPIRATION_HOURS environment variable (defaults to 24 hours). Args: user_info: User information from the database team_id: Team ID for the user (optional, uses user's team if available) team_alias: Team alias for the selected team, if available Returns: Encrypted JWT token string """ from datetime import timedelta from litellm.proxy.common_utils.encrypt_decrypt_utils import ( encrypt_value_helper, ) if user_info.user_role is None: raise Exception("User role is required for CLI JWT login") # Calculate expiration time (configurable via LITELLM_CLI_JWT_EXPIRATION_HOURS env var) expiration_time = get_utc_datetime() + timedelta(hours=CLI_JWT_EXPIRATION_HOURS) # Format the expiration time as ISO 8601 string expires = expiration_time.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "+00:00" # Use provided team_id, or fall back to user's teams if available _team_id = team_id if _team_id is None and hasattr(user_info, "teams") and user_info.teams: # Use first team if user has teams _team_id = user_info.teams[0] if len(user_info.teams) > 0 else None valid_token = UserAPIKeyAuth( token=CLI_JWT_TOKEN_NAME, key_name=CLI_JWT_TOKEN_NAME, key_alias=CLI_JWT_TOKEN_NAME, max_budget=litellm.max_ui_session_budget, expires=expires, user_id=user_info.user_id, team_id=_team_id, team_alias=team_alias, models=user_info.models, max_parallel_requests=None, user_role=LitellmUserRoles(user_info.user_role), ) return encrypt_value_helper(valid_token.model_dump_json(exclude_none=True)) @staticmethod def get_key_object_from_ui_hash_key( hashed_token: str, ) -> Optional[UserAPIKeyAuth]: import json from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, ) decrypted_token = decrypt_value_helper( hashed_token, key="ui_hash_key", exception_type="debug" ) if decrypted_token is None: return None try: return UserAPIKeyAuth(**json.loads(decrypted_token)) except Exception as e: raise Exception( f"Invalid hash key. Hash key={hashed_token}. Decrypted token={decrypted_token}. Error: {e}" ) async def _fetch_key_object_from_db_with_reconnect( hashed_token: str, prisma_client: PrismaClient, parent_otel_span: Optional[Span], proxy_logging_obj: Optional[ProxyLogging], ) -> Optional[BaseModel]: """ Fetch key object from DB and retry once if a DB connection error can be healed. """ try: return await prisma_client.get_data( token=hashed_token, table_name="combined_view", parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, ) except Exception as e: if PrismaDBExceptionHandler.is_database_transport_error(e): did_reconnect = False if hasattr(prisma_client, "attempt_db_reconnect"): auth_reconnect_timeout = getattr( prisma_client, "_db_auth_reconnect_timeout_seconds", 2.0 ) if not isinstance(auth_reconnect_timeout, (int, float)): auth_reconnect_timeout = 2.0 auth_reconnect_lock_timeout = getattr( prisma_client, "_db_auth_reconnect_lock_timeout_seconds", 0.1 ) if not isinstance(auth_reconnect_lock_timeout, (int, float)): auth_reconnect_lock_timeout = 0.1 did_reconnect = await prisma_client.attempt_db_reconnect( reason="auth_get_key_object_lookup_failure", timeout_seconds=auth_reconnect_timeout, lock_timeout_seconds=auth_reconnect_lock_timeout, ) if did_reconnect: return await prisma_client.get_data( token=hashed_token, table_name="combined_view", parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, ) raise @log_db_metrics async def get_jwt_key_mapping_object( jwt_claim_name: str, jwt_claim_value: str, prisma_client: PrismaClient, ) -> Optional[str]: """ Lookup a JWT-to-virtual-key mapping from the database. Returns the hashed token (str) if a matching active mapping is found, else None. """ mapping = await JWTKeyMappingRepository(prisma_client).table.find_first( where={ "jwt_claim_name": jwt_claim_name, "jwt_claim_value": jwt_claim_value, "is_active": True, } ) if mapping is not None: return mapping.token return None @log_db_metrics async def get_key_object( hashed_token: str, prisma_client: Optional[PrismaClient], user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span] = None, proxy_logging_obj: Optional[ProxyLogging] = None, check_cache_only: Optional[bool] = None, ) -> UserAPIKeyAuth: """ - Check if team id in proxy Team Table - if valid, return LiteLLM_TeamTable object with defined limits - if not, then raise an error """ if prisma_client is None: raise Exception( "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys" ) # check if in cache key = hashed_token # Same flow as before: use cache only when we have a hit we can turn into UserAPIKeyAuth # (dict from Redis / model_dump, or UserAPIKeyAuth from in-memory). Otherwise fall through to DB. user_api_key_auth = await user_api_key_cache.async_get_cache( key=key, model_type=UserAPIKeyAuth, ) if user_api_key_auth is not None: return _copy_user_api_key_auth_for_cache(user_api_key_obj=user_api_key_auth) if check_cache_only: raise Exception( f"Key doesn't exist in cache + check_cache_only=True. key={key}." ) # else, check db _valid_token: Optional[BaseModel] = await _fetch_key_object_from_db_with_reconnect( hashed_token=hashed_token, prisma_client=prisma_client, parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, ) if _valid_token is None: raise ProxyException( message="Authentication Error, Invalid proxy server token passed. key={}, not found in db. Create key via `/key/generate` call.".format( hashed_token ), type=ProxyErrorTypes.token_not_found_in_db, param="key", code=status.HTTP_401_UNAUTHORIZED, ) _response = UserAPIKeyAuth(**_valid_token.model_dump(exclude_none=True)) # Load object_permission if object_permission_id exists but object_permission is not loaded if _response.object_permission_id and not _response.object_permission: try: _response.object_permission = await get_object_permission( object_permission_id=_response.object_permission_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, ) except Exception as e: verbose_proxy_logger.debug( f"Failed to load object_permission for key with object_permission_id={_response.object_permission_id}: {e}" ) # save the key object to cache await _cache_key_object( hashed_token=hashed_token, user_api_key_obj=_response, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) return _response def _copy_user_api_key_auth_for_cache( user_api_key_obj: UserAPIKeyAuth, ) -> UserAPIKeyAuth: copied_key_obj = user_api_key_obj.model_copy() copied_key_obj.budget_reservation = None copied_key_obj.parent_otel_span = None copied_key_obj.request_route = None return copied_key_obj @log_db_metrics async def get_object_permission( object_permission_id: str, prisma_client: Optional[PrismaClient], user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span] = None, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> Optional[LiteLLM_ObjectPermissionTable]: """ - Check if object permission id in proxy ObjectPermissionTable - if valid, return LiteLLM_ObjectPermissionTable object - if not, then raise an error """ if prisma_client is None: raise Exception( "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys" ) # check if in cache key = "object_permission_id:{}".format(object_permission_id) deserialized_perm = await user_api_key_cache.async_get_cache( key=key, model_type=LiteLLM_ObjectPermissionTable, ) if deserialized_perm is not None: return deserialized_perm # else, check db try: response = await ObjectPermissionRepository(prisma_client).table.find_unique( where={"object_permission_id": object_permission_id} ) if response is None: return None _perm_obj = LiteLLM_ObjectPermissionTable(**response.dict()) await user_api_key_cache.async_set_cache( key=key, value=_perm_obj, model_type=LiteLLM_ObjectPermissionTable, ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) return _perm_obj except Exception: return None @log_db_metrics async def get_managed_vector_store_rows_by_uuids( uuids: List[str], prisma_client: Optional[PrismaClient], user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span] = None, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> List[LiteLLM_ManagedVectorStoresTable]: """ Fetch managed vector store rows by their internal UUIDs. Follows the get_team_object / get_key_object / get_object_permission pattern: cache-first lookup (in-memory / Redis), DB fallback only on cache miss. Critical-path DB access must go through this helper to avoid raw Prisma calls on the hot request path. """ if not uuids or prisma_client is None: return [] result: List[LiteLLM_ManagedVectorStoresTable] = [] cache_misses: List[str] = [] for uuid in uuids: key = "managed_vector_store_id:{}".format(uuid) deserialized_vs = await user_api_key_cache.async_get_cache( key=key, model_type=LiteLLM_ManagedVectorStoresTable, ) if deserialized_vs is not None: result.append(deserialized_vs) else: cache_misses.append(uuid) if not cache_misses: return result rows = await ManagedVectorStoresRepository(prisma_client).table.find_many( where={"vector_store_id": {"in": cache_misses}}, take=len(cache_misses), ) for row in rows: row_dict = ( row.model_dump() if hasattr(row, "model_dump") else (row.dict() if hasattr(row, "dict") else None) ) if not isinstance(row_dict, dict) or not row_dict: row_dict = dict(row) if hasattr(row, "__dict__") else {} if not row_dict: continue cached_obj = LiteLLM_ManagedVectorStoresTable(**row_dict) key = "managed_vector_store_id:{}".format(cached_obj.vector_store_id) await user_api_key_cache.async_set_cache( key=key, value=cached_obj, model_type=LiteLLM_ManagedVectorStoresTable, ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) result.append(cached_obj) return result @log_db_metrics async def get_org_object( org_id: str, prisma_client: Optional[PrismaClient], user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span] = None, proxy_logging_obj: Optional[ProxyLogging] = None, include_budget_table: bool = False, ) -> Optional[LiteLLM_OrganizationTable]: """ - Check if org id in proxy Org Table - if valid, return LiteLLM_OrganizationTable object - if not, then raise an error Args: org_id: Organization ID to look up prisma_client: Database client user_api_key_cache: Cache for storing results parent_otel_span: Optional OpenTelemetry span proxy_logging_obj: Optional proxy logging object include_budget_table: If True, includes litellm_budget_table in the query """ if prisma_client is None: raise Exception( "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys" ) if not isinstance(org_id, str): return None # Use different cache key if budget table is included cache_key = "org_id:{}".format(org_id) if include_budget_table: cache_key = "org_id:{}:with_budget".format(org_id) # check if in cache deserialized_org = await user_api_key_cache.async_get_cache( key=cache_key, model_type=LiteLLM_OrganizationTable, ) if deserialized_org is not None: return deserialized_org # else, check db try: query_kwargs: Dict[str, Any] = {"where": {"organization_id": org_id}} if include_budget_table: query_kwargs["include"] = {"litellm_budget_table": True} response = await OrganizationRepository(prisma_client).table.find_unique( **query_kwargs ) if response is None: raise Exception _org_obj = LiteLLM_OrganizationTable(**response.model_dump()) # Cache the result await user_api_key_cache.async_set_cache( key=cache_key, value=_org_obj, model_type=LiteLLM_OrganizationTable, ttl=DEFAULT_IN_MEMORY_TTL, ) return _org_obj except Exception: raise Exception( f"Organization doesn't exist in db. Organization={org_id}. Create organization via `/organization/new` call." ) async def _get_resources_from_access_groups( access_group_ids: List[str], resource_field: Literal[ "access_model_names", "access_mcp_server_ids", "access_agent_ids" ], prisma_client: Optional[PrismaClient] = None, user_api_key_cache: Optional[UserApiKeyCache] = None, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> List[str]: """ Fetch access groups by their IDs (from cache or DB) and collect the specified resource field across all of them. Args: access_group_ids: List of access group IDs to fetch resource_field: Which resource list to extract from each access group - "access_model_names": model names (for model access checks) - "access_mcp_server_ids": MCP server IDs (for MCP access checks) - "access_agent_ids": agent IDs (for agent access checks) prisma_client: Optional PrismaClient (lazy-imported from proxy_server if None) user_api_key_cache: Optional DualCache (lazy-imported from proxy_server if None) proxy_logging_obj: Optional ProxyLogging (lazy-imported from proxy_server if None) Returns: Deduplicated list of resource identifiers from all resolved access groups. """ if not access_group_ids: return [] # Lazy import to avoid circular imports if prisma_client is None or user_api_key_cache is None: from litellm.proxy.proxy_server import prisma_client as _prisma_client from litellm.proxy.proxy_server import proxy_logging_obj as _proxy_logging_obj from litellm.proxy.proxy_server import user_api_key_cache as _user_api_key_cache prisma_client = prisma_client or _prisma_client user_api_key_cache = user_api_key_cache or _user_api_key_cache proxy_logging_obj = proxy_logging_obj or _proxy_logging_obj if user_api_key_cache is None: return [] resources: List[str] = [] for ag_id in access_group_ids: try: ag = await get_access_object( access_group_id=ag_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) resources.extend(getattr(ag, resource_field, [])) except Exception: verbose_proxy_logger.debug( "Could not fetch access group %s for resource field %s", ag_id, resource_field, ) return list(set(resources)) async def _get_models_from_access_groups( access_group_ids: List[str], prisma_client: Optional[PrismaClient] = None, user_api_key_cache: Optional[UserApiKeyCache] = None, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> List[str]: """ Collect model names from unified access groups. Models are matched by model name for backwards compatibility. """ return await _get_resources_from_access_groups( access_group_ids=access_group_ids, resource_field="access_model_names", prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) async def _get_mcp_server_ids_from_access_groups( access_group_ids: List[str], prisma_client: Optional[PrismaClient] = None, user_api_key_cache: Optional[UserApiKeyCache] = None, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> List[str]: """ Collect MCP server IDs from unified access groups. MCPs are matched by server ID. """ return await _get_resources_from_access_groups( access_group_ids=access_group_ids, resource_field="access_mcp_server_ids", prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) async def _get_agent_ids_from_access_groups( access_group_ids: List[str], prisma_client: Optional[PrismaClient] = None, user_api_key_cache: Optional[UserApiKeyCache] = None, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> List[str]: """ Collect agent IDs from unified access groups. Agents are matched by agent ID. """ return await _get_resources_from_access_groups( access_group_ids=access_group_ids, resource_field="access_agent_ids", prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) def _check_model_access_helper( model: str, llm_router: Optional[Router], models: List[str], team_model_aliases: Optional[Dict[str, str]] = None, team_id: Optional[str] = None, ) -> bool: ## check if model in allowed model names from collections import defaultdict access_groups: Dict[str, List[str]] = defaultdict(list) if llm_router: access_groups = llm_router.get_model_access_groups( model_name=model, team_id=team_id ) if ( len(access_groups) > 0 and llm_router is not None ): # check if token contains any model access groups for idx, m in enumerate( models ): # loop token models, if any of them are an access group add the access group if m in access_groups: return True # Filter out models that are access_groups filtered_models = [m for m in models if m not in access_groups] if _model_in_team_aliases(model=model, team_model_aliases=team_model_aliases): return True if _model_matches_any_wildcard_pattern_in_list( model=model, allowed_model_list=filtered_models ): return True all_model_access: bool = False if (len(filtered_models) == 0 and len(models) == 0) or "*" in filtered_models: all_model_access = True if SpecialModelNames.all_proxy_models.value in filtered_models: all_model_access = True if model is not None and model not in filtered_models and all_model_access is False: return False return True def _can_object_call_model( model: Union[str, List[str]], llm_router: Optional[Router], models: List[str], team_model_aliases: Optional[Dict[str, str]] = None, team_id: Optional[str] = None, object_type: Literal["user", "team", "key", "org", "project"] = "user", fallback_depth: int = 0, ) -> Literal[True]: """ Checks if token can call a given model Args: - model: str - llm_router: Optional[Router] - models: List[str] - team_model_aliases: Optional[Dict[str, str]] - object_type: Literal["user", "team", "key", "org"]. We use the object type to raise the correct exception type Returns: - True: if token allowed to call model Raises: - Exception: If token not allowed to call model """ if fallback_depth >= DEFAULT_MAX_RECURSE_DEPTH: raise Exception( "Unable to parse model, max fallback depth exceeded - received model: {}".format( model ) ) if isinstance(model, list): for m in model: _can_object_call_model( model=m, llm_router=llm_router, models=models, team_model_aliases=team_model_aliases, team_id=team_id, object_type=object_type, fallback_depth=fallback_depth + 1, ) return True potential_models = [model] if model in litellm.model_alias_map: potential_models.append(litellm.model_alias_map[model]) elif llm_router and model in llm_router.model_group_alias: _model = llm_router._get_model_from_alias(model) if _model: potential_models.append(_model) ## check model access for alias + underlying model - allow if either is in allowed models for m in potential_models: if _check_model_access_helper( model=m, llm_router=llm_router, models=models, team_model_aliases=team_model_aliases, team_id=team_id, ): return True raise ProxyException( message=f"{object_type} not allowed to access model. This {object_type} can only access models={models}. Tried to access {model}", type=ProxyErrorTypes.get_model_access_error_type_for_object( object_type=object_type ), param="model", code=status.HTTP_403_FORBIDDEN, ) def _model_in_team_aliases( model: str, team_model_aliases: Optional[Dict[str, str]] = None ) -> bool: """ Returns True if `model` being accessed is an alias of a team model - `model=gpt-4o` - `team_model_aliases={"gpt-4o": "gpt-4o-team-1"}` - returns True - `model=gp-4o` - `team_model_aliases={"o-3": "o3-preview"}` - returns False """ if team_model_aliases: if model in team_model_aliases: return True return False def _resolve_key_models_for_auth_check(valid_token: UserAPIKeyAuth) -> List[str]: """ Expand key model sentinels before auth checks. ``all-team-models`` means inherit the parent team's allowlist — same semantics as ``get_key_models`` in ``model_checks.py``. If the key has no team_id the sentinel cannot be resolved, so the original model list (still containing the sentinel string) is returned unchanged. That string won't match any real model, so access is denied rather than silently falling through to unrestricted access. """ models = list(valid_token.models or []) if SpecialModelNames.all_team_models.value in models: if valid_token.team_id is None: return models return list(valid_token.team_models or []) return models async def can_key_call_model( model: Union[str, List[str]], llm_model_list: Optional[list], valid_token: UserAPIKeyAuth, llm_router: Optional[litellm.Router], ) -> Literal[True]: """ Checks if token can call a given model 1. First checks native key-level model permissions (current implementation) 2. If not allowed natively, falls back to access_group_ids on the key Returns: - True: if token allowed to call model Raises: - Exception: If token not allowed to call model """ key_models = _resolve_key_models_for_auth_check(valid_token=valid_token) try: return _can_object_call_model( model=model, llm_router=llm_router, models=key_models, team_model_aliases=valid_token.team_model_aliases, team_id=valid_token.team_id, object_type="key", ) except ProxyException: # Fallback: check key's access_group_ids key_access_group_ids = valid_token.access_group_ids or [] if key_access_group_ids: models_from_groups = await _get_models_from_access_groups( access_group_ids=key_access_group_ids, ) if models_from_groups: return _can_object_call_model( model=model, llm_router=llm_router, models=models_from_groups, team_model_aliases=valid_token.team_model_aliases, team_id=valid_token.team_id, object_type="key", ) raise async def can_key_call_resolved_model( model: str, llm_model_list: Optional[list], valid_token: UserAPIKeyAuth, llm_router: Optional[litellm.Router], ) -> None: from litellm.proxy.proxy_server import ( prisma_client, proxy_logging_obj, user_api_key_cache, ) skip_key_model_check = valid_token.config or ( isinstance(valid_token.models, list) and SpecialModelNames.all_team_models.value in valid_token.models ) if not skip_key_model_check: await can_key_call_model( model=model, llm_model_list=llm_model_list, valid_token=valid_token, llm_router=llm_router, ) team_object: Optional[LiteLLM_TeamTableCachedObj] = None team_object_from_lookup = False if valid_token.team_id is not None: try: team_object = await get_team_object( team_id=valid_token.team_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, parent_otel_span=valid_token.parent_otel_span, proxy_logging_obj=proxy_logging_obj, ) team_object_from_lookup = True except Exception: team_object = LiteLLM_TeamTableCachedObj( team_id=valid_token.team_id, models=valid_token.team_models, blocked=valid_token.team_blocked, team_alias=valid_token.team_alias, metadata=valid_token.team_metadata, object_permission_id=valid_token.team_object_permission_id, object_permission=valid_token.team_object_permission, ) if team_object is not None: try: await can_team_access_model( model=model, team_object=team_object, llm_router=llm_router, team_model_aliases=valid_token.team_model_aliases, ) except ProxyException as team_denial: if team_denial.type != ProxyErrorTypes.team_model_access_denied: raise if not await _key_access_group_grants_model( model=model, valid_token=valid_token, team_object=team_object, llm_router=llm_router, ): raise if valid_token.user_id is not None and team_object_from_lookup: await _check_team_member_model_access( model=model, team_object=team_object, valid_token=valid_token, llm_router=llm_router, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) if valid_token.project_id is not None: project_object = await get_project_object( project_id=valid_token.project_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) if project_object is not None and len(project_object.models) > 0: can_project_access_model( model=model, project_object=project_object, llm_router=llm_router, ) def can_org_access_model( model: str, org_object: Optional[LiteLLM_OrganizationTable], llm_router: Optional[Router], team_model_aliases: Optional[Dict[str, str]] = None, ) -> Literal[True]: """ Returns True if the team can access a specific model. """ return _can_object_call_model( model=model, llm_router=llm_router, models=org_object.models if org_object else [], team_model_aliases=team_model_aliases, object_type="org", ) async def can_team_access_model( model: Union[str, List[str]], team_object: Optional[LiteLLM_TeamTable], llm_router: Optional[Router], team_model_aliases: Optional[Dict[str, str]] = None, ) -> Literal[True]: """ Returns True if the team can access a specific model. 1. First checks native team-level model permissions (current implementation) 2. If not allowed natively, falls back to access_group_ids on the team """ try: return _can_object_call_model( model=model, llm_router=llm_router, models=team_object.models if team_object else [], team_model_aliases=team_model_aliases, team_id=team_object.team_id if team_object else None, object_type="team", ) except ProxyException: # Fallback: check team's access_group_ids team_access_group_ids = ( (team_object.access_group_ids or []) if team_object else [] ) if team_access_group_ids: models_from_groups = await _get_models_from_access_groups( access_group_ids=team_access_group_ids, ) if models_from_groups: return _can_object_call_model( model=model, llm_router=llm_router, models=models_from_groups, team_model_aliases=team_model_aliases, team_id=team_object.team_id if team_object else None, object_type="team", ) raise async def get_authorized_resources_from_key_access_groups( valid_token: Optional[UserAPIKeyAuth], team_object: Optional[LiteLLM_TeamTable], resource_field: Literal[ "access_model_names", "access_mcp_server_ids", "access_agent_ids" ], ) -> List[str]: """ For each access_group_id on the key, fetch the LiteLLM_AccessGroupTable row and contribute its `resource_field` only if the group authorizes the caller as an owner — that is, the group's `assigned_team_ids` includes the key's `team_id`, or the group's `assigned_key_ids` includes the key's token. This preserves the team-as-owner boundary while still letting a group reach the key without first being added to the team's `access_group_ids` list. """ if valid_token is None: return [] key_access_group_ids = list(valid_token.access_group_ids or []) if not key_access_group_ids: return [] from litellm.proxy.proxy_server import prisma_client as _prisma_client from litellm.proxy.proxy_server import proxy_logging_obj as _proxy_logging_obj from litellm.proxy.proxy_server import user_api_key_cache as _user_api_key_cache if _prisma_client is None or _user_api_key_cache is None: return [] key_team_id = valid_token.team_id or ( team_object.team_id if team_object is not None else None ) key_token = valid_token.token authorized_resources: List[str] = [] for ag_id in key_access_group_ids: try: ag = await get_access_object( access_group_id=ag_id, prisma_client=_prisma_client, user_api_key_cache=_user_api_key_cache, proxy_logging_obj=_proxy_logging_obj, ) except Exception: continue team_authorized = bool( key_team_id and key_team_id in (ag.assigned_team_ids or []) ) key_authorized = bool(key_token and key_token in (ag.assigned_key_ids or [])) if team_authorized or key_authorized: authorized_resources.extend(getattr(ag, resource_field, []) or []) return list(set(authorized_resources)) async def _key_access_group_grants_model( model: Union[str, List[str]], valid_token: Optional[UserAPIKeyAuth], team_object: Optional[LiteLLM_TeamTable], llm_router: Optional[Router], ) -> bool: """ Returns True if the key's `access_group_ids` expand to models that grant access to `model`. Used to let a key's access group override a team's model restriction in `common_checks`. """ authorized_models = await get_authorized_resources_from_key_access_groups( valid_token=valid_token, team_object=team_object, resource_field="access_model_names", ) if not authorized_models: return False try: _can_object_call_model( model=model, llm_router=llm_router, models=authorized_models, team_model_aliases=valid_token.team_model_aliases if valid_token else None, team_id=valid_token.team_id if valid_token else None, object_type="key", ) return True except ProxyException: return False def can_project_access_model( model: Union[str, List[str]], project_object: LiteLLM_ProjectTableCachedObj, llm_router: Optional[Router], ) -> Literal[True]: """ Returns True if the project can access a specific model. Raises ProxyException if access is denied. """ return _can_object_call_model( model=model, llm_router=llm_router, models=project_object.models if project_object else [], object_type="project", ) async def can_user_call_model( model: Union[str, List[str]], llm_router: Optional[Router], user_object: Optional[LiteLLM_UserTable], ) -> Literal[True]: if user_object is None: return True if SpecialModelNames.no_default_models.value in user_object.models: raise ProxyException( message=f"User not allowed to access model. No default model access, only team models allowed. Tried to access {model}", type=ProxyErrorTypes.key_model_access_denied, param="model", code=status.HTTP_403_FORBIDDEN, ) return _can_object_call_model( model=model, llm_router=llm_router, models=user_object.models, object_type="user", ) def _search_tool_names_from_object_permission( object_permission: Optional[LiteLLM_ObjectPermissionTable], ) -> List[str]: """Return allowlisted search tool names from object_permission (empty = unrestricted).""" if object_permission is None: return [] raw = object_permission.search_tools if not raw: return [] return list(raw) def _can_object_call_search_tool( search_tool_name: str, allowed_search_tools: List[str], object_type: Literal["key", "team", "project"], ) -> Literal[True]: """ Check if an object (key/team/project) can access a specific search tool. Similar to _can_object_call_model but for search tools. Args: search_tool_name: The search tool being requested allowed_search_tools: List of allowed search tool names for this object object_type: Type of object for error messaging Returns: True if access is allowed Raises: ProxyException if access is denied """ # Empty list means all search tools are allowed if not allowed_search_tools: return True # Check if the search tool is in the allowlist if search_tool_name in allowed_search_tools: return True # Access denied raise ProxyException( message=f"{object_type.capitalize()} not allowed to access search tool: {search_tool_name}. " f"Allowed search tools: {allowed_search_tools}", type=ProxyErrorTypes.key_model_access_denied, param="search_tool_name", code=status.HTTP_403_FORBIDDEN, ) async def can_key_call_search_tool( search_tool_name: str, valid_token: UserAPIKeyAuth, ) -> Literal[True]: """ Check if a key can access a specific search tool. Similar to can_key_call_model but for search tools. Args: search_tool_name: The search tool being requested valid_token: The authenticated key Returns: True if access is allowed Raises: ProxyException if access is denied """ return _can_object_call_search_tool( search_tool_name=search_tool_name, allowed_search_tools=_search_tool_names_from_object_permission( valid_token.object_permission ), object_type="key", ) async def can_team_call_search_tool( search_tool_name: str, team_object: Optional[LiteLLM_TeamTable], ) -> Literal[True]: """ Check if a team can access a specific search tool. Similar to can_team_access_model but for search tools. Args: search_tool_name: The search tool being requested team_object: The team object Returns: True if access is allowed Raises: ProxyException if access is denied """ if team_object is None: return True return _can_object_call_search_tool( search_tool_name=search_tool_name, allowed_search_tools=_search_tool_names_from_object_permission( team_object.object_permission ), object_type="team", ) async def can_user_view_search_tool( search_tool_name: str, valid_token: UserAPIKeyAuth, team_object: Optional[LiteLLM_TeamTable], ) -> bool: """ Boolean variant of the key + team authorization enforced on /search, used to scope /search_tools/list so a non-admin caller only sees tools it may invoke. """ try: await can_key_call_search_tool( search_tool_name=search_tool_name, valid_token=valid_token, ) await can_team_call_search_tool( search_tool_name=search_tool_name, team_object=team_object, ) except ProxyException: return False return True async def is_valid_fallback_model( model: str, llm_router: Optional[Router], user_model: Optional[str], ) -> Literal[True]: """ Try to route the fallback model request. Validate if it can't be routed. Help catch invalid fallback models. """ await route_request( data={ "model": model, "messages": [{"role": "user", "content": "Who was Alexander?"}], }, llm_router=llm_router, user_model=user_model, route_type="acompletion", # route type shouldn't affect the fallback model check ) return True async def _virtual_key_max_budget_check( valid_token: UserAPIKeyAuth, proxy_logging_obj: ProxyLogging, user_obj: Optional[LiteLLM_UserTable] = None, ): """ Raises: BudgetExceededError if the token is over it's max budget. Triggers a budget alert if the token is over it's max budget. """ if valid_token.max_budget is not None: from litellm.proxy.proxy_server import get_current_spend fallback_spend = valid_token.spend or 0.0 counter_key = f"spend:key:{valid_token.token}" # Read spend from cross-pod counter (Redis-first) or cached object (fallback) spend = await get_current_spend( counter_key=counter_key, fallback_spend=fallback_spend, max_budget=valid_token.max_budget, ) #################################### # collect information for alerting # #################################### user_email = None # Check if the token has any user id information if user_obj is not None: user_email = user_obj.user_email call_info = CallInfo( token=valid_token.token, spend=spend, max_budget=valid_token.max_budget, soft_budget=valid_token.soft_budget, user_id=valid_token.user_id, team_id=valid_token.team_id, organization_id=valid_token.org_id, user_email=user_email, key_alias=valid_token.key_alias, event_group=Litellm_EntityType.KEY, ) asyncio.create_task( proxy_logging_obj.budget_alerts( type="token_budget", user_info=call_info, ) ) #################################### # collect information for alerting # #################################### # Defense-in-depth (GHSA-2rv4-xv66-fpjg): spend >= NaN is always False, # so a NaN max_budget would silently disable enforcement. Treat a # non-finite max_budget as "no configured limit" rather than as a bypass. if math.isfinite(valid_token.max_budget) and spend >= valid_token.max_budget: raise litellm.BudgetExceededError( current_cost=spend, max_budget=valid_token.max_budget, ) async def _virtual_key_multi_budget_check( valid_token: UserAPIKeyAuth, ): """ Raises BudgetExceededError if any budget window in valid_token.budget_limits is exceeded. Each window has its own Redis counter keyed by spend:key:{token}:window:{budget_duration}. Using budget_duration (not list index) keeps counters stable when windows are reordered or removed during a key update. Note: counters are not seeded from DB on Redis cold-start. After a Redis flush, per-window spend resets to zero within the current window period. This is an acceptable trade-off: the DB stores reset_at timestamps but not per-window accumulated spend. """ if not valid_token.budget_limits: return from litellm.proxy.proxy_server import get_current_spend for window in valid_token.budget_limits: w: dict = window if isinstance(window, dict) else window.model_dump() counter_key = f"spend:key:{valid_token.token}:window:{w['budget_duration']}" window_spend = await get_current_spend( counter_key=counter_key, fallback_spend=0.0, max_budget=w["max_budget"], window_entity_type="Key", window_entity_id=valid_token.token, window_start=get_budget_window_start(w), ) if math.isfinite(w["max_budget"]) and window_spend >= w["max_budget"]: raise litellm.BudgetExceededError( current_cost=window_spend, max_budget=w["max_budget"], message=( f"ExceededBudget: Key over {w['budget_duration']} budget. " f"Spend=${window_spend:.4f}, Limit=${w['max_budget']:.2f}" ), ) async def _virtual_key_soft_budget_check( valid_token: UserAPIKeyAuth, proxy_logging_obj: ProxyLogging, user_obj: Optional[LiteLLM_UserTable] = None, ): """ Triggers a budget alert if the token is over it's soft budget. """ if valid_token.soft_budget and valid_token.spend >= valid_token.soft_budget: verbose_proxy_logger.debug( "Crossed Soft Budget for token %s, spend %s, soft_budget %s", valid_token.token, valid_token.spend, valid_token.soft_budget, ) call_info = CallInfo( token=valid_token.token, spend=valid_token.spend, max_budget=valid_token.max_budget, soft_budget=valid_token.soft_budget, user_id=valid_token.user_id, team_id=valid_token.team_id, team_alias=valid_token.team_alias, organization_id=valid_token.org_id, user_email=user_obj.user_email if user_obj else None, key_alias=valid_token.key_alias, event_group=Litellm_EntityType.KEY, ) asyncio.create_task( proxy_logging_obj.budget_alerts( type="soft_budget", user_info=call_info, ) ) def _parse_email_list(raw: Any) -> List[str]: """Parse emails from a list or comma-separated string.""" if isinstance(raw, list): return [e.strip() for e in raw if isinstance(e, str) and e.strip()] elif isinstance(raw, str): return [e.strip() for e in raw.split(",") if e.strip()] return [] def _normalize_alert_emails( cfg: Optional[Dict[str, Any]], ) -> Dict[str, List[str]]: """Coerce user-supplied threshold→recipients mapping to Dict[str, List[str]]. Values may legitimately arrive as list, comma-separated string, or None from YAML/metadata; _parse_email_list tolerates all three. """ if not cfg: return {} return {k: _parse_email_list(v) for k, v in cfg.items()} def _merge_budget_alert_email_configs( global_cfg: Optional[Dict[str, Any]], per_key_cfg: Optional[Dict[str, Any]], ) -> Optional[Dict[str, List[str]]]: """ Per-threshold additive merge: each threshold's recipient list is the union of global + per-key entries (deduped, global-first ordering). Missing thresholds on one side are inherited from the other. """ global_cfg_normalized = _normalize_alert_emails(global_cfg) per_key_cfg_normalized = _normalize_alert_emails(per_key_cfg) if not global_cfg_normalized and not per_key_cfg_normalized: return None thresholds = set(global_cfg_normalized) | set(per_key_cfg_normalized) return { t: list( dict.fromkeys( global_cfg_normalized.get(t, []) + per_key_cfg_normalized.get(t, []) ) ) for t in thresholds } async def _virtual_key_max_budget_alert_check( valid_token: UserAPIKeyAuth, proxy_logging_obj: ProxyLogging, user_obj: Optional[LiteLLM_UserTable] = None, ): """ Triggers a budget alert if the token has reached EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE (default 80%) of its max budget. This is a warning alert before the token actually exceeds the max budget. """ if ( valid_token.max_budget is not None and valid_token.spend is not None and valid_token.spend > 0 ): owner_email = user_obj.user_email if user_obj else None alert_email_config: Optional[Dict[str, List[str]]] = ( _merge_budget_alert_email_configs( global_cfg=litellm.default_key_max_budget_alert_emails, per_key_cfg=(valid_token.metadata or {}).get("max_budget_alert_emails"), ) ) if isinstance(alert_email_config, dict) and alert_email_config: # New path: only create task if spend has crossed the lowest threshold min_pct = min( (int(k) for k in alert_email_config if k.isdigit()), default=None, ) if min_pct is None or valid_token.spend < valid_token.max_budget * ( min_pct / 100.0 ): return call_info = CallInfo( token=valid_token.token, spend=valid_token.spend, max_budget=valid_token.max_budget, soft_budget=valid_token.soft_budget, user_id=valid_token.user_id, team_id=valid_token.team_id, team_alias=valid_token.team_alias, organization_id=valid_token.org_id, user_email=owner_email, key_alias=valid_token.key_alias, event_group=Litellm_EntityType.KEY, max_budget_alert_emails=alert_email_config, ) asyncio.create_task( proxy_logging_obj.budget_alerts( type="max_budget_alert", user_info=call_info, ) ) else: # Old path: existing single 80% threshold — completely unchanged alert_threshold = ( valid_token.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE ) if ( valid_token.spend >= alert_threshold and valid_token.spend < valid_token.max_budget ): verbose_proxy_logger.debug( "Reached Max Budget Alert Threshold for token %s, spend %s, max_budget %s, alert_threshold %s", valid_token.token, valid_token.spend, valid_token.max_budget, alert_threshold, ) call_info = CallInfo( token=valid_token.token, spend=valid_token.spend, max_budget=valid_token.max_budget, soft_budget=valid_token.soft_budget, user_id=valid_token.user_id, team_id=valid_token.team_id, team_alias=valid_token.team_alias, organization_id=valid_token.org_id, user_email=owner_email, key_alias=valid_token.key_alias, event_group=Litellm_EntityType.KEY, ) asyncio.create_task( proxy_logging_obj.budget_alerts( type="max_budget_alert", user_info=call_info, ) ) async def _check_team_member_budget( team_object: Optional[LiteLLM_TeamTable], user_object: Optional[LiteLLM_UserTable], valid_token: Optional[UserAPIKeyAuth], prisma_client: Optional[PrismaClient], user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, ): """Check if team member is over their max budget within the team.""" if ( team_object is not None and team_object.team_id is not None and valid_token is not None and valid_token.user_id is not None ): team_membership = await get_team_membership( user_id=valid_token.user_id, team_id=team_object.team_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) # Per-member override wins; otherwise fall back to the team-level # default configured via team.metadata["team_member_budget_id"]. team_member_budget: Optional[float] = None if ( team_membership is not None and team_membership.litellm_budget_table is not None and team_membership.litellm_budget_table.max_budget is not None ): team_member_budget = team_membership.litellm_budget_table.max_budget else: default_budget_id = (team_object.metadata or {}).get( "team_member_budget_id" ) if isinstance(default_budget_id, str): default_budget = await get_team_member_default_budget( budget_id=default_budget_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, ) # Treat 0 on the team default as "no cap". # Per-member rows still respect 0 as an explicit admin disable. if ( default_budget is not None and default_budget.max_budget is not None and default_budget.max_budget > 0 ): team_member_budget = default_budget.max_budget if team_member_budget is not None: team_member_spend = ( team_membership.spend if team_membership is not None else 0.0 ) or 0.0 # Read from cross-pod counter (Redis-first) if available from litellm.proxy.proxy_server import get_current_spend team_member_spend = await get_current_spend( counter_key=f"spend:team_member:{valid_token.user_id}:{team_object.team_id}", fallback_spend=team_member_spend, max_budget=team_member_budget, ) if ( math.isfinite(team_member_budget) and team_member_spend >= team_member_budget ): raise litellm.BudgetExceededError( current_cost=team_member_spend, max_budget=team_member_budget, 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}", ) async def _check_team_member_model_access( model: Union[str, List[str]], team_object: LiteLLM_TeamTable, valid_token: UserAPIKeyAuth, llm_router: Optional[Router], prisma_client: Optional["PrismaClient"], user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, ) -> None: """ Check if a team member's per-member model scope allows access to the requested model. Only enforced when the member's budget table has a non-empty allowed_models list. If allowed_models is empty or absent, the team-level models list applies (no extra restriction). """ if valid_token.user_id is None or team_object.team_id is None: return team_membership = await get_team_membership( user_id=valid_token.user_id, team_id=team_object.team_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) if ( team_membership is None or team_membership.litellm_budget_table is None or not team_membership.litellm_budget_table.allowed_models ): return # no per-member restriction — inherit team-level check member_allowed_models: List[str] = ( team_membership.litellm_budget_table.allowed_models ) try: _can_object_call_model( model=model, llm_router=llm_router, models=member_allowed_models, object_type="team", team_id=team_object.team_id, ) except ProxyException: raise ProxyException( message=f"Team member not allowed to access model. User={valid_token.user_id}, Team={team_object.team_id}, Model={model}. Allowed member models = {member_allowed_models}", type=ProxyErrorTypes.team_model_access_denied, param="model", code=status.HTTP_403_FORBIDDEN, ) async def _team_max_budget_check( team_object: Optional[LiteLLM_TeamTable], valid_token: Optional[UserAPIKeyAuth], proxy_logging_obj: ProxyLogging, ): """ Check if the team is over it's max budget. Raises: BudgetExceededError if the team is over it's max budget. Triggers a budget alert if the team is over it's max budget. """ if team_object is not None and team_object.max_budget is not None: from litellm.proxy.proxy_server import get_current_spend # Read spend from cross-pod counter (Redis-first) or cached object (fallback) spend = await get_current_spend( counter_key=f"spend:team:{team_object.team_id}", fallback_spend=team_object.spend or 0.0, max_budget=team_object.max_budget, ) if math.isfinite(team_object.max_budget) and spend > team_object.max_budget: if valid_token: call_info = CallInfo( token=valid_token.token, spend=spend, max_budget=team_object.max_budget, user_id=valid_token.user_id, team_id=valid_token.team_id, team_alias=valid_token.team_alias, organization_id=valid_token.org_id, event_group=Litellm_EntityType.TEAM, ) asyncio.create_task( proxy_logging_obj.budget_alerts( type="team_budget", user_info=call_info, ) ) raise litellm.BudgetExceededError( current_cost=spend, max_budget=team_object.max_budget, message=f"Budget has been exceeded! Team={team_object.team_id} Current cost: {spend}, Max budget: {team_object.max_budget}", ) async def _team_multi_budget_check( team_object: Optional[LiteLLM_TeamTable], ): """ Raises BudgetExceededError if any budget window in team_object.budget_limits is exceeded. Each window has its own Redis counter keyed by spend:team:{team_id}:window:{budget_duration}. Using budget_duration (not list index) keeps counters stable when windows are reordered or removed during a team update. """ if team_object is None or not team_object.budget_limits: return from litellm.proxy.proxy_server import get_current_spend for window in team_object.budget_limits: w: dict = window if isinstance(window, dict) else window.model_dump() counter_key = f"spend:team:{team_object.team_id}:window:{w['budget_duration']}" window_spend = await get_current_spend( counter_key=counter_key, fallback_spend=0.0, max_budget=w["max_budget"], window_entity_type="Team", window_entity_id=team_object.team_id, window_start=get_budget_window_start(w), ) if math.isfinite(w["max_budget"]) and window_spend >= w["max_budget"]: raise litellm.BudgetExceededError( current_cost=window_spend, max_budget=w["max_budget"], message=( f"ExceededBudget: Team={team_object.team_id} over {w['budget_duration']} budget. " f"Spend=${window_spend:.4f}, Limit=${w['max_budget']:.2f}" ), ) async def _team_soft_budget_check( team_object: Optional[LiteLLM_TeamTable], valid_token: Optional[UserAPIKeyAuth], proxy_logging_obj: ProxyLogging, ): """ Triggers a budget alert if the team is over it's soft budget. """ if ( team_object is not None and team_object.soft_budget is not None and team_object.spend is not None and team_object.spend >= team_object.soft_budget ): verbose_proxy_logger.debug( "Crossed Soft Budget for team %s, spend %s, soft_budget %s", team_object.team_id, team_object.spend, team_object.soft_budget, ) if valid_token: # Extract alert emails from team metadata alert_emails: Optional[List[str]] = None if team_object.metadata is not None and isinstance( team_object.metadata, dict ): soft_budget_alert_emails = team_object.metadata.get( "soft_budget_alerting_emails" ) if soft_budget_alert_emails is not None: if isinstance(soft_budget_alert_emails, list): alert_emails = [ email for email in soft_budget_alert_emails if isinstance(email, str) and email.strip() ] elif isinstance(soft_budget_alert_emails, str): # Handle comma-separated string alert_emails = [ email.strip() for email in soft_budget_alert_emails.split(",") if email.strip() ] # Filter out empty strings if alert_emails: alert_emails = [email for email in alert_emails if email] else: alert_emails = None # Only send team soft budget alerts if alert_emails are configured # Team soft budget alerts are sent via metadata.soft_budget_alerting_emails, not global alerting if alert_emails is None or len(alert_emails) == 0: verbose_proxy_logger.debug( "Skipping team soft budget alert for team %s: no alert_emails configured in metadata.soft_budget_alerting_emails", team_object.team_id, ) return call_info = CallInfo( token=valid_token.token, spend=team_object.spend, max_budget=team_object.max_budget, soft_budget=team_object.soft_budget, user_id=valid_token.user_id, team_id=valid_token.team_id, team_alias=valid_token.team_alias, organization_id=valid_token.org_id, user_email=None, # Team-level alert, no specific user email key_alias=valid_token.key_alias, event_group=Litellm_EntityType.TEAM, alert_emails=alert_emails, ) asyncio.create_task( proxy_logging_obj.budget_alerts( type="soft_budget", user_info=call_info, ) ) async def _project_max_budget_check( project_object: Optional[LiteLLM_ProjectTableCachedObj], valid_token: Optional[UserAPIKeyAuth], proxy_logging_obj: ProxyLogging, ): """ Check if the project is over its max budget. Raises: BudgetExceededError if the project is over its max budget. Triggers a budget alert if the project is over its max budget. """ if project_object is None: return max_budget = None if project_object.litellm_budget_table is not None: max_budget = project_object.litellm_budget_table.max_budget if ( max_budget is not None and project_object.spend is not None and math.isfinite(max_budget) and project_object.spend > max_budget ): if valid_token: call_info = CallInfo( token=valid_token.token, spend=project_object.spend, max_budget=max_budget, user_id=valid_token.user_id, team_id=valid_token.team_id, team_alias=valid_token.team_alias, organization_id=valid_token.org_id, event_group=Litellm_EntityType.PROJECT, ) asyncio.create_task( proxy_logging_obj.budget_alerts( type="project_budget", user_info=call_info, ) ) raise litellm.BudgetExceededError( current_cost=project_object.spend, max_budget=max_budget, message=f"Budget has been exceeded! Project={project_object.project_id} Current cost: {project_object.spend}, Max budget: {max_budget}", ) async def _project_soft_budget_check( project_object: Optional[LiteLLM_ProjectTableCachedObj], valid_token: Optional[UserAPIKeyAuth], proxy_logging_obj: ProxyLogging, ): """ Triggers a budget alert if the project is over its soft budget. Mirrors _team_soft_budget_check() pattern. """ if project_object is None: return soft_budget = None if project_object.litellm_budget_table is not None: soft_budget = project_object.litellm_budget_table.soft_budget if ( soft_budget is not None and project_object.spend is not None and project_object.spend >= soft_budget ): verbose_proxy_logger.debug( "Crossed Soft Budget for project %s, spend %s, soft_budget %s", project_object.project_id, project_object.spend, soft_budget, ) if valid_token: call_info = CallInfo( token=valid_token.token, spend=project_object.spend, max_budget=None, soft_budget=soft_budget, user_id=valid_token.user_id, team_id=valid_token.team_id, team_alias=valid_token.team_alias, organization_id=valid_token.org_id, event_group=Litellm_EntityType.PROJECT, ) asyncio.create_task( proxy_logging_obj.budget_alerts( type="soft_budget", user_info=call_info, ) ) async def get_project_object( project_id: str, prisma_client: Optional[PrismaClient], user_api_key_cache: UserApiKeyCache, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> Optional[LiteLLM_ProjectTableCachedObj]: """ Fetch project object from cache or DB. Follows get_team_object() caching pattern with TTL and last_refreshed_at. Returns LiteLLM_ProjectTableCachedObj or None if not found. """ if prisma_client is None: return None # Check cache first cache_key = "project_id:{}".format(project_id) deserialized_project = await user_api_key_cache.async_get_cache( key=cache_key, model_type=LiteLLM_ProjectTableCachedObj, ) if deserialized_project is not None: return deserialized_project # Fetch from DB project_row = await ProjectRepository(prisma_client).table.find_unique( where={"project_id": project_id}, include={"litellm_budget_table": True}, ) if project_row is None: return None project_obj = LiteLLM_ProjectTableCachedObj(**project_row.model_dump()) # Cache with TTL following _cache_management_object pattern project_obj.last_refreshed_at = time.time() await _cache_management_object( key=cache_key, value=project_obj, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, model_type=LiteLLM_ProjectTableCachedObj, ) return project_obj async def _organization_max_budget_check( valid_token: Optional[UserAPIKeyAuth], team_object: Optional[LiteLLM_TeamTable], prisma_client: Optional[PrismaClient], user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, ): """ Check if the organization is over its max budget. This function checks the organization budget using: 1. First, tries to use valid_token.org_id (if key has organization_id set) 2. Falls back to team_object.organization_id (if key doesn't have org_id but team does) This ensures organization budget checks work even when keys don't have organization_id set directly, as long as their team belongs to an organization. Raises: BudgetExceededError if the organization is over its max budget. Triggers a budget alert if the organization is over its max budget. """ if valid_token is None or prisma_client is None: return # Determine organization_id: first try from token, then fallback to team org_id: Optional[str] = None if valid_token.org_id is not None: org_id = valid_token.org_id elif team_object is not None and team_object.organization_id is not None: org_id = team_object.organization_id # If no organization_id found, skip the check if org_id is None: return # Get organization object with budget table - use get_org_object so it can be mocked in tests try: org_table = await get_org_object( org_id=org_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, include_budget_table=True, ) except Exception: # If organization lookup fails, skip the check return if org_table is None: return # Get max_budget from organization's budget table org_max_budget: Optional[float] = None if org_table.litellm_budget_table is not None: org_max_budget = org_table.litellm_budget_table.max_budget # Only check if organization has a valid max_budget set if org_max_budget is None or org_max_budget <= 0: return # Read spend from cross-pod counter (Redis-first) or cached object (fallback) from litellm.proxy.proxy_server import get_current_spend org_spend = await get_current_spend( counter_key=f"spend:org:{org_id}", fallback_spend=org_table.spend or 0.0, max_budget=org_max_budget, ) # Check if organization spend exceeds max budget if math.isfinite(org_max_budget) and org_spend >= org_max_budget: # Trigger budget alert call_info = CallInfo( token=valid_token.token, spend=org_spend, max_budget=org_max_budget, user_id=valid_token.user_id, team_id=valid_token.team_id, team_alias=valid_token.team_alias, organization_id=org_id, event_group=Litellm_EntityType.ORGANIZATION, ) asyncio.create_task( proxy_logging_obj.budget_alerts( type="organization_budget", user_info=call_info, ) ) raise litellm.BudgetExceededError( current_cost=org_spend, max_budget=org_max_budget, message=f"Budget has been exceeded! Organization={org_id} Current cost: {org_spend}, Max budget: {org_max_budget}", ) async def _tag_max_budget_check( request_body: dict, prisma_client: Optional[PrismaClient], user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, valid_token: Optional[UserAPIKeyAuth], ): """ Check if any tags in the request are over their max budget. Raises: BudgetExceededError if any tag is over its max budget. Triggers a budget alert if any tag is over its max budget. """ from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body if prisma_client is None: return # Get tags from request metadata tags = get_tags_from_request_body(request_body=request_body) if not tags: return # Batch fetch all tags in one go tag_objects = await get_tag_objects_batch( tag_names=tags, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) # Check budget for each tag for tag_name in tags: tag_object = tag_objects.get(tag_name) if tag_object is None: continue # Check if tag has budget limits if ( tag_object.litellm_budget_table is not None and tag_object.litellm_budget_table.max_budget is not None ): from litellm.proxy.proxy_server import get_current_spend tag_spend = await get_current_spend( counter_key=f"spend:tag:{tag_name}", fallback_spend=tag_object.spend or 0.0, max_budget=tag_object.litellm_budget_table.max_budget, fallback_authoritative=True, ) if tag_spend <= tag_object.litellm_budget_table.max_budget: continue raise litellm.BudgetExceededError( current_cost=tag_spend, max_budget=tag_object.litellm_budget_table.max_budget, message=f"Budget has been exceeded! Tag={tag_name} Current cost: {tag_spend}, Max budget: {tag_object.litellm_budget_table.max_budget}", ) def is_model_allowed_by_pattern(model: str, allowed_model_pattern: str) -> bool: """ Check if a model matches an allowed pattern. Handles exact matches and wildcard patterns. Args: model (str): The model to check (e.g., "bedrock/anthropic.claude-3-5-sonnet-20240620") allowed_model_pattern (str): The allowed pattern (e.g., "bedrock/*", "*", "openai/*") Returns: bool: True if model matches the pattern, False otherwise """ if "*" in allowed_model_pattern: pattern = f"^{allowed_model_pattern.replace('*', '.*')}$" return bool(re.match(pattern, model)) return False def _model_matches_any_wildcard_pattern_in_list( model: str, allowed_model_list: list ) -> bool: """ Returns True if a model matches any wildcard pattern in a list. eg. - model=`bedrock/us.amazon.nova-micro-v1:0`, allowed_models=`bedrock/*` returns True - model=`bedrock/us.amazon.nova-micro-v1:0`, allowed_models=`bedrock/us.*` returns True - model=`bedrockzzzz/us.amazon.nova-micro-v1:0`, allowed_models=`bedrock/*` returns False """ if any( _is_wildcard_pattern(allowed_model_pattern) and is_model_allowed_by_pattern( model=model, allowed_model_pattern=allowed_model_pattern ) for allowed_model_pattern in allowed_model_list ): return True if any( _is_wildcard_pattern(allowed_model_pattern) and _model_custom_llm_provider_matches_wildcard_pattern( model=model, allowed_model_pattern=allowed_model_pattern ) for allowed_model_pattern in allowed_model_list ): return True return False def _model_custom_llm_provider_matches_wildcard_pattern( model: str, allowed_model_pattern: str ) -> bool: """ Returns True for this scenario: - `model=gpt-4o` - `allowed_model_pattern=openai/*` or - `model=claude-3-5-sonnet-20240620` - `allowed_model_pattern=anthropic/*` """ try: model, custom_llm_provider, _, _ = get_llm_provider(model=model) except Exception: return False return is_model_allowed_by_pattern( model=f"{custom_llm_provider}/{model}", allowed_model_pattern=allowed_model_pattern, ) def _is_wildcard_pattern(allowed_model_pattern: str) -> bool: """ Returns True if the pattern is a wildcard pattern. Checks if `*` is in the pattern. """ return "*" in allowed_model_pattern async def vector_store_access_check( request_body: dict, team_object: Optional[LiteLLM_TeamTable], valid_token: Optional[UserAPIKeyAuth], ): """ Checks if the object (key, team, org) has access to the vector store. Raises ProxyException if the object (key, team, org) cannot access the specific vector store. """ from litellm.proxy.proxy_server import prisma_client ######################################################### # Get the vector store the user is trying to access ######################################################### if prisma_client is None: verbose_proxy_logger.debug( "Prisma client not found, skipping vector store access check" ) return True if litellm.vector_store_registry is None: verbose_proxy_logger.debug( "Vector store registry not found, skipping vector store access check" ) return True vector_store_ids_to_run = litellm.vector_store_registry.get_vector_store_ids_to_run( non_default_params=request_body, tools=request_body.get("tools", None) ) if vector_store_ids_to_run is None: verbose_proxy_logger.debug( "Vector store to run not found, skipping vector store access check" ) return True ######################################################### # Check if the object (key, team, org) has access to the vector store ######################################################### # Check if the key can access the vector store if valid_token is not None and valid_token.object_permission_id is not None: key_object_permission = await ObjectPermissionRepository( prisma_client ).table.find_unique( where={"object_permission_id": valid_token.object_permission_id}, ) if key_object_permission is not None: _can_object_call_vector_stores( object_type="key", vector_store_ids_to_run=vector_store_ids_to_run, object_permissions=key_object_permission, ) # Check if the team can access the vector store if team_object is not None and team_object.object_permission_id is not None: team_object_permission = await ObjectPermissionRepository( prisma_client ).table.find_unique( where={"object_permission_id": team_object.object_permission_id}, ) if team_object_permission is not None: _can_object_call_vector_stores( object_type="team", vector_store_ids_to_run=vector_store_ids_to_run, object_permissions=team_object_permission, ) return True def _can_object_call_vector_stores( object_type: Literal["key", "team", "org"], vector_store_ids_to_run: List[str], object_permissions: Optional[LiteLLM_ObjectPermissionTable], ): """ Raises ProxyException if the object (key, team, org) cannot access the specific vector store. """ if object_permissions is None: return True if object_permissions.vector_stores is None: return True # If length is 0, then the object has access to all vector stores. if len(object_permissions.vector_stores) == 0: return True for vector_store_id in vector_store_ids_to_run: if vector_store_id not in object_permissions.vector_stores: raise ProxyException( message=f"User not allowed to access vector store. Tried to access {vector_store_id}. Only allowed to access {object_permissions.vector_stores}", type=ProxyErrorTypes.get_vector_store_access_error_type_for_object( object_type ), param="vector_store", code=status.HTTP_401_UNAUTHORIZED, ) return True