From 838ab59a84f76e0eb3903b9d36c83dbbd578e73b Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 31 May 2024 10:31:19 -0700 Subject: [PATCH] return correct fields in NewUserResponse --- litellm/integrations/slack_alerting.py | 254 +++++++++++++++---------- litellm/proxy/_types.py | 14 ++ litellm/proxy/proxy_config.yaml | 2 +- litellm/proxy/proxy_server.py | 69 +++++-- 4 files changed, 227 insertions(+), 112 deletions(-) diff --git a/litellm/integrations/slack_alerting.py b/litellm/integrations/slack_alerting.py index 49a8d0e2c..0b62d6c69 100644 --- a/litellm/integrations/slack_alerting.py +++ b/litellm/integrations/slack_alerting.py @@ -68,6 +68,67 @@ class SlackAlertingArgsEnum(Enum): max_outage_alert_list_size: int = 1 * 10 +KEY_CREATED_EMAIL_TEMPLATE = """ + LiteLLM Logo + +

Hi {recipient_email},
+ + I'm happy to provide you with an OpenAI Proxy API Key, loaded with ${key_budget} per month.

+ + + Key: 

{key_token}

+ + +

Usage Example

+ + Detailed Documentation on Usage with OpenAI Python SDK, Langchain, LlamaIndex, Curl + + 
+
+                    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"
+                            }}
+                        ]
+                    )
+
+
+ + + If you have any questions, please send an email to {email_support_contact}

+ + Best,
+ The LiteLLM team
+""" + + +USER_INVITED_EMAIL_TEMPLATE = """ + LiteLLM Logo + +

Hi {recipient_email},
+ + You were invited to use OpenAI Proxy API for team {team_name}

+ +

+ + + + If you have any questions, please send an email to {email_support_contact}

+ + Best,
+ The LiteLLM team
+""" + + 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""" - LiteLLM Logo - -

Hi {recipient_email},
- - I'm happy to provide you with an OpenAI Proxy API Key, loaded with ${key_budget} per month.

- - - Key: 

{key_token}

- - -

Usage Example

- - Detailed Documentation on Usage with OpenAI Python SDK, Langchain, LlamaIndex, Curl - - 
-
-            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"
-                    }}
-                ]
-            )
-
-
- - - If you have any questions, please send an email to {email_support_contact}

- - Best,
- The LiteLLM team
- """ - - 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: """ diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c22ae02f3..b86df4e9b 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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"] diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index ff7343c3a..1d6306d0c 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -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 \ No newline at end of file diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f2af60bff..108124d2b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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 ) )