add pub sub

litellm/integrations/gcs_pubsub/pub_sub.py

@@ -41,6 +41,10 @@ def __init__(
             topic_id (str): Pub/Sub topic ID
             credentials_path (str, optional): Path to Google Cloud credentials JSON file
         """
+        from litellm.proxy.utils import _premium_user_check
+
+        _premium_user_check()
+
         self.async_httpx_client = get_async_httpx_client(
             llm_provider=httpxSpecialProvider.LoggingCallback
         )
@@ -102,6 +106,9 @@ async def async_log_success_event(self, kwargs, response_obj, start_time, end_ti
         from litellm.proxy.spend_tracking.spend_tracking_utils import (
             get_logging_payload,
         )
+        from litellm.proxy.utils import _premium_user_check
+
+        _premium_user_check()
 
         try:
             verbose_logger.debug(

litellm/proxy/management_endpoints/team_endpoints.py

@@ -58,7 +58,7 @@
     add_new_member,
     management_endpoint_wrapper,
 )
-from litellm.proxy.utils import PrismaClient
+from litellm.proxy.utils import PrismaClient, _premium_user_check
 
 router = APIRouter()
 
@@ -1527,15 +1527,3 @@ def _set_team_metadata_field(
         _premium_user_check()
     team_data.metadata = team_data.metadata or {}
     team_data.metadata[field_name] = value
-
-
-def _premium_user_check():
-    from litellm.proxy.proxy_server import premium_user
-
-    if not premium_user:
-        raise HTTPException(
-            status_code=403,
-            detail={
-                "error": f"This feature is only available for LiteLLM Enterprise users. {CommonProxyErrors.not_premium_user.value}"
-            },
-        )

litellm/proxy/utils.py

@@ -16,6 +16,7 @@
 from litellm.litellm_core_utils.duration_parser import duration_in_seconds
 from litellm.proxy._types import (
     DB_CONNECTION_ERROR_TYPES,
+    CommonProxyErrors,
     ProxyErrorTypes,
     ProxyException,
 )
@@ -2959,3 +2960,18 @@ def handle_exception_on_proxy(e: Exception) -> ProxyException:
         param=getattr(e, "param", "None"),
         code=status.HTTP_500_INTERNAL_SERVER_ERROR,
     )
+
+
+def _premium_user_check():
+    """
+    Raises an HTTPException if the user is not a premium user
+    """
+    from litellm.proxy.proxy_server import premium_user
+
+    if not premium_user:
+        raise HTTPException(
+            status_code=403,
+            detail={
+                "error": f"This feature is only available for LiteLLM Enterprise users. {CommonProxyErrors.not_premium_user.value}"
+            },
+        )