forked from phoenix/litellm-mirror
Merge branch 'main' into litellm_ui_invite_for_spend_viewer
This commit is contained in:
Ishaan Jaff
GitHub
commit
6623c1e017
8 changed files with 299 additions and 126 deletions
|
|
@ -3247,6 +3247,9 @@ def _duration_in_seconds(duration: str):
|
|||
|
||||
|
||||
async def generate_key_helper_fn(
|
||||
request_type: Literal[
|
||||
"user", "key"
|
||||
], # identifies if this request is from /user/new or /key/generate
|
||||
duration: Optional[str],
|
||||
models: list,
|
||||
aliases: dict,
|
||||
|
|
@ -3278,6 +3281,7 @@ async def generate_key_helper_fn(
|
|||
teams: Optional[list] = None,
|
||||
organization_id: Optional[str] = None,
|
||||
table_name: Optional[Literal["key", "user"]] = None,
|
||||
send_invite_email: Optional[bool] = None,
|
||||
):
|
||||
global prisma_client, custom_db_client, user_api_key_cache, litellm_proxy_admin_name, premium_user
|
||||
|
||||
|
|
@ -3312,7 +3316,7 @@ async def generate_key_helper_fn(
|
|||
permissions_json = json.dumps(permissions)
|
||||
metadata_json = json.dumps(metadata)
|
||||
model_max_budget_json = json.dumps(model_max_budget)
|
||||
user_role = user_role or "app_user"
|
||||
user_role = user_role
|
||||
tpm_limit = tpm_limit
|
||||
rpm_limit = rpm_limit
|
||||
allowed_cache_controls = allowed_cache_controls
|
||||
|
|
@ -3378,7 +3382,7 @@ async def generate_key_helper_fn(
|
|||
"get_spend_routes" in saved_token["permissions"]
|
||||
and premium_user != True
|
||||
):
|
||||
raise Exception(
|
||||
raise ValueError(
|
||||
"get_spend_routes permission is only available for LiteLLM Enterprise users"
|
||||
)
|
||||
|
||||
|
|
@ -3435,6 +3439,10 @@ async def generate_key_helper_fn(
|
|||
|
||||
# Add budget related info in key_data - this ensures it's returned
|
||||
key_data["budget_id"] = budget_id
|
||||
|
||||
if request_type == "user":
|
||||
# if this is a /user/new request update the key_date with user_data fields
|
||||
key_data.update(user_data)
|
||||
return key_data
|
||||
|
||||
|
||||
|
|
@ -3782,6 +3790,7 @@ async def startup_event():
|
|||
)
|
||||
asyncio.create_task(
|
||||
generate_key_helper_fn(
|
||||
request_type="user",
|
||||
duration=None,
|
||||
models=[],
|
||||
aliases={},
|
||||
|
|
@ -3804,6 +3813,7 @@ async def startup_event():
|
|||
# add proxy budget to db in the user table
|
||||
asyncio.create_task(
|
||||
generate_key_helper_fn(
|
||||
request_type="user",
|
||||
user_id=litellm_proxy_budget_name,
|
||||
duration=None,
|
||||
models=[],
|
||||
|
|
@ -3826,7 +3836,13 @@ async def startup_event():
|
|||
if custom_db_client is not None and master_key is not None:
|
||||
# add master key to db
|
||||
await generate_key_helper_fn(
|
||||
duration=None, models=[], aliases={}, config={}, spend=0, token=master_key
|
||||
request_type="key",
|
||||
duration=None,
|
||||
models=[],
|
||||
aliases={},
|
||||
config={},
|
||||
spend=0,
|
||||
token=master_key,
|
||||
)
|
||||
|
||||
### CHECK IF VIEW EXISTS ###
|
||||
|
|
@ -6034,6 +6050,7 @@ async def generate_key_fn(
|
|||
- aliases: Optional[dict] - Any alias mappings, on top of anything in the config.yaml model list. - https://docs.litellm.ai/docs/proxy/virtual_keys#managing-auth---upgradedowngrade-models
|
||||
- config: Optional[dict] - any key-specific configs, overrides config in config.yaml
|
||||
- spend: Optional[int] - Amount spent by key. Default is 0. Will be updated by proxy whenever key is used. https://docs.litellm.ai/docs/proxy/virtual_keys#managing-auth---tracking-spend
|
||||
- send_invite_email: Optional[bool] - Whether to send an invite email to the user_id, with the generate key
|
||||
- max_budget: Optional[float] - Specify max budget for a given key.
|
||||
- max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x.
|
||||
- metadata: Optional[dict] - Metadata for key, store information for key. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" }
|
||||
|
|
@ -6162,29 +6179,37 @@ async def generate_key_fn(
|
|||
if "budget_duration" in data_json:
|
||||
data_json["key_budget_duration"] = data_json.pop("budget_duration", None)
|
||||
|
||||
response = await generate_key_helper_fn(**data_json, table_name="key")
|
||||
response = await generate_key_helper_fn(
|
||||
request_type="key", **data_json, table_name="key"
|
||||
)
|
||||
|
||||
response["soft_budget"] = (
|
||||
data.soft_budget
|
||||
) # include the user-input soft budget in the response
|
||||
event = WebhookEvent(
|
||||
event="key_created",
|
||||
event_group="key",
|
||||
event_message=f"API Key Created",
|
||||
token=response.get("token", ""),
|
||||
spend=response.get("spend", 0.0),
|
||||
max_budget=response.get("max_budget", 0.0),
|
||||
user_id=response.get("user_id", None),
|
||||
team_id=response.get("team_id", "Default Team"),
|
||||
key_alias=response.get("key_alias", None),
|
||||
)
|
||||
|
||||
# If user configured email alerting - send an Email letting their end-user know the key was created
|
||||
asyncio.create_task(
|
||||
proxy_logging_obj.slack_alerting_instance.send_key_created_email(
|
||||
webhook_event=event,
|
||||
if data.send_invite_email is True:
|
||||
if "email" not in general_settings.get("alerting", []):
|
||||
raise ValueError(
|
||||
"Email alerting not setup on config.yaml. Please set `alerting=['email']. \nDocs: https://docs.litellm.ai/docs/proxy/email`"
|
||||
)
|
||||
event = WebhookEvent(
|
||||
event="key_created",
|
||||
event_group="key",
|
||||
event_message=f"API Key Created",
|
||||
token=response.get("token", ""),
|
||||
spend=response.get("spend", 0.0),
|
||||
max_budget=response.get("max_budget", 0.0),
|
||||
user_id=response.get("user_id", None),
|
||||
team_id=response.get("team_id", "Default Team"),
|
||||
key_alias=response.get("key_alias", None),
|
||||
)
|
||||
|
||||
# If user configured email alerting - send an Email letting their end-user know the key was created
|
||||
asyncio.create_task(
|
||||
proxy_logging_obj.slack_alerting_instance.send_key_created_or_user_invited_email(
|
||||
webhook_event=event,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
return GenerateKeyResponse(**response)
|
||||
except Exception as e:
|
||||
|
|
@ -8146,6 +8171,7 @@ async def new_user(data: NewUserRequest):
|
|||
- teams: Optional[list] - specify a list of team id's a user belongs to.
|
||||
- organization_id: Optional[str] - specify the org a user belongs to.
|
||||
- user_email: Optional[str] - Specify a user email.
|
||||
- send_invite_email: Optional[bool] - Specify if an invite email should be sent.
|
||||
- user_role: Optional[str] - Specify a user role - "admin", "app_owner", "app_user"
|
||||
- max_budget: Optional[float] - Specify max budget for a given user.
|
||||
- models: Optional[list] - Model_name's a user is allowed to call. (if empty, key is allowed to call all models)
|
||||
|
|
@ -8167,7 +8193,7 @@ async def new_user(data: NewUserRequest):
|
|||
data_json["table_name"] = (
|
||||
"user" # only create a user, don't create key if 'auto_create_key' set to False
|
||||
)
|
||||
response = await generate_key_helper_fn(**data_json)
|
||||
response = await generate_key_helper_fn(request_type="user", **data_json)
|
||||
|
||||
# Admin UI Logic
|
||||
# if team_id passed add this user to the team
|
||||
|
|
@ -8182,11 +8208,42 @@ async def new_user(data: NewUserRequest):
|
|||
),
|
||||
)
|
||||
)
|
||||
|
||||
if data.send_invite_email is True:
|
||||
# check if user has setup email alerting
|
||||
if "email" not in general_settings.get("alerting", []):
|
||||
raise ValueError(
|
||||
"Email alerting not setup on config.yaml. Please set `alerting=['email']. \nDocs: https://docs.litellm.ai/docs/proxy/email`"
|
||||
)
|
||||
|
||||
event = WebhookEvent(
|
||||
event="internal_user_created",
|
||||
event_group="internal_user",
|
||||
event_message=f"Welcome to LiteLLM Proxy",
|
||||
token=response.get("token", ""),
|
||||
spend=response.get("spend", 0.0),
|
||||
max_budget=response.get("max_budget", 0.0),
|
||||
user_id=response.get("user_id", None),
|
||||
user_email=response.get("user_email", None),
|
||||
team_id=response.get("team_id", "Default Team"),
|
||||
key_alias=response.get("key_alias", None),
|
||||
)
|
||||
|
||||
# If user configured email alerting - send an Email letting their end-user know the key was created
|
||||
asyncio.create_task(
|
||||
proxy_logging_obj.slack_alerting_instance.send_key_created_or_user_invited_email(
|
||||
webhook_event=event,
|
||||
)
|
||||
)
|
||||
|
||||
return NewUserResponse(
|
||||
key=response.get("token", ""),
|
||||
expires=response.get("expires", None),
|
||||
max_budget=response["max_budget"],
|
||||
user_id=response["user_id"],
|
||||
user_role=response.get("user_role", None),
|
||||
user_email=response.get("user_email", None),
|
||||
teams=response.get("teams", None),
|
||||
team_id=response.get("team_id", None),
|
||||
metadata=response.get("metadata", None),
|
||||
models=response.get("models", None),
|
||||
|
|
@ -8243,11 +8300,13 @@ async def user_auth(request: Request):
|
|||
if response is not None:
|
||||
user_id = response.user_id
|
||||
response = await generate_key_helper_fn(
|
||||
**{"duration": "24hr", "models": [], "aliases": {}, "config": {}, "spend": 0, "user_id": user_id} # type: ignore
|
||||
request_type="key",
|
||||
**{"duration": "24hr", "models": [], "aliases": {}, "config": {}, "spend": 0, "user_id": user_id}, # type: ignore
|
||||
)
|
||||
else: ### else - create new user
|
||||
response = await generate_key_helper_fn(
|
||||
**{"duration": "24hr", "models": [], "aliases": {}, "config": {}, "spend": 0, "user_email": user_email} # type: ignore
|
||||
request_type="key",
|
||||
**{"duration": "24hr", "models": [], "aliases": {}, "config": {}, "spend": 0, "user_email": user_email}, # type: ignore
|
||||
)
|
||||
|
||||
base_url = os.getenv("LITELLM_HOSTED_UI", "https://dashboard.litellm.ai/")
|
||||
|
|
@ -11739,7 +11798,8 @@ async def login(request: Request):
|
|||
)
|
||||
if os.getenv("DATABASE_URL") is not None:
|
||||
response = await generate_key_helper_fn(
|
||||
**{"user_role": LitellmUserRoles.PROXY_ADMIN, "duration": "2hr", "key_max_budget": 5, "models": [], "aliases": {}, "config": {}, "spend": 0, "user_id": key_user_id, "team_id": "litellm-dashboard"} # type: ignore
|
||||
request_type="key",
|
||||
**{"user_role": LitellmUserRoles.PROXY_ADMIN, "duration": "2hr", "key_max_budget": 5, "models": [], "aliases": {}, "config": {}, "spend": 0, "user_id": key_user_id, "team_id": "litellm-dashboard"}, # type: ignore
|
||||
)
|
||||
else:
|
||||
raise ProxyException(
|
||||
|
|
@ -12155,8 +12215,11 @@ async def auth_callback(request: Request):
|
|||
verbose_proxy_logger.info(
|
||||
f"user_defined_values for creating ui key: {user_defined_values}"
|
||||
)
|
||||
|
||||
default_ui_key_values.update(user_defined_values)
|
||||
default_ui_key_values["request_type"] = "key"
|
||||
response = await generate_key_helper_fn(
|
||||
**default_ui_key_values, **user_defined_values # type: ignore
|
||||
**default_ui_key_values, # type: ignore
|
||||
)
|
||||
key = response["token"] # type: ignore
|
||||
user_id = response["user_id"] # type: ignore
|
||||
|
|
@ -13261,7 +13324,7 @@ async def health_services_endpoint(
|
|||
|
||||
# use create task - this can take 10 seconds. don't keep ui users waiting for notification to check their email
|
||||
asyncio.create_task(
|
||||
proxy_logging_obj.slack_alerting_instance.send_key_created_email(
|
||||
proxy_logging_obj.slack_alerting_instance.send_key_created_or_user_invited_email(
|
||||
webhook_event=webhook_event
|
||||
)
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue