diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index ac487fb06d0..b7b63aef4e3 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -2071,6 +2071,14 @@ async def test_custom_code_guardrail( detail="Admin access required to test custom code guardrails", ) + if os.getenv("LITELLM_ENABLE_CUSTOM_CODE_GUARDRAILS", "").lower() != "true": + return TestCustomCodeGuardrailResponse( + success=False, + error="Custom code guardrails are disabled by default. " + "Set LITELLM_ENABLE_CUSTOM_CODE_GUARDRAILS=true to enable.", + error_type="compilation", + ) + EXECUTION_TIMEOUT_SECONDS = 5 try: diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py index 58502e309ef..7486f18c4e4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py @@ -35,6 +35,7 @@ Example: block when response rejects the user (input_type response only): """ import asyncio +import os import threading from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Type, cast @@ -143,8 +144,17 @@ class CustomCodeGuardrail(CustomGuardrail): """Returns the config model for the UI.""" return CustomCodeGuardrailConfigModel + @staticmethod + def _require_custom_code_enabled() -> None: + if os.getenv("LITELLM_ENABLE_CUSTOM_CODE_GUARDRAILS", "").lower() != "true": + raise CustomCodeCompilationError( + "Custom code guardrails are disabled by default. " + "Set LITELLM_ENABLE_CUSTOM_CODE_GUARDRAILS=true to enable." + ) + def _do_compile(self) -> None: """Internal compilation method without lock. Expected to run inside _compile_lock.""" + self._require_custom_code_enabled() exec_globals = build_sandbox_globals() compiled = compile_sandboxed(self.custom_code) exec(compiled, exec_globals) # noqa: S102 diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_response_rejection_guardrail_code.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_response_rejection_guardrail_code.py index eeba2e49728..9ac1c883d4f 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_response_rejection_guardrail_code.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_response_rejection_guardrail_code.py @@ -9,6 +9,11 @@ from litellm.proxy.guardrails.guardrail_hooks.custom_code import ( ) +@pytest.fixture(autouse=True) +def enable_custom_code_guardrails(monkeypatch): + monkeypatch.setenv("LITELLM_ENABLE_CUSTOM_CODE_GUARDRAILS", "true") + + @pytest.fixture def response_rejection_guardrail(): """Guardrail instance using the response-rejection custom code.""" diff --git a/tests/test_litellm/proxy/guardrails/test_custom_code_security.py b/tests/test_litellm/proxy/guardrails/test_custom_code_security.py index 00cf3f317c9..d5177e1a1be 100644 --- a/tests/test_litellm/proxy/guardrails/test_custom_code_security.py +++ b/tests/test_litellm/proxy/guardrails/test_custom_code_security.py @@ -1,3 +1,5 @@ +import os + import pytest from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import ( @@ -6,6 +8,11 @@ from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail ) +@pytest.fixture(autouse=True) +def enable_custom_code_guardrails(monkeypatch): + monkeypatch.setenv("LITELLM_ENABLE_CUSTOM_CODE_GUARDRAILS", "true") + + # str.mro() + generator gi_code + code.replace(co_names=...) + __setattr__ # to swap a function's bytecode and read http_get's real builtins dict. BYTECODE_REWRITE_PAYLOAD = (