diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 97724f6e14c..e72a3219f72 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2592,7 +2592,20 @@ async def can_team_access_model( ) if effective_models: - models_to_check = effective_models + # Defense-in-depth: intersect with team.models so that + # misconfigured default_models can never grant access + # beyond the team's allowed model list. + if ( + team_object + and team_object.models + and SpecialModelNames.all_proxy_models.value + not in team_object.models + ): + models_to_check = list( + set(effective_models) & set(team_object.models) + ) + else: + models_to_check = effective_models elif team_object and team_object.models: # Fallback: effective models empty but team has models configured. # Graceful degradation prevents misconfiguration cliff when feature diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 2ff78c60877..573232fd437 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -756,19 +756,25 @@ async def new_team( # noqa: PLR0915 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}"} + detail={ + "error": f"max_budget cannot be negative. Received: {data.max_budget}" + }, ) if data.team_member_budget is not None and data.team_member_budget < 0: raise HTTPException( status_code=400, - detail={"error": f"team_member_budget cannot be negative. Received: {data.team_member_budget}"} + detail={ + "error": f"team_member_budget cannot be negative. Received: {data.team_member_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}"} + detail={ + "error": f"soft_budget cannot be negative. Received: {data.soft_budget}" + }, ) - + if data.soft_budget is not None: if data.max_budget is not None: # If max_budget is set, soft_budget must be strictly lower than max_budget @@ -777,7 +783,22 @@ async def new_team( # noqa: PLR0915 status_code=400, detail={ "error": f"soft_budget ({data.soft_budget}) must be strictly lower than max_budget ({data.max_budget})" - } + }, + ) + + # Validate default_models ⊆ team.models (prevent privilege escalation) + if data.default_models: + if ( + data.models + and SpecialModelNames.all_proxy_models.value not in data.models + ): + invalid = set(data.default_models) - set(data.models) + if invalid: + raise HTTPException( + status_code=400, + detail={ + "error": f"default_models {list(invalid)} not in team's allowed models: {data.models}" + }, ) # Check if license is over limit @@ -937,12 +958,16 @@ async def new_team( # noqa: PLR0915 complete_team_data.members_with_roles = [] complete_team_data_dict = complete_team_data.model_dump(exclude_none=True) - + # Serialize router_settings to JSON (matching key creation pattern) router_settings_value = getattr(data, "router_settings", None) - router_settings_json = safe_dumps(router_settings_value) if router_settings_value is not None else safe_dumps({}) + router_settings_json = ( + safe_dumps(router_settings_value) + if router_settings_value is not None + else safe_dumps({}) + ) complete_team_data_dict["router_settings"] = router_settings_json - + complete_team_data_dict = prisma_client.jsonify_team_object( db_data=complete_team_data_dict ) @@ -1118,7 +1143,9 @@ async def fetch_and_validate_organization( validate_team_org_change( team=LiteLLM_TeamTable(**existing_team_row.model_dump()), - organization=LiteLLM_OrganizationTableWithMembers(**organization_row.model_dump()), + organization=LiteLLM_OrganizationTableWithMembers( + **organization_row.model_dump() + ), llm_router=llm_router, ) @@ -1126,7 +1153,9 @@ async def fetch_and_validate_organization( def validate_team_org_change( - team: LiteLLM_TeamTable, organization: LiteLLM_OrganizationTableWithMembers, llm_router: Router + team: LiteLLM_TeamTable, + organization: LiteLLM_OrganizationTableWithMembers, + llm_router: Router, ) -> bool: """ Validate that a team can be moved to an organization. @@ -1177,7 +1206,9 @@ def validate_team_org_change( # Check if the team's user_id is a member of the org team_members = [m.user_id for m in team.members_with_roles] - org_members = [m.user_id for m in organization.members] if organization.members else [] + org_members = ( + [m.user_id for m in organization.members] if organization.members else [] + ) not_in_org = [ m for m in team_members @@ -1223,7 +1254,7 @@ def validate_team_org_change( "/team/update", tags=["team management"], dependencies=[Depends(user_api_key_auth)] ) @management_endpoint_wrapper -async def update_team( # noqa: PLR0915 +async def update_team( # noqa: PLR0915 data: UpdateTeamRequest, http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), @@ -1309,24 +1340,32 @@ async def update_team( # noqa: PLR0915 ) if data.team_id is None: - raise HTTPException(status_code=400, detail={"error": "No team id passed in"}) + raise HTTPException( + status_code=400, detail={"error": "No team id passed in"} + ) verbose_proxy_logger.debug("/team/update - %s", data) # Validate budget values are not negative 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}"} + detail={ + "error": f"max_budget cannot be negative. Received: {data.max_budget}" + }, ) if data.team_member_budget is not None and data.team_member_budget < 0: raise HTTPException( status_code=400, - detail={"error": f"team_member_budget cannot be negative. Received: {data.team_member_budget}"} + detail={ + "error": f"team_member_budget cannot be negative. Received: {data.team_member_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}"} + detail={ + "error": f"soft_budget cannot be negative. Received: {data.soft_budget}" + }, ) existing_team_row = await prisma_client.db.litellm_teamtable.find_unique( @@ -1338,28 +1377,38 @@ async def update_team( # noqa: PLR0915 status_code=404, detail={"error": f"Team not found, passed team_id={data.team_id}"}, ) - + if data.soft_budget is not None: - max_budget_to_check = data.max_budget if data.max_budget is not None else existing_team_row.max_budget + max_budget_to_check = ( + data.max_budget + if data.max_budget is not None + else existing_team_row.max_budget + ) if max_budget_to_check is not None: if data.soft_budget >= max_budget_to_check: raise HTTPException( status_code=400, detail={ "error": f"soft_budget ({data.soft_budget}) must be strictly lower than max_budget ({max_budget_to_check})" - } + }, ) - + if data.max_budget is not None: - existing_soft_budget = getattr(existing_team_row, 'soft_budget', None) - soft_budget_to_check = data.soft_budget if data.soft_budget is not None else existing_soft_budget - if soft_budget_to_check is not None and isinstance(soft_budget_to_check, (int, float)): + existing_soft_budget = getattr(existing_team_row, "soft_budget", None) + soft_budget_to_check = ( + data.soft_budget + if data.soft_budget is not None + else existing_soft_budget + ) + if soft_budget_to_check is not None and isinstance( + soft_budget_to_check, (int, float) + ): if data.max_budget <= soft_budget_to_check: raise HTTPException( status_code=400, detail={ "error": f"max_budget ({data.max_budget}) must be strictly greater than soft_budget ({soft_budget_to_check})" - } + }, ) if ( @@ -1414,6 +1463,24 @@ async def update_team( # noqa: PLR0915 user_api_key_cache=user_api_key_cache, ) + # Validate default_models ⊆ team.models (prevent privilege escalation) + if data.default_models is not None: + team_models = ( + data.models if data.models is not None else existing_team_row.models + ) + if ( + team_models + and SpecialModelNames.all_proxy_models.value not in team_models + ): + invalid = set(data.default_models) - set(team_models) + if invalid: + raise HTTPException( + status_code=400, + detail={ + "error": f"default_models {list(invalid)} not in team's allowed models: {team_models}" + }, + ) + updated_kv = data.json(exclude_unset=True) # Check budget_duration and budget_reset_at @@ -1460,16 +1527,19 @@ async def update_team( # noqa: PLR0915 updated_kv["model_id"] = _model_id # Serialize router_settings to JSON if present (matching key update pattern) - if "router_settings" in updated_kv and updated_kv["router_settings"] is not None: + if ( + "router_settings" in updated_kv + and updated_kv["router_settings"] is not None + ): updated_kv["router_settings"] = safe_dumps(updated_kv["router_settings"]) updated_kv = prisma_client.jsonify_team_object(db_data=updated_kv) - team_row: Optional[LiteLLM_TeamTable] = ( - await prisma_client.db.litellm_teamtable.update( - where={"team_id": data.team_id}, - data=updated_kv, - include={"litellm_model_table": True}, # type: ignore - ) + team_row: Optional[ + LiteLLM_TeamTable + ] = await prisma_client.db.litellm_teamtable.update( + where={"team_id": data.team_id}, + data=updated_kv, + include={"litellm_model_table": True}, # type: ignore ) if team_row is None or team_row.team_id is None: @@ -1478,7 +1548,9 @@ async def update_team( # noqa: PLR0915 detail={"error": "Team doesn't exist. Got={}".format(team_row)}, ) - verbose_proxy_logger.info("Successfully updated team - %s, info", team_row.team_id) + verbose_proxy_logger.info( + "Successfully updated team - %s, info", team_row.team_id + ) await _cache_team_object( team_id=team_row.team_id, team_table=LiteLLM_TeamTableCachedObj(**team_row.model_dump()), @@ -1848,14 +1920,14 @@ async def _validate_and_populate_member_user_info( ) -> Member: """ Validate and populate user_email/user_id for a member. - + Logic: 1. If both user_email and user_id are provided, verify they belong to the same user (use user_email as source of truth) 2. If only user_email is provided, populate user_id from DB 3. If only user_id is provided, populate user_email from DB (if user exists) 4. If only user_id is provided and doesn't exist, allow it to pass with user_email as None (will be upserted later) 5. If user_email and user_id mismatch, throw error - + Returns a Member with user_email and user_id populated (user_email may be None if only user_id provided and user doesn't exist). """ if member.user_email is None and member.user_id is None: @@ -1863,7 +1935,7 @@ async def _validate_and_populate_member_user_info( status_code=400, detail={"error": "Either user_id or user_email must be provided"}, ) - + # Case 1: Both user_email and user_id provided - verify they match if member.user_email is not None and member.user_id is not None: # Use user_email as source of truth @@ -1873,13 +1945,13 @@ async def _validate_and_populate_member_user_info( table_name="user", query_type="find_all", ) - + if users_by_email is None or ( isinstance(users_by_email, list) and len(users_by_email) == 0 ): # User doesn't exist yet - this is fine, will be created later return member - + if isinstance(users_by_email, list) and len(users_by_email) > 1: raise HTTPException( status_code=400, @@ -1887,10 +1959,10 @@ async def _validate_and_populate_member_user_info( "error": f"Multiple users found with email '{member.user_email}'. Please use 'user_id' instead." }, ) - + # Get the single user user_by_email = users_by_email[0] - + # Verify the user_id matches if user_by_email.user_id != member.user_id: raise HTTPException( @@ -1899,56 +1971,61 @@ async def _validate_and_populate_member_user_info( "error": f"user_email '{member.user_email}' and user_id '{member.user_id}' do not belong to the same user." }, ) - + # Both match, return as is return member - + # Case 2: Only user_email provided - populate user_id from DB if member.user_email is not None and member.user_id is None: user_by_email = await prisma_client.db.litellm_usertable.find_first( where={"user_email": {"equals": member.user_email, "mode": "insensitive"}} ) - + if user_by_email is None: # User doesn't exist yet - this is fine, will be created later return member - + # Check for multiple users with same email users_by_email = await prisma_client.get_data( key_val={"user_email": member.user_email}, table_name="user", query_type="find_all", ) - - if users_by_email and isinstance(users_by_email, list) and len(users_by_email) > 1: + + if ( + users_by_email + and isinstance(users_by_email, list) + and len(users_by_email) > 1 + ): raise HTTPException( status_code=400, detail={ "error": f"Multiple users found with email '{member.user_email}'. Please use 'user_id' instead." }, ) - + # Populate user_id member.user_id = user_by_email.user_id return member - + # Case 3: Only user_id provided - populate user_email from DB if user exists if member.user_id is not None and member.user_email is None: user_by_id = await prisma_client.db.litellm_usertable.find_unique( where={"user_id": member.user_id} ) - + if user_by_id is None: # User doesn't exist yet - allow it to pass with user_email as None # Will be upserted later with just user_id and null email return member - + # Populate user_email member.user_email = user_by_id.user_email return member - + return member + @router.post( "/team/member_add", tags=["team management"], @@ -2037,14 +2114,16 @@ async def team_member_add( prisma_client=prisma_client, ) - updated_team, updated_users, updated_team_memberships = ( - await _add_team_members_to_team( - data=data, - complete_team_data=complete_team_data, - prisma_client=prisma_client, - user_api_key_dict=user_api_key_dict, - litellm_proxy_admin_name=litellm_proxy_admin_name, - ) + ( + updated_team, + updated_users, + updated_team_memberships, + ) = await _add_team_members_to_team( + data=data, + complete_team_data=complete_team_data, + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, ) # Check if updated_team is None @@ -2223,15 +2302,15 @@ async def team_member_delete( ) # Fetch keys before deletion to persist them - keys_to_delete: List[LiteLLM_VerificationToken] = ( - await prisma_client.db.litellm_verificationtoken.find_many( - where={ - "user_id": {"in": list(user_ids_to_delete)}, - "team_id": data.team_id, - } - ) + keys_to_delete: List[ + LiteLLM_VerificationToken + ] = await prisma_client.db.litellm_verificationtoken.find_many( + where={ + "user_id": {"in": list(user_ids_to_delete)}, + "team_id": data.team_id, + } ) - + if keys_to_delete: await _persist_deleted_verification_tokens( keys=keys_to_delete, @@ -2633,10 +2712,10 @@ async def delete_team( team_rows: List[LiteLLM_TeamTable] = [] for team_id in data.team_ids: try: - team_row_base: Optional[BaseModel] = ( - await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id} - ) + team_row_base: Optional[ + BaseModel + ] = await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id} ) if team_row_base is None: raise Exception @@ -2695,10 +2774,10 @@ async def delete_team( _persist_deleted_verification_tokens, ) - keys_to_delete: List[LiteLLM_VerificationToken] = ( - await prisma_client.db.litellm_verificationtoken.find_many( - where={"team_id": {"in": data.team_ids}} - ) + keys_to_delete: List[ + LiteLLM_VerificationToken + ] = await prisma_client.db.litellm_verificationtoken.find_many( + where={"team_id": {"in": data.team_ids}} ) if keys_to_delete: @@ -2737,7 +2816,6 @@ async def delete_team( return deleted_teams - def _transform_teams_to_deleted_records( teams: List[LiteLLM_TeamTable], user_api_key_dict: UserAPIKeyAuth, @@ -2760,7 +2838,13 @@ def _transform_teams_to_deleted_records( ) record = deleted_record.model_dump() - for json_field in ["members_with_roles", "metadata", "model_spend", "model_max_budget", "router_settings"]: + for json_field in [ + "members_with_roles", + "metadata", + "model_spend", + "model_max_budget", + "router_settings", + ]: if json_field in record and record[json_field] is not None: record[json_field] = json.dumps(record[json_field]) @@ -2779,9 +2863,7 @@ async def _save_deleted_team_records( """Save deleted team records to the database.""" if not records: return - await prisma_client.db.litellm_deletedteamtable.create_many( - data=records - ) + await prisma_client.db.litellm_deletedteamtable.create_many(data=records) async def _persist_deleted_team_records( @@ -2801,6 +2883,7 @@ async def _persist_deleted_team_records( prisma_client=prisma_client, ) + def validate_membership( user_api_key_dict: UserAPIKeyAuth, team_table: LiteLLM_TeamTable ): @@ -2924,11 +3007,11 @@ async def team_info( ) try: - team_info: Optional[BaseModel] = ( - await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id}, - include={"object_permission": True}, - ) + team_info: Optional[ + BaseModel + ] = await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id}, + include={"object_permission": True}, ) if team_info is None: raise Exception @@ -3386,7 +3469,9 @@ async def list_team_v2( order=order_by if order_by else {"created_at": "desc"}, # Default sort ) # Get total count for pagination - total_count = await prisma_client.db.litellm_teamtable.count(where=where_conditions) + total_count = await prisma_client.db.litellm_teamtable.count( + where=where_conditions + ) # Calculate total pages total_pages = -(-total_count // page_size) # Ceiling division diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 11a5d6e39e2..de175218a9b 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -23,23 +23,31 @@ from typing import ( ) from litellm import _custom_logger_compatible_callbacks_literal -from litellm.constants import (DEFAULT_MODEL_CREATED_AT_TIME, - MAX_TEAM_LIST_LIMIT) -from litellm.proxy._types import (DB_CONNECTION_ERROR_TYPES, CommonProxyErrors, - ProxyErrorTypes, ProxyException, - SpendLogsMetadata, SpendLogsPayload) +from litellm.constants import DEFAULT_MODEL_CREATED_AT_TIME, MAX_TEAM_LIST_LIMIT +from litellm.proxy._types import ( + DB_CONNECTION_ERROR_TYPES, + CommonProxyErrors, + ProxyErrorTypes, + ProxyException, + SpendLogsMetadata, + SpendLogsPayload, +) from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import CallTypes, CallTypesLiteral try: - from litellm_enterprise.enterprise_callbacks.send_emails.base_email import \ - BaseEmailLogger - from litellm_enterprise.enterprise_callbacks.send_emails.resend_email import \ - ResendEmailLogger - from litellm_enterprise.enterprise_callbacks.send_emails.sendgrid_email import \ - SendGridEmailLogger - from litellm_enterprise.enterprise_callbacks.send_emails.smtp_email import \ - SMTPEmailLogger + from litellm_enterprise.enterprise_callbacks.send_emails.base_email import ( + BaseEmailLogger, + ) + from litellm_enterprise.enterprise_callbacks.send_emails.resend_email import ( + ResendEmailLogger, + ) + from litellm_enterprise.enterprise_callbacks.send_emails.sendgrid_email import ( + SendGridEmailLogger, + ) + from litellm_enterprise.enterprise_callbacks.send_emails.smtp_email import ( + SMTPEmailLogger, + ) except ImportError: BaseEmailLogger = None # type: ignore SendGridEmailLogger = None # type: ignore @@ -58,56 +66,70 @@ from fastapi import HTTPException, status import litellm import litellm.litellm_core_utils import litellm.litellm_core_utils.litellm_logging -from litellm import (EmbeddingResponse, ImageResponse, ModelResponse, - ModelResponseStream, Router) +from litellm import ( + EmbeddingResponse, + ImageResponse, + ModelResponse, + ModelResponseStream, + Router, +) from litellm._logging import verbose_proxy_logger from litellm._service_logger import ServiceLogging, ServiceTypes from litellm.caching.caching import DualCache, RedisCache from litellm.caching.dual_cache import LimitedSizeOrderedDict from litellm.exceptions import RejectedRequestError -from litellm.integrations.custom_guardrail import (CustomGuardrail, - ModifyResponseException) +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + ModifyResponseException, +) from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting -from litellm.integrations.SlackAlerting.utils import \ - _add_langfuse_trace_id_to_alert +from litellm.integrations.SlackAlerting.utils import _add_langfuse_trace_id_to_alert from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -from litellm.proxy._types import (AlertType, CallInfo, - LiteLLM_VerificationTokenView, Member, - UserAPIKeyAuth) +from litellm.proxy._types import ( + AlertType, + CallInfo, + LiteLLM_VerificationTokenView, + Member, + UserAPIKeyAuth, +) from litellm.proxy.auth.route_checks import RouteChecks -from litellm.proxy.db.create_views import (create_missing_views, - should_create_missing_views) +from litellm.proxy.db.create_views import ( + create_missing_views, + should_create_missing_views, +) from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.db.log_db_metrics import log_db_metrics from litellm.proxy.db.prisma_client import PrismaWrapper -from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import \ - UnifiedLLMGuardrails +from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, +) from litellm.proxy.hooks import PROXY_HOOKS, get_proxy_hook from litellm.proxy.hooks.cache_control_check import _PROXY_CacheControlCheck from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter -from litellm.proxy.hooks.parallel_request_limiter import \ - _PROXY_MaxParallelRequestsHandler +from litellm.proxy.hooks.parallel_request_limiter import ( + _PROXY_MaxParallelRequestsHandler, +) from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor from litellm.secret_managers.main import str_to_bool from litellm.types.integrations.slack_alerting import DEFAULT_ALERT_TYPES -from litellm.types.mcp import (MCPDuringCallResponseObject, - MCPPreCallRequestObject, - MCPPreCallResponseObject) -from litellm.types.proxy.policy_engine.pipeline_types import \ - PipelineExecutionResult +from litellm.types.mcp import ( + MCPDuringCallResponseObject, + MCPPreCallRequestObject, + MCPPreCallResponseObject, +) +from litellm.types.proxy.policy_engine.pipeline_types import PipelineExecutionResult from litellm.types.utils import LLMResponseTypes, LoggedLiteLLMParams if TYPE_CHECKING: from opentelemetry.trace import Span as _Span - from litellm.litellm_core_utils.litellm_logging import \ - Logging as LiteLLMLoggingObj + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj Span = Union[_Span, Any] else: @@ -1050,9 +1072,10 @@ class ProxyLogging: """Process prompt template if applicable.""" from litellm.proxy.prompts.prompt_endpoints import ( - construct_versioned_prompt_id, get_latest_version_prompt_id) - from litellm.proxy.prompts.prompt_registry import \ - IN_MEMORY_PROMPT_REGISTRY + construct_versioned_prompt_id, + get_latest_version_prompt_id, + ) + from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY from litellm.utils import get_non_default_completion_params if prompt_version is None: @@ -1102,8 +1125,9 @@ class ProxyLogging: def _process_guardrail_metadata(self, data: dict) -> None: """Process guardrails from metadata and add to applied_guardrails.""" - from litellm.proxy.common_utils.callback_utils import \ - add_guardrail_to_applied_guardrails_header + from litellm.proxy.common_utils.callback_utils import ( + add_guardrail_to_applied_guardrails_header, + ) metadata_standard = data.get("metadata") or {} metadata_litellm = data.get("litellm_metadata") or {} @@ -2000,8 +2024,7 @@ class ProxyLogging: if isinstance(response, (ModelResponse, ModelResponseStream)): response_str = litellm.get_response_string(response_obj=response) elif isinstance(response, dict) and self.is_a2a_streaming_response(response): - from litellm.llms.a2a.common_utils import \ - extract_text_from_a2a_response + from litellm.llms.a2a.common_utils import extract_text_from_a2a_response response_str = extract_text_from_a2a_response(response) if response_str is not None: @@ -2010,8 +2033,7 @@ class ProxyLogging: _callback: Optional[CustomLogger] = None if isinstance(callback, CustomGuardrail): # Main - V2 Guardrails implementation - from litellm.types.guardrails import \ - GuardrailEventHooks + from litellm.types.guardrails import GuardrailEventHooks ## CHECK FOR MODEL-LEVEL GUARDRAILS modified_data = _check_and_merge_model_level_guardrails( @@ -2824,16 +2846,28 @@ class PrismaClient: detail={"error": f"No token passed in. Token={token}"}, ) + # Only include override columns when feature is enabled; + # avoids SQL errors if the columns haven't been migrated yet. + _override_cols = "" + if ( + os.getenv("LITELLM_TEAM_MODEL_OVERRIDES", "false").lower() + == "true" + ): + _override_cols = ( + "t.default_models AS team_default_models,\n" + " tm.models AS team_member_models," + ) + sql_query = f""" - SELECT + SELECT v.*, - t.spend AS team_spend, + t.spend AS team_spend, t.max_budget AS team_max_budget, t.soft_budget AS team_soft_budget, t.tpm_limit AS team_tpm_limit, t.rpm_limit AS team_rpm_limit, t.models AS team_models, - t.default_models AS team_default_models, + {_override_cols} t.metadata AS team_metadata, t.blocked AS team_blocked, t.team_alias AS team_alias, @@ -2842,7 +2876,6 @@ class PrismaClient: t.object_permission_id AS team_object_permission_id, t.organization_id as org_id, tm.spend AS team_member_spend, - tm.models AS team_member_models, m.aliases AS team_model_aliases, -- Added comma to separate b.* columns b.max_budget AS litellm_budget_table_max_budget, @@ -3602,13 +3635,15 @@ class PrismaClient: probe_pid, _ = os.waitpid(pid, os.WNOHANG) except ChildProcessError: verbose_proxy_logger.debug( - "PID %s is not a child process; skipping waitpid watch.", pid, + "PID %s is not a child process; skipping waitpid watch.", + pid, ) return False if probe_pid == pid: verbose_proxy_logger.warning( - "prisma-query-engine PID %s already dead at watch start.", pid, + "prisma-query-engine PID %s already dead at watch start.", + pid, ) self._engine_confirmed_dead = True self._reap_all_zombies() @@ -3785,11 +3820,17 @@ class PrismaClient: waitpid thread nor pidfd are available. """ - if self._watching_engine or self._engine_pidfd >= 0 or self._engine_wait_thread is not None: + if ( + self._watching_engine + or self._engine_pidfd >= 0 + or self._engine_wait_thread is not None + ): return pid = self._get_engine_pid() if pid == 0: - verbose_proxy_logger.debug("Could not find prisma-query-engine PID; engine death detection unavailable.") + verbose_proxy_logger.debug( + "Could not find prisma-query-engine PID; engine death detection unavailable." + ) return self._engine_pid = pid self._engine_confirmed_dead = False @@ -3798,15 +3839,18 @@ class PrismaClient: pidfd_ok = False if waitpid_ok else self._try_pidfd_watch(pid) if waitpid_ok: verbose_proxy_logger.info( - "Watching engine PID %s via waitpid thread.", pid, + "Watching engine PID %s via waitpid thread.", + pid, ) elif pidfd_ok: verbose_proxy_logger.info( - "Watching engine PID %s via pidfd.", pid, + "Watching engine PID %s via pidfd.", + pid, ) else: verbose_proxy_logger.info( - "Watching engine PID %s via os.kill polling.", pid, + "Watching engine PID %s via os.kill polling.", + pid, ) self._watching_engine = True asyncio.create_task(self._poll_engine_proc()) @@ -3829,7 +3873,9 @@ class PrismaClient: blip -- disconnect, connect, SELECT 1). """ effective_timeout = ( - timeout_seconds if timeout_seconds is not None else self._db_watchdog_reconnect_timeout_seconds + timeout_seconds + if timeout_seconds is not None + else self._db_watchdog_reconnect_timeout_seconds ) engine_is_dead = self._engine_confirmed_dead or ( @@ -3849,14 +3895,18 @@ class PrismaClient: async def _do_heavy_reconnect() -> None: db_url = os.getenv("DATABASE_URL", "") if not db_url: - verbose_proxy_logger.error("DATABASE_URL not set; cannot recreate Prisma client.") + verbose_proxy_logger.error( + "DATABASE_URL not set; cannot recreate Prisma client." + ) raise RuntimeError("DATABASE_URL not set") await self.db.recreate_prisma_client(db_url) await self._start_engine_watcher() await asyncio.wait_for(_do_heavy_reconnect(), timeout=effective_timeout) else: - verbose_proxy_logger.debug("Performing Prisma DB reconnect (engine alive or unknown).") + verbose_proxy_logger.debug( + "Performing Prisma DB reconnect (engine alive or unknown)." + ) async def _do_direct_reconnect() -> None: try: @@ -3939,7 +3989,9 @@ class PrismaClient: if lock_timeout_seconds is None: async with self._db_reconnect_lock: - return await self._attempt_reconnect_inside_lock(force, reason, timeout_seconds) + return await self._attempt_reconnect_inside_lock( + force, reason, timeout_seconds + ) lock_acquired_by_timeout_task = False @@ -3988,14 +4040,17 @@ class PrismaClient: return False try: - return await self._attempt_reconnect_inside_lock(force, reason, timeout_seconds) + return await self._attempt_reconnect_inside_lock( + force, reason, timeout_seconds + ) finally: self._db_reconnect_lock.release() async def start_db_health_watchdog_task(self) -> None: """Start background tasks that monitor DB health: - A periodic SELECT 1 probe that triggers reconnect on network/connection failure. - - A process-level watcher that detects engine death via waitpid thread, pidfd, or os.kill polling.""" + - A process-level watcher that detects engine death via waitpid thread, pidfd, or os.kill polling. + """ if self._db_health_watchdog_enabled is not True: verbose_proxy_logger.debug( "Prisma DB health watchdog disabled via PRISMA_HEALTH_WATCHDOG_ENABLED" @@ -4455,9 +4510,9 @@ class ProxyUpdateSpend: :MAX_LOGS_PER_INTERVAL ] # Remove the logs we're about to process - prisma_client.spend_log_transactions = prisma_client.spend_log_transactions[ - len(logs_to_process) : - ] + prisma_client.spend_log_transactions = ( + prisma_client.spend_log_transactions[len(logs_to_process) :] + ) popped_batch = True if len(logs_to_process) > 0: verbose_proxy_logger.info( @@ -4611,9 +4666,7 @@ async def update_spend_logs_job( return async with prisma_client._spend_log_transactions_lock: - logs_to_process = prisma_client.spend_log_transactions[ - :MAX_LOGS_PER_INTERVAL - ] + logs_to_process = prisma_client.spend_log_transactions[:MAX_LOGS_PER_INTERVAL] prisma_client.spend_log_transactions = prisma_client.spend_log_transactions[ len(logs_to_process) : ] @@ -4628,8 +4681,10 @@ async def update_spend_logs_job( # Guardrail/policy usage tracking (same batch, outside spend-logs update) try: - from litellm.proxy.guardrails.usage_tracking import \ - process_spend_logs_guardrail_usage + from litellm.proxy.guardrails.usage_tracking import ( + process_spend_logs_guardrail_usage, + ) + await process_spend_logs_guardrail_usage( prisma_client=prisma_client, logs_to_process=logs_to_process, @@ -4655,8 +4710,10 @@ async def _monitor_spend_logs_queue( db_writer_client: Optional HTTP handler for external spend logs endpoint proxy_logging_obj: Proxy logging object """ - from litellm.constants import (SPEND_LOG_QUEUE_POLL_INTERVAL, - SPEND_LOG_QUEUE_SIZE_THRESHOLD) + from litellm.constants import ( + SPEND_LOG_QUEUE_POLL_INTERVAL, + SPEND_LOG_QUEUE_SIZE_THRESHOLD, + ) threshold = SPEND_LOG_QUEUE_SIZE_THRESHOLD base_interval = SPEND_LOG_QUEUE_POLL_INTERVAL @@ -5177,11 +5234,12 @@ async def get_available_models_for_user( List of model names available to the user """ from litellm.proxy.auth.auth_checks import get_team_object - from litellm.proxy.auth.model_checks import (get_complete_model_list, - get_key_models, - get_team_models) - from litellm.proxy.management_endpoints.team_endpoints import \ - validate_membership + from litellm.proxy.auth.model_checks import ( + get_complete_model_list, + get_key_models, + get_team_models, + ) + from litellm.proxy.management_endpoints.team_endpoints import validate_membership # Get proxy model list and access groups if llm_router is None: