diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py index a530682ae4..28f0d830f1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py @@ -59,7 +59,7 @@ class lakeraAI_Moderation(CustomGuardrail): self.async_handler = get_async_httpx_client( llm_provider=httpxSpecialProvider.GuardrailCallback ) - self.lakera_api_key = api_key or os.environ["LAKERA_API_KEY"] + self.lakera_api_key = api_key or os.environ.get("LAKERA_API_KEY") or "" self.moderation_check = moderation_check self.category_thresholds = category_thresholds self.api_base = ( diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index 732331349e..738827b7ad 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -54,7 +54,7 @@ class LakeraAIGuardrail(CustomGuardrail): self.async_handler = get_async_httpx_client( llm_provider=httpxSpecialProvider.GuardrailCallback ) - self.lakera_api_key = api_key or os.environ["LAKERA_API_KEY"] + self.lakera_api_key = api_key or os.environ.get("LAKERA_API_KEY") or "" self.project_id = project_id self.api_base = ( api_base or get_secret_str("LAKERA_API_BASE") or "https://api.lakera.ai" diff --git a/tests/proxy_unit_tests/test_proxy_setting_guardrails.py b/tests/proxy_unit_tests/test_proxy_setting_guardrails.py index b845f86b6e..d5dac59b3c 100644 --- a/tests/proxy_unit_tests/test_proxy_setting_guardrails.py +++ b/tests/proxy_unit_tests/test_proxy_setting_guardrails.py @@ -30,7 +30,8 @@ from litellm.proxy.proxy_server import ( # Replace with the actual module where def client(): filepath = os.path.dirname(os.path.abspath(__file__)) config_fp = f"{filepath}/test_configs/test_guardrails_config.yaml" - asyncio.run(initialize(config=config_fp)) + with mock.patch("litellm.proxy.proxy_server.premium_user", True): + asyncio.run(initialize(config=config_fp)) from litellm.proxy.proxy_server import app return TestClient(app)