diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/internal_user_endpoints.py b/enterprise/litellm_enterprise/proxy/management_endpoints/internal_user_endpoints.py index d17946171b..e60b4d6990 100644 --- a/enterprise/litellm_enterprise/proxy/management_endpoints/internal_user_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/internal_user_endpoints.py @@ -2,6 +2,8 @@ Enterprise internal user management endpoints """ +import os + from fastapi import APIRouter, Depends, HTTPException from litellm.proxy._types import UserAPIKeyAuth @@ -21,7 +23,7 @@ async def available_enterprise_users( """ For keys with `max_users` set, return the list of users that are allowed to use the key. """ - from litellm.proxy._types import CommonProxyErrors + from litellm.proxy._types import CommonProxyErrors, EnterpriseLicenseData from litellm.proxy.proxy_server import ( premium_user, premium_user_data, @@ -34,10 +36,14 @@ async def available_enterprise_users( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - if premium_user is None: - raise HTTPException( - status_code=500, detail={"error": CommonProxyErrors.not_premium_user.value} - ) + if not premium_user: + # check if SSO is enabled - show 5 user limit + from litellm.proxy.auth.auth_utils import _has_user_setup_sso + + if _has_user_setup_sso(): + premium_user_data = EnterpriseLicenseData( + max_users=5, + ) # Count number of rows in LiteLLM_UserTable user_count = await prisma_client.db.litellm_usertable.count() diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index e8a528845f..202e7f0759 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -17691,22 +17691,6 @@ }, "supports_tool_choice": true }, - "databricks/databricks-meta-llama-3-1-70b-instruct": { - "max_tokens": 128000, - "max_input_tokens": 128000, - "max_output_tokens": 128000, - "input_cost_per_token": 1.00002e-06, - "input_dbu_cost_per_token": 1.4286e-05, - "output_cost_per_token": 2.99999e-06, - "output_dbu_cost_per_token": 4.2857e-05, - "litellm_provider": "databricks", - "mode": "chat", - "source": "https://www.databricks.com/product/pricing/foundation-model-serving", - "metadata": { - "notes": "Input/output cost per token is dbu cost * $0.070, based on databricks Llama 3.1 70B conversion. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." - }, - "supports_tool_choice": true - }, "databricks/databricks-meta-llama-3-3-70b-instruct": { "max_tokens": 128000, "max_input_tokens": 128000, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 3e155135d7..2239526a31 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -530,7 +530,7 @@ class LiteLLMRoutes(enum.Enum): # Routes accessible by Admin Viewer (read-only admin access) admin_viewer_routes = [ "/user/list", - "/user/available_users", + "/user/available_users", "/user/available_roles", "/user/daily/activity", "/team/daily/activity", @@ -540,7 +540,10 @@ class LiteLLMRoutes(enum.Enum): # All routes accesible by an Org Admin org_admin_allowed_routes = ( - org_admin_only_routes + management_routes + self_managed_routes + admin_viewer_routes + org_admin_only_routes + + management_routes + + self_managed_routes + + admin_viewer_routes ) @@ -585,13 +588,14 @@ class LiteLLMPromptInjectionParams(LiteLLMPydanticObjectBase): ######### Request Class Definition ###### class ProxyChatCompletionRequest(LiteLLMPydanticObjectBase): """ - Pydantic model for chat completion requests that includes both OpenAI standard fields + Pydantic model for chat completion requests that includes both OpenAI standard fields and LiteLLM-specific parameters. This replaces the previous TypedDict version. """ + # Required fields (from ChatCompletionRequest) model: str messages: List[AllMessageValues] - + # Standard OpenAI completion parameters (all optional) frequency_penalty: Optional[float] = None logit_bias: Optional[Dict[str, float]] = None @@ -614,10 +618,10 @@ class ProxyChatCompletionRequest(LiteLLMPydanticObjectBase): functions: Optional[List[Dict[str, Any]]] = None user: Optional[str] = None stream: Optional[bool] = None - + # LiteLLM-specific metadata param (from original ChatCompletionRequest) metadata: Optional[Dict[str, Any]] = None - + # Optional LiteLLM params guardrails: Optional[List[str]] = None caching: Optional[bool] = None @@ -1873,7 +1877,8 @@ class UserAPIKeyAuth( key_alias=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, team_alias=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, ) - + + class UserInfoResponse(LiteLLMPydanticObjectBase): user_id: Optional[str] user_info: Optional[Union[dict, BaseModel]] @@ -2120,7 +2125,6 @@ class TokenCountRequest(LiteLLMPydanticObjectBase): Anthropic token counting endpoint uses /messages """ - contents: Optional[List[dict]] = None """ Google /countTokens endpoint expects contents to be a list of dicts with the following structure: @@ -2265,7 +2269,7 @@ class AllCallbacks(LiteLLMPydanticObjectBase): braintrust: CallbackOnUI = CallbackOnUI( litellm_callback_name="braintrust", - litellm_callback_params=["BRAINTRUST_API_KEY","BRAINTRUST_API_BASE"], + litellm_callback_params=["BRAINTRUST_API_KEY", "BRAINTRUST_API_BASE"], ui_callback_name="Braintrust", ) @@ -2319,7 +2323,9 @@ class SpendLogsMetadata(TypedDict): error_information: Optional[StandardLoggingPayloadErrorInformation] usage_object: Optional[dict] model_map_information: Optional[StandardLoggingModelInformation] - cold_storage_object_key: Optional[str] # S3/GCS object key for cold storage retrieval + cold_storage_object_key: Optional[ + str + ] # S3/GCS object key for cold storage retrieval class SpendLogsPayload(TypedDict): @@ -2646,7 +2652,7 @@ class LiteLLM_TeamMembership(LiteLLMPydanticObjectBase): if self.litellm_budget_table is not None: return self.litellm_budget_table.rpm_limit return None - + def safe_get_team_member_tpm_limit(self) -> Optional[int]: if self.litellm_budget_table is not None: return self.litellm_budget_table.tpm_limit @@ -2763,14 +2769,11 @@ class TeamMemberUpdateRequest(TeamMemberDeleteRequest): max_budget_in_team: Optional[float] = None role: Optional[Literal["admin", "user"]] = None tpm_limit: Optional[int] = Field( - default=None, - description="Tokens per minute limit for this team member" + default=None, description="Tokens per minute limit for this team member" ) rpm_limit: Optional[int] = Field( - default=None, - description="Requests per minute limit for this team member" + default=None, description="Requests per minute limit for this team member" ) - class TeamMemberUpdateResponse(MemberUpdateResponse): diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 451f110a89..98165a5a0c 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -77,7 +77,9 @@ router = APIRouter() @router.get("/sso/key/generate", tags=["experimental"], include_in_schema=False) -async def google_login(request: Request, source: Optional[str] = None, key: Optional[str] = None): # noqa: PLR0915 +async def google_login( + request: Request, source: Optional[str] = None, key: Optional[str] = None +): # noqa: PLR0915 """ Create Proxy API Keys using Google Workspace SSO. Requires setting PROXY_BASE_URL in .env PROXY_BASE_URL should be the your deployed proxy endpoint, e.g. PROXY_BASE_URL="https://litellm-production-7002.up.railway.app/" @@ -85,6 +87,7 @@ async def google_login(request: Request, source: Optional[str] = None, key: Opti """ from litellm.proxy.proxy_server import ( premium_user, + prisma_client, user_custom_ui_sso_sign_in_handler, ) @@ -106,12 +109,23 @@ async def google_login(request: Request, source: Optional[str] = None, key: Opti or generic_client_id is not None ): if premium_user is not True: - raise ProxyException( - message="You must be a LiteLLM Enterprise user to use SSO. If you have a license please set `LITELLM_LICENSE` in your env. If you want to obtain a license meet with us here: https://calendly.com/d/4mp-gd3-k5k/litellm-1-1-onboarding-chat You are seeing this error message because You set one of `MICROSOFT_CLIENT_ID`, `GOOGLE_CLIENT_ID`, or `GENERIC_CLIENT_ID` in your env. Please unset this", - type=ProxyErrorTypes.auth_error, - param="premium_user", - code=status.HTTP_403_FORBIDDEN, - ) + # Check if under 'free SSO user' limit + if prisma_client is not None: + total_users = await prisma_client.db.litellm_usertable.count() + if total_users and total_users > 5: + raise ProxyException( + message="You must be a LiteLLM Enterprise user to use SSO for more than 5 users. If you have a license please set `LITELLM_LICENSE` in your env. If you want to obtain a license meet with us here: https://calendly.com/d/4mp-gd3-k5k/litellm-1-1-onboarding-chat You are seeing this error message because You set one of `MICROSOFT_CLIENT_ID`, `GOOGLE_CLIENT_ID`, or `GENERIC_CLIENT_ID` in your env. Please unset this", + type=ProxyErrorTypes.auth_error, + param="premium_user", + code=status.HTTP_403_FORBIDDEN, + ) + else: + raise ProxyException( + message=CommonProxyErrors.db_not_connected_error.value, + type=ProxyErrorTypes.auth_error, + param="premium_user", + code=status.HTTP_403_FORBIDDEN, + ) ####### Detect DB + MASTER KEY in .env ####### missing_env_vars = show_missing_vars_in_env() @@ -124,7 +138,7 @@ async def google_login(request: Request, source: Optional[str] = None, key: Opti request=request, sso_callback_route="sso/callback", ) - + # Store CLI key in state for OAuth flow cli_state: Optional[str] = SSOAuthenticationHandler._get_cli_state( source=source, @@ -137,11 +151,14 @@ async def google_login(request: Request, source: Optional[str] = None, key: Opti from litellm_enterprise.proxy.auth.custom_sso_handler import ( EnterpriseCustomSSOHandler, ) + return await EnterpriseCustomSSOHandler.handle_custom_ui_sso_sign_in( request=request, ) except ImportError: - raise ValueError("Enterprise features are not available. Custom UI SSO sign-in requires LiteLLM Enterprise.") + raise ValueError( + "Enterprise features are not available. Custom UI SSO sign-in requires LiteLLM Enterprise." + ) # Check if we should use SSO handler if ( @@ -525,15 +542,16 @@ async def check_and_update_if_proxy_admin_id( async def auth_callback(request: Request, state: Optional[str] = None): # noqa: PLR0915 """Verify login""" verbose_proxy_logger.info(f"Starting SSO callback with state: {state}") - + # Check if this is a CLI login (state starts with our CLI prefix) from litellm.constants import LITELLM_CLI_SESSION_TOKEN_PREFIX + if state and state.startswith(f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:"): # Extract the key ID from the state key_id = state.split(":", 1)[1] verbose_proxy_logger.info(f"CLI SSO callback detected for key: {key_id}") return await cli_sso_callback(request, key=key_id) - + from litellm.proxy._types import LiteLLM_JWTAuth from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.proxy_server import ( @@ -608,7 +626,7 @@ async def auth_callback(request: Request, state: Optional[str] = None): # noqa: status_code=401, detail="Result not returned by SSO provider.", ) - + return await SSOAuthenticationHandler.get_redirect_response_from_openid( result=result, request=request, @@ -618,28 +636,26 @@ async def auth_callback(request: Request, state: Optional[str] = None): # noqa: ) - - async def cli_sso_callback(request: Request, key: Optional[str] = None): """CLI SSO callback - generates the key with pre-specified ID""" verbose_proxy_logger.info(f"CLI SSO callback for key: {key}") - + from litellm.proxy.management_endpoints.key_management_endpoints import ( generate_key_helper_fn, ) from litellm.proxy.proxy_server import prisma_client - - if not key or not key.startswith('sk-'): + + if not key or not key.startswith("sk-"): raise HTTPException( status_code=400, - detail="Invalid key parameter. Must be a valid key ID starting with 'sk-'" + detail="Invalid key parameter. Must be a valid key ID starting with 'sk-'", ) - + if prisma_client is None: raise HTTPException( status_code=500, detail=CommonProxyErrors.db_not_connected_error.value ) - + # Generate a simple key for CLI usage with the pre-specified key ID try: await generate_key_helper_fn( @@ -653,63 +669,57 @@ async def cli_sso_callback(request: Request, key: Optional[str] = None): table_name="key", token=key, # Use the pre-specified key ID ) - + verbose_proxy_logger.info(f"Generated CLI key: {key}") - + # Return success page from fastapi.responses import HTMLResponse from litellm.proxy.common_utils.html_forms.cli_sso_success import ( render_cli_sso_success_page, ) - + html_content = render_cli_sso_success_page() return HTMLResponse(content=html_content, status_code=200) - + except Exception as e: verbose_proxy_logger.error(f"Error generating CLI key: {e}") - raise HTTPException( - status_code=500, - detail=f"Failed to generate key: {str(e)}" - ) + raise HTTPException(status_code=500, detail=f"Failed to generate key: {str(e)}") @router.get("/sso/cli/poll/{key_id}", tags=["experimental"], include_in_schema=False) async def cli_poll_key(key_id: str): """CLI polling endpoint - checks if key exists in DB""" from litellm.proxy.proxy_server import prisma_client - - if not key_id.startswith('sk-'): - raise HTTPException( - status_code=400, - detail="Invalid key ID format" - ) - + + if not key_id.startswith("sk-"): + raise HTTPException(status_code=400, detail="Invalid key ID format") + if prisma_client is None: raise HTTPException( status_code=500, detail=CommonProxyErrors.db_not_connected_error.value ) - + try: # Check if key exists in database from litellm.proxy.utils import hash_token + hashed_token = hash_token(key_id) - + key_obj = await prisma_client.db.litellm_verificationtoken.find_unique( where={"token": hashed_token} ) - + if key_obj: verbose_proxy_logger.info(f"CLI key found: {key_id}") return {"status": "ready", "key": key_id} else: return {"status": "pending"} - + except Exception as e: verbose_proxy_logger.error(f"Error polling for CLI key: {e}") raise HTTPException( - status_code=500, - detail=f"Error checking key status: {str(e)}" + status_code=500, detail=f"Error checking key status: {str(e)}" ) @@ -811,6 +821,7 @@ class SSOAuthenticationHandler: """ Handler for SSO Authentication across all SSO providers """ + @staticmethod async def get_sso_login_redirect( redirect_url: str, @@ -1163,7 +1174,6 @@ class SSOAuthenticationHandler: _new_team_request.update(_default_team_params) team_request = NewTeamRequest(**_new_team_request) return team_request - @staticmethod def _get_cli_state(source: Optional[str], key: Optional[str]) -> Optional[str]: @@ -1176,13 +1186,15 @@ class SSOAuthenticationHandler: LITELLM_CLI_SESSION_TOKEN_PREFIX, LITELLM_CLI_SOURCE_IDENTIFIER, ) - return f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:{key}" if source == LITELLM_CLI_SOURCE_IDENTIFIER and key else None - - + return ( + f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:{key}" + if source == LITELLM_CLI_SOURCE_IDENTIFIER and key + else None + ) @staticmethod - async def get_redirect_response_from_openid( # noqa: PLR0915 + async def get_redirect_response_from_openid( # noqa: PLR0915 result: Union[OpenID, dict, CustomOpenID], request: Request, received_response: Optional[dict] = None, @@ -1202,14 +1214,18 @@ class SSOAuthenticationHandler: ) from litellm.proxy.utils import get_prisma_client_or_throw from litellm.types.proxy.ui_sso import ReturnedUITokenObject - prisma_client = get_prisma_client_or_throw("Prisma client is None, connect a database to your proxy") + prisma_client = get_prisma_client_or_throw( + "Prisma client is None, connect a database to your proxy" + ) # User is Authe'd in - generate key for the UI to access Proxy verbose_proxy_logger.info(f"SSO callback result: {result}") user_email: Optional[str] = getattr(result, "email", None) - user_id: Optional[str] = getattr(result, "id", None) if result is not None else None + user_id: Optional[str] = ( + getattr(result, "id", None) if result is not None else None + ) if user_email is not None and os.getenv("ALLOWED_EMAIL_DOMAINS") is not None: email_domain = user_email.split("@")[1] @@ -1394,7 +1410,8 @@ class SSOAuthenticationHandler: redirect_response = RedirectResponse(url=litellm_dashboard_ui, status_code=303) redirect_response.set_cookie(key="token", value=jwt_token) return redirect_response - + + class MicrosoftSSOHandler: """ Handles Microsoft SSO callback response and returns a CustomOpenID object diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index e8a528845f..202e7f0759 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -17691,22 +17691,6 @@ }, "supports_tool_choice": true }, - "databricks/databricks-meta-llama-3-1-70b-instruct": { - "max_tokens": 128000, - "max_input_tokens": 128000, - "max_output_tokens": 128000, - "input_cost_per_token": 1.00002e-06, - "input_dbu_cost_per_token": 1.4286e-05, - "output_cost_per_token": 2.99999e-06, - "output_dbu_cost_per_token": 4.2857e-05, - "litellm_provider": "databricks", - "mode": "chat", - "source": "https://www.databricks.com/product/pricing/foundation-model-serving", - "metadata": { - "notes": "Input/output cost per token is dbu cost * $0.070, based on databricks Llama 3.1 70B conversion. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." - }, - "supports_tool_choice": true - }, "databricks/databricks-meta-llama-3-3-70b-instruct": { "max_tokens": 128000, "max_input_tokens": 128000,