scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 152

fix: team scoped models overrides req changes

Browse source
This commit is contained in:
Harshit28j 2026-03-05 19:55:39 +05:30
parent f755a5e6aa
commit 956daa2760
3 changed files with 318 additions and 162 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
litellm/proxy/auth/auth_checks.py
View file

@ -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

255
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -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

210
litellm/proxy/utils.py
View file

@ -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: