fix(proxy_server.py): fix proxy_server.py premium user check for encrypted license key

This commit is contained in:
Krrish Dholakia 2024-07-10 12:25:31 -07:00
parent d5d782f844
commit aace0b22a3
5 changed files with 36 additions and 9 deletions

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

View file

@ -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 2. _verify: checks if license is valid calling litellm API. This is the old way we were generating/validating license
""" """
try: 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: if self.license_str is None:
return False return False
elif ( elif (

View file

@ -2165,7 +2165,7 @@ async def initialize(
generate_feedback_box() generate_feedback_box()
user_model = model user_model = model
user_debug = debug 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 import logging
from litellm._logging import ( from litellm._logging import (
@ -2424,7 +2424,7 @@ def giveup(e):
@router.on_event("startup") @router.on_event("startup")
async def startup_event(): 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 import json
### LOAD MASTER KEY ### ### LOAD MASTER KEY ###
@ -2454,6 +2454,21 @@ async def startup_event():
worker_config = json.loads(os.getenv("WORKER_CONFIG")) worker_config = json.loads(os.getenv("WORKER_CONFIG"))
await initialize(**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 ##
cost_tracking() cost_tracking()
@ -7645,7 +7660,7 @@ async def login(request: Request):
litellm_dashboard_ui += "/ui/" litellm_dashboard_ui += "/ui/"
import jwt import jwt
jwt_token = jwt.encode( jwt_token = jwt.encode( # type: ignore
{ {
"user_id": user_id, "user_id": user_id,
"key": key, "key": key,
@ -7709,7 +7724,7 @@ async def login(request: Request):
litellm_dashboard_ui += "/ui/" litellm_dashboard_ui += "/ui/"
import jwt import jwt
jwt_token = jwt.encode( jwt_token = jwt.encode( # type: ignore
{ {
"user_id": user_id, "user_id": user_id,
"key": key, "key": key,
@ -7844,7 +7859,7 @@ async def onboarding(invite_link: str):
litellm_dashboard_ui += "/ui/onboarding" litellm_dashboard_ui += "/ui/onboarding"
import jwt import jwt
jwt_token = jwt.encode( jwt_token = jwt.encode( # type: ignore
{ {
"user_id": user_obj.user_id, "user_id": user_obj.user_id,
"key": key, "key": key,
@ -8261,7 +8276,7 @@ async def auth_callback(request: Request):
import jwt import jwt
jwt_token = jwt.encode( jwt_token = jwt.encode( # type: ignore
{ {
"user_id": user_id, "user_id": user_id,
"key": key, "key": key,