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 overrides

Browse source
This commit is contained in:
Harshit28j 2026-03-05 19:25:36 +05:30
parent a5ed93e8b7
commit f755a5e6aa
5 changed files with 74 additions and 58 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
litellm/constants.py
View file

@ -52,9 +52,7 @@ DEFAULT_IMAGE_TOKEN_COUNT = int(os.getenv("DEFAULT_IMAGE_TOKEN_COUNT", 250))
# Maximum number of base64 characters to keep in logging payloads.
# Data URIs exceeding this are replaced with a size placeholder.
# Set to 0 to disable truncation.
MAX_BASE64_LENGTH_FOR_LOGGING = int(
os.getenv("MAX_BASE64_LENGTH_FOR_LOGGING", 64)
)
MAX_BASE64_LENGTH_FOR_LOGGING = int(os.getenv("MAX_BASE64_LENGTH_FOR_LOGGING", 64))
# When true, adds detailed per-phase timing breakdown headers to responses.
# Headers: x-litellm-timing-{pre-processing,llm-api,post-processing,message-copy}-ms
@ -183,11 +181,6 @@ RUNWAYML_POLLING_TIMEOUT = int(
########## Networking constants ##############################################################
_DEFAULT_TTL_FOR_HTTPX_CLIENTS = 3600 # 1 hour, re-use the same httpx client for 1 hour
# Team-scoped model overrides feature flag
LITELLM_TEAM_MODEL_OVERRIDES: bool = (
os.getenv("LITELLM_TEAM_MODEL_OVERRIDES", "false").lower() == "true"
)
# Aiohttp connection pooling - prevents memory leaks from unbounded connection growth
# Set to 0 for unlimited (not recommended for production)
AIOHTTP_CONNECTOR_LIMIT = int(os.getenv("AIOHTTP_CONNECTOR_LIMIT", 300))
@ -1406,9 +1399,7 @@ SPECIAL_LITELLM_AUTH_TOKEN = ["ui-token"]
DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL = int(
os.getenv("DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL", 60)
)
DEFAULT_ACCESS_GROUP_CACHE_TTL = int(
os.getenv("DEFAULT_ACCESS_GROUP_CACHE_TTL", 600)
)
DEFAULT_ACCESS_GROUP_CACHE_TTL = int(os.getenv("DEFAULT_ACCESS_GROUP_CACHE_TTL", 600))
# Sentry Scrubbing Configuration
SENTRY_DENYLIST = [

47
litellm/proxy/_types.py
View file

@ -1106,7 +1106,9 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase):
raise ValueError("args is required for stdio transport")
elif transport in [MCPTransport.http, MCPTransport.sse]:
if not values.get("url") and not values.get("spec_path"):
raise ValueError("url or spec_path is required for HTTP/SSE transport")
raise ValueError(
"url or spec_path is required for HTTP/SSE transport"
)
return values
@model_validator(mode="before")
@ -1158,7 +1160,9 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase):
raise ValueError("args is required for stdio transport")
elif transport in [MCPTransport.http, MCPTransport.sse]:
if not values.get("url") and not values.get("spec_path"):
raise ValueError("url or spec_path is required for HTTP/SSE transport")
raise ValueError(
"url or spec_path is required for HTTP/SSE transport"
)
return values
@ -1409,12 +1413,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
object_permission: Optional[LiteLLM_ObjectPermissionBase] = None
@model_validator(mode="before")
@ -1437,12 +1441,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
object_permission: Optional[LiteLLM_ObjectPermissionBase] = None
@ -2253,8 +2257,8 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken):
team_max_budget: Optional[float] = None
team_soft_budget: Optional[float] = None
team_models: List = []
team_default_models: List[str] = []
team_member_models: List[str] = []
team_default_models: Optional[List[str]] = None
team_member_models: Optional[List[str]] = None
team_blocked: bool = False
soft_budget: Optional[float] = None
team_model_aliases: Optional[Dict] = None
@ -3060,7 +3064,9 @@ class SpendLogsMetadata(TypedDict):
str
] # S3/GCS object key for cold storage retrieval
litellm_overhead_time_ms: Optional[float] # LiteLLM overhead time in milliseconds
attempted_retries: Optional[int] # Number of retries attempted (0 = first attempt succeeded)
attempted_retries: Optional[
int
] # Number of retries attempted (0 = first attempt succeeded)
max_retries: Optional[int] # Max retries configured for this request
cost_breakdown: Optional[
CostBreakdown
@ -4125,10 +4131,10 @@ class SpendUpdateQueueItem(TypedDict, total=False):
class ToolDiscoveryQueueItem(TypedDict, total=False):
tool_name: str
origin: Optional[str] # MCP server name or "user_defined"
origin: Optional[str] # MCP server name or "user_defined"
created_by: Optional[str]
key_hash: Optional[str] # hash of virtual key that triggered discovery
team_id: Optional[str] # team that triggered discovery
key_hash: Optional[str] # hash of virtual key that triggered discovery
team_id: Optional[str] # team that triggered discovery
key_alias: Optional[str] # human-readable key alias
@ -4152,6 +4158,7 @@ class LiteLLM_ManagedObjectTable(LiteLLMPydanticObjectBase):
class LiteLLM_ManagedVectorStoreTable(LiteLLMPydanticObjectBase):
"""Table for managing vector stores with target_model_names support."""
unified_resource_id: str
resource_object: Optional[Any] = None # VectorStoreCreateResponse
model_mappings: Dict[str, str]

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

@ -65,7 +65,9 @@ from litellm.utils import get_utc_datetime
from .auth_checks_organization import organization_role_based_access_check
from .auth_utils import get_model_from_request
from litellm.constants import LITELLM_TEAM_MODEL_OVERRIDES
from litellm.proxy.management_endpoints.common_utils import (
_is_team_model_overrides_enabled,
)
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
@ -2068,13 +2070,11 @@ async def get_key_object(
)
# else, check db
_valid_token: Optional[BaseModel] = (
await _fetch_key_object_from_db_with_reconnect(
hashed_token=hashed_token,
prisma_client=prisma_client,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
_valid_token: Optional[BaseModel] = await _fetch_key_object_from_db_with_reconnect(
hashed_token=hashed_token,
prisma_client=prisma_client,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
if _valid_token is None:
@ -2560,11 +2560,11 @@ def can_org_access_model(
def compute_effective_team_models(
team_default_models: List[str],
team_member_models: List[str],
team_default_models: Optional[List[str]],
team_member_models: Optional[List[str]],
) -> List[str]:
"""Union of team defaults and per-user overrides, deduplicated."""
return list(set(team_default_models) | set(team_member_models))
return list(set(team_default_models or []) | set(team_member_models or []))
async def can_team_access_model(
@ -2581,7 +2581,7 @@ async def can_team_access_model(
2. If not allowed natively, falls back to access_group_ids on the team
"""
models_to_check: List[str] = team_object.models if team_object else []
if LITELLM_TEAM_MODEL_OVERRIDES and valid_token:
if _is_team_model_overrides_enabled() and valid_token:
# Only apply override logic when overrides are actually configured.
# Teams that haven't set default_models or member models continue
# to use the original team_object.models / access_group path unchanged.

30
litellm/proxy/management_endpoints/common_utils.py
View file

@ -1,3 +1,4 @@
import os
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
from litellm._logging import verbose_proxy_logger
@ -49,10 +50,7 @@ def _team_member_has_permission(
if permission not in team_obj.team_member_permissions:
return False
for member in team_obj.members_with_roles:
if (
member.user_id is not None
and member.user_id == user_api_key_dict.user_id
):
if member.user_id is not None and member.user_id == user_api_key_dict.user_id:
return True
return False
@ -332,10 +330,19 @@ async def _upsert_budget_and_membership(
"""
if max_budget is None and tpm_limit is None and rpm_limit is None:
if models is not None:
# Only updating models — do not touch the existing budget
await tx.litellm_teammembership.update(
# Only updating models — do not touch the existing budget.
# Use upsert so that the membership row is created if it
# doesn't exist yet (e.g. member added without budget/models).
await tx.litellm_teammembership.upsert(
where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}},
data={"models": models},
data={
"create": {
"user_id": user_id,
"team_id": team_id,
"models": models,
},
"update": {"models": models},
},
)
return
@ -451,6 +458,11 @@ def _update_metadata_fields(updated_kv: dict) -> None:
def _is_team_model_overrides_enabled() -> bool:
from litellm.constants import LITELLM_TEAM_MODEL_OVERRIDES
"""
Check if team-scoped model overrides feature is enabled.
return LITELLM_TEAM_MODEL_OVERRIDES
Reads os.getenv at call time (not import time) so that environment
variables set via the YAML ``environment_variables`` section — which
are applied after module import — are respected.
"""
return os.getenv("LITELLM_TEAM_MODEL_OVERRIDES", "false").lower() == "true"

18
litellm/proxy/management_helpers/utils.py
View file

@ -2,7 +2,7 @@
## Helper utils for the management endpoints (keys/users/teams)
from datetime import datetime
from functools import wraps
from typing import Optional, Tuple
from typing import List, Optional, Tuple
from fastapi import HTTPException, Request
@ -64,19 +64,19 @@ async def handle_budget_for_entity(
) -> Optional[str]:
"""
Common helper to handle budget creation/updates for entities (organizations, tags, etc).
This function:
1. Creates a new budget if budget_id is None but budget fields are provided
2. Updates an existing budget if budget fields are provided and budget_id exists
3. Returns the budget_id to use (existing or newly created)
Args:
data: The request object (e.g., TagNewRequest, NewOrganizationRequest, etc.) containing budget fields
existing_budget_id: The existing budget_id if updating an entity, None if creating new
user_api_key_dict: User authentication info
prisma_client: Database client
litellm_proxy_admin_name: Admin name for audit trail
Returns:
Optional[str]: The budget_id to use, or None if no budget was created/updated
"""
@ -88,7 +88,9 @@ async def handle_budget_for_entity(
budget_params = LiteLLM_BudgetTable.model_fields.keys()
# Extract budget fields from data
_json_data = data.model_dump(exclude_none=True) if hasattr(data, "model_dump") else data
_json_data = (
data.model_dump(exclude_none=True) if hasattr(data, "model_dump") else data
)
_budget_data = {k: v for k, v in _json_data.items() if k in budget_params}
# Check if budget_id is explicitly provided in the data
@ -221,7 +223,11 @@ async def add_new_member(
else:
_budget_id = default_team_budget_id
if (_budget_id or models) and returned_user is not None and returned_user.user_id is not None:
if (
(_budget_id or models)
and returned_user is not None
and returned_user.user_id is not None
):
create_data: Dict[str, Any] = {
"team_id": team_id,
"user_id": returned_user.user_id,