3562 lines
118 KiB
Python
3562 lines
118 KiB
Python
# 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 re
|
|
import time
|
|
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cast
|
|
|
|
from fastapi import HTTPException, Request, status
|
|
from pydantic import BaseModel
|
|
|
|
import litellm
|
|
from litellm._logging import verbose_proxy_logger
|
|
from litellm.caching.caching import DualCache
|
|
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.get_llm_provider_logic import get_llm_provider
|
|
from litellm.proxy._types import (
|
|
RBAC_ROLES,
|
|
CallInfo,
|
|
LiteLLM_AccessGroupTable,
|
|
LiteLLM_BudgetTable,
|
|
LiteLLM_EndUserTable,
|
|
Litellm_EntityType,
|
|
LiteLLM_JWTAuth,
|
|
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.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.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 _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
|
|
|
|
for model_name in model_list:
|
|
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"
|
|
)
|
|
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"
|
|
)
|
|
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})"
|
|
)
|
|
return False
|
|
|
|
# This model has zero cost explicitly configured
|
|
verbose_proxy_logger.debug(
|
|
f"Model {model_name} has zero cost explicitly configured (input: {input_cost}, output: {output_cost})"
|
|
)
|
|
|
|
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
|
|
|
|
|
|
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 global_proxy_spend > litellm.max_budget:
|
|
raise litellm.BudgetExceededError(
|
|
current_cost=global_proxy_spend, max_budget=litellm.max_budget
|
|
)
|
|
|
|
|
|
def _guardrail_modification_check(
|
|
request_body: dict, team_object: Optional[LiteLLM_TeamTable]
|
|
) -> None:
|
|
_request_metadata: dict = request_body.get("metadata", {}) or {}
|
|
if not _request_metadata.get("guardrails"):
|
|
return
|
|
|
|
from litellm.proxy.guardrails.guardrail_helpers import can_modify_guardrails
|
|
|
|
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,
|
|
)
|
|
|
|
|
|
async def common_checks( # noqa: PLR0915
|
|
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_body, route
|
|
)
|
|
|
|
# 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 your admin."
|
|
)
|
|
|
|
# 2. If team can call model
|
|
if _model and team_object:
|
|
if not 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,
|
|
):
|
|
raise ProxyException(
|
|
message=f"Team not allowed to access model. Team={team_object.team_id}, Model={_model}. Allowed team models = {team_object.models}",
|
|
type=ProxyErrorTypes.team_model_access_denied,
|
|
param="model",
|
|
code=status.HTTP_401_UNAUTHORIZED,
|
|
)
|
|
|
|
# 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:
|
|
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)
|
|
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,
|
|
)
|
|
|
|
# If this is a free model, skip all budget checks
|
|
if not skip_budget_checks:
|
|
# 3. If team is in budget
|
|
await _team_max_budget_check(
|
|
team_object=team_object,
|
|
proxy_logging_obj=proxy_logging_obj,
|
|
valid_token=valid_token,
|
|
)
|
|
|
|
# 3.0.5. If team is over soft budget (alert only, doesn't block)
|
|
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
|
|
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,
|
|
)
|
|
|
|
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
|
|
if user_budget < user_object.spend:
|
|
raise litellm.BudgetExceededError(
|
|
current_cost=user_object.spend,
|
|
max_budget=user_budget,
|
|
message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_object.spend}, Budget={user_budget}",
|
|
)
|
|
|
|
## 4.2 check team member budget, if team key
|
|
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
|
|
):
|
|
end_user_budget = end_user_object.litellm_budget_table.max_budget
|
|
if end_user_budget is not None and end_user_object.spend > end_user_budget:
|
|
raise litellm.BudgetExceededError(
|
|
current_cost=end_user_object.spend,
|
|
max_budget=end_user_budget,
|
|
message=f"ExceededBudget: End User={end_user_object.user_id} over budget. Spend={end_user_object.spend}, Budget={end_user_budget}",
|
|
)
|
|
|
|
_enforce_user_param_check(general_settings, request, request_body, route)
|
|
_reject_clientside_metadata_tags_check(general_settings, 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
|
|
)
|
|
|
|
token_team = getattr(valid_token, "team_id", None)
|
|
token_type: Literal["ui", "api"] = (
|
|
"ui" if token_team is not None and token_team == "litellm-dashboard" else "api"
|
|
)
|
|
_is_route_allowed = _is_allowed_route(
|
|
route=route,
|
|
token_type=token_type,
|
|
user_obj=user_object,
|
|
request=request,
|
|
request_data=request_body,
|
|
valid_token=valid_token,
|
|
)
|
|
|
|
# 11. [OPTIONAL] Vector store checks - is the object allowed to access the vector store
|
|
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)
|
|
await check_tools_allowlist(
|
|
request_body=request_body,
|
|
valid_token=valid_token,
|
|
team_object=team_object,
|
|
route=route,
|
|
)
|
|
|
|
return True
|
|
|
|
|
|
def _is_ui_route(
|
|
route: str,
|
|
user_obj: Optional[LiteLLM_UserTable] = None,
|
|
) -> bool:
|
|
"""
|
|
- Check if the route is a UI used route
|
|
"""
|
|
# this token is only used for managing the ui
|
|
allowed_routes = LiteLLMRoutes.ui_routes.value
|
|
# check if the current route startswith any of the allowed routes
|
|
if (
|
|
route is not None
|
|
and isinstance(route, str)
|
|
and any(route.startswith(allowed_route) for allowed_route in allowed_routes)
|
|
):
|
|
# Do something if the current route starts with any of the allowed routes
|
|
return True
|
|
elif any(
|
|
RouteChecks._route_matches_pattern(route=route, pattern=allowed_route)
|
|
for allowed_route in allowed_routes
|
|
):
|
|
return True
|
|
return False
|
|
|
|
|
|
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 _is_allowed_route(
|
|
route: str,
|
|
token_type: Literal["ui", "api"],
|
|
request: Request,
|
|
request_data: dict,
|
|
valid_token: Optional[UserAPIKeyAuth],
|
|
user_obj: Optional[LiteLLM_UserTable] = None,
|
|
) -> bool:
|
|
"""
|
|
- Route b/w ui token check and normal token check
|
|
"""
|
|
|
|
if token_type == "ui" and _is_ui_route(route=route, user_obj=user_obj):
|
|
return True
|
|
else:
|
|
return _is_api_route_allowed(
|
|
route=route,
|
|
request=request,
|
|
request_data=request_data,
|
|
valid_token=valid_token,
|
|
user_obj=user_obj,
|
|
)
|
|
|
|
|
|
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: DualCache,
|
|
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)
|
|
if cached_budget is not None:
|
|
return LiteLLM_BudgetTable(**cached_budget)
|
|
|
|
# Fetch from database
|
|
try:
|
|
budget_record = await prisma_client.db.litellm_budgettable.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
|
|
|
|
# Cache the budget for 60 seconds
|
|
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 as e:
|
|
verbose_proxy_logger.error(f"Error fetching default end user budget: {str(e)}")
|
|
return None
|
|
|
|
|
|
async def _apply_default_budget_to_end_user(
|
|
end_user_obj: LiteLLM_EndUserTable,
|
|
prisma_client: PrismaClient,
|
|
user_api_key_cache: DualCache,
|
|
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
|
|
|
|
|
|
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 not None and end_user_obj.spend > end_user_budget:
|
|
raise litellm.BudgetExceededError(
|
|
current_cost=end_user_obj.spend,
|
|
max_budget=end_user_budget,
|
|
message=f"ExceededBudget: End User={end_user_obj.user_id} over budget. Spend={end_user_obj.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: DualCache,
|
|
route: 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)
|
|
if cached_user_obj is not None:
|
|
return_obj = LiteLLM_EndUserTable(**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,
|
|
)
|
|
|
|
# Check budget limits
|
|
_check_end_user_budget(end_user_obj=return_obj, route=route)
|
|
|
|
return return_obj
|
|
|
|
# Fetch from database
|
|
try:
|
|
response = await prisma_client.db.litellm_endusertable.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 (always store as dict for consistency)
|
|
await user_api_key_cache.async_set_cache(
|
|
key="end_user_id:{}".format(end_user_id), value=_response.dict()
|
|
)
|
|
|
|
# Check budget limits
|
|
_check_end_user_budget(end_user_obj=_response, route=route)
|
|
|
|
return _response
|
|
|
|
except Exception as e:
|
|
if isinstance(e, litellm.BudgetExceededError):
|
|
raise e
|
|
return None
|
|
|
|
|
|
@log_db_metrics
|
|
async def get_tag_objects_batch(
|
|
tag_names: List[str],
|
|
prisma_client: Optional[PrismaClient],
|
|
user_api_key_cache: DualCache,
|
|
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)
|
|
if cached_tag is not None:
|
|
if isinstance(cached_tag, dict):
|
|
tag_objects[tag_name] = LiteLLM_TagTable(**cached_tag)
|
|
else:
|
|
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 prisma_client.db.litellm_tagtable.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}"
|
|
# Cache with default TTL (same as end_user objects)
|
|
await user_api_key_cache.async_set_cache(
|
|
key=cache_key, value=db_tag.dict()
|
|
)
|
|
tag_objects[tag_name] = LiteLLM_TagTable(**db_tag.dict())
|
|
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: DualCache,
|
|
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: DualCache,
|
|
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)
|
|
if cached_membership_obj is not None:
|
|
return LiteLLM_TeamMembership(**cached_membership_obj)
|
|
|
|
# else, check db
|
|
try:
|
|
response = await prisma_client.db.litellm_teammembership.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
|
|
|
|
# save the team membership object to cache (store as dict)
|
|
await user_api_key_cache.async_set_cache(key=_key, value=response.dict())
|
|
|
|
_response = LiteLLM_TeamMembership(**response.dict())
|
|
|
|
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 prisma_client.db.litellm_usertable.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 prisma_client.db.litellm_usertable.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
|
|
prisma_client.db.litellm_usertable.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: DualCache,
|
|
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)
|
|
if cached_user_obj is not None:
|
|
if isinstance(cached_user_obj, dict):
|
|
return LiteLLM_UserTable(**cached_user_obj)
|
|
elif isinstance(cached_user_obj, LiteLLM_UserTable):
|
|
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 prisma_client.db.litellm_usertable.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 prisma_client.db.litellm_usertable.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_dict,
|
|
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: BaseModel,
|
|
user_api_key_cache: DualCache,
|
|
proxy_logging_obj: Optional[ProxyLogging],
|
|
):
|
|
await user_api_key_cache.async_set_cache(
|
|
key=key,
|
|
value=value,
|
|
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
|
|
)
|
|
|
|
|
|
async def _cache_team_object(
|
|
team_id: str,
|
|
team_table: LiteLLM_TeamTableCachedObj,
|
|
user_api_key_cache: DualCache,
|
|
proxy_logging_obj: Optional[ProxyLogging],
|
|
):
|
|
key = "team_id:{}".format(team_id)
|
|
|
|
## CACHE REFRESH TIME!
|
|
team_table.last_refreshed_at = time.time()
|
|
|
|
await _cache_management_object(
|
|
key=key,
|
|
value=team_table,
|
|
user_api_key_cache=user_api_key_cache,
|
|
proxy_logging_obj=proxy_logging_obj,
|
|
)
|
|
|
|
|
|
async def _cache_key_object(
|
|
hashed_token: str,
|
|
user_api_key_obj: UserAPIKeyAuth,
|
|
user_api_key_cache: DualCache,
|
|
proxy_logging_obj: Optional[ProxyLogging],
|
|
):
|
|
key = hashed_token
|
|
|
|
## CACHE REFRESH TIME
|
|
user_api_key_obj.last_refreshed_at = time.time()
|
|
|
|
await _cache_management_object(
|
|
key=key,
|
|
value=user_api_key_obj,
|
|
user_api_key_cache=user_api_key_cache,
|
|
proxy_logging_obj=proxy_logging_obj,
|
|
)
|
|
|
|
|
|
async def _delete_cache_key_object(
|
|
hashed_token: str,
|
|
user_api_key_cache: DualCache,
|
|
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 prisma_client.db.litellm_teamtable.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 prisma_client.db.litellm_teamtable.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: DualCache,
|
|
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: DualCache,
|
|
parent_otel_span: Optional[Span],
|
|
) -> Optional[LiteLLM_TeamTableCachedObj]:
|
|
cached_team_obj: Optional[LiteLLM_TeamTableCachedObj] = None
|
|
|
|
## CHECK REDIS CACHE ##
|
|
if (
|
|
proxy_logging_obj is not None
|
|
and proxy_logging_obj.internal_usage_cache.dual_cache
|
|
):
|
|
cached_team_obj = (
|
|
await proxy_logging_obj.internal_usage_cache.dual_cache.async_get_cache(
|
|
key=key, parent_otel_span=parent_otel_span
|
|
)
|
|
)
|
|
|
|
if cached_team_obj is None:
|
|
cached_team_obj = await user_api_key_cache.async_get_cache(key=key)
|
|
|
|
if cached_team_obj is not None:
|
|
if isinstance(cached_team_obj, dict):
|
|
return LiteLLM_TeamTableCachedObj(**cached_team_obj)
|
|
elif isinstance(cached_team_obj, LiteLLM_TeamTableCachedObj):
|
|
return cached_team_obj
|
|
|
|
return None
|
|
|
|
|
|
async def get_team_object(
|
|
team_id: str,
|
|
prisma_client: Optional[PrismaClient],
|
|
user_api_key_cache: DualCache,
|
|
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: DualCache,
|
|
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,
|
|
ttl=DEFAULT_ACCESS_GROUP_CACHE_TTL,
|
|
)
|
|
|
|
|
|
async def _delete_cache_access_object(
|
|
access_group_id: str,
|
|
user_api_key_cache: DualCache,
|
|
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: DualCache,
|
|
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)
|
|
|
|
# Always check cache first
|
|
cached_access_obj = await user_api_key_cache.async_get_cache(key=key)
|
|
if cached_access_obj is not None:
|
|
if isinstance(cached_access_obj, dict):
|
|
return LiteLLM_AccessGroupTable(**cached_access_obj)
|
|
elif isinstance(cached_access_obj, LiteLLM_AccessGroupTable):
|
|
return cached_access_obj
|
|
|
|
# Not in cache - fetch from DB
|
|
try:
|
|
response = await prisma_client.db.litellm_accessgrouptable.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: DualCache,
|
|
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 prisma_client.db.litellm_teamtable.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,
|
|
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,
|
|
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: DualCache,
|
|
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)
|
|
if cached_org_obj is not None:
|
|
if isinstance(cached_org_obj, dict):
|
|
return LiteLLM_OrganizationTable(**cached_org_obj)
|
|
elif isinstance(cached_org_obj, LiteLLM_OrganizationTable):
|
|
return cached_org_obj
|
|
|
|
# Query database by organization_alias
|
|
try:
|
|
orgs = await prisma_client.db.litellm_organizationtable.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_dump(),
|
|
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_dump(),
|
|
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
|
|
) -> 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)
|
|
|
|
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,
|
|
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 prisma_client.db.litellm_jwtkeymapping.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: DualCache,
|
|
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
|
|
|
|
cached_key_obj: Optional[UserAPIKeyAuth] = await user_api_key_cache.async_get_cache(
|
|
key=key
|
|
)
|
|
|
|
if cached_key_obj is not None:
|
|
if isinstance(cached_key_obj, dict):
|
|
return UserAPIKeyAuth(**cached_key_obj)
|
|
elif isinstance(cached_key_obj, UserAPIKeyAuth):
|
|
return cached_key_obj
|
|
|
|
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
|
|
|
|
|
|
@log_db_metrics
|
|
async def get_object_permission(
|
|
object_permission_id: str,
|
|
prisma_client: Optional[PrismaClient],
|
|
user_api_key_cache: DualCache,
|
|
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)
|
|
cached_obj_permission = await user_api_key_cache.async_get_cache(key=key)
|
|
if cached_obj_permission is not None:
|
|
if isinstance(cached_obj_permission, dict):
|
|
return LiteLLM_ObjectPermissionTable(**cached_obj_permission)
|
|
elif isinstance(cached_obj_permission, LiteLLM_ObjectPermissionTable):
|
|
return cached_obj_permission
|
|
|
|
# else, check db
|
|
try:
|
|
response = await prisma_client.db.litellm_objectpermissiontable.find_unique(
|
|
where={"object_permission_id": object_permission_id}
|
|
)
|
|
|
|
if response is None:
|
|
return None
|
|
|
|
# save the object permission to cache
|
|
await user_api_key_cache.async_set_cache(
|
|
key=key,
|
|
value=response.model_dump(),
|
|
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
|
|
)
|
|
|
|
return LiteLLM_ObjectPermissionTable(**response.dict())
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
@log_db_metrics
|
|
async def get_org_object(
|
|
org_id: str,
|
|
prisma_client: Optional[PrismaClient],
|
|
user_api_key_cache: DualCache,
|
|
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
|
|
cached_org_obj = user_api_key_cache.async_get_cache(key=cache_key)
|
|
if cached_org_obj is not None:
|
|
if isinstance(cached_org_obj, dict):
|
|
return LiteLLM_OrganizationTable(**cached_org_obj)
|
|
elif isinstance(cached_org_obj, LiteLLM_OrganizationTable):
|
|
return cached_org_obj
|
|
# 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 prisma_client.db.litellm_organizationtable.find_unique(
|
|
**query_kwargs
|
|
)
|
|
|
|
if response is None:
|
|
raise Exception
|
|
|
|
# Cache the result
|
|
await user_api_key_cache.async_set_cache(
|
|
key=cache_key,
|
|
value=(
|
|
response.model_dump() if hasattr(response, "model_dump") else response
|
|
),
|
|
ttl=DEFAULT_IN_MEMORY_TTL,
|
|
)
|
|
|
|
return response
|
|
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[DualCache] = 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[DualCache] = 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[DualCache] = 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[DualCache] = 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_401_UNAUTHORIZED,
|
|
)
|
|
|
|
|
|
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
|
|
|
|
|
|
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
|
|
"""
|
|
try:
|
|
return _can_object_call_model(
|
|
model=model,
|
|
llm_router=llm_router,
|
|
models=valid_token.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
|
|
|
|
|
|
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
|
|
|
|
|
|
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_401_UNAUTHORIZED,
|
|
)
|
|
|
|
return _can_object_call_model(
|
|
model=model,
|
|
llm_router=llm_router,
|
|
models=user_object.models,
|
|
object_type="user",
|
|
)
|
|
|
|
|
|
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.spend is not None and valid_token.max_budget is not None:
|
|
####################################
|
|
# 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=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,
|
|
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 #
|
|
####################################
|
|
|
|
if valid_token.spend >= valid_token.max_budget:
|
|
raise litellm.BudgetExceededError(
|
|
current_cost=valid_token.spend,
|
|
max_budget=valid_token.max_budget,
|
|
)
|
|
|
|
|
|
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,
|
|
)
|
|
)
|
|
|
|
|
|
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
|
|
):
|
|
alert_threshold = (
|
|
valid_token.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE
|
|
)
|
|
|
|
# Only alert if we've crossed the threshold but haven't exceeded max_budget yet
|
|
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=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="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: DualCache,
|
|
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 user_object 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,
|
|
)
|
|
|
|
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
|
|
team_member_spend = team_membership.spend or 0.0
|
|
|
|
if 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 _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
|
|
and team_object.spend is not None
|
|
and team_object.spend > team_object.max_budget
|
|
):
|
|
if valid_token:
|
|
call_info = CallInfo(
|
|
token=valid_token.token,
|
|
spend=team_object.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=team_object.spend,
|
|
max_budget=team_object.max_budget,
|
|
message=f"Budget has been exceeded! Team={team_object.team_id} Current cost: {team_object.spend}, Max budget: {team_object.max_budget}",
|
|
)
|
|
|
|
|
|
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 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: DualCache,
|
|
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)
|
|
cached_obj = await user_api_key_cache.async_get_cache(key=cache_key)
|
|
if cached_obj is not None:
|
|
if isinstance(cached_obj, dict):
|
|
return LiteLLM_ProjectTableCachedObj(**cached_obj)
|
|
elif isinstance(cached_obj, LiteLLM_ProjectTableCachedObj):
|
|
return cached_obj
|
|
|
|
# Fetch from DB
|
|
project_row = await prisma_client.db.litellm_projecttable.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,
|
|
)
|
|
|
|
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: DualCache,
|
|
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
|
|
|
|
# Check if organization spend exceeds max budget
|
|
if org_table.spend >= org_max_budget:
|
|
# Trigger budget alert
|
|
call_info = CallInfo(
|
|
token=valid_token.token,
|
|
spend=org_table.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_table.spend,
|
|
max_budget=org_max_budget,
|
|
message=f"Budget has been exceeded! Organization={org_id} Current cost: {org_table.spend}, Max budget: {org_max_budget}",
|
|
)
|
|
|
|
|
|
async def _tag_max_budget_check(
|
|
request_body: dict,
|
|
prisma_client: Optional[PrismaClient],
|
|
user_api_key_cache: DualCache,
|
|
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
|
|
and tag_object.spend is not None
|
|
and tag_object.spend > tag_object.litellm_budget_table.max_budget
|
|
):
|
|
raise litellm.BudgetExceededError(
|
|
current_cost=tag_object.spend,
|
|
max_budget=tag_object.litellm_budget_table.max_budget,
|
|
message=f"Budget has been exceeded! Tag={tag_name} Current cost: {tag_object.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 prisma_client.db.litellm_objectpermissiontable.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 prisma_client.db.litellm_objectpermissiontable.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
|