From aace0b22a3c2c20867f105da149d6d6fd8ea381b Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 10 Jul 2024 12:25:31 -0700 Subject: [PATCH] fix(proxy_server.py): fix proxy_server.py premium user check for encrypted license key --- litellm/proxy/_experimental/out/404.html | 1 - .../proxy/_experimental/out/model_hub.html | 1 - .../proxy/_experimental/out/onboarding.html | 1 - litellm/proxy/auth/litellm_license.py | 15 +++++++++++ litellm/proxy/proxy_server.py | 27 ++++++++++++++----- 5 files changed, 36 insertions(+), 9 deletions(-) delete mode 100644 litellm/proxy/_experimental/out/404.html delete mode 100644 litellm/proxy/_experimental/out/model_hub.html delete mode 100644 litellm/proxy/_experimental/out/onboarding.html diff --git a/litellm/proxy/_experimental/out/404.html b/litellm/proxy/_experimental/out/404.html deleted file mode 100644 index 47104dc832..0000000000 --- a/litellm/proxy/_experimental/out/404.html +++ /dev/null @@ -1 +0,0 @@ -404: This page could not be found.LiteLLM Dashboard

404

This page could not be found.

\ No newline at end of file diff --git a/litellm/proxy/_experimental/out/model_hub.html b/litellm/proxy/_experimental/out/model_hub.html deleted file mode 100644 index d239941bb8..0000000000 --- a/litellm/proxy/_experimental/out/model_hub.html +++ /dev/null @@ -1 +0,0 @@ -LiteLLM Dashboard \ No newline at end of file diff --git a/litellm/proxy/_experimental/out/onboarding.html b/litellm/proxy/_experimental/out/onboarding.html deleted file mode 100644 index 6ae2f18a74..0000000000 --- a/litellm/proxy/_experimental/out/onboarding.html +++ /dev/null @@ -1 +0,0 @@ -LiteLLM Dashboard \ No newline at end of file diff --git a/litellm/proxy/auth/litellm_license.py b/litellm/proxy/auth/litellm_license.py index ab45985620..18c20b35ad 100644 --- a/litellm/proxy/auth/litellm_license.py +++ b/litellm/proxy/auth/litellm_license.py @@ -71,6 +71,21 @@ class LicenseCheck: 2. _verify: checks if license is valid calling litellm API. This is the old way we were generating/validating license """ try: + verbose_proxy_logger.debug( + "litellm.proxy.auth.litellm_license.py::is_premium() - ENTERING 'IS_PREMIUM' - {}".format( + self.license_str + ) + ) + + if self.license_str is None: + self.license_str = os.getenv("LITELLM_LICENSE", None) + + verbose_proxy_logger.debug( + "litellm.proxy.auth.litellm_license.py::is_premium() - Updated 'self.license_str' - {}".format( + self.license_str + ) + ) + if self.license_str is None: return False elif ( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 2a580f8e7b..a600ac0189 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2165,7 +2165,7 @@ async def initialize( generate_feedback_box() user_model = model user_debug = debug - if debug == True: # this needs to be first, so users can see Router init debugg + if debug is True: # this needs to be first, so users can see Router init debugg import logging from litellm._logging import ( @@ -2424,7 +2424,7 @@ def giveup(e): @router.on_event("startup") async def startup_event(): - global prisma_client, master_key, use_background_health_checks, llm_router, llm_model_list, general_settings, proxy_budget_rescheduler_min_time, proxy_budget_rescheduler_max_time, litellm_proxy_admin_name, db_writer_client, store_model_in_db + global prisma_client, master_key, use_background_health_checks, llm_router, llm_model_list, general_settings, proxy_budget_rescheduler_min_time, proxy_budget_rescheduler_max_time, litellm_proxy_admin_name, db_writer_client, store_model_in_db, premium_user, _license_check import json ### LOAD MASTER KEY ### @@ -2454,6 +2454,21 @@ async def startup_event(): worker_config = json.loads(os.getenv("WORKER_CONFIG")) await initialize(**worker_config) + ## CHECK PREMIUM USER + verbose_proxy_logger.debug( + "litellm.proxy.proxy_server.py::startup() - CHECKING PREMIUM USER - {}".format( + premium_user + ) + ) + if premium_user is False: + premium_user = _license_check.is_premium() + + verbose_proxy_logger.debug( + "litellm.proxy.proxy_server.py::startup() - PREMIUM USER value - {}".format( + premium_user + ) + ) + ## COST TRACKING ## cost_tracking() @@ -7645,7 +7660,7 @@ async def login(request: Request): litellm_dashboard_ui += "/ui/" import jwt - jwt_token = jwt.encode( + jwt_token = jwt.encode( # type: ignore { "user_id": user_id, "key": key, @@ -7709,7 +7724,7 @@ async def login(request: Request): litellm_dashboard_ui += "/ui/" import jwt - jwt_token = jwt.encode( + jwt_token = jwt.encode( # type: ignore { "user_id": user_id, "key": key, @@ -7844,7 +7859,7 @@ async def onboarding(invite_link: str): litellm_dashboard_ui += "/ui/onboarding" import jwt - jwt_token = jwt.encode( + jwt_token = jwt.encode( # type: ignore { "user_id": user_obj.user_id, "key": key, @@ -8261,7 +8276,7 @@ async def auth_callback(request: Request): import jwt - jwt_token = jwt.encode( + jwt_token = jwt.encode( # type: ignore { "user_id": user_id, "key": key,