diff --git a/litellm/constants.py b/litellm/constants.py index a25aac293fd..3df2fb9c2bc 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -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 = [ diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c39d9b5e359..1ba7bb7a929 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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] diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 2bcb6dcb59c..97724f6e14c 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -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. diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index 29b5b31a486..e38dbf366ec 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -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" diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index 76f1ad50d5f..1edb4fe9df1 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -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,