From 5f87e3bd28bf7a74df8d3c0fa14ca9c05f61ec46 Mon Sep 17 00:00:00 2001 From: Harshit Jain Date: Fri, 13 Feb 2026 18:07:27 +0530 Subject: [PATCH] fix: add project managemenet api with proper working --- litellm/proxy/_types.py | 115 +++-- litellm/proxy/auth/auth_checks.py | 284 +++++++++++- litellm/proxy/auth/user_api_key_auth.py | 43 +- .../key_management_endpoints.py | 79 ++++ .../management_endpoints/project_endpoints.py | 253 +++++++++-- .../test_project_endpoints_prisma.py | 415 ++++++++++++++++-- 6 files changed, 1025 insertions(+), 164 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 2cd67d8108..739a9d2736 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -198,6 +198,7 @@ class Litellm_EntityType(enum.Enum): TEAM = "team" TEAM_MEMBER = "team_member" ORGANIZATION = "organization" + PROJECT = "project" TAG = "tag" # global proxy level entity @@ -851,9 +852,9 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase): allowed_cache_controls: Optional[list] = [] config: Optional[dict] = {} permissions: Optional[dict] = {} - model_max_budget: Optional[dict] = ( - {} - ) # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {} + model_max_budget: Optional[ + dict + ] = {} # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {} model_config = ConfigDict(protected_namespaces=()) model_rpm_limit: Optional[dict] = None @@ -1405,12 +1406,12 @@ class NewCustomerRequest(BudgetNewRequest): blocked: bool = False # allow/disallow requests for this end-user budget_id: Optional[str] = None # give either a budget_id or max_budget spend: Optional[float] = None - allowed_model_region: Optional[AllowedModelRegion] = ( - None # require all user requests to use models in this specific region - ) - default_model: Optional[str] = ( - None # if no equivalent model in allowed region - default all requests to this model - ) + allowed_model_region: Optional[ + AllowedModelRegion + ] = None # require all user requests to use models in this specific region + default_model: Optional[ + str + ] = None # if no equivalent model in allowed region - default all requests to this model @model_validator(mode="before") @classmethod @@ -1432,12 +1433,12 @@ class UpdateCustomerRequest(LiteLLMPydanticObjectBase): blocked: bool = False # allow/disallow requests for this end-user max_budget: Optional[float] = None budget_id: Optional[str] = None # give either a budget_id or max_budget - allowed_model_region: Optional[AllowedModelRegion] = ( - None # require all user requests to use models in this specific region - ) - default_model: Optional[str] = ( - None # if no equivalent model in allowed region - default all requests to this model - ) + allowed_model_region: Optional[ + AllowedModelRegion + ] = None # require all user requests to use models in this specific region + default_model: Optional[ + str + ] = None # if no equivalent model in allowed region - default all requests to this model class DeleteCustomerRequest(LiteLLMPydanticObjectBase): @@ -1526,15 +1527,15 @@ class NewTeamRequest(TeamBase): ] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating tpm model_tpm_limit: Optional[Dict[str, int]] = None - team_member_budget: Optional[float] = ( - None # allow user to set a budget for all team members - ) - team_member_rpm_limit: Optional[int] = ( - None # allow user to set RPM limit for all team members - ) - team_member_tpm_limit: Optional[int] = ( - None # allow user to set TPM limit for all team members - ) + team_member_budget: Optional[ + float + ] = None # allow user to set a budget for all team members + team_member_rpm_limit: Optional[ + int + ] = None # allow user to set RPM limit for all team members + team_member_tpm_limit: Optional[ + int + ] = None # allow user to set TPM limit for all team members team_member_key_duration: Optional[str] = None # e.g. "1d", "1w", "1m" allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None @@ -1626,9 +1627,9 @@ class BlockKeyRequest(LiteLLMPydanticObjectBase): class AddTeamCallback(LiteLLMPydanticObjectBase): callback_name: str - callback_type: Optional[Literal["success", "failure", "success_and_failure"]] = ( - "success_and_failure" - ) + callback_type: Optional[ + Literal["success", "failure", "success_and_failure"] + ] = "success_and_failure" callback_vars: Dict[str, str] @model_validator(mode="before") @@ -1960,9 +1961,9 @@ class ConfigList(LiteLLMPydanticObjectBase): stored_in_db: Optional[bool] field_default_value: Any premium_field: bool = False - nested_fields: Optional[List[FieldDetail]] = ( - None # For nested dictionary or Pydantic fields - ) + nested_fields: Optional[ + List[FieldDetail] + ] = None # For nested dictionary or Pydantic fields class UserHeaderMapping(LiteLLMPydanticObjectBase): @@ -2403,9 +2404,9 @@ class LiteLLM_OrganizationMembershipTable(LiteLLMPydanticObjectBase): budget_id: Optional[str] = None created_at: datetime updated_at: datetime - user: Optional[Any] = ( - None # You might want to replace 'Any' with a more specific type if available - ) + user: Optional[ + Any + ] = None # You might want to replace 'Any' with a more specific type if available litellm_budget_table: Optional[LiteLLM_BudgetTable] = None model_config = ConfigDict(protected_namespaces=()) @@ -2615,25 +2616,10 @@ class NewProjectResponse(LiteLLM_ProjectTable): updated_at: datetime -class LiteLLM_ProjectTableCachedObj(LiteLLMPydanticObjectBase): - """Cached version for auth checks""" +class LiteLLM_ProjectTableCachedObj(LiteLLM_ProjectTable): + """Cached version for auth checks. Mirrors LiteLLM_TeamTableCachedObj pattern.""" - project_id: str - project_alias: Optional[str] = None - description: Optional[str] = None - team_id: Optional[str] = None - budget_id: Optional[str] = None - metadata: Optional[dict] = None - models: List[str] = [] - spend: float = 0.0 - model_spend: Optional[dict] = None - model_rpm_limit: Optional[dict] = None - model_tpm_limit: Optional[dict] = None - blocked: bool = False - object_permission_id: Optional[str] = None - litellm_budget_table: Optional[LiteLLM_BudgetTable] = None - created_at: Optional[datetime] = None - updated_at: Optional[datetime] = None + last_refreshed_at: Optional[float] = None class LiteLLM_UserTableFiltered(BaseModel): # done to avoid exposing sensitive data @@ -3228,6 +3214,11 @@ class ProxyErrorTypes(str, enum.Enum): Organization does not have access to the model """ + project_model_access_denied = "project_model_access_denied" + """ + Project does not have access to the model + """ + expired_key = "expired_key" """ Key has expired @@ -3290,7 +3281,7 @@ class ProxyErrorTypes(str, enum.Enum): @classmethod def get_model_access_error_type_for_object( - cls, object_type: Literal["key", "user", "team", "org"] + cls, object_type: Literal["key", "user", "team", "org", "project"] ) -> "ProxyErrorTypes": """ Get the model access error type for object_type @@ -3303,6 +3294,8 @@ class ProxyErrorTypes(str, enum.Enum): return cls.user_model_access_denied elif object_type == "org": return cls.org_model_access_denied + elif object_type == "project": + return cls.project_model_access_denied @classmethod def get_vector_store_access_error_type_for_object( @@ -3507,9 +3500,9 @@ class TeamModelDeleteRequest(BaseModel): # Organization Member Requests class OrganizationMemberAddRequest(OrgMemberAddRequest): organization_id: str - max_budget_in_organization: Optional[float] = ( - None # Users max budget within the organization - ) + max_budget_in_organization: Optional[ + float + ] = None # Users max budget within the organization class OrganizationMemberDeleteRequest(MemberDeleteRequest): @@ -3727,9 +3720,9 @@ class ProviderBudgetResponse(LiteLLMPydanticObjectBase): Maps provider names to their budget configs. """ - providers: Dict[str, ProviderBudgetResponseObject] = ( - {} - ) # Dictionary mapping provider names to their budget configurations + providers: Dict[ + str, ProviderBudgetResponseObject + ] = {} # Dictionary mapping provider names to their budget configurations class ProxyStateVariables(TypedDict): @@ -3872,9 +3865,9 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): enforce_rbac: bool = False roles_jwt_field: Optional[str] = None # v2 on role mappings role_mappings: Optional[List[RoleMapping]] = None - object_id_jwt_field: Optional[str] = ( - None # can be either user / team, inferred from the role mapping - ) + object_id_jwt_field: Optional[ + str + ] = None # can be either user / team, inferred from the role mapping scope_mappings: Optional[List[ScopeMapping]] = None enforce_scope_based_access: bool = False enforce_team_based_model_access: bool = False diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 76ec67ab10..dc5d5d749b 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -39,6 +39,7 @@ from litellm.proxy._types import ( LiteLLM_ObjectPermissionTable, LiteLLM_OrganizationMembershipTable, LiteLLM_OrganizationTable, + LiteLLM_ProjectTableCachedObj, LiteLLM_TagTable, LiteLLM_TeamMembership, LiteLLM_TeamTable, @@ -75,6 +76,7 @@ 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. @@ -92,38 +94,41 @@ def _log_budget_lookup_failure(entity: str, error: Exception) -> None: 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." + 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 @@ -131,42 +136,87 @@ def _is_model_cost_zero( 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, + ) + + async def common_checks( request_body: dict, team_object: Optional[LiteLLM_TeamTable], @@ -180,13 +230,18 @@ async def common_checks( 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 @@ -231,6 +286,16 @@ async def common_checks( 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 @@ -290,7 +355,10 @@ async def common_checks( ) # 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: + 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( @@ -1368,7 +1436,7 @@ async def _get_team_object_from_user_api_key_cache( 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: @@ -1383,7 +1451,7 @@ async def _get_team_object_from_user_api_key_cache( 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, @@ -2069,7 +2137,7 @@ def _can_object_call_model( models: List[str], team_model_aliases: Optional[Dict[str, str]] = None, team_id: Optional[str] = None, - object_type: Literal["user", "team", "key", "org"] = "user", + object_type: Literal["user", "team", "key", "org", "project"] = "user", fallback_depth: int = 0, ) -> Literal[True]: """ @@ -2220,6 +2288,24 @@ def can_team_access_model( ) +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], @@ -2520,14 +2606,26 @@ async def _team_soft_budget_check( 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 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()] + 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()] + 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] @@ -2566,6 +2664,150 @@ async def _team_soft_budget_check( ) +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], diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 42f10ff859..fc61f2b05c 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -36,6 +36,7 @@ from litellm.proxy.auth.auth_checks import ( common_checks, get_end_user_object, get_key_object, + get_project_object, get_team_object, get_user_object, is_valid_fallback_model, @@ -120,12 +121,12 @@ def _get_bearer_token_or_received_api_key(api_key: str) -> str: # Handle AWS Signature V4 format from LangChain # Format: AWS4-HMAC-SHA256 Credential=Bearer sk-12345/date/region/service/aws4_request, SignedHeaders=..., Signature=... # Extract the Bearer token from the Credential field - match = re.search(r'Credential=Bearer\s+([^/\s,]+)', api_key) + match = re.search(r"Credential=Bearer\s+([^/\s,]+)", api_key) if match: api_key = match.group(1) else: # If no Bearer token found in Credential, try to extract just the credential value - match = re.search(r'Credential=([^/\s,]+)', api_key) + match = re.search(r"Credential=([^/\s,]+)", api_key) if match: api_key = match.group(1) @@ -145,12 +146,12 @@ def _get_bearer_token( # Handle AWS Signature V4 format from LangChain # Format: AWS4-HMAC-SHA256 Credential=Bearer sk-12345/date/region/service/aws4_request, SignedHeaders=..., Signature=... # Extract the Bearer token from the Credential field - match = re.search(r'Credential=Bearer\s+([^/\s,]+)', api_key) + match = re.search(r"Credential=Bearer\s+([^/\s,]+)", api_key) if match: api_key = match.group(1) else: # If no Bearer token found in Credential, try to extract just the credential value - match = re.search(r'Credential=([^/\s,]+)', api_key) + match = re.search(r"Credential=([^/\s,]+)", api_key) if match: api_key = match.group(1) else: @@ -274,7 +275,9 @@ async def get_global_proxy_spend( proxy_logging_obj: ProxyLogging, ) -> Optional[float]: global_proxy_spend = None - if litellm.max_budget > 0 and prisma_client is not None: # user set proxy max budget + if ( + litellm.max_budget > 0 and prisma_client is not None + ): # user set proxy max budget # Use event-driven coordination to prevent cache stampede cache_key = "{}:spend".format(litellm_proxy_admin_name) global_proxy_spend = await _fetch_global_spend_with_event_coordination( @@ -631,13 +634,13 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 if team_object is not None else None, ) - + # Check if model has zero cost - if so, skip all budget checks model = get_model_from_request(request_data, route) skip_budget_checks = False if model is not None and llm_router is not None: from litellm.proxy.auth.auth_checks import _is_model_cost_zero - + skip_budget_checks = _is_model_cost_zero( model=model, llm_router=llm_router ) @@ -645,7 +648,17 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 verbose_proxy_logger.info( f"Skipping all budget checks for zero-cost model: {model}" ) - + + # Fetch project object for JWT path if project_id is set + _jwt_project_obj = None + if valid_token.project_id is not None: + _jwt_project_obj = await get_project_object( + project_id=valid_token.project_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + # run through common checks _ = await common_checks( request=request, @@ -660,6 +673,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 proxy_logging_obj=proxy_logging_obj, valid_token=valid_token, skip_budget_checks=skip_budget_checks, + project_object=_jwt_project_obj, ) # return UserAPIKeyAuth object @@ -1056,7 +1070,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 skip_budget_checks = False if model is not None and llm_router is not None: from litellm.proxy.auth.auth_checks import _is_model_cost_zero - + skip_budget_checks = _is_model_cost_zero( model=model, llm_router=llm_router ) @@ -1192,6 +1206,16 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 key=valid_token.team_id, value=_team_obj ) # save team table in cache - used for tpm/rpm limiting - tpm_rpm_limiter.py + # Fetch project object if key belongs to a project + _project_obj = None + if valid_token.project_id is not None: + _project_obj = await get_project_object( + project_id=valid_token.project_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + global_proxy_spend = None if ( litellm.max_budget > 0 and prisma_client is not None @@ -1231,6 +1255,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 proxy_logging_obj=proxy_logging_obj, valid_token=valid_token, skip_budget_checks=skip_budget_checks, + project_object=_project_obj, ) # Token passed all checks if valid_token is None: diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index cbb271f860..c6c03a119c 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -43,6 +43,7 @@ from litellm.proxy.auth.auth_checks import ( can_team_access_model, get_key_object, get_org_object, + get_project_object, get_team_object, ) from litellm.proxy.auth.auth_utils import abbreviate_api_key @@ -881,6 +882,61 @@ async def _check_team_key_limits( ) +async def _check_project_key_limits( + project_id: str, + data: Union[GenerateKeyRequest, UpdateKeyRequest], + prisma_client: PrismaClient, + user_api_key_cache: DualCache, +) -> None: + """ + Validate that key's models and budget respect its project's limits. + + - Key models must be a subset of project models + - Key max_budget must be <= project max_budget + """ + project_obj = await get_project_object( + project_id=project_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + + if project_obj is None: + raise HTTPException( + status_code=404, + detail={"error": f"Project not found, project_id={project_id}"}, + ) + + # Validate key models are a subset of project models + if data.models and len(project_obj.models) > 0: + for m in data.models: + if m not in project_obj.models: + raise HTTPException( + status_code=400, + detail={ + "error": f"Model '{m}' not in project's allowed models. Project allowed models={project_obj.models}. Project: {project_id}" + }, + ) + + # Validate key max_budget <= project max_budget + project_max_budget = None + if project_obj.litellm_budget_table is not None: + project_max_budget = getattr( + project_obj.litellm_budget_table, "max_budget", None + ) + + if ( + data.max_budget is not None + and project_max_budget is not None + and data.max_budget > project_max_budget + ): + raise HTTPException( + status_code=400, + detail={ + "error": f"Key max_budget ({data.max_budget}) exceeds project's max_budget ({project_max_budget}). Project: {project_id}" + }, + ) + + def check_org_key_model_specific_limits( keys: List[LiteLLM_VerificationToken], org_table: LiteLLM_OrganizationTable, @@ -1135,6 +1191,15 @@ async def generate_key_fn( prisma_client=prisma_client, ) + # Validate key against project limits if project_id is set + if data.project_id is not None: + await _check_project_key_limits( + project_id=data.project_id, + data=data, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + return await _common_key_generation_helper( data=data, user_api_key_dict=user_api_key_dict, @@ -1808,6 +1873,20 @@ async def update_key_fn( prisma_client=prisma_client, ) + # Validate key against project limits if project_id is being set + _project_id_to_check = getattr(data, "project_id", None) or getattr( + existing_key_row, "project_id", None + ) + if _project_id_to_check is not None and ( + data.models is not None or data.max_budget is not None + ): + await _check_project_key_limits( + project_id=_project_id_to_check, + data=data, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + # if team change - check if this is possible if is_different_team(data=data, existing_key_row=existing_key_row): if llm_router is None: diff --git a/litellm/proxy/management_endpoints/project_endpoints.py b/litellm/proxy/management_endpoints/project_endpoints.py index b6c51d3422..ba3238ebfd 100644 --- a/litellm/proxy/management_endpoints/project_endpoints.py +++ b/litellm/proxy/management_endpoints/project_endpoints.py @@ -11,23 +11,19 @@ Endpoints for /project operations #### PROJECT MANAGEMENT #### import json -from typing import List, Optional +from typing import List, Optional, Union -from fastapi import APIRouter, Depends, HTTPException, Request, status +from fastapi import APIRouter, Depends, HTTPException, Request from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.management_endpoints.common_utils import _set_object_metadata_field -from litellm.proxy.management_helpers.object_permission_utils import ( - handle_update_object_permission_common, -) from litellm.proxy.management_helpers.utils import ( management_endpoint_wrapper, ) from litellm.proxy.utils import PrismaClient, handle_exception_on_proxy -from litellm.utils import _update_dictionary router = APIRouter() @@ -37,31 +33,37 @@ async def _check_user_permission_for_project( team_id: Optional[str], prisma_client: PrismaClient, require_admin: bool = False, + team_object: Optional[LiteLLM_TeamTable] = None, ) -> bool: """ Check if user has permission to manage a project. - + Returns True if user is proxy admin or team admin (when team_id provided). If require_admin=True, only proxy admins are allowed. + + If team_object is provided, it will be used instead of fetching from DB + (avoids duplicate DB queries when team was already fetched for validation). """ is_proxy_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN - + if require_admin: return is_proxy_admin - + if is_proxy_admin: return True - + if not team_id or not user_api_key_dict.user_id: return False - - team = await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id} - ) - + + team = team_object + if team is None: + team = await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id} + ) + if team and team.admins: return user_api_key_dict.user_id in team.admins - + return False @@ -69,11 +71,11 @@ async def _validate_team_exists( team_id: str, prisma_client: PrismaClient, ): - """Validate that a team exists.""" + """Validate that a team exists. Returns the team row.""" team = await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id} + where={"team_id": team_id}, ) - + if team is None: raise ProxyException( message=f"Team not found, team_id={team_id}", @@ -81,10 +83,110 @@ async def _validate_team_exists( code=404, param="team_id", ) - + return team +def _check_team_project_limits( + team_object: LiteLLM_TeamTable, + data: Union[NewProjectRequest, UpdateProjectRequest], +) -> None: + """ + Check that project limits respect its parent Team's limits. + + Mirrors _check_org_team_limits() from team_endpoints.py. + + Validates: + - Project models are a subset of Team models + - Project max_budget <= Team max_budget + - Project tpm_limit <= Team tpm_limit + - Project rpm_limit <= Team rpm_limit + - Budget values are non-negative + - soft_budget < max_budget + """ + # --- Budget non-negativity checks --- + if data.max_budget is not None and data.max_budget < 0: + raise HTTPException( + status_code=400, + detail={ + "error": f"max_budget cannot be negative. Received: {data.max_budget}" + }, + ) + if data.soft_budget is not None and data.soft_budget < 0: + raise HTTPException( + status_code=400, + detail={ + "error": f"soft_budget cannot be negative. Received: {data.soft_budget}" + }, + ) + + # --- soft_budget < max_budget --- + if data.soft_budget is not None and data.max_budget is not None: + if data.soft_budget >= data.max_budget: + raise HTTPException( + status_code=400, + detail={ + "error": f"soft_budget ({data.soft_budget}) must be strictly lower than max_budget ({data.max_budget})" + }, + ) + + # --- Validate project models are a subset of team models --- + project_models = getattr(data, "models", None) + team_models = team_object.models or [] + if project_models and len(team_models) > 0: + # If team has 'all-proxy-models', skip validation as it allows all models + if SpecialModelNames.all_proxy_models.value not in team_models: + for m in project_models: + if m not in team_models: + raise HTTPException( + status_code=400, + detail={ + "error": f"Model '{m}' not in team's allowed models. Team allowed models={team_models}. Team: {team_object.team_id}" + }, + ) + + # --- Validate project max_budget <= team max_budget --- + # Team stores budget fields directly (max_budget, tpm_limit, rpm_limit) + # unlike Project which uses a separate LiteLLM_BudgetTable relation + if ( + data.max_budget is not None + and team_object.max_budget is not None + and data.max_budget > team_object.max_budget + ): + raise HTTPException( + status_code=400, + detail={ + "error": f"Project max_budget ({data.max_budget}) exceeds team's max_budget ({team_object.max_budget}). Team: {team_object.team_id}" + }, + ) + + # --- Validate project tpm_limit <= team tpm_limit --- + if ( + data.tpm_limit is not None + and team_object.tpm_limit is not None + and data.tpm_limit > team_object.tpm_limit + ): + raise HTTPException( + status_code=400, + detail={ + "error": f"Project tpm_limit ({data.tpm_limit}) exceeds team's tpm_limit ({team_object.tpm_limit}). Team: {team_object.team_id}" + }, + ) + + # --- Validate project rpm_limit <= team rpm_limit --- + if ( + data.rpm_limit is not None + and team_object.rpm_limit is not None + and data.rpm_limit > team_object.rpm_limit + ): + raise HTTPException( + status_code=400, + detail={ + "error": f"Project rpm_limit ({data.rpm_limit}) exceeds team's rpm_limit ({team_object.rpm_limit}). Team: {team_object.team_id}" + }, + ) + + async def _create_budget_for_project( data: NewProjectRequest, user_id: Optional[str], @@ -96,9 +198,9 @@ async def _create_budget_for_project( _json_data = data.json(exclude_none=True) _budget_data = {k: v for k, v in _json_data.items() if k in budget_params} budget_row = LiteLLM_BudgetTable(**_budget_data) - + new_budget = prisma_client.jsonify_object(budget_row.json(exclude_none=True)) - + _budget = await prisma_client.db.litellm_budgettable.create( data={ **new_budget, @@ -106,7 +208,7 @@ async def _create_budget_for_project( "updated_by": user_id or litellm_proxy_admin_name, } ) - + return _budget.budget_id @@ -137,7 +239,7 @@ def _remove_budget_fields_from_project_data(project_data: dict) -> dict: Remove budget fields from project data. Budget fields belong to LiteLLM_BudgetTable, not LiteLLM_ProjectTable. Keep budget_id as it's a foreign key. - + Following the pattern from organization_endpoints.py """ budget_fields = LiteLLM_BudgetTable.model_fields.keys() @@ -232,18 +334,36 @@ async def new_project( """ from litellm.proxy.proxy_server import ( litellm_proxy_admin_name, + premium_user, prisma_client, ) try: + if not premium_user: + raise HTTPException( + status_code=403, + detail={ + "error": "Project management is an enterprise feature. " + + CommonProxyErrors.not_premium_user.value + }, + ) + if prisma_client is None: raise HTTPException( status_code=500, detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - # Validate team exists - await _validate_team_exists(team_id=data.team_id, prisma_client=prisma_client) + # Validate team exists and get team object with budget + team_object = await _validate_team_exists( + team_id=data.team_id, prisma_client=prisma_client + ) + + # Validate project limits against team limits + _check_team_project_limits( + team_object=LiteLLM_TeamTable(**team_object.model_dump()), + data=data, + ) # Check if user has permission to create projects for this team # only team admins can create projects for their team @@ -251,8 +371,9 @@ async def new_project( user_api_key_dict=user_api_key_dict, team_id=data.team_id, prisma_client=prisma_client, + team_object=LiteLLM_TeamTable(**team_object.model_dump()), ) - + if not has_permission: raise HTTPException( status_code=403, @@ -311,10 +432,10 @@ async def new_project( new_project_row = prisma_client.jsonify_object( project_row.json(exclude_none=True) ) - + # Remove budget fields (following organization_endpoints.py pattern) new_project_row = _remove_budget_fields_from_project_data(new_project_row) - + verbose_proxy_logger.info( f"new_project_row: {json.dumps(new_project_row, indent=2)}" ) @@ -388,10 +509,20 @@ async def update_project( """ from litellm.proxy.proxy_server import ( litellm_proxy_admin_name, + premium_user, prisma_client, ) try: + if not premium_user: + raise HTTPException( + status_code=403, + detail={ + "error": "Project management is an enterprise feature. " + + CommonProxyErrors.not_premium_user.value + }, + ) + if prisma_client is None: raise HTTPException( status_code=500, @@ -417,23 +548,43 @@ async def update_project( param="project_id", ) + # Validate team exists and get team object for limit + permission checks + team_id_to_check = data.team_id or existing_project.team_id + team_obj_for_checks = None + if team_id_to_check is not None: + team_obj_for_checks = await _validate_team_exists( + team_id=team_id_to_check, prisma_client=prisma_client + ) + # Check if user has permission to update this project has_permission = await _check_user_permission_for_project( user_api_key_dict=user_api_key_dict, team_id=existing_project.team_id, prisma_client=prisma_client, + team_object=LiteLLM_TeamTable(**team_obj_for_checks.model_dump()) + if team_obj_for_checks + else None, ) - + if not has_permission: raise HTTPException( status_code=403, detail={"error": "Only admins or team admins can update projects"}, ) + # Validate project limits against team limits + if team_obj_for_checks is not None: + _check_team_project_limits( + team_object=LiteLLM_TeamTable(**team_obj_for_checks.model_dump()), + data=data, + ) + # Prepare update data update_data = data.json(exclude_none=True, exclude={"project_id"}) update_data = prisma_client.jsonify_object(update_data) - update_data["updated_by"] = user_api_key_dict.user_id or litellm_proxy_admin_name + update_data["updated_by"] = ( + user_api_key_dict.user_id or litellm_proxy_admin_name + ) # Handle budget updates budget_fields = LiteLLM_BudgetTable.model_fields.keys() @@ -459,15 +610,21 @@ async def update_project( if existing_project.object_permission_id: # Update existing permission await prisma_client.db.litellm_objectpermissiontable.update( - where={"object_permission_id": existing_project.object_permission_id}, + where={ + "object_permission_id": existing_project.object_permission_id + }, data=object_permission_data, ) else: # Create new permission - created_permission = await prisma_client.db.litellm_objectpermissiontable.create( - data=object_permission_data, + created_permission = ( + await prisma_client.db.litellm_objectpermissiontable.create( + data=object_permission_data, + ) ) - update_data["object_permission_id"] = created_permission.object_permission_id + update_data[ + "object_permission_id" + ] = created_permission.object_permission_id # Handle metadata fields for field in LiteLLM_ManagementEndpoint_MetadataFields: @@ -524,9 +681,18 @@ async def delete_project( }' ``` """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import premium_user, prisma_client try: + if not premium_user: + raise HTTPException( + status_code=403, + detail={ + "error": "Project management is an enterprise feature. " + + CommonProxyErrors.not_premium_user.value + }, + ) + if prisma_client is None: raise HTTPException( status_code=500, @@ -540,7 +706,7 @@ async def delete_project( prisma_client=prisma_client, require_admin=True, ) - + if not has_permission: raise HTTPException( status_code=403, @@ -564,8 +730,10 @@ async def delete_project( ) # Check if there are any keys associated with this project - associated_keys = await prisma_client.db.litellm_verificationtoken.find_many( - where={"project_id": project_id} + associated_keys = ( + await prisma_client.db.litellm_verificationtoken.find_many( + where={"project_id": project_id} + ) ) if len(associated_keys) > 0: @@ -641,15 +809,15 @@ async def project_info( # Check if user has access to this project (admin or team member) is_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN is_team_member = False - + if project.team_id and user_api_key_dict.user_id: team = await prisma_client.db.litellm_teamtable.find_unique( where={"team_id": project.team_id} ) if team: is_team_member = ( - user_api_key_dict.user_id in team.admins or - user_api_key_dict.user_id in team.members + user_api_key_dict.user_id in team.admins + or user_api_key_dict.user_id in team.members ) if not (is_admin or is_team_member): @@ -726,4 +894,3 @@ async def list_projects( ) ) raise handle_exception_on_proxy(e) - diff --git a/tests/proxy_unit_tests/test_project_endpoints_prisma.py b/tests/proxy_unit_tests/test_project_endpoints_prisma.py index daa3af2d67..c2366139b0 100644 --- a/tests/proxy_unit_tests/test_project_endpoints_prisma.py +++ b/tests/proxy_unit_tests/test_project_endpoints_prisma.py @@ -2,23 +2,15 @@ import os import sys import traceback from litellm._uuid import uuid -from datetime import datetime, timezone from unittest import mock from dotenv import load_dotenv from fastapi import Request -from fastapi.routing import APIRoute -import httpx load_dotenv() -import io -import os import time -sys.path.insert( - 0, os.path.abspath("../..") -) -import asyncio +sys.path.insert(0, os.path.abspath("../..")) import logging import pytest @@ -27,8 +19,6 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.proxy.management_endpoints.team_endpoints import ( new_team, - team_info, - update_team, ) from litellm.proxy.management_endpoints.project_endpoints import ( new_project, @@ -38,13 +28,11 @@ from litellm.proxy.management_endpoints.project_endpoints import ( ) from litellm.proxy.proxy_server import ( LitellmUserRoles, - user_api_key_auth, ) -from litellm.proxy.utils import PrismaClient, ProxyLogging, hash_token +from litellm.proxy.utils import PrismaClient, ProxyLogging verbose_proxy_logger.setLevel(level=logging.DEBUG) -from starlette.datastructures import URL from litellm.caching.caching import DualCache from litellm.proxy._types import ( @@ -52,7 +40,6 @@ from litellm.proxy._types import ( UpdateProjectRequest, DeleteProjectRequest, NewTeamRequest, - ProxyException, UserAPIKeyAuth, ) @@ -80,6 +67,9 @@ def prisma_client(): ) litellm.proxy.proxy_server.user_custom_key_generate = None + # Enable premium_user for project management tests + setattr(litellm.proxy.proxy_server, "premium_user", True) + return prisma_client @@ -115,10 +105,7 @@ async def test_new_project(prisma_client): project_alias="test-project", description="Test project for unit testing", team_id=_team_id, - metadata={ - "use_case_id": "TEST-001", - "responsible_ai_id": "RAI-001" - }, + metadata={"use_case_id": "TEST-001", "responsible_ai_id": "RAI-001"}, models=["gpt-4", "gpt-3.5-turbo"], max_budget=100.0, model_rpm_limit={"gpt-4": 100}, @@ -136,7 +123,7 @@ async def test_new_project(prisma_client): ) print("New project response:", response) - + # Assertions assert response.project_id is not None assert response.project_alias == "test-project" @@ -216,7 +203,7 @@ async def test_update_project(prisma_client): description="Updated description", metadata={ "use_case_id": "TEST-002-UPDATED", - "additional_field": "new_value" + "additional_field": "new_value", }, models=["gpt-4", "gpt-3.5-turbo", "claude-3"], max_budget=200.0, @@ -244,8 +231,14 @@ async def test_update_project(prisma_client): # model_rpm_limit and model_tpm_limit are stored in metadata assert update_response.metadata["use_case_id"] == "TEST-002-UPDATED" assert update_response.metadata["additional_field"] == "new_value" - assert update_response.metadata["model_rpm_limit"] == {"gpt-4": 200, "claude-3": 50} - assert update_response.metadata["model_tpm_limit"] == {"gpt-4": 2000, "claude-3": 500} + assert update_response.metadata["model_rpm_limit"] == { + "gpt-4": 200, + "claude-3": 50, + } + assert update_response.metadata["model_tpm_limit"] == { + "gpt-4": 2000, + "claude-3": 500, + } assert update_response.litellm_budget_table is not None assert update_response.litellm_budget_table.max_budget == 200.0 @@ -304,9 +297,7 @@ async def test_delete_project(prisma_client): project_id = create_response.project_id # Delete the project - delete_data = DeleteProjectRequest( - project_ids=[project_id] - ) + delete_data = DeleteProjectRequest(project_ids=[project_id]) delete_response = await delete_project( data=delete_data, @@ -378,10 +369,7 @@ async def test_project_info(prisma_client): project_alias="test-project-info", description="Test project info endpoint", team_id=_team_id, - metadata={ - "use_case_id": "TEST-003", - "cost_center": "engineering" - }, + metadata={"use_case_id": "TEST-003", "cost_center": "engineering"}, models=["gpt-4", "claude-3"], max_budget=150.0, model_rpm_limit={"gpt-4": 150}, @@ -432,3 +420,370 @@ async def test_project_info(prisma_client): traceback.print_exc() pytest.fail(f"Got exception {e}") + +### VALIDATION TESTS ### + + +def test_check_team_project_limits_models_not_in_team(): + """ + Test that creating a project with models not in the team raises an error. + """ + from litellm.proxy.management_endpoints.project_endpoints import ( + _check_team_project_limits, + ) + from litellm.proxy._types import LiteLLM_TeamTable + + team = LiteLLM_TeamTable( + team_id="test-team", + models=["gpt-4", "gpt-3.5-turbo"], + ) + + data = NewProjectRequest( + team_id="test-team", + models=["gpt-4", "claude-3"], # claude-3 not in team + ) + + with pytest.raises(Exception) as exc_info: + _check_team_project_limits(team_object=team, data=data) + + assert "claude-3" in str(exc_info.value.detail) + assert "not in team's allowed models" in str(exc_info.value.detail) + + +def test_check_team_project_limits_budget_exceeds_team(): + """ + Test that creating a project with budget > team budget raises an error. + """ + from litellm.proxy.management_endpoints.project_endpoints import ( + _check_team_project_limits, + ) + from litellm.proxy._types import LiteLLM_TeamTable + + team = LiteLLM_TeamTable( + team_id="test-team", + models=["gpt-4"], + max_budget=100.0, + ) + + data = NewProjectRequest( + team_id="test-team", + models=["gpt-4"], + max_budget=150.0, # exceeds team's 100.0 + ) + + with pytest.raises(Exception) as exc_info: + _check_team_project_limits(team_object=team, data=data) + + assert "exceeds team's max_budget" in str(exc_info.value.detail) + + +def test_check_team_project_limits_valid_subset(): + """ + Test that a valid project (models subset, budget within limit) passes. + """ + from litellm.proxy.management_endpoints.project_endpoints import ( + _check_team_project_limits, + ) + from litellm.proxy._types import LiteLLM_TeamTable + + team = LiteLLM_TeamTable( + team_id="test-team", + models=["gpt-4", "gpt-3.5-turbo", "claude-3"], + max_budget=1000.0, + ) + + data = NewProjectRequest( + team_id="test-team", + models=["gpt-4", "gpt-3.5-turbo"], + max_budget=500.0, + ) + + # Should not raise + _check_team_project_limits(team_object=team, data=data) + + +def test_check_team_project_limits_all_proxy_models(): + """ + Test that team with 'all-proxy-models' allows any project models. + """ + from litellm.proxy.management_endpoints.project_endpoints import ( + _check_team_project_limits, + ) + from litellm.proxy._types import LiteLLM_TeamTable + + team = LiteLLM_TeamTable( + team_id="test-team", + models=["all-proxy-models"], + ) + + data = NewProjectRequest( + team_id="test-team", + models=["gpt-4", "claude-3", "anything-goes"], + ) + + # Should not raise - team allows all models + _check_team_project_limits(team_object=team, data=data) + + +def test_check_team_project_limits_tpm_exceeds_team(): + """ + Test that project tpm_limit exceeding team tpm_limit raises an error. + """ + from litellm.proxy.management_endpoints.project_endpoints import ( + _check_team_project_limits, + ) + from litellm.proxy._types import LiteLLM_TeamTable + + team = LiteLLM_TeamTable( + team_id="test-team", + models=["gpt-4"], + tpm_limit=10000, + ) + + data = NewProjectRequest( + team_id="test-team", + models=["gpt-4"], + tpm_limit=20000, # exceeds team's 10000 + ) + + with pytest.raises(Exception) as exc_info: + _check_team_project_limits(team_object=team, data=data) + + assert "exceeds team's tpm_limit" in str(exc_info.value.detail) + + +def test_check_team_project_limits_negative_budget(): + """ + Test that negative budget values raise an error. + """ + from litellm.proxy.management_endpoints.project_endpoints import ( + _check_team_project_limits, + ) + from litellm.proxy._types import LiteLLM_TeamTable + + team = LiteLLM_TeamTable( + team_id="test-team", + models=["gpt-4"], + ) + + data = NewProjectRequest( + team_id="test-team", + models=["gpt-4"], + max_budget=-10.0, + ) + + with pytest.raises(Exception) as exc_info: + _check_team_project_limits(team_object=team, data=data) + + assert "cannot be negative" in str(exc_info.value.detail) + + +def test_check_team_project_limits_soft_budget_gte_max(): + """ + Test that soft_budget >= max_budget raises an error. + """ + from litellm.proxy.management_endpoints.project_endpoints import ( + _check_team_project_limits, + ) + from litellm.proxy._types import LiteLLM_TeamTable + + team = LiteLLM_TeamTable( + team_id="test-team", + models=["gpt-4"], + ) + + data = NewProjectRequest( + team_id="test-team", + models=["gpt-4"], + max_budget=100.0, + soft_budget=100.0, # equal to max, should fail + ) + + with pytest.raises(Exception) as exc_info: + _check_team_project_limits(team_object=team, data=data) + + assert "must be strictly lower" in str(exc_info.value.detail) + + +def test_premium_user_gate(): + """ + Test that project endpoints require premium_user=True. + """ + + # This test just validates the premium_user check exists + # The actual endpoint test would need prisma, but we can verify + # the import path works + setattr(litellm.proxy.proxy_server, "premium_user", False) + + # Verify that CommonProxyErrors.not_premium_user exists + from litellm.proxy._types import CommonProxyErrors + + assert hasattr(CommonProxyErrors, "not_premium_user") + + # Reset + setattr(litellm.proxy.proxy_server, "premium_user", True) + + +def test_project_model_access_denied_error_type(): + """ + Test that ProxyErrorTypes.project_model_access_denied exists. + """ + from litellm.proxy._types import ProxyErrorTypes + + assert hasattr(ProxyErrorTypes, "project_model_access_denied") + assert ( + ProxyErrorTypes.project_model_access_denied.value + == "project_model_access_denied" + ) + + # Test the classmethod resolves correctly + result = ProxyErrorTypes.get_model_access_error_type_for_object("project") + assert result == ProxyErrorTypes.project_model_access_denied + + +def test_project_cached_obj_has_last_refreshed_at(): + """ + Test that LiteLLM_ProjectTableCachedObj has last_refreshed_at field + matching LiteLLM_TeamTableCachedObj pattern. + """ + from litellm.proxy._types import ( + LiteLLM_ProjectTableCachedObj, + LiteLLM_ProjectTable, + ) + + # Verify inheritance + assert issubclass(LiteLLM_ProjectTableCachedObj, LiteLLM_ProjectTable) + + # Verify last_refreshed_at field exists and defaults to None + obj = LiteLLM_ProjectTableCachedObj( + project_id="test", + created_by="admin", + updated_by="admin", + ) + assert obj.last_refreshed_at is None + + # Verify it can be set + obj.last_refreshed_at = 1234567890.0 + assert obj.last_refreshed_at == 1234567890.0 + + +@pytest.mark.asyncio +async def test_project_max_budget_check_fires_alert(): + """ + Test that _project_max_budget_check fires a budget alert + when project exceeds its max budget (matches _team_max_budget_check pattern). + """ + from litellm.proxy.auth.auth_checks import _project_max_budget_check + from litellm.proxy._types import ( + LiteLLM_BudgetTable, + LiteLLM_ProjectTableCachedObj, + ) + + project = LiteLLM_ProjectTableCachedObj( + project_id="test-project", + spend=150.0, + created_by="admin", + updated_by="admin", + litellm_budget_table=LiteLLM_BudgetTable(max_budget=100.0), + ) + + valid_token = UserAPIKeyAuth( + token="test-token", + user_id="user-1", + team_id="team-1", + ) + + mock_proxy_logging = mock.AsyncMock(spec=ProxyLogging) + mock_proxy_logging.budget_alerts = mock.AsyncMock() + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _project_max_budget_check( + project_object=project, + valid_token=valid_token, + proxy_logging_obj=mock_proxy_logging, + ) + + assert "Project=test-project" in str(exc_info.value) + assert "150.0" in str(exc_info.value) + + +@pytest.mark.asyncio +async def test_project_soft_budget_check(): + """ + Test that _project_soft_budget_check triggers alert when soft budget is exceeded. + """ + from litellm.proxy.auth.auth_checks import _project_soft_budget_check + from litellm.proxy._types import ( + LiteLLM_BudgetTable, + LiteLLM_ProjectTableCachedObj, + ) + + project = LiteLLM_ProjectTableCachedObj( + project_id="test-project", + spend=80.0, + created_by="admin", + updated_by="admin", + litellm_budget_table=LiteLLM_BudgetTable(soft_budget=75.0), + ) + + valid_token = UserAPIKeyAuth( + token="test-token", + user_id="user-1", + team_id="team-1", + ) + + mock_proxy_logging = mock.AsyncMock(spec=ProxyLogging) + mock_proxy_logging.budget_alerts = mock.AsyncMock() + + # Should not raise (soft budget only alerts, doesn't block) + await _project_soft_budget_check( + project_object=project, + valid_token=valid_token, + proxy_logging_obj=mock_proxy_logging, + ) + + +@pytest.mark.asyncio +async def test_project_soft_budget_check_no_alert_under_budget(): + """ + Test that _project_soft_budget_check does NOT trigger alert when under soft budget. + """ + from litellm.proxy.auth.auth_checks import _project_soft_budget_check + from litellm.proxy._types import ( + LiteLLM_BudgetTable, + LiteLLM_ProjectTableCachedObj, + ) + + project = LiteLLM_ProjectTableCachedObj( + project_id="test-project", + spend=50.0, + created_by="admin", + updated_by="admin", + litellm_budget_table=LiteLLM_BudgetTable(soft_budget=75.0), + ) + + valid_token = UserAPIKeyAuth( + token="test-token", + user_id="user-1", + team_id="team-1", + ) + + mock_proxy_logging = mock.AsyncMock(spec=ProxyLogging) + mock_proxy_logging.budget_alerts = mock.AsyncMock() + + # Should not raise and should not alert + await _project_soft_budget_check( + project_object=project, + valid_token=valid_token, + proxy_logging_obj=mock_proxy_logging, + ) + + +def test_litellm_entity_type_has_project(): + """ + Test that Litellm_EntityType has PROJECT member for budget alerts. + """ + from litellm.proxy._types import Litellm_EntityType + + assert hasattr(Litellm_EntityType, "PROJECT") + assert Litellm_EntityType.PROJECT.value == "project"