phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

return correct fields in NewUserResponse

Browse source
This commit is contained in:
Ishaan Jaff 2024-05-31 10:31:19 -07:00
parent f3687f68e7
commit 838ab59a84
4 changed files with 227 additions and 112 deletions
Download patch file Download diff file Expand all files Collapse all files

254
litellm/integrations/slack_alerting.py
View file

@ -68,6 +68,67 @@ class SlackAlertingArgsEnum(Enum):
max_outage_alert_list_size: int = 1 * 10
KEY_CREATED_EMAIL_TEMPLATE = """
<img src="{email_logo_url}" alt="LiteLLM Logo" width="150" height="50" />
<p> Hi {recipient_email}, <br/>
I'm happy to provide you with an OpenAI Proxy API Key, loaded with ${key_budget} per month. <br /> <br />
<b>
Key: <pre>{key_token}</pre> <br>
</b>
<h2>Usage Example</h2>
Detailed Documentation on <a href="https://docs.litellm.ai/docs/proxy/user_keys">Usage with OpenAI Python SDK, Langchain, LlamaIndex, Curl</a>
<pre>
import openai
client = openai.OpenAI(
api_key="{key_token}",
base_url={{base_url}}
)
response = client.chat.completions.create(
model="gpt-3.5-turbo", # model to send to the proxy
messages = [
{{
"role": "user",
"content": "this is a test request, write a short poem"
}}
]
)
</pre>
If you have any questions, please send an email to {email_support_contact} <br /> <br />
Best, <br />
The LiteLLM team <br />
"""
USER_INVITED_EMAIL_TEMPLATE = """
<img src="{email_logo_url}" alt="LiteLLM Logo" width="150" height="50" />
<p> Hi {recipient_email}, <br/>
You were invited to use OpenAI Proxy API for team {team_name} <br /> <br />
<button><a href="{base_url}">Get Started here</a></button> <br /> <br />
If you have any questions, please send an email to {email_support_contact} <br /> <br />
Best, <br />
The LiteLLM team <br />
"""
class SlackAlertingArgs(LiteLLMBase):
daily_report_frequency: int = Field(
default=int(
@ -1190,105 +1251,106 @@ Model Info:
raise ValueError(
f"Trying to Customize Email Alerting\n {CommonProxyErrors.not_premium_user.value}"
)
return
async def send_key_created_email(self, webhook_event: WebhookEvent) -> bool:
from litellm.proxy.utils import send_email
async def send_key_created_or_user_invited_email(
self, webhook_event: WebhookEvent
) -> bool:
try:
from litellm.proxy.utils import send_email
if self.alerting is None or "email" not in self.alerting:
# do nothing if user does not want email alerts
if self.alerting is None or "email" not in self.alerting:
# do nothing if user does not want email alerts
return False
from litellm.proxy.proxy_server import premium_user, prisma_client
email_logo_url = os.getenv("SMTP_SENDER_LOGO", None)
email_support_contact = os.getenv("EMAIL_SUPPORT_CONTACT", None)
await self._check_if_using_premium_email_feature(
premium_user, email_logo_url, email_support_contact
)
if email_logo_url is None:
email_logo_url = LITELLM_LOGO_URL
if email_support_contact is None:
email_support_contact = LITELLM_SUPPORT_CONTACT
event_name = webhook_event.event_message
recipient_email = webhook_event.user_email
recipient_user_id = webhook_event.user_id
if (
recipient_email is None
and recipient_user_id is not None
and prisma_client is not None
):
user_row = await prisma_client.db.litellm_usertable.find_unique(
where={"user_id": recipient_user_id}
)
if user_row is not None:
recipient_email = user_row.user_email
key_name = webhook_event.key_alias
key_token = webhook_event.token
key_budget = webhook_event.max_budget
base_url = os.getenv("PROXY_BASE_URL", "http://0.0.0.0:4000")
email_html_content = "Alert from LiteLLM Server"
if recipient_email is None:
verbose_proxy_logger.error(
"Trying to send email alert to no recipient",
extra=webhook_event.dict(),
)
if webhook_event.event == "key_created":
email_html_content = KEY_CREATED_EMAIL_TEMPLATE.format(
email_logo_url=email_logo_url,
recipient_email=recipient_email,
key_budget=key_budget,
key_token=key_token,
base_url=base_url,
email_support_contact=email_support_contact,
)
elif webhook_event.event == "internal_user_created":
# GET TEAM NAME
team_id = webhook_event.team_id
team_name = "Default Team"
if team_id is not None and prisma_client is not None:
team_row = await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": team_id}
)
if team_row is not None:
team_name = team_row.team_alias or "-"
email_html_content = USER_INVITED_EMAIL_TEMPLATE.format(
email_logo_url=email_logo_url,
recipient_email=recipient_email,
team_name=team_name,
base_url=base_url,
email_support_contact=email_support_contact,
)
else:
verbose_proxy_logger.error(
"Trying to send email alert on unknown webhook event",
extra=webhook_event.model_dump(),
)
payload = webhook_event.model_dump_json()
email_event = {
"to": recipient_email,
"subject": f"LiteLLM: {event_name}",
"html": email_html_content,
}
response = await send_email(
receiver_email=email_event["to"],
subject=email_event["subject"],
html=email_event["html"],
)
return True
except Exception as e:
verbose_proxy_logger.error("Error sending email alert %s", str(e))
return False
from litellm.proxy.proxy_server import premium_user, prisma_client
email_logo_url = os.getenv("SMTP_SENDER_LOGO", None)
email_support_contact = os.getenv("EMAIL_SUPPORT_CONTACT", None)
await self._check_if_using_premium_email_feature(
premium_user, email_logo_url, email_support_contact
)
if email_logo_url is None:
email_logo_url = LITELLM_LOGO_URL
if email_support_contact is None:
email_support_contact = LITELLM_SUPPORT_CONTACT
event_name = webhook_event.event_message
recipient_email = webhook_event.user_email
recipient_user_id = webhook_event.user_id
if (
recipient_email is None
and recipient_user_id is not None
and prisma_client is not None
):
user_row = await prisma_client.db.litellm_usertable.find_unique(
where={"user_id": recipient_user_id}
)
if user_row is not None:
recipient_email = user_row.user_email
key_name = webhook_event.key_alias
key_token = webhook_event.token
key_budget = webhook_event.max_budget
email_html_content = "Alert from LiteLLM Server"
if recipient_email is None:
verbose_proxy_logger.error(
"Trying to send email alert to no recipient", extra=webhook_event.dict()
)
email_html_content = f"""
<img src="{email_logo_url}" alt="LiteLLM Logo" width="150" height="50" />
<p> Hi {recipient_email}, <br/>
I'm happy to provide you with an OpenAI Proxy API Key, loaded with ${key_budget} per month. <br /> <br />
<b>
Key: <pre>{key_token}</pre> <br>
</b>
<h2>Usage Example</h2>
Detailed Documentation on <a href="https://docs.litellm.ai/docs/proxy/user_keys">Usage with OpenAI Python SDK, Langchain, LlamaIndex, Curl</a>
<pre>
import openai
client = openai.OpenAI(
api_key="{key_token}",
base_url={os.getenv("PROXY_BASE_URL", "http://0.0.0.0:4000")}
)
response = client.chat.completions.create(
model="gpt-3.5-turbo", # model to send to the proxy
messages = [
{{
"role": "user",
"content": "this is a test request, write a short poem"
}}
]
)
</pre>
If you have any questions, please send an email to {email_support_contact} <br /> <br />
Best, <br />
The LiteLLM team <br />
"""
payload = webhook_event.model_dump_json()
email_event = {
"to": recipient_email,
"subject": f"LiteLLM: {event_name}",
"html": email_html_content,
}
response = await send_email(
receiver_email=email_event["to"],
subject=email_event["subject"],
html=email_event["html"],
)
return False
async def send_email_alert_using_smtp(self, webhook_event: WebhookEvent) -> bool:
"""

14
litellm/proxy/_types.py
View file

@ -587,6 +587,19 @@ class NewUserRequest(GenerateKeyRequest):
class NewUserResponse(GenerateKeyResponse):
max_budget: Optional[float] = None
user_email: Optional[str] = None
user_role: Optional[
Literal[
LitellmUserRoles.PROXY_ADMIN,
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
LitellmUserRoles.INTERNAL_USER,
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
LitellmUserRoles.TEAM,
LitellmUserRoles.CUSTOMER,
]
] = None
teams: Optional[list] = None
organization_id: Optional[str] = None
class UpdateUserRequest(GenerateRequestBase):
@ -1292,6 +1305,7 @@ class WebhookEvent(CallInfo):
"threshold_crossed",
"projected_limit_exceeded",
"key_created",
"internal_user_created",
"spend_tracked",
]
event_group: Literal["internal_user", "key", "team", "proxy", "customer"]

2
litellm/proxy/proxy_config.yaml
View file

@ -21,7 +21,7 @@ model_list:
general_settings:
master_key: sk-1234
alerting: ["slack"]
alerting: ["slack", "email"]
litellm_settings:
callbacks: custom_callbacks1.proxy_handler_instance

69
litellm/proxy/proxy_server.py
View file

@ -3209,6 +3209,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,
@ -3240,6 +3243,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
@ -3340,7 +3344,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"
)
@ -3397,6 +3401,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
@ -3744,6 +3752,7 @@ async def startup_event():
)
asyncio.create_task(
generate_key_helper_fn(
request_type="user",
duration=None,
models=[],
aliases={},
@ -3766,6 +3775,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=[],
@ -3788,7 +3798,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 ###
@ -6125,13 +6141,19 @@ 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
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",
@ -6146,7 +6168,7 @@ async def generate_key_fn(
# 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(
proxy_logging_obj.slack_alerting_instance.send_key_created_or_user_invited_email(
webhook_event=event,
)
)
@ -8133,7 +8155,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
@ -8150,21 +8172,28 @@ 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="key_created",
event_group="key",
event_message=f"API Key Created",
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_email(
proxy_logging_obj.slack_alerting_instance.send_key_created_or_user_invited_email(
webhook_event=event,
)
)
@ -8174,6 +8203,9 @@ async def new_user(data: NewUserRequest):
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),
@ -8230,11 +8262,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/")
@ -11726,7 +11760,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(
@ -11822,7 +11857,8 @@ async def onboarding(invite_link: str):
user_email = user_obj.user_email
response = await generate_key_helper_fn(
**{"user_role": user_obj.user_role or "app_owner", "duration": "2hr", "key_max_budget": 5, "models": [], "aliases": {}, "config": {}, "spend": 0, "user_id": user_obj.user_id, "team_id": "litellm-dashboard"} # type: ignore
request_type="key",
**{"user_role": user_obj.user_role or "app_owner", "duration": "2hr", "key_max_budget": 5, "models": [], "aliases": {}, "config": {}, "spend": 0, "user_id": user_obj.user_id, "team_id": "litellm-dashboard"}, # type: ignore
)
key = response["token"] # type: ignore
@ -12125,8 +12161,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
@ -13231,7 +13270,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
)
)