From cf8fb58bcd3011914be5a89f7686ceee053c7ba1 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Wed, 6 Nov 2024 14:29:28 +0530
Subject: [PATCH] refactor(proxy_server.py): add debug logging around license
 check event (refactor position in startup_event logic)

---
 litellm/proxy/auth/litellm_license.py | 38 ++++++++++++++++++++++-----
 litellm/proxy/proxy_cli.py            |  3 +++
 litellm/proxy/proxy_server.py         | 24 +++++++----------
 3 files changed, 44 insertions(+), 21 deletions(-)

diff --git a/litellm/proxy/auth/litellm_license.py b/litellm/proxy/auth/litellm_license.py
index 784b4274e5e..a736a1f5e81 100644
--- a/litellm/proxy/auth/litellm_license.py
+++ b/litellm/proxy/auth/litellm_license.py
@@ -5,6 +5,9 @@ import json
 import os
 import traceback
 from datetime import datetime
+from typing import Optional
+
+import httpx
 
 from litellm._logging import verbose_proxy_logger
 from litellm.llms.custom_httpx.http_handler import HTTPHandler
@@ -44,23 +47,46 @@ class LicenseCheck:
             verbose_proxy_logger.error(f"Error reading public key: {str(e)}")
 
     def _verify(self, license_str: str) -> bool:
+
+        verbose_proxy_logger.debug(
+            "litellm.proxy.auth.litellm_license.py::_verify - Checking license against {}/verify_license - {}".format(
+                self.base_url, license_str
+            )
+        )
         url = "{}/verify_license/{}".format(self.base_url, license_str)
 
+        response: Optional[httpx.Response] = None
         try:  # don't impact user, if call fails
-            response = self.http_handler.get(url=url)
+            num_retries = 3
+            for i in range(num_retries):
+                try:
+                    response = self.http_handler.get(url=url)
+                    if response is None:
+                        raise Exception("No response from license server")
+                    response.raise_for_status()
+                except httpx.HTTPStatusError:
+                    if i == num_retries - 1:
+                        raise
 
-            response.raise_for_status()
+            if response is None:
+                raise Exception("No response from license server")
 
             response_json = response.json()
 
             premium = response_json["verify"]
 
             assert isinstance(premium, bool)
+
+            verbose_proxy_logger.debug(
+                "litellm.proxy.auth.litellm_license.py::_verify - License={} is premium={}".format(
+                    license_str, premium
+                )
+            )
             return premium
         except Exception as e:
-            verbose_proxy_logger.error(
-                "litellm.proxy.auth.litellm_license.py::_verify - Unable to verify License via api. - {}".format(
-                    str(e)
+            verbose_proxy_logger.exception(
+                "litellm.proxy.auth.litellm_license.py::_verify - Unable to verify License={} via api. - {}".format(
+                    license_str, str(e)
                 )
             )
             return False
@@ -72,7 +98,7 @@ class LicenseCheck:
         """
         try:
             verbose_proxy_logger.debug(
-                "litellm.proxy.auth.litellm_license.py::is_premium() - ENTERING 'IS_PREMIUM' - {}".format(
+                "litellm.proxy.auth.litellm_license.py::is_premium() - ENTERING 'IS_PREMIUM' - LiteLLM License={}".format(
                     self.license_str
                 )
             )
diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py
index 1fb628a809b..f9f8276c7bd 100644
--- a/litellm/proxy/proxy_cli.py
+++ b/litellm/proxy/proxy_cli.py
@@ -694,6 +694,9 @@ def run_server(  # noqa: PLR0915
 
         import litellm
 
+        if detailed_debug is True:
+            litellm._turn_on_debug()
+
         # DO NOT DELETE - enables global variables to work across files
         from litellm.proxy.proxy_server import app  # noqa
 
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 9f657924277..eb7a3f4f004 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -3060,6 +3060,15 @@ async def startup_event():
             user_api_key_cache=user_api_key_cache,
         )
 
+    ## 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()
+
     ### LOAD CONFIG ###
     worker_config: Optional[Union[str, dict]] = get_secret("WORKER_CONFIG")  # type: ignore
     env_config_yaml: Optional[str] = get_secret_str("CONFIG_FILE_PATH")
@@ -3107,21 +3116,6 @@ async def startup_event():
             if isinstance(worker_config, dict):
                 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
-        )
-    )
-
     ProxyStartupEvent._initialize_startup_logging(
         llm_router=llm_router,
         proxy_logging_obj=proxy_logging_obj,