From 17831f45a6d5cd0c26830faf57412782ee9315d9 Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Mon, 16 Mar 2026 13:29:06 +0530 Subject: [PATCH] fix: guardrails working template policy --- litellm/integrations/custom_guardrail.py | 8 +- .../proxy/guardrails/guardrail_endpoints.py | 51 +- .../proxy/guardrails/guardrail_registry.py | 26 +- .../policy_endpoints/endpoints.py | 46 +- .../guardrails/test_guardrail_endpoints.py | 938 +++++++++++------- 5 files changed, 697 insertions(+), 372 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index aa2a8121ee8..6301622371b 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -805,7 +805,13 @@ class CustomGuardrail(CustomLogger): """ Update the guardrails litellm params in memory """ - for key, value in vars(litellm_params).items(): + # Handle both dict and Pydantic model/object + items = ( + litellm_params.items() + if isinstance(litellm_params, dict) + else vars(litellm_params).items() + ) + for key, value in items: setattr(self, key, value) def get_guardrails_messages_for_call_type( diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 2b20876ba22..97fa4e46bc8 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -335,11 +335,25 @@ async def create_guardrail( f"Immediate sync: Successfully initialized guardrail '{guardrail_name}' (ID: {guardrail_id})" ) except Exception as init_error: - verbose_proxy_logger.warning( + verbose_proxy_logger.error( f"Immediate sync: Failed to initialize guardrail '{guardrail_name}' (ID: {guardrail_id}) in memory: {init_error}" ) + # Rollback: remove the ghost row from DB + try: + await GUARDRAIL_REGISTRY.delete_guardrail_from_db( + guardrail_id=guardrail_id, prisma_client=prisma_client + ) + except Exception: + verbose_proxy_logger.error("Failed to rollback guardrail DB entry") + + raise HTTPException( + status_code=422, + detail=f"Guardrail saved but failed to initialize: {init_error}", + ) return result + except HTTPException as e: + raise e except Exception as e: verbose_proxy_logger.exception(f"Error adding guardrail to db: {e}") raise HTTPException(status_code=500, detail=str(e)) @@ -444,9 +458,30 @@ async def update_guardrail( f"Immediate sync: Successfully updated guardrail '{guardrail_name}' (ID: {guardrail_id})" ) except Exception as update_error: - verbose_proxy_logger.warning( + verbose_proxy_logger.error( f"Immediate sync: Failed to update '{guardrail_name}' (ID: {guardrail_id}) in memory: {update_error}" ) + # Rollback: restore previous guardrail data in DB + try: + await GUARDRAIL_REGISTRY.update_guardrail_in_db( + guardrail_id=guardrail_id, + guardrail=cast(Guardrail, existing_guardrail), + prisma_client=prisma_client, + ) + # Re-initialize the old guardrail in memory + IN_MEMORY_GUARDRAIL_HANDLER.update_in_memory_guardrail( + guardrail_id=guardrail_id, + guardrail=cast(Guardrail, existing_guardrail), + ) + except Exception: + verbose_proxy_logger.error( + "Failed to rollback guardrail DB entry after update failure" + ) + + raise HTTPException( + status_code=422, + detail=f"Guardrail update failed, rolled back: {update_error}", + ) return result except HTTPException as e: @@ -518,9 +553,13 @@ async def delete_guardrail( f"Immediate sync: Successfully removed guardrail '{guardrail_name}' (ID: {guardrail_id}) from memory" ) except Exception as delete_error: - verbose_proxy_logger.warning( + verbose_proxy_logger.error( f"Immediate sync: Failed to remove guardrail '{guardrail_name}' (ID: {guardrail_id}) from memory: {delete_error}" ) + raise HTTPException( + status_code=422, + detail=f"Guardrail deleted from DB but failed to remove from memory: {delete_error}", + ) return result except HTTPException as e: @@ -1101,9 +1140,13 @@ async def patch_guardrail( f"Immediate sync: Successfully updated guardrail '{guardrail_name}' (ID: {guardrail_id})" ) except Exception as update_error: - verbose_proxy_logger.warning( + verbose_proxy_logger.error( f"Immediate sync: Failed to update '{guardrail_name}' (ID: {guardrail_id}) in memory: {update_error}" ) + raise HTTPException( + status_code=422, + detail=f"Guardrail patched in DB but failed to update in memory: {update_error}", + ) return result except HTTPException as e: diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index d41be370f7b..4765f03f783 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -544,23 +544,21 @@ class InMemoryGuardrailHandler: self, guardrail_id: str, guardrail: Guardrail ) -> None: """ - Update a guardrail in memory + Update a guardrail in memory by deleting and re-initializing. - - updates the guardrail in memory - - updates the guardrail params in litellm.callback_manager + Re-initialization is necessary because guardrails like + ``ContentFilterGuardrail`` compile patterns at init time. + Simply patching attributes via ``setattr`` would leave stale + compiled state and skip validation of new patterns. """ - self.IN_MEMORY_GUARDRAILS[guardrail_id] = guardrail + # Delete old callback and in-memory references + self.delete_in_memory_guardrail(guardrail_id) - custom_guardrail_callback = self.guardrail_id_to_custom_guardrail.get( - guardrail_id - ) - if custom_guardrail_callback: - updated_litellm_params = cast( - LitellmParams, guardrail.get("litellm_params", {}) - ) - custom_guardrail_callback.update_in_memory_litellm_params( - litellm_params=updated_litellm_params - ) + # Ensure the guardrail_id is set so initialize_guardrail uses it + guardrail["guardrail_id"] = guardrail_id + + # Re-initialize (validates patterns, compiles regexes, registers callback) + self.initialize_guardrail(guardrail=guardrail) def delete_in_memory_guardrail(self, guardrail_id: str) -> None: """ diff --git a/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py b/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py index 57578d98b75..a34166ba12c 100644 --- a/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py +++ b/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py @@ -659,8 +659,10 @@ async def get_policy_templates( "yes", ) if use_local: - return _load_policy_templates_from_local_backup() + templates = _load_policy_templates_from_local_backup() + return _filter_templates_by_available_patterns(templates) + templates = [] try: from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.llms.custom_http import httpxSpecialProvider @@ -671,13 +673,51 @@ async def get_policy_templates( ) response = await async_client.get(POLICY_TEMPLATES_GITHUB_URL) if response.status_code == 200: - return response.json() + templates = response.json() except Exception as e: verbose_proxy_logger.debug( "Failed to fetch policy templates from GitHub, using local backup: %s", e ) - return _load_policy_templates_from_local_backup() + if not templates: + templates = _load_policy_templates_from_local_backup() + + return _filter_templates_by_available_patterns(templates) + + +def _filter_templates_by_available_patterns(templates: list) -> list: + """Filter out templates that reference unavailable prebuilt patterns.""" + from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.patterns import ( + PREBUILT_PATTERNS, + ) + + available_patterns = set(PREBUILT_PATTERNS.keys()) + filtered_templates = [] + + for template in templates: + is_valid = True + guardrail_definitions = template.get("guardrailDefinitions", []) + for gd in guardrail_definitions: + litellm_params = gd.get("litellm_params", {}) + if litellm_params.get("guardrail") == "litellm_content_filter": + patterns = litellm_params.get("patterns", []) + for p in patterns: + if p.get("pattern_type") == "prebuilt": + pattern_name = p.get("pattern_name") + if pattern_name and pattern_name not in available_patterns: + is_valid = False + break + if not is_valid: + break + + if is_valid: + filtered_templates.append(template) + else: + verbose_proxy_logger.debug( + f"Filtering out policy template '{template.get('id')}' because it references unavailable patterns" + ) + + return filtered_templates class EnrichTemplateRequest(BaseModel): diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index ca224726361..8a9a80b8426 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -2,7 +2,7 @@ import json import os import sys from datetime import datetime -from typing import Dict, List, Optional +from typing import List, Optional from unittest.mock import AsyncMock import pytest @@ -35,7 +35,6 @@ from litellm.proxy.guardrails.guardrail_endpoints import ( MOCK_ADMIN_USER = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) from litellm.proxy.guardrails.guardrail_registry import ( - IN_MEMORY_GUARDRAIL_HANDLER, InMemoryGuardrailHandler, ) from litellm.types.guardrails import ( @@ -72,7 +71,7 @@ MOCK_CONFIG_GUARDRAIL = { MOCK_GUARDRAIL = Guardrail( guardrail_name=MOCK_CONFIG_GUARDRAIL["guardrail_name"], litellm_params=LitellmParams(**MOCK_CONFIG_GUARDRAIL["litellm_params"]), - guardrail_info=MOCK_CONFIG_GUARDRAIL["guardrail_info"] + guardrail_info=MOCK_CONFIG_GUARDRAIL["guardrail_info"], ) MOCK_CREATE_REQUEST = CreateGuardrailRequest(guardrail=MOCK_GUARDRAIL) @@ -80,7 +79,7 @@ MOCK_UPDATE_REQUEST = UpdateGuardrailRequest(guardrail=MOCK_GUARDRAIL) MOCK_PATCH_REQUEST = PatchGuardrailRequest( guardrail_name="Updated Test Guardrail", litellm_params={"guardrail": "updated.guardrail", "mode": "post_call"}, - guardrail_info={"description": "Updated test guardrail"} + guardrail_info={"description": "Updated test guardrail"}, ) @@ -111,19 +110,22 @@ def mock_in_memory_handler(mocker): mock_handler.delete_in_memory_guardrail = mocker.Mock() return mock_handler + @pytest.fixture def mock_guardrail_registry(mocker): """Mock GuardrailRegistry for testing""" mock_registry = mocker.Mock() - mock_registry.add_guardrail_to_db = AsyncMock(return_value={ - **MOCK_DB_GUARDRAIL, - "guardrail_id": "new-test-guardrail-id" - }) + mock_registry.add_guardrail_to_db = AsyncMock( + return_value={**MOCK_DB_GUARDRAIL, "guardrail_id": "new-test-guardrail-id"} + ) mock_registry.delete_guardrail_from_db = AsyncMock(return_value=MOCK_DB_GUARDRAIL) - mock_registry.get_guardrail_by_id_from_db = AsyncMock(return_value=MOCK_DB_GUARDRAIL) + mock_registry.get_guardrail_by_id_from_db = AsyncMock( + return_value=MOCK_DB_GUARDRAIL + ) mock_registry.update_guardrail_in_db = AsyncMock(return_value=MOCK_DB_GUARDRAIL) return mock_registry + @pytest.mark.asyncio async def test_list_guardrails_v2_with_db_and_config( mocker, mock_prisma_client, mock_in_memory_handler @@ -199,7 +201,11 @@ async def test_list_guardrails_v2_masks_sensitive_data_in_db_guardrails(mocker): if isinstance(litellm_params, dict): params = litellm_params else: - params = litellm_params.model_dump() if hasattr(litellm_params, "model_dump") else dict(litellm_params) + params = ( + litellm_params.model_dump() + if hasattr(litellm_params, "model_dump") + else dict(litellm_params) + ) # Sensitive keys (containing "key", "secret", "token", etc.) should be masked assert params["api_key"] != "sk-1234567890abcdef" @@ -228,9 +234,7 @@ async def test_list_guardrails_v2_masks_sensitive_data_in_config_guardrails(mock mock_prisma_client = mocker.Mock() mock_prisma_client.db = mocker.Mock() mock_prisma_client.db.litellm_guardrailstable = mocker.Mock() - mock_prisma_client.db.litellm_guardrailstable.find_many = AsyncMock( - return_value=[] - ) + mock_prisma_client.db.litellm_guardrailstable.find_many = AsyncMock(return_value=[]) mock_in_memory_handler = mocker.Mock() mock_in_memory_handler.list_in_memory_guardrails.return_value = [ @@ -251,7 +255,11 @@ async def test_list_guardrails_v2_masks_sensitive_data_in_config_guardrails(mock if isinstance(litellm_params, dict): params = litellm_params else: - params = litellm_params.model_dump() if hasattr(litellm_params, "model_dump") else dict(litellm_params) + params = ( + litellm_params.model_dump() + if hasattr(litellm_params, "model_dump") + else dict(litellm_params) + ) # Sensitive keys should be masked assert params["api_key"] != "my-secret-bedrock-key" @@ -336,7 +344,6 @@ def test_get_provider_specific_params(): pytest.skip("Azure config model not available") fields = _get_fields_from_model(config_model) - print("FIELDS", fields) # Test that we get the expected nested structure assert isinstance(fields, dict) @@ -352,12 +359,12 @@ def test_get_provider_specific_params(): fields["api_key"]["description"] == "API key for the Azure Content Safety Prompt Shield guardrail" ) - assert fields["api_key"]["required"] == False + assert fields["api_key"]["required"] is False assert fields["api_key"]["type"] == "string" # Should be string, not None # Check the structure of the nested optional_params field assert fields["optional_params"]["type"] == "nested" - assert fields["optional_params"]["required"] == True + assert fields["optional_params"]["required"] is True assert "fields" in fields["optional_params"] # Check nested fields within optional_params @@ -374,7 +381,7 @@ def test_get_provider_specific_params(): nested_fields["severity_threshold"]["description"] == "Severity threshold for the Azure Content Safety Text Moderation guardrail across all categories" ) - assert nested_fields["severity_threshold"]["required"] == False + assert nested_fields["severity_threshold"]["required"] is False assert ( nested_fields["severity_threshold"]["type"] == "number" ) # Should be number, not None @@ -390,9 +397,8 @@ def test_get_provider_specific_params(): def test_optional_params_not_returned_when_not_overridden(): """Test that optional_params is not returned when the config model doesn't override it""" - from typing import Optional - from pydantic import BaseModel, Field + from pydantic import Field from litellm.proxy.guardrails.guardrail_endpoints import _get_fields_from_model from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel @@ -413,13 +419,11 @@ def test_optional_params_not_returned_when_not_overridden(): # Get fields from the model fields = _get_fields_from_model(TestGuardrailConfig) - print("FIELDS", fields) assert "optional_params" not in fields def test_optional_params_returned_when_properly_overridden(): """Test that optional_params IS returned when the config model properly overrides it""" - from typing import Optional from pydantic import BaseModel, Field @@ -451,14 +455,13 @@ def test_optional_params_returned_when_properly_overridden(): # Get fields from the model fields = _get_fields_from_model(TestGuardrailConfigWithOptionalParams) - print("FIELDS", fields) assert "optional_params" in fields @pytest.mark.asyncio async def test_bedrock_guardrail_prepare_request_with_api_key(): """Test _prepare_request method uses Bearer token when api_key is provided in data""" - from unittest.mock import Mock, patch + from unittest.mock import Mock from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockGuardrail, @@ -466,27 +469,23 @@ async def test_bedrock_guardrail_prepare_request_with_api_key(): # Setup guardrail hook guardrail_hook = BedrockGuardrail( - guardrailIdentifier="test-guardrail-id", - guardrailVersion="1" + guardrailIdentifier="test-guardrail-id", guardrailVersion="1" ) mock_credentials = Mock() - test_data = { - "source": "INPUT", - "content": [{"text": {"text": "test content"}}] - } - + test_data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]} + prepared_request = guardrail_hook._prepare_request( credentials=mock_credentials, data=test_data, optional_params={}, aws_region_name="us-east-1", - api_key="test-bearer-token-123" + api_key="test-bearer-token-123", ) - + # Verify Bearer token is used in Authorization header assert "Authorization" in prepared_request.headers assert prepared_request.headers["Authorization"] == "Bearer test-bearer-token-123" - + # Verify URL is correct expected_url = "https://bedrock-runtime.us-east-1.amazonaws.com/guardrail/test-guardrail-id/version/1/apply" assert prepared_request.url == expected_url @@ -503,45 +502,44 @@ async def test_bedrock_guardrail_prepare_request_without_api_key(): # Setup guardrail hook guardrail_hook = BedrockGuardrail( - guardrailIdentifier="test-guardrail-id", - guardrailVersion="1" + guardrailIdentifier="test-guardrail-id", guardrailVersion="1" ) - + # Mock credentials mock_credentials = Mock() - + # Test data without api_key - test_data = { - "source": "INPUT", - "content": [{"text": {"text": "test content"}}] - } - - with patch("litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.get_secret_str") as mock_get_secret, \ - patch("botocore.auth.SigV4Auth") as mock_sigv4_auth, \ - patch("botocore.awsrequest.AWSRequest") as mock_aws_request: - + test_data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]} + + with patch( + "litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.get_secret_str" + ) as mock_get_secret, patch("botocore.auth.SigV4Auth") as mock_sigv4_auth, patch( + "botocore.awsrequest.AWSRequest" + ) as mock_aws_request: # Mock no AWS_BEARER_TOKEN_BEDROCK mock_get_secret.return_value = None - + # Mock SigV4Auth mock_sigv4_instance = Mock() mock_sigv4_auth.return_value = mock_sigv4_instance - + # Mock AWSRequest mock_request_instance = Mock() mock_request_instance.prepare.return_value = Mock() mock_aws_request.return_value = mock_request_instance - + # Call _prepare_request - prepared_request = guardrail_hook._prepare_request( + guardrail_hook._prepare_request( credentials=mock_credentials, data=test_data, optional_params={}, - aws_region_name="us-east-1" + aws_region_name="us-east-1", ) - + # Verify SigV4 auth was used - mock_sigv4_auth.assert_called_once_with(mock_credentials, "bedrock", "us-east-1") + mock_sigv4_auth.assert_called_once_with( + mock_credentials, "bedrock", "us-east-1" + ) mock_sigv4_instance.add_auth.assert_called_once() @@ -556,34 +554,30 @@ async def test_bedrock_guardrail_prepare_request_with_bearer_token_env(): # Setup guardrail hook guardrail_hook = BedrockGuardrail( - guardrailIdentifier="test-guardrail-id", - guardrailVersion="1" + guardrailIdentifier="test-guardrail-id", guardrailVersion="1" ) - + # Mock credentials mock_credentials = Mock() - + # Test data without api_key - test_data = { - "source": "INPUT", - "content": [{"text": {"text": "test content"}}] - } - - with patch("litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.get_secret_str") as mock_get_secret, \ - patch("botocore.awsrequest.AWSRequest") as mock_aws_request: - + test_data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]} + + with patch( + "litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.get_secret_str" + ) as mock_get_secret, patch("botocore.awsrequest.AWSRequest") as mock_aws_request: mock_get_secret.return_value = "env-bearer-token-456" mock_request_instance = Mock() mock_request_instance.prepare.return_value = Mock() mock_aws_request.return_value = mock_request_instance - - prepared_request = guardrail_hook._prepare_request( + + guardrail_hook._prepare_request( credentials=mock_credentials, data=test_data, optional_params={}, - aws_region_name="us-east-1" + aws_region_name="us-east-1", ) - + # Verify Bearer token from environment is used mock_aws_request.assert_called_once() call_args = mock_aws_request.call_args @@ -599,45 +593,51 @@ async def test_bedrock_guardrail_make_api_request_passes_api_key(): from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockGuardrail, ) - + guardrail_hook = BedrockGuardrail( - guardrailIdentifier="test-guardrail-id", - guardrailVersion="1" + guardrailIdentifier="test-guardrail-id", guardrailVersion="1" ) - + guardrail_hook.async_handler = Mock() mock_response = Mock() mock_response.status_code = 200 mock_response.json.return_value = {"action": "NONE", "outputs": []} - - test_request_data = { - "api_key": "test-api-key-789" - } - - with patch.object(guardrail_hook.async_handler, "post", AsyncMock(return_value=mock_response)), \ - patch.object(guardrail_hook, "_load_credentials") as mock_load_creds, \ - patch.object(guardrail_hook, "convert_to_bedrock_format") as mock_convert, \ - patch.object(guardrail_hook, "get_guardrail_dynamic_request_body_params") as mock_get_params, \ - patch.object(guardrail_hook, "add_standard_logging_guardrail_information_to_request_data"), \ - patch("botocore.awsrequest.AWSRequest") as mock_aws_request: - + + test_request_data = {"api_key": "test-api-key-789"} + + with patch.object( + guardrail_hook.async_handler, "post", AsyncMock(return_value=mock_response) + ), patch.object( + guardrail_hook, "_load_credentials" + ) as mock_load_creds, patch.object( + guardrail_hook, "convert_to_bedrock_format" + ) as mock_convert, patch.object( + guardrail_hook, "get_guardrail_dynamic_request_body_params" + ) as mock_get_params, patch.object( + guardrail_hook, "add_standard_logging_guardrail_information_to_request_data" + ), patch( + "botocore.awsrequest.AWSRequest" + ) as mock_aws_request: mock_load_creds.return_value = (Mock(), "us-east-1") mock_convert.return_value = {"source": "INPUT", "content": []} mock_get_params.return_value = {} - + mock_request_instance = Mock() mock_request_instance.url = "test-url" mock_request_instance.body = b"test-body" - mock_request_instance.headers = {"Content-Type": "application/json", "Authorization": "Bearer test-api-key-789"} + mock_request_instance.headers = { + "Content-Type": "application/json", + "Authorization": "Bearer test-api-key-789", + } mock_request_instance.prepare.return_value = Mock() mock_aws_request.return_value = mock_request_instance - + await guardrail_hook.make_bedrock_api_request( source="INPUT", messages=[{"role": "user", "content": "test"}], - request_data=test_request_data + request_data=test_request_data, ) - + # Verify _prepare_request was invoked and used the api_key mock_aws_request.assert_called_once() call_args = mock_aws_request.call_args @@ -645,340 +645,401 @@ async def test_bedrock_guardrail_make_api_request_passes_api_key(): assert headers["Authorization"] == "Bearer test-api-key-789" -@pytest.mark.parametrize("scenario,expected_result,expected_exception", [ - ( - "success_with_sync", - "new-test-guardrail-id", - None - ), - ( - "success_sync_fails", - "new-test-guardrail-id", - None - ), - ( - "database_failure", - None, - HTTPException - ), - ( - "no_prisma_client", - None, - HTTPException - ), -], ids=[ - "success_with_immediate_sync", - "success_but_sync_fails", - "database_error", - "missing_prisma_client" -]) +@pytest.mark.parametrize( + "scenario,expected_result,expected_exception", + [ + ("success_with_sync", "new-test-guardrail-id", None), + ("success_sync_fails", None, HTTPException), + ("database_failure", None, HTTPException), + ("no_prisma_client", None, HTTPException), + ], + ids=[ + "success_with_immediate_sync", + "success_but_sync_fails", + "database_error", + "missing_prisma_client", + ], +) @pytest.mark.asyncio async def test_create_guardrail_endpoint( - scenario, expected_result, expected_exception, - mocker, mock_guardrail_registry, mock_in_memory_handler + scenario, + expected_result, + expected_exception, + mocker, + mock_guardrail_registry, + mock_in_memory_handler, ): """Test create_guardrail endpoint with different scenarios""" - + # Configure mocks based on scenario - mock_logger = None if scenario == "success_with_sync": mock_prisma_client = mocker.Mock() mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_guardrail_registry) - mocker.patch("litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", mock_in_memory_handler) - + mocker.patch( + "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", + mock_guardrail_registry, + ) + mocker.patch( + "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", + mock_in_memory_handler, + ) + elif scenario == "success_sync_fails": mock_prisma_client = mocker.Mock() - mock_in_memory_handler.initialize_guardrail.side_effect = Exception("Sync failed") - mock_logger = mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger") - + mock_in_memory_handler.initialize_guardrail.side_effect = Exception( + "Sync failed" + ) + mocker.patch( + "litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger" + ) + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_guardrail_registry) - mocker.patch("litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", mock_in_memory_handler) - + mocker.patch( + "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", + mock_guardrail_registry, + ) + mocker.patch( + "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", + mock_in_memory_handler, + ) + elif scenario == "database_failure": mock_prisma_client = mocker.Mock() - mock_guardrail_registry.add_guardrail_to_db.side_effect = Exception("Database error") - - mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_guardrail_registry) - + mock_guardrail_registry.add_guardrail_to_db.side_effect = Exception( + "Database error" + ) + + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mocker.patch( + "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", + mock_guardrail_registry, + ) + elif scenario == "no_prisma_client": mocker.patch("litellm.proxy.proxy_server.prisma_client", None) - + # Run the test if expected_exception: with pytest.raises(expected_exception) as exc_info: - await create_guardrail(MOCK_CREATE_REQUEST, user_api_key_dict=MOCK_ADMIN_USER) + await create_guardrail( + MOCK_CREATE_REQUEST, user_api_key_dict=MOCK_ADMIN_USER + ) if scenario == "database_failure": assert "Database error" in str(exc_info.value.detail) elif scenario == "no_prisma_client": assert "Prisma client not initialized" in str(exc_info.value.detail) - + elif scenario == "success_sync_fails": + assert exc_info.value.status_code == 422 + assert "failed to initialize" in str(exc_info.value.detail) + # Verify rollback + mock_guardrail_registry.delete_guardrail_from_db.assert_called_once() else: - result = await create_guardrail(MOCK_CREATE_REQUEST, user_api_key_dict=MOCK_ADMIN_USER) - + result = await create_guardrail( + MOCK_CREATE_REQUEST, user_api_key_dict=MOCK_ADMIN_USER + ) + assert result["guardrail_id"] == expected_result assert result["guardrail_name"] == "Test DB Guardrail" - - mock_guardrail_registry.add_guardrail_to_db.assert_called_once_with( - guardrail=MOCK_CREATE_REQUEST.guardrail, - prisma_client=mocker.ANY - ) - - mock_in_memory_handler.initialize_guardrail.assert_called_once() - - if scenario == "success_sync_fails": - assert mock_logger is not None - mock_logger.warning.assert_called_once() - assert "Failed to initialize guardrail" in str(mock_logger.warning.call_args) -@pytest.mark.parametrize("scenario,expected_result,expected_exception", [ - ( - "success_with_sync", - "test-db-guardrail", - None - ), - ( - "success_sync_fails", - "test-db-guardrail", - None - ), - ( - "database_failure", - None, - HTTPException - ), - ( - "no_prisma_client", - None, - HTTPException - ), -], ids=[ - "success_with_immediate_sync", - "success_but_sync_fails", - "database_error", - "missing_prisma_client" -]) + mock_guardrail_registry.add_guardrail_to_db.assert_called_once_with( + guardrail=MOCK_CREATE_REQUEST.guardrail, prisma_client=mocker.ANY + ) + + mock_in_memory_handler.initialize_guardrail.assert_called_once() + + +@pytest.mark.parametrize( + "scenario,expected_result,expected_exception", + [ + ("success_with_sync", "test-db-guardrail", None), + ("success_sync_fails", None, HTTPException), + ("database_failure", None, HTTPException), + ("no_prisma_client", None, HTTPException), + ], + ids=[ + "success_with_immediate_sync", + "success_but_sync_fails", + "database_error", + "missing_prisma_client", + ], +) @pytest.mark.asyncio async def test_update_guardrail_endpoint( - scenario, expected_result, expected_exception, - mocker, mock_guardrail_registry, mock_in_memory_handler + scenario, + expected_result, + expected_exception, + mocker, + mock_guardrail_registry, + mock_in_memory_handler, ): """Test update_guardrail endpoint with different scenarios""" - + # Configure mocks based on scenario - mock_logger = None if scenario == "success_with_sync": mock_prisma_client = mocker.Mock() mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_guardrail_registry) - mocker.patch("litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", mock_in_memory_handler) - + mocker.patch( + "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", + mock_guardrail_registry, + ) + mocker.patch( + "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", + mock_in_memory_handler, + ) + elif scenario == "success_sync_fails": mock_prisma_client = mocker.Mock() - mock_in_memory_handler.update_in_memory_guardrail.side_effect = Exception("Sync failed") - mock_logger = mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger") - + mock_in_memory_handler.update_in_memory_guardrail.side_effect = Exception( + "Sync failed" + ) + mocker.patch( + "litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger" + ) + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_guardrail_registry) - mocker.patch("litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", mock_in_memory_handler) - + mocker.patch( + "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", + mock_guardrail_registry, + ) + mocker.patch( + "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", + mock_in_memory_handler, + ) + elif scenario == "database_failure": mock_prisma_client = mocker.Mock() - mock_guardrail_registry.update_guardrail_in_db.side_effect = Exception("Database error") - - mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_guardrail_registry) - + mock_guardrail_registry.update_guardrail_in_db.side_effect = Exception( + "Database error" + ) + + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mocker.patch( + "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", + mock_guardrail_registry, + ) + elif scenario == "no_prisma_client": mocker.patch("litellm.proxy.proxy_server.prisma_client", None) - + # Run the test if expected_exception: with pytest.raises(expected_exception) as exc_info: - await update_guardrail("test-guardrail-id", MOCK_UPDATE_REQUEST, user_api_key_dict=MOCK_ADMIN_USER) + await update_guardrail( + "test-guardrail-id", + MOCK_UPDATE_REQUEST, + user_api_key_dict=MOCK_ADMIN_USER, + ) if scenario == "database_failure": assert "Database error" in str(exc_info.value.detail) elif scenario == "no_prisma_client": assert "Prisma client not initialized" in str(exc_info.value.detail) - + elif scenario == "success_sync_fails": + assert exc_info.value.status_code == 422 + assert "rolled back" in str(exc_info.value.detail) else: - result = await update_guardrail("test-guardrail-id", MOCK_UPDATE_REQUEST, user_api_key_dict=MOCK_ADMIN_USER) - + result = await update_guardrail( + "test-guardrail-id", MOCK_UPDATE_REQUEST, user_api_key_dict=MOCK_ADMIN_USER + ) + assert result["guardrail_id"] == expected_result assert result["guardrail_name"] == "Test DB Guardrail" - + mock_guardrail_registry.update_guardrail_in_db.assert_called_once_with( guardrail_id="test-guardrail-id", guardrail=MOCK_UPDATE_REQUEST.guardrail, - prisma_client=mocker.ANY + prisma_client=mocker.ANY, ) - - mock_in_memory_handler.update_in_memory_guardrail.assert_called_once_with( - guardrail_id="test-guardrail-id", - guardrail=mocker.ANY - ) - - if scenario == "success_sync_fails": - assert mock_logger is not None - mock_logger.warning.assert_called_once() - assert "Failed to update" in str(mock_logger.warning.call_args) -@pytest.mark.parametrize("scenario,expected_result,expected_exception", [ - ( - "success_with_sync", - "test-db-guardrail", - None - ), - ( - "success_sync_fails", - "test-db-guardrail", - None - ), - ( - "database_failure", - None, - HTTPException - ), - ( - "no_prisma_client", - None, - HTTPException - ), -], ids=[ - "success_with_immediate_sync", - "success_but_sync_fails", - "database_error", - "missing_prisma_client" -]) + mock_in_memory_handler.update_in_memory_guardrail.assert_called_once_with( + guardrail_id="test-guardrail-id", guardrail=mocker.ANY + ) + + +@pytest.mark.parametrize( + "scenario,expected_result,expected_exception", + [ + ("success_with_sync", "test-db-guardrail", None), + ("success_sync_fails", None, HTTPException), + ("database_failure", None, HTTPException), + ("no_prisma_client", None, HTTPException), + ], + ids=[ + "success_with_immediate_sync", + "success_but_sync_fails", + "database_error", + "missing_prisma_client", + ], +) @pytest.mark.asyncio async def test_patch_guardrail_endpoint( - scenario, expected_result, expected_exception, - mocker, mock_guardrail_registry, mock_in_memory_handler + scenario, + expected_result, + expected_exception, + mocker, + mock_guardrail_registry, + mock_in_memory_handler, ): """Test patch_guardrail endpoint with different scenarios""" - + # Configure mocks based on scenario mock_logger = None if scenario == "success_with_sync": mock_prisma_client = mocker.Mock() mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock() mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_guardrail_registry) - mocker.patch("litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", mock_in_memory_handler) - + mocker.patch( + "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", + mock_guardrail_registry, + ) + mocker.patch( + "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", + mock_in_memory_handler, + ) + elif scenario == "success_sync_fails": mock_prisma_client = mocker.Mock() - mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock(side_effect=Exception("Sync failed")) - mock_logger = mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger") - + mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock( + side_effect=Exception("Sync failed") + ) + mock_logger = mocker.patch( + "litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger" + ) + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_guardrail_registry) - mocker.patch("litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", mock_in_memory_handler) - + mocker.patch( + "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", + mock_guardrail_registry, + ) + mocker.patch( + "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", + mock_in_memory_handler, + ) + elif scenario == "database_failure": mock_prisma_client = mocker.Mock() - mock_guardrail_registry.update_guardrail_in_db.side_effect = Exception("Database error") - - mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_guardrail_registry) - + mock_guardrail_registry.update_guardrail_in_db.side_effect = Exception( + "Database error" + ) + + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mocker.patch( + "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", + mock_guardrail_registry, + ) + elif scenario == "no_prisma_client": mocker.patch("litellm.proxy.proxy_server.prisma_client", None) - + # Run the test if expected_exception: with pytest.raises(expected_exception) as exc_info: - await patch_guardrail("test-guardrail-id", MOCK_PATCH_REQUEST, user_api_key_dict=MOCK_ADMIN_USER) - - if scenario == "database_failure": - assert "Database error" in str(exc_info.value.detail) - elif scenario == "no_prisma_client": - assert "Prisma client not initialized" in str(exc_info.value.detail) + await patch_guardrail( + "test-guardrail-id", + MOCK_PATCH_REQUEST, + user_api_key_dict=MOCK_ADMIN_USER, + ) + if scenario == "success_sync_fails": + assert exc_info.value.status_code == 422 + assert "failed to update in memory" in str(exc_info.value.detail) else: - result = await patch_guardrail("test-guardrail-id", MOCK_PATCH_REQUEST, user_api_key_dict=MOCK_ADMIN_USER) - + result = await patch_guardrail( + "test-guardrail-id", MOCK_PATCH_REQUEST, user_api_key_dict=MOCK_ADMIN_USER + ) + assert result["guardrail_id"] == expected_result assert result["guardrail_name"] == "Test DB Guardrail" - + mock_guardrail_registry.update_guardrail_in_db.assert_called_once() - + mock_in_memory_handler.sync_guardrail_from_db.assert_called_once_with( guardrail=mocker.ANY ) - + if scenario == "success_sync_fails": assert mock_logger is not None mock_logger.warning.assert_called_once() assert "Failed to update" in str(mock_logger.warning.call_args) -@pytest.mark.parametrize("scenario,expected_result,expected_exception", [ - ( - "success_with_sync", - "test-db-guardrail", - None - ), - ( - "success_sync_fails", - "test-db-guardrail", - None - ), -], ids=[ - "success_with_immediate_sync", - "success_but_sync_fails" -]) + +@pytest.mark.parametrize( + "scenario,expected_result,expected_exception", + [ + ("success_with_sync", "test-db-guardrail", None), + ("success_sync_fails", "test-db-guardrail", HTTPException), + ], + ids=["success_with_immediate_sync", "success_but_sync_fails"], +) @pytest.mark.asyncio async def test_delete_guardrail_endpoint( - scenario, expected_result, expected_exception, - mocker, mock_guardrail_registry, mock_in_memory_handler + scenario, + expected_result, + expected_exception, + mocker, + mock_guardrail_registry, + mock_in_memory_handler, ): """Test delete_guardrail endpoint with different scenarios""" - + # Configure mocks based on scenario mock_prisma_client = mocker.Mock() - mock_logger = None - + if scenario == "success_with_sync": mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_guardrail_registry) - mocker.patch("litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", mock_in_memory_handler) - + mocker.patch( + "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", + mock_guardrail_registry, + ) + mocker.patch( + "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", + mock_in_memory_handler, + ) + elif scenario == "success_sync_fails": - mock_in_memory_handler.delete_in_memory_guardrail.side_effect = Exception("Sync failed") - mock_logger = mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger") + mock_in_memory_handler.delete_in_memory_guardrail.side_effect = Exception( + "Sync failed" + ) + mocker.patch( + "litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger" + ) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_guardrail_registry) - mocker.patch("litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", mock_in_memory_handler) - + mocker.patch( + "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", + mock_guardrail_registry, + ) + mocker.patch( + "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", + mock_in_memory_handler, + ) + if expected_exception: - with pytest.raises(expected_exception): - await delete_guardrail(guardrail_id=expected_result, user_api_key_dict=MOCK_ADMIN_USER) + with pytest.raises(expected_exception) as exc_info: + await delete_guardrail( + guardrail_id=expected_result, user_api_key_dict=MOCK_ADMIN_USER + ) + + if scenario == "success_sync_fails": + assert exc_info.value.status_code == 422 + assert "failed to remove from memory" in str(exc_info.value.detail) else: - result = await delete_guardrail(guardrail_id=expected_result, user_api_key_dict=MOCK_ADMIN_USER) - + result = await delete_guardrail( + guardrail_id=expected_result, user_api_key_dict=MOCK_ADMIN_USER + ) + assert result == MOCK_DB_GUARDRAIL - + mock_guardrail_registry.get_guardrail_by_id_from_db.assert_called_once_with( - guardrail_id=expected_result, - prisma_client=mock_prisma_client + guardrail_id=expected_result, prisma_client=mock_prisma_client ) mock_guardrail_registry.delete_guardrail_from_db.assert_called_once_with( - guardrail_id=expected_result, - prisma_client=mock_prisma_client + guardrail_id=expected_result, prisma_client=mock_prisma_client ) - + mock_in_memory_handler.delete_in_memory_guardrail.assert_called_once_with( guardrail_id=expected_result ) - - if scenario == "success_sync_fails": - assert mock_logger is not None - mock_logger.warning.assert_called_once() - assert "Failed to remove guardrail" in str(mock_logger.warning.call_args) @pytest.mark.asyncio @@ -991,21 +1052,22 @@ async def test_apply_guardrail_not_found(mocker): # Mock the GUARDRAIL_REGISTRY to return None (guardrail not found) mock_registry = mocker.Mock() mock_registry.get_initialized_guardrail_callback.return_value = None - mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry) - + mocker.patch( + "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry + ) + # Create request request = ApplyGuardrailRequest( - guardrail_name="non-existent-guardrail", - text="Test input text" + guardrail_name="non-existent-guardrail", text="Test input text" ) - + # Mock user auth mock_user_auth = UserAPIKeyAuth() - + # Call endpoint and expect ProxyException with pytest.raises(ProxyException) as exc_info: await apply_guardrail(request=request, user_api_key_dict=mock_user_auth) - + # Verify error details assert str(exc_info.value.code) == "404" assert "not found" in str(exc_info.value.message).lower() @@ -1023,28 +1085,30 @@ async def test_apply_guardrail_execution_error(mocker): mock_guardrail.apply_guardrail = AsyncMock( side_effect=Exception("Bedrock guardrail failed: Violated guardrail policy") ) - + # Mock the GUARDRAIL_REGISTRY mock_registry = mocker.Mock() mock_registry.get_initialized_guardrail_callback.return_value = mock_guardrail - mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry) - + mocker.patch( + "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry + ) + # Create request request = ApplyGuardrailRequest( - guardrail_name="test-guardrail", - text="Test input text with forbidden content" + guardrail_name="test-guardrail", text="Test input text with forbidden content" ) - + # Mock user auth mock_user_auth = UserAPIKeyAuth() - + # Call endpoint and expect ProxyException with pytest.raises(ProxyException) as exc_info: await apply_guardrail(request=request, user_api_key_dict=mock_user_auth) - + # Verify error is properly handled assert "Bedrock guardrail failed" in str(exc_info.value.message) + @pytest.mark.asyncio async def test_get_guardrail_info_endpoint_config_guardrail(mocker): """ @@ -1059,17 +1123,22 @@ async def test_get_guardrail_info_endpoint_config_guardrail(mocker): # Mock the GUARDRAIL_REGISTRY to return None from DB (so it checks config) mock_registry = mocker.Mock() mock_registry.get_guardrail_by_id_from_db = AsyncMock(return_value=None) - mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry) + mocker.patch( + "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry + ) # Mock IN_MEMORY_GUARDRAIL_HANDLER at its source to return config guardrail mock_in_memory_handler = mocker.Mock() mock_in_memory_handler.get_guardrail_by_id.return_value = MOCK_CONFIG_GUARDRAIL - mocker.patch("litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", mock_in_memory_handler) + mocker.patch( + "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", + mock_in_memory_handler, + ) # Mock _get_masked_values to return values as-is mocker.patch( "litellm.litellm_core_utils.litellm_logging._get_masked_values", - side_effect=lambda x, **kwargs: x + side_effect=lambda x, **kwargs: x, ) # Call endpoint and expect GuardrailInfoResponse @@ -1081,6 +1150,7 @@ async def test_get_guardrail_info_endpoint_config_guardrail(mocker): assert result.guardrail_name == "Test Config Guardrail" assert result.guardrail_definition_location == "config" + @pytest.mark.asyncio async def test_get_guardrail_info_endpoint_db_guardrail(mocker): """ @@ -1094,13 +1164,20 @@ async def test_get_guardrail_info_endpoint_db_guardrail(mocker): # Mock the GUARDRAIL_REGISTRY to return a guardrail from DB mock_registry = mocker.Mock() - mock_registry.get_guardrail_by_id_from_db = AsyncMock(return_value=MOCK_DB_GUARDRAIL) - mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry) + mock_registry.get_guardrail_by_id_from_db = AsyncMock( + return_value=MOCK_DB_GUARDRAIL + ) + mocker.patch( + "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry + ) # Mock IN_MEMORY_GUARDRAIL_HANDLER to return None mock_in_memory_handler = mocker.Mock() mock_in_memory_handler.get_guardrail_by_id.return_value = None - mocker.patch("litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", mock_in_memory_handler) + mocker.patch( + "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", + mock_in_memory_handler, + ) # Call endpoint and expect GuardrailInfoResponse result = await get_guardrail_info(guardrail_id="test-db-guardrail") @@ -1122,7 +1199,10 @@ class TestBuildFieldDict: from litellm.proxy.guardrails.guardrail_endpoints import _build_field_dict field = MagicMock() - field.json_schema_extra = {"ui_type": "multiselect", "options": ["python", "javascript"]} + field.json_schema_extra = { + "ui_type": "multiselect", + "options": ["python", "javascript"], + } result = _build_field_dict( field=field, @@ -1153,6 +1233,8 @@ class TestBuildFieldDict: assert result["type"] == "bool" assert result["required"] is True + + # --- Team guardrail registration (register / submissions) --- MOCK_REGISTER_REQUEST = RegisterGuardrailRequest( @@ -1198,7 +1280,11 @@ async def test_register_guardrail_rejects_non_generic_api(mocker): mocker.patch("litellm.proxy.proxy_server.prisma_client", mocker.Mock()) req = RegisterGuardrailRequest( guardrail_name="other-guard", - litellm_params={"guardrail": "bedrock", "mode": "pre_call", "api_base": "https://x.com"}, + litellm_params={ + "guardrail": "bedrock", + "mode": "pre_call", + "api_base": "https://x.com", + }, ) user = UserAPIKeyAuth(user_id="u1", user_email="a@b.com", team_id="team-1") @@ -1256,7 +1342,10 @@ async def test_list_guardrail_submissions_success(mocker): guardrail_name="pending-guard", status="pending_review", team_id="t1", - litellm_params={"guardrail": "generic_guardrail_api", "api_base": "https://x.com"}, + litellm_params={ + "guardrail": "generic_guardrail_api", + "api_base": "https://x.com", + }, guardrail_info={ "description": "A guard", "submitted_by_user_id": "u1", @@ -1362,7 +1451,11 @@ async def test_approve_guardrail_submission_success(mocker): guardrail_id="approve-me", guardrail_name="my-guard", status="pending_review", - litellm_params={"guardrail": "generic_guardrail_api", "mode": "pre_call", "api_base": "https://g.com"}, + litellm_params={ + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "api_base": "https://g.com", + }, guardrail_info={}, ) mock_prisma.db.litellm_guardrailstable.find_unique = AsyncMock(return_value=row) @@ -1421,7 +1514,9 @@ async def test_reject_guardrail_submission_success(mocker): async def test_reject_guardrail_submission_not_pending(mocker): """Reject returns 400 when status is not pending_review (e.g. already active).""" mock_prisma = mocker.Mock() - row = mocker.Mock(guardrail_id="already-active", guardrail_name="g", status="active") + row = mocker.Mock( + guardrail_id="already-active", guardrail_name="g", status="active" + ) mock_prisma.db.litellm_guardrailstable.find_unique = AsyncMock(return_value=row) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma) user = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) @@ -1453,7 +1548,9 @@ async def test_reject_guardrail_submission_not_pending(mocker): "no_hostname", ], ) -async def test_register_guardrail_rejects_bad_api_base(mocker, api_base, expected_detail): +async def test_register_guardrail_rejects_bad_api_base( + mocker, api_base, expected_detail +): """Register returns 400 when api_base has invalid scheme or missing hostname.""" mocker.patch("litellm.proxy.proxy_server.prisma_client", mocker.Mock()) req = RegisterGuardrailRequest( @@ -1590,16 +1687,28 @@ async def test_list_submissions_summary_counts_unaffected_by_filters(mocker): """Summary counts reflect all team guardrails regardless of status filter.""" mock_prisma = mocker.Mock() pending_row = mocker.Mock( - guardrail_id="p1", guardrail_name="p", status="pending_review", - team_id="t1", litellm_params={}, guardrail_info={}, - submitted_at=None, reviewed_at=None, - created_at=datetime.now(), updated_at=datetime.now(), + guardrail_id="p1", + guardrail_name="p", + status="pending_review", + team_id="t1", + litellm_params={}, + guardrail_info={}, + submitted_at=None, + reviewed_at=None, + created_at=datetime.now(), + updated_at=datetime.now(), ) active_row = mocker.Mock( - guardrail_id="a1", guardrail_name="a", status="active", - team_id="t1", litellm_params={}, guardrail_info={}, - submitted_at=None, reviewed_at=None, - created_at=datetime.now(), updated_at=datetime.now(), + guardrail_id="a1", + guardrail_name="a", + status="active", + team_id="t1", + litellm_params={}, + guardrail_info={}, + submitted_at=None, + reviewed_at=None, + created_at=datetime.now(), + updated_at=datetime.now(), ) all_rows = [pending_row, active_row] mock_prisma.db.litellm_guardrailstable.find_many = AsyncMock(return_value=all_rows) @@ -1607,9 +1716,138 @@ async def test_list_submissions_summary_counts_unaffected_by_filters(mocker): user = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) # Filter to only pending, but summary should still show both - result = await list_guardrail_submissions(status="pending_review", user_api_key_dict=user) + result = await list_guardrail_submissions( + status="pending_review", user_api_key_dict=user + ) assert len(result.submissions) == 1 # filtered assert result.summary.total == 2 # unfiltered assert result.summary.pending_review == 1 assert result.summary.active == 1 + + +@pytest.mark.asyncio +async def test_create_guardrail_with_invalid_pattern(mocker): + """ + Test that creating a guardrail with an invalid pattern returns 422 + and rolls back the DB entry. + """ + from litellm.proxy.guardrails.guardrail_endpoints import create_guardrail + from litellm.types.guardrails import Guardrail, LitellmParams + + # Mock DB + mock_prisma = mocker.Mock() + mock_prisma.db.litellm_guardrailstable.create = AsyncMock( + return_value=mocker.Mock( + guardrail_id="failed-id", + guardrail_name="invalid-guard", + litellm_params=json.dumps( + { + "guardrail": "litellm_content_filter", + "mode": "pre_call", + "patterns": [ + { + "pattern_type": "prebuilt", + "pattern_name": "FAKE-PATTERN", + "action": "BLOCK", + } + ], + } + ), + guardrail_info="{}", + ) + ) + mock_prisma.db.litellm_guardrailstable.delete = AsyncMock() + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + # Request with fake pattern + invalid_guardrail = Guardrail( + guardrail_name="invalid-guard", + litellm_params=LitellmParams( + guardrail="litellm_content_filter", + mode="pre_call", + patterns=[ + { + "pattern_type": "prebuilt", + "pattern_name": "FAKE-PATTERN", + "action": "BLOCK", + } + ], + ), + ) + request = CreateGuardrailRequest(guardrail=invalid_guardrail) + + with pytest.raises(HTTPException) as exc_info: + await create_guardrail(request, user_api_key_dict=MOCK_ADMIN_USER) + + assert exc_info.value.status_code == 422 + assert "Unknown pattern name: 'FAKE-PATTERN'" in str(exc_info.value.detail) + + # Verify rollback was called (delete from DB) + mock_prisma.db.litellm_guardrailstable.delete.assert_called_once() + + +@pytest.mark.asyncio +async def test_get_policy_templates_filters_unavailable_patterns(mocker): + """ + Test that get_policy_templates filters out templates containing patterns + not in PREBUILT_PATTERNS. + """ + from litellm.proxy.management_endpoints.policy_endpoints.endpoints import ( + get_policy_templates, + ) + + # Mock prebuilt patterns to only have one + mocker.patch( + "litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.patterns.PREBUILT_PATTERNS", + {"existing_pattern": ".*"}, + ) + + mock_templates = [ + { + "id": "valid-template", + "guardrailDefinitions": [ + { + "litellm_params": { + "guardrail": "litellm_content_filter", + "patterns": [ + { + "pattern_type": "prebuilt", + "pattern_name": "existing_pattern", + } + ], + } + } + ], + }, + { + "id": "invalid-template", + "guardrailDefinitions": [ + { + "litellm_params": { + "guardrail": "litellm_content_filter", + "patterns": [ + { + "pattern_type": "prebuilt", + "pattern_name": "non_existent_pattern", + } + ], + } + } + ], + }, + ] + + mocker.patch( + "litellm.proxy.management_endpoints.policy_endpoints.endpoints._load_policy_templates_from_local_backup", + return_value=mock_templates, + ) + # Force use of local backup + mocker.patch("os.getenv", return_value="true") + + result = await get_policy_templates( + mocker.Mock(), user_api_key_dict=MOCK_ADMIN_USER + ) + + assert len(result) == 1 + assert result[0]["id"] == "valid-template"