forked from phoenix/litellm-mirror
Merge branch 'main' into litellm_google_text_moderation
This commit is contained in:
Krish Dholakia
GitHub
commit
b41cdf598b
29 changed files with 489 additions and 62 deletions
|
|
@ -93,6 +93,7 @@ from litellm.proxy.utils import (
|
|||
html_form,
|
||||
_read_request_body,
|
||||
_is_valid_team_configs,
|
||||
_is_user_proxy_admin,
|
||||
)
|
||||
from litellm.proxy.secret_managers.google_kms import load_google_kms
|
||||
import pydantic
|
||||
|
|
@ -379,6 +380,11 @@ async def user_api_key_auth(
|
|||
# 3. If 'user' passed to /chat/completions, /embeddings endpoint is in budget
|
||||
# 4. If token is expired
|
||||
# 5. If token spend is under Budget for the token
|
||||
# 6. If token spend per model is under budget per model
|
||||
|
||||
request_data = await _read_request_body(
|
||||
request=request
|
||||
) # request data, used across all checks. Making this easily available
|
||||
|
||||
# Check 1. If token can call model
|
||||
litellm.model_alias_map = valid_token.aliases
|
||||
|
|
@ -453,7 +459,6 @@ async def user_api_key_auth(
|
|||
if (
|
||||
litellm.max_user_budget is not None
|
||||
): # Check if 'user' passed in /chat/completions is in budget, only checked if litellm.max_user_budget is set
|
||||
request_data = await _read_request_body(request=request)
|
||||
user_passed_to_chat_completions = request_data.get("user", None)
|
||||
if user_passed_to_chat_completions is not None:
|
||||
user_id_list.append(user_passed_to_chat_completions)
|
||||
|
|
@ -499,11 +504,7 @@ async def user_api_key_auth(
|
|||
continue
|
||||
assert isinstance(_user, dict)
|
||||
# check if user is admin #
|
||||
if (
|
||||
_user.get("user_role", None) is not None
|
||||
and _user.get("user_role") == "proxy_admin"
|
||||
):
|
||||
return UserAPIKeyAuth(api_key=master_key)
|
||||
|
||||
# Token exists, not expired now check if its in budget for the user
|
||||
user_max_budget = _user.get("max_budget", None)
|
||||
user_current_spend = _user.get("spend", None)
|
||||
|
|
@ -590,6 +591,25 @@ async def user_api_key_auth(
|
|||
f"ExceededTokenBudget: Current spend for token: {valid_token.spend}; Max Budget for Token: {valid_token.max_budget}"
|
||||
)
|
||||
|
||||
# Check 5. Token Model Spend is under Model budget
|
||||
max_budget_per_model = valid_token.model_max_budget
|
||||
spend_per_model = valid_token.model_spend
|
||||
|
||||
if max_budget_per_model is not None and spend_per_model is not None:
|
||||
current_model = request_data.get("model")
|
||||
if current_model is not None:
|
||||
current_model_spend = spend_per_model.get(current_model, None)
|
||||
current_model_budget = max_budget_per_model.get(current_model, None)
|
||||
|
||||
if (
|
||||
current_model_spend is not None
|
||||
and current_model_budget is not None
|
||||
):
|
||||
if current_model_spend > current_model_budget:
|
||||
raise Exception(
|
||||
f"ExceededModelBudget: Current spend for model: {current_model_spend}; Max Budget for Model: {current_model_budget}"
|
||||
)
|
||||
|
||||
# Token passed all checks
|
||||
api_key = valid_token.token
|
||||
|
||||
|
|
@ -619,11 +639,15 @@ async def user_api_key_auth(
|
|||
)
|
||||
)
|
||||
if (
|
||||
route.startswith("/key/")
|
||||
or route.startswith("/user/")
|
||||
or route.startswith("/model/")
|
||||
or route.startswith("/spend/")
|
||||
) and (not is_master_key_valid):
|
||||
(
|
||||
route.startswith("/key/")
|
||||
or route.startswith("/user/")
|
||||
or route.startswith("/model/")
|
||||
or route.startswith("/spend/")
|
||||
)
|
||||
and (not is_master_key_valid)
|
||||
and (not _is_user_proxy_admin(user_id_information))
|
||||
):
|
||||
allow_user_auth = False
|
||||
if (
|
||||
general_settings.get("allow_user_auth", False) == True
|
||||
|
|
@ -715,9 +739,12 @@ async def user_api_key_auth(
|
|||
# Do something if the current route starts with any of the allowed routes
|
||||
pass
|
||||
else:
|
||||
raise Exception(
|
||||
f"This key is made for LiteLLM UI, Tried to access route: {route}. Not allowed"
|
||||
)
|
||||
if _is_user_proxy_admin(user_id_information):
|
||||
pass
|
||||
else:
|
||||
raise Exception(
|
||||
f"This key is made for LiteLLM UI, Tried to access route: {route}. Not allowed"
|
||||
)
|
||||
return UserAPIKeyAuth(api_key=api_key, **valid_token_dict)
|
||||
except Exception as e:
|
||||
# verbose_proxy_logger.debug(f"An exception occurred - {traceback.format_exc()}")
|
||||
|
|
@ -3141,6 +3168,19 @@ async def generate_key_fn(
|
|||
- permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false}
|
||||
- model_max_budget: Optional[dict] - key-specific model budget in USD. Example - {"text-davinci-002": 0.5, "gpt-3.5-turbo": 0.5}. IF null or {} then no model specific budget.
|
||||
|
||||
Examples:
|
||||
|
||||
1. Allow users to turn on/off pii masking
|
||||
|
||||
```bash
|
||||
curl --location 'http://0.0.0.0:8000/key/generate' \
|
||||
--header 'Authorization: Bearer sk-1234' \
|
||||
--header 'Content-Type: application/json' \
|
||||
--data '{
|
||||
"permissions": {"allow_pii_controls": true}
|
||||
}'
|
||||
```
|
||||
|
||||
Returns:
|
||||
- key: (str) The generated api key
|
||||
- expires: (datetime) Datetime object for when key expires.
|
||||
|
|
@ -4933,7 +4973,7 @@ async def auth_callback(request: Request):
|
|||
if user_id is None:
|
||||
user_id = getattr(result, "first_name", "") + getattr(result, "last_name", "")
|
||||
response = await generate_key_helper_fn(
|
||||
**{"duration": "1hr", "key_max_budget": 0, "models": [], "aliases": {}, "config": {}, "spend": 0, "user_id": user_id, "team_id": "litellm-dashboard", "user_email": user_email} # type: ignore
|
||||
**{"duration": "1hr", "key_max_budget": 0.01, "models": [], "aliases": {}, "config": {}, "spend": 0, "user_id": user_id, "team_id": "litellm-dashboard", "user_email": user_email} # type: ignore
|
||||
)
|
||||
key = response["token"] # type: ignore
|
||||
user_id = response["user_id"] # type: ignore
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue