From a44bc9f61398c940e4f123e4e775c2b7f8f3046b Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Sat, 24 May 2025 08:59:02 -0700 Subject: [PATCH] Remove + Check for unsafe `enterprise/` folder imports (#11107) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 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 * test: cleanup redundant test * fix: fix import * fix: fix linting errors --------- Co-authored-by: Ishaan Jaff Co-authored-by: JuHyun Bae --- .circleci/config.yml | 1 + docs/my-website/docs/proxy/config_settings.md | 2 + .../send_emails/base_email.py | 11 ++- .../send_emails/endpoints.py | 10 +- .../enterprise_callbacks/send_emails.py | 0 litellm/proxy/_types.py | 4 + litellm/proxy/common_utils/callback_utils.py | 84 ++++++++++++---- .../guardrails/guardrail_initializers.py | 13 ++- .../proxy/hooks/key_management_event_hooks.py | 26 +++-- .../hooks/user_management_event_hooks.py | 35 ++++--- .../customer_endpoints.py | 15 ++- .../spend_management_endpoints.py | 27 +++--- .../check_unsafe_enterprise_import.py | 95 +++++++++++++++++++ .../send_emails/test_base_email.py | 6 +- .../send_emails/test_endpoints.py | 3 +- 15 files changed, 257 insertions(+), 75 deletions(-) rename {litellm/types/enterprise => enterprise/litellm_enterprise/types}/enterprise_callbacks/send_emails.py (100%) create mode 100644 tests/code_coverage_tests/check_unsafe_enterprise_import.py diff --git a/.circleci/config.yml b/.circleci/config.yml index 2319cedbdf..779f302c94 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -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: diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 6abe507e27..d9830279f7 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -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 diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py index 779d4f2eb3..b4b128b624 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py @@ -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 diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/endpoints.py b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/endpoints.py index cc6f0be80f..61681c27ee 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/endpoints.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/endpoints.py @@ -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() diff --git a/litellm/types/enterprise/enterprise_callbacks/send_emails.py b/enterprise/litellm_enterprise/types/enterprise_callbacks/send_emails.py similarity index 100% rename from litellm/types/enterprise/enterprise_callbacks/send_emails.py rename to enterprise/litellm_enterprise/types/enterprise_callbacks/send_emails.py diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 06ac41d6a1..1802b1f25d 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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): diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index 9f4f370f33..39498202b6 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -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( diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index 468df07ea4..6685dc1393 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -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, diff --git a/litellm/proxy/hooks/key_management_event_hooks.py b/litellm/proxy/hooks/key_management_event_hooks.py index dad1cf3910..0289a92961 100644 --- a/litellm/proxy/hooks/key_management_event_hooks.py +++ b/litellm/proxy/hooks/key_management_event_hooks.py @@ -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", ""), diff --git a/litellm/proxy/hooks/user_management_event_hooks.py b/litellm/proxy/hooks/user_management_event_hooks.py index eb0b003d23..cc33f63837 100644 --- a/litellm/proxy/hooks/user_management_event_hooks.py +++ b/litellm/proxy/hooks/user_management_event_hooks.py @@ -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 ################ diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index b75e3644e1..a0c090653d 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -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) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 453d09a200..49e2ae300e 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -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}} diff --git a/tests/code_coverage_tests/check_unsafe_enterprise_import.py b/tests/code_coverage_tests/check_unsafe_enterprise_import.py new file mode 100644 index 0000000000..dee6b6eeb1 --- /dev/null +++ b/tests/code_coverage_tests/check_unsafe_enterprise_import.py @@ -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.") diff --git a/tests/litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py b/tests/litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py index d591a38dee..71608065c8 100644 --- a/tests/litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py +++ b/tests/litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py @@ -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(): diff --git a/tests/litellm/enterprise/enterprise_callbacks/send_emails/test_endpoints.py b/tests/litellm/enterprise/enterprise_callbacks/send_emails/test_endpoints.py index 8169c457c2..d1e8f37184 100644 --- a/tests/litellm/enterprise/enterprise_callbacks/send_emails/test_endpoints.py +++ b/tests/litellm/enterprise/enterprise_callbacks/send_emails/test_endpoints.py @@ -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,