Remove + Check for unsafe enterprise/ folder imports (#11107)

* refactor(litellm/): wrap enterprise package usage in safe try-excepts prevents missing package error for OSS users Fixes https://github.com/BerriAI/litellm/issues/10795 * ci(config.yml): add script to check for unsafe imports to ci/cd Closes https://github.com/BerriAI/litellm/issues/10795 * fix: add missing docs * test: test_openai_image_edit_litellm_sdk * use n 4 for mapped tests (#11109) * Fix/background health check (#10887) * fix: improve health check logic by deep copying model list on each iteration * test: add async test for background health check reflecting model list changes * fix: validate health check interval before executing background health check * fix: specify type for health check results dictionary * fix(user_api_key_auth.py): handle user custom auth set with no custom settings * bump: version 0.1.21 → 0.2.0 * ci(config.yml): run enterprise and litellm tests separately * fix: fix linting error * docs: add missing docs * [Feat] Add content policy violation error mapping for image editd (#11113) * feat: add image edit mapping for content policy violations * test fix * Expose `/list` and `/info` endpoints for Audit Log events (#11102) * feat(audit_logging_endpoints.py): expose list endpoint to show all audit logs make it easier for user to retrieve individual endpoints * feat(enterprise/): add audit logging endpoint * feat(audit_logging_endpoints.py): expose new GET `/audit/{id}` endpoint make it easier to retrieve view individual audit logs * feat(key_management_event_hooks.py): correctly show the key of the user who initiated the change * fix(key_management_event_hooks.py): add key rotations as an audit log event ' * test(test_audit_logging_endpoints.py): add simple unit testing for audit log endpoint * fix: testing fixes * fix: fix ruff check * [Feat] Use aiohttp transport by default - 97% lower median latency (#11097) * fix: add flag for disabling use_aiohttp_transport * feat: add _create_async_transport * feat: fixes for transport * add httpx-aiohttp * feat: fixes for transport * refactor: fixes for transport * build: fix deps * fixes: test fixes * fix: ensure aiohttp does not auto set content type * test: test fixes * feat: add LiteLLMAiohttpTransport * fix: fixes for responses API handling * test: fixes for responses API handling * test: fixes for responses API handling * feat: fixes for transport * fix: base embedding handler * test: test_async_http_handler_force_ipv4 * test: fix failing deepeval test * fix: add YARL for bedrock urls * fix: issues with transport * fix: comment out linting issues * test fix * test: XAI is unstable * test: fixes for using respx * test: XAI fixes * test: XAI fixes * test: infinity testing fixes * docs(config_settings.md): document param * test: test_openai_image_edit_litellm_sdk * test: remove deprecated test * bump respx==0.22.0 * test: test_xai_message_name_filtering * test: fix anthropic test after bumping httpx * use n 4 for mapped tests (#11109) * fix: use 1 session per event loop * test: test_client_session_helper * fix: linting error * fix: resolving GET requests on httpx 0.28.1 * test fixes proxy unit tests * fix: add ssl verify settings * fix: proxy unit tests * fix: refactor * tests: basic unit tests for aiohttp transports * tests: fixes xai --------- Co-authored-by: Krrish Dholakia <krrishdholakia@gmail.com> * test: cleanup redundant test * fix: fix import * fix: fix linting errors --------- Co-authored-by: Ishaan Jaff <ishaanjaffer0324@gmail.com> Co-authored-by: JuHyun Bae <jhyun0408@nate.com>
2026-08-05 02:23:42 +00:00 · 2025-05-24 08:59:02 -07:00
parent dc79424e6b
commit a44bc9f613
15 changed files with 257 additions and 75 deletions
@@ -1349,6 +1349,7 @@ jobs:
      - run: python ./tests/code_coverage_tests/enforce_llms_folder_style.py
      - run: python ./tests/documentation_tests/test_circular_imports.py
      - run: python ./tests/code_coverage_tests/prevent_key_leaks_in_exceptions.py
+      - run: python ./tests/code_coverage_tests/check_unsafe_enterprise_import.py
      - run: helm lint ./deploy/charts/litellm-helm

  db_migration_disable_update_check:
@@ -352,6 +352,7 @@ router_settings:
 | CONFIG_FILE_PATH | File path for configuration file
 | CONFIDENT_API_KEY | API key for DeepEval integration
 | CUSTOM_TIKTOKEN_CACHE_DIR | Custom directory for Tiktoken cache
+| CONFIDENT_API_KEY | API key for Confident AI (Deepeval) Logging service
 | DATABASE_HOST | Hostname for the database server
 | DATABASE_NAME | Name of the database
 | DATABASE_PASSWORD | Password for the database user
@@ -507,6 +508,7 @@ router_settings:
 | LITELLM_GLOBAL_MAX_PARALLEL_REQUEST_RETRY_TIMEOUT | Timeout for retries of parallel requests in LiteLLM
 | LITELLM_MIGRATION_DIR | Custom migrations directory for prisma migrations, used for baselining db in read-only file systems.
 | LITELLM_HOSTED_UI | URL of the hosted UI for LiteLLM
+| LITELM_ENVIRONMENT | Environment of LiteLLM Instance, used by logging services. Currently only used by DeepEval.
 | LITELLM_LICENSE | License key for LiteLLM usage
 | LITELLM_LOCAL_MODEL_COST_MAP | Local configuration for model cost mapping in LiteLLM
 | LITELLM_LOG | Enable detailed logging for LiteLLM
@@ -7,6 +7,12 @@ import json
 import os
 from typing import List, Optional

+from litellm_enterprise.types.enterprise_callbacks.send_emails import (
+    EmailEvent,
+    EmailParams,
+    SendKeyCreatedEmailEvent,
+)
+
 from litellm._logging import verbose_proxy_logger
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.integrations.email_templates.email_footer import EMAIL_FOOTER
@@ -17,11 +23,6 @@ from litellm.integrations.email_templates.user_invitation_email import (
    USER_INVITATION_EMAIL_TEMPLATE,
 )
 from litellm.proxy._types import WebhookEvent
-from litellm.types.enterprise.enterprise_callbacks.send_emails import (
-    EmailEvent,
-    EmailParams,
-    SendKeyCreatedEmailEvent,
-)
 from litellm.types.integrations.slack_alerting import LITELLM_LOGO_URL


@@ -6,11 +6,7 @@ import json
 from typing import Dict

 from fastapi import APIRouter, Depends, HTTPException
-
-from litellm._logging import verbose_proxy_logger
-from litellm.proxy._types import UserAPIKeyAuth
-from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
-from litellm.types.enterprise.enterprise_callbacks.send_emails import (
+from litellm_enterprise.types.enterprise_callbacks.send_emails import (
    DefaultEmailSettings,
    EmailEvent,
    EmailEventSettings,
@@ -18,6 +14,10 @@ from litellm.types.enterprise.enterprise_callbacks.send_emails import (
    EmailEventSettingsUpdateRequest,
 )

+from litellm._logging import verbose_proxy_logger
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+
 router = APIRouter()


@@ -2167,6 +2167,10 @@ class CommonProxyErrors(str, enum.Enum):
    max_parallel_request_limit_reached = (
        "Crossed TPM / RPM / Max Parallel Request Limit"
    )
+    missing_enterprise_package = "Missing litellm-enterprise package. Please install it to use this feature. Run `pip install litellm-enterprise`"
+    missing_enterprise_package_docker = (
+        "This uses the enterprise folder - only available on the Docker image."
+    )


 class SpendCalculateRequest(LiteLLMPydanticObjectBase):
@@ -56,9 +56,15 @@ def initialize_callbacks_on_proxy(  # noqa: PLR0915
                pii_masking_object = _OPTIONAL_PresidioPIIMasking(**params)
                imported_list.append(pii_masking_object)
            elif isinstance(callback, str) and callback == "llamaguard_moderations":
-                from litellm_enterprise.enterprise_callbacks.llama_guard import (
-                    _ENTERPRISE_LlamaGuard,
-                )
+                try:
+                    from litellm_enterprise.enterprise_callbacks.llama_guard import (
+                        _ENTERPRISE_LlamaGuard,
+                    )
+                except ImportError:
+                    raise Exception(
+                        "MissingTrying to use Llama Guard"
+                        + CommonProxyErrors.missing_enterprise_package.value
+                    )

                if premium_user is not True:
                    raise Exception(
@@ -69,9 +75,15 @@ def initialize_callbacks_on_proxy(  # noqa: PLR0915
                llama_guard_object = _ENTERPRISE_LlamaGuard()
                imported_list.append(llama_guard_object)
            elif isinstance(callback, str) and callback == "hide_secrets":
-                from litellm_enterprise.enterprise_callbacks.secret_detection import (
-                    _ENTERPRISE_SecretDetection,
-                )
+                try:
+                    from litellm_enterprise.enterprise_callbacks.secret_detection import (
+                        _ENTERPRISE_SecretDetection,
+                    )
+                except ImportError:
+                    raise Exception(
+                        "Trying to use Secret Detection"
+                        + CommonProxyErrors.missing_enterprise_package.value
+                    )

                if premium_user is not True:
                    raise Exception(
@@ -82,9 +94,15 @@ def initialize_callbacks_on_proxy(  # noqa: PLR0915
                _secret_detection_object = _ENTERPRISE_SecretDetection()
                imported_list.append(_secret_detection_object)
            elif isinstance(callback, str) and callback == "openai_moderations":
-                from enterprise.enterprise_hooks.openai_moderation import (
-                    _ENTERPRISE_OpenAI_Moderation,
-                )
+                try:
+                    from enterprise.enterprise_hooks.openai_moderation import (
+                        _ENTERPRISE_OpenAI_Moderation,
+                    )
+                except ImportError:
+                    raise Exception(
+                        "Trying to use OpenAI Moderations Check,"
+                        + CommonProxyErrors.missing_enterprise_package_docker.value
+                    )

                if premium_user is not True:
                    raise Exception(
@@ -112,9 +130,15 @@ def initialize_callbacks_on_proxy(  # noqa: PLR0915
                aporia_guardrail_object = AporiaGuardrail()
                imported_list.append(aporia_guardrail_object)
            elif isinstance(callback, str) and callback == "google_text_moderation":
-                from enterprise.enterprise_hooks.google_text_moderation import (
-                    _ENTERPRISE_GoogleTextModeration,
-                )
+                try:
+                    from enterprise.enterprise_hooks.google_text_moderation import (
+                        _ENTERPRISE_GoogleTextModeration,
+                    )
+                except ImportError:
+                    raise Exception(
+                        "Trying to use Google Text Moderation,"
+                        + CommonProxyErrors.missing_enterprise_package_docker.value
+                    )

                if premium_user is not True:
                    raise Exception(
@@ -125,9 +149,15 @@ def initialize_callbacks_on_proxy(  # noqa: PLR0915
                google_text_moderation_obj = _ENTERPRISE_GoogleTextModeration()
                imported_list.append(google_text_moderation_obj)
            elif isinstance(callback, str) and callback == "llmguard_moderations":
-                from litellm_enterprise.enterprise_callbacks.llm_guard import (
-                    _ENTERPRISE_LLMGuard,
-                )
+                try:
+                    from litellm_enterprise.enterprise_callbacks.llm_guard import (
+                        _ENTERPRISE_LLMGuard,
+                    )
+                except ImportError:
+                    raise Exception(
+                        "Trying to use Llm Guard"
+                        + CommonProxyErrors.missing_enterprise_package.value
+                    )

                if premium_user is not True:
                    raise Exception(
@@ -138,9 +168,15 @@ def initialize_callbacks_on_proxy(  # noqa: PLR0915
                llm_guard_moderation_obj = _ENTERPRISE_LLMGuard()
                imported_list.append(llm_guard_moderation_obj)
            elif isinstance(callback, str) and callback == "blocked_user_check":
-                from enterprise.enterprise_hooks.blocked_user_list import (
-                    _ENTERPRISE_BlockedUserList,
-                )
+                try:
+                    from enterprise.enterprise_hooks.blocked_user_list import (
+                        _ENTERPRISE_BlockedUserList,
+                    )
+                except ImportError:
+                    raise Exception(
+                        "Trying to use Blocked User List"
+                        + CommonProxyErrors.missing_enterprise_package_docker.value
+                    )

                if premium_user is not True:
                    raise Exception(
@@ -153,9 +189,15 @@ def initialize_callbacks_on_proxy(  # noqa: PLR0915
                )
                imported_list.append(blocked_user_list)
            elif isinstance(callback, str) and callback == "banned_keywords":
-                from enterprise.enterprise_hooks.banned_keywords import (
-                    _ENTERPRISE_BannedKeywords,
-                )
+                try:
+                    from enterprise.enterprise_hooks.banned_keywords import (
+                        _ENTERPRISE_BannedKeywords,
+                    )
+                except ImportError:
+                    raise Exception(
+                        "Trying to use Banned Keywords"
+                        + CommonProxyErrors.missing_enterprise_package_docker.value
+                    )

                if premium_user is not True:
                    raise Exception(
@@ -1,5 +1,6 @@
 # litellm/proxy/guardrails/guardrail_initializers.py
 import litellm
+from litellm.proxy._types import CommonProxyErrors
 from litellm.types.guardrails import *


@@ -131,9 +132,15 @@ def initialize_presidio(litellm_params: LitellmParams, guardrail: Guardrail):


 def initialize_hide_secrets(litellm_params: LitellmParams, guardrail: Guardrail):
-    from litellm_enterprise.enterprise_callbacks.secret_detection import (
-        _ENTERPRISE_SecretDetection,
-    )
+    try:
+        from litellm_enterprise.enterprise_callbacks.secret_detection import (
+            _ENTERPRISE_SecretDetection,
+        )
+    except ImportError:
+        raise Exception(
+            "Trying to use Secret Detection"
+            + CommonProxyErrors.missing_enterprise_package.value
+        )

    _secret_detection_object = _ENTERPRISE_SecretDetection(
        detect_secrets_config=litellm_params.detect_secrets_config,
@@ -7,6 +7,7 @@ from typing import Any, List, Optional
 import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm.proxy._types import (
+    CommonProxyErrors,
    GenerateKeyRequest,
    GenerateKeyResponse,
    KeyRequest,
@@ -309,14 +310,27 @@ class KeyManagementEventHooks:

    @staticmethod
    async def _send_key_created_email(response: dict):
-        from litellm_enterprise.enterprise_callbacks.send_emails.base_email import (
-            BaseEmailLogger,
-        )
+        try:
+            from litellm_enterprise.enterprise_callbacks.send_emails.base_email import (
+                BaseEmailLogger,
+            )
+        except ImportError:
+            raise Exception(
+                "Trying to use Email Hooks"
+                + CommonProxyErrors.missing_enterprise_package.value
+            )

        from litellm.proxy.proxy_server import general_settings, proxy_logging_obj
-        from litellm.types.enterprise.enterprise_callbacks.send_emails import (
-            SendKeyCreatedEmailEvent,
-        )
+
+        try:
+            from litellm_enterprise.types.enterprise_callbacks.send_emails import (
+                SendKeyCreatedEmailEvent,
+            )
+        except ImportError:
+            raise Exception(
+                "Trying to use Email Hooks"
+                + CommonProxyErrors.missing_enterprise_package.value
+            )

        event = SendKeyCreatedEmailEvent(
            virtual_key=response.get("key", ""),
@@ -7,9 +7,6 @@ import uuid
 from datetime import datetime, timezone
 from typing import Optional

-from litellm_enterprise.enterprise_callbacks.send_emails.base_email import (
-    BaseEmailLogger,
-)
 from pydantic import BaseModel

 import litellm
@@ -111,17 +108,29 @@ class UserManagementEventHooks:
        #########################################################
        ########## V2 USER INVITATION EMAIL ################
        #########################################################
-        initialized_email_loggers = (
-            litellm.logging_callback_manager.get_custom_loggers_for_type(
-                callback_type=BaseEmailLogger
+        try:
+            from litellm_enterprise.enterprise_callbacks.send_emails.base_email import (
+                BaseEmailLogger,
            )
-        )
-        if len(initialized_email_loggers) > 0:
-            for email_logger in initialized_email_loggers:
-                if isinstance(email_logger, BaseEmailLogger):
-                    await email_logger.send_user_invitation_email(
-                        event=event,
-                    )
+
+            use_enterprise_email_hooks = True
+        except ImportError:
+            verbose_proxy_logger.warning(
+                "Defaulting to using Legacy Email Hooks."
+                + CommonProxyErrors.missing_enterprise_package.value
+            )
+            use_enterprise_email_hooks = False
+
+        if use_enterprise_email_hooks:
+            initialized_email_loggers = litellm.logging_callback_manager.get_custom_loggers_for_type(
+                callback_type=BaseEmailLogger  # type: ignore
+            )
+            if len(initialized_email_loggers) > 0:
+                for email_logger in initialized_email_loggers:
+                    if isinstance(email_logger, BaseEmailLogger):  # type: ignore
+                        await email_logger.send_user_invitation_email(  # type: ignore
+                            event=event,
+                        )

        #########################################################
        ########## LEGACY V1 USER INVITATION EMAIL ################
@@ -102,9 +102,18 @@ async def unblock_user(data: BlockUsers):
    }'
    ```
    """
-    from enterprise.enterprise_hooks.blocked_user_list import (
-        _ENTERPRISE_BlockedUserList,
-    )
+    try:
+        from enterprise.enterprise_hooks.blocked_user_list import (
+            _ENTERPRISE_BlockedUserList,
+        )
+    except ImportError:
+        raise HTTPException(
+            status_code=400,
+            detail={
+                "error": "Blocked user check was never set. This call has no effect."
+                + CommonProxyErrors.missing_enterprise_package_docker.value
+            },
+        )

    if (
        not any(isinstance(x, _ENTERPRISE_BlockedUserList) for x in litellm.callbacks)
@@ -149,7 +149,13 @@ async def view_spend_tags(
    ```
    """

-    from enterprise.utils import get_spend_by_tags
+    try:
+        from enterprise.utils import get_spend_by_tags
+    except ImportError:
+        raise Exception(
+            "Trying to use Spend by Tags"
+            + CommonProxyErrors.missing_enterprise_package_docker.value
+        )
    from litellm.proxy.proxy_server import prisma_client

    try:
@@ -1655,12 +1661,10 @@ async def ui_view_spend_logs(  # noqa: PLR0915
    ),
    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
    status_filter: Optional[str] = fastapi.Query(
-        default=None,
-        description="Filter logs by status (e.g., success, failure)"
+        default=None, description="Filter logs by status (e.g., success, failure)"
    ),
    model: Optional[str] = fastapi.Query(
-        default=None,
-        description="Filter logs by model"
+        default=None, description="Filter logs by model"
    ),
 ):
    """
@@ -2922,22 +2926,17 @@ async def ui_view_session_spend_logs(
 def _build_status_filter_condition(status_filter: Optional[str]) -> Dict[str, Any]:
    """
    Helper function to build the status filter condition for database queries.
-    
+
    Args:
        status_filter (Optional[str]): The status to filter by. Can be "success" or "failure".
-        
+
    Returns:
        Dict[str, Any]: A dictionary containing the status filter condition.
    """
    if status_filter is None:
        return {}
-        
+
    if status_filter == "success":
-        return {
-            "OR": [
-                {"status": {"equals": "success"}},
-                {"status": None}
-            ]
-        }
+        return {"OR": [{"status": {"equals": "success"}}, {"status": None}]}
    else:
        return {"status": {"equals": status_filter}}
@@ -0,0 +1,95 @@
+import ast
+import os
+
+class EnterpriseImportFinder(ast.NodeVisitor):
+    def __init__(self):
+        self.unsafe_imports = []
+        self.current_file = None
+        self.in_try_block = False
+        self.try_blocks = []
+
+    def visit_Try(self, node):
+        # Track that we're entering a try block
+        self.in_try_block = True
+        self.try_blocks.append(node)
+        # Visit all nodes in the try block
+        for item in node.body:
+            self.visit(item)
+        # Visit except blocks
+        for handler in node.handlers:
+            for item in handler.body:
+                self.visit(item)
+        # Visit else block if it exists
+        for item in node.orelse:
+            self.visit(item)
+        # Visit finally block if it exists
+        for item in node.finalbody:
+            self.visit(item)
+        # We're leaving the try block
+        self.try_blocks.pop()
+        self.in_try_block = len(self.try_blocks) > 0
+
+    def visit_Import(self, node):
+        # Check for direct imports of litellm_enterprise
+        for name in node.names:
+            if "litellm_enterprise" in name.name or "enterprise" in name.name:
+                if not self.in_try_block:
+                    self.unsafe_imports.append({
+                        "file": self.current_file,
+                        "line": node.lineno,
+                        "import": name.name,
+                        "context": "direct import"
+                    })
+        self.generic_visit(node)
+
+    def visit_ImportFrom(self, node):
+        # Check for from litellm_enterprise imports
+        if node.module and ("litellm_enterprise" in node.module or "enterprise" in node.module):
+            if not self.in_try_block:
+                self.unsafe_imports.append({
+                    "file": self.current_file,
+                    "line": node.lineno,
+                    "import": f"from {node.module}",
+                    "context": "from import"
+                })
+        self.generic_visit(node)
+
+def find_unsafe_enterprise_imports_in_file(file_path):
+    with open(file_path, "r") as file:
+        tree = ast.parse(file.read(), filename=file_path)
+    finder = EnterpriseImportFinder()
+    finder.current_file = file_path
+    finder.visit(tree)
+    return finder.unsafe_imports
+
+def find_unsafe_enterprise_imports_in_directory(directory):
+    unsafe_imports = []
+    for root, _, files in os.walk(directory):
+        for file in files:
+            if file.endswith(".py"):
+                file_path = os.path.join(root, file)
+                imports = find_unsafe_enterprise_imports_in_file(file_path)
+                if imports:
+                    unsafe_imports.extend(imports)
+    return unsafe_imports
+
+if __name__ == "__main__":
+    # Check for unsafe enterprise imports in the litellm directory
+    directory_path = "./litellm"
+    unsafe_imports = find_unsafe_enterprise_imports_in_directory(directory_path)
+    
+    if unsafe_imports:
+        print("🚨 UNSAFE ENTERPRISE IMPORTS FOUND (not in try-except blocks):")
+        for imp in unsafe_imports:
+            print(f"File: {imp['file']}")
+            print(f"Line: {imp['line']}")
+            print(f"Import: {imp['import']}")
+            print(f"Context: {imp['context']}")
+            print("---")
+        
+        # Raise exception to fail CI/CD
+        raise Exception(
+            "🚨 Unsafe enterprise imports found. All enterprise imports must be wrapped in try-except blocks."
+        )
+    else:
+        print("✅ No unsafe enterprise imports found.")
@@ -11,13 +11,13 @@ sys.path.insert(0, os.path.abspath("../../.."))
 from litellm_enterprise.enterprise_callbacks.send_emails.base_email import (
    BaseEmailLogger,
 )
-
-from litellm.proxy._types import Litellm_EntityType, WebhookEvent
-from litellm.types.enterprise.enterprise_callbacks.send_emails import (
+from litellm_enterprise.types.enterprise_callbacks.send_emails import (
    EmailEvent,
    SendKeyCreatedEmailEvent,
 )

+from litellm.proxy._types import Litellm_EntityType, WebhookEvent
+

@pytest.fixture
 def base_email_logger():
@@ -17,8 +17,7 @@ from litellm_enterprise.enterprise_callbacks.send_emails.endpoints import (
    router,
    update_event_settings,
 )
-
-from litellm.types.enterprise.enterprise_callbacks.send_emails import (
+from litellm_enterprise.types.enterprise_callbacks.send_emails import (
    DefaultEmailSettings,
    EmailEvent,
    EmailEventSettings,