From 725671395ac6ef8697ed6ff150dad72b7f59cb3c Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 31 Oct 2025 16:29:24 -0700 Subject: [PATCH] [Feature] UI - Config Guardrails should not be editable and guardrail info fix (#16142) * UI Config Guardrails should not be editable and guardrails info definition location fix * Remove unused import * Added literals for guardrail definition location --- .../proxy/guardrails/guardrail_endpoints.py | 4 + litellm/types/guardrails.py | 5 +- .../guardrails/test_guardrail_endpoints.py | 68 ++++++++++- .../guardrails/guardrail_info.test.tsx | 113 ++++++++++++++++++ .../components/guardrails/guardrail_info.tsx | 14 ++- 5 files changed, 200 insertions(+), 4 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/guardrails/guardrail_info.test.tsx diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index f2f63778e4..0ad1a496a9 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -630,11 +630,13 @@ async def get_guardrail_info(guardrail_id: str): from litellm.litellm_core_utils.litellm_logging import _get_masked_values from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER from litellm.proxy.proxy_server import prisma_client + from litellm.types.guardrails import GUARDRAIL_DEFINITION_LOCATION if prisma_client is None: raise HTTPException(status_code=500, detail="Prisma client not initialized") try: + guardrail_definition_location: GUARDRAIL_DEFINITION_LOCATION = GUARDRAIL_DEFINITION_LOCATION.DB result = await GUARDRAIL_REGISTRY.get_guardrail_by_id_from_db( guardrail_id=guardrail_id, prisma_client=prisma_client ) @@ -642,6 +644,7 @@ async def get_guardrail_info(guardrail_id: str): result = IN_MEMORY_GUARDRAIL_HANDLER.get_guardrail_by_id( guardrail_id=guardrail_id ) + guardrail_definition_location: GUARDRAIL_DEFINITION_LOCATION = GUARDRAIL_DEFINITION_LOCATION.CONFIG if result is None: raise HTTPException( @@ -669,6 +672,7 @@ async def get_guardrail_info(guardrail_id: str): guardrail_info=dict(result.get("guardrail_info") or {}), created_at=result.get("created_at"), updated_at=result.get("updated_at"), + guardrail_definition_location=guardrail_definition_location, ) except HTTPException as e: raise e diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 366ddba935..0fe7a2c417 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -589,6 +589,9 @@ class GuardrailEventHooks(str, Enum): class DynamicGuardrailParams(TypedDict): extra_body: Dict[str, Any] +class GUARDRAIL_DEFINITION_LOCATION(str, Enum): + DB = "db" + CONFIG = "config" class GuardrailInfoResponse(BaseModel): guardrail_id: Optional[str] = None @@ -597,7 +600,7 @@ class GuardrailInfoResponse(BaseModel): guardrail_info: Optional[Dict] = None created_at: Optional[datetime] = None updated_at: Optional[datetime] = None - guardrail_definition_location: Literal["config", "db"] = "config" + guardrail_definition_location: GUARDRAIL_DEFINITION_LOCATION = GUARDRAIL_DEFINITION_LOCATION.CONFIG def __init__(self, **kwargs): super().__init__(**kwargs) diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index 75daf54016..a2d072dc7d 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -929,4 +929,70 @@ async def test_apply_guardrail_execution_error(mocker): 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) \ No newline at end of file + assert "Bedrock guardrail failed" in str(exc_info.value.message) + +@pytest.mark.asyncio +async def test_get_guardrail_info_endpoint_config_guardrail(mocker): + """ + Test get_guardrail_info endpoint returns proper response when guardrail is found in config. + """ + from litellm.proxy.guardrails.guardrail_endpoints import get_guardrail_info + + # Mock prisma_client to not be None (patch at the source where it's imported from) + mock_prisma = mocker.Mock() + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + # 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) + + # 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) + + # 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 + ) + + # Call endpoint and expect GuardrailInfoResponse + result = await get_guardrail_info(guardrail_id="test-config-guardrail") + + # Verify the response is of the correct type + assert isinstance(result, GuardrailInfoResponse) + assert result.guardrail_id == "test-config-guardrail" + 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): + """ + Test get_guardrail_info endpoint returns proper response when guardrail is found in DB. + """ + from litellm.proxy.guardrails.guardrail_endpoints import get_guardrail_info + + # Mock prisma_client to not be None (patch at the source where it's imported from) + mock_prisma = mocker.Mock() + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + # 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 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) + + # Call endpoint and expect GuardrailInfoResponse + result = await get_guardrail_info(guardrail_id="test-db-guardrail") + + # Verify the response is of the correct type + assert isinstance(result, GuardrailInfoResponse) + assert result.guardrail_id == "test-db-guardrail" + assert result.guardrail_name == "Test DB Guardrail" + assert result.guardrail_definition_location == "db" \ No newline at end of file diff --git a/ui/litellm-dashboard/src/components/guardrails/guardrail_info.test.tsx b/ui/litellm-dashboard/src/components/guardrails/guardrail_info.test.tsx new file mode 100644 index 0000000000..301045103b --- /dev/null +++ b/ui/litellm-dashboard/src/components/guardrails/guardrail_info.test.tsx @@ -0,0 +1,113 @@ +import * as networking from "@/components/networking"; +import { fireEvent, render, waitFor } from "@testing-library/react"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import GuardrailInfoView from "./guardrail_info"; + +// Mock the networking module +vi.mock("@/components/networking", () => ({ + getGuardrailInfo: vi.fn(), + getGuardrailUISettings: vi.fn(), + getGuardrailProviderSpecificParams: vi.fn(), + updateGuardrailCall: vi.fn(), +})); + +describe("Guardrail Info", () => { + afterEach(() => { + vi.clearAllMocks(); + }); + + it("should render the guardrail info after loading", async () => { + // Mock the network responses + vi.mocked(networking.getGuardrailInfo).mockResolvedValue({ + guardrail_id: "123", + guardrail_name: "Test Guardrail", + litellm_params: { + guardrail: "presidio", + mode: "pre_call", + default_on: true, + }, + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + guardrail_definition_location: "database", + }); + + vi.mocked(networking.getGuardrailUISettings).mockResolvedValue({ + supported_entities: ["PERSON", "EMAIL"], + supported_actions: ["MASK", "REDACT"], + pii_entity_categories: [], + supported_modes: ["pre_call", "post_call"], + }); + + vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); + + const { getAllByText, getByText } = render( + {}} accessToken="123" isAdmin={true} />, + ); + + // Wait for the loading to complete and data to be rendered + await waitFor(() => { + // The guardrail name appears in multiple places (title and settings tab) + const elements = getAllByText("Test Guardrail"); + expect(elements.length).toBeGreaterThan(0); + }); + + // Verify other key elements are present + expect(getByText("Back to Guardrails")).toBeInTheDocument(); + expect(getByText("Overview")).toBeInTheDocument(); + expect(getByText("Settings")).toBeInTheDocument(); + }); + + it("should not render the edit button for config guardrails", async () => { + // Mock the network responses + vi.mocked(networking.getGuardrailInfo).mockResolvedValue({ + guardrail_id: "123", + guardrail_name: "Test Guardrail", + litellm_params: { + guardrail: "presidio", + mode: "pre_call", + default_on: true, + }, + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + guardrail_definition_location: "config", + }); + + vi.mocked(networking.getGuardrailUISettings).mockResolvedValue({ + supported_entities: ["PERSON", "EMAIL"], + supported_actions: ["MASK", "REDACT"], + pii_entity_categories: [], + supported_modes: ["pre_call", "post_call"], + }); + + vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); + + const { getByText, container } = render( + {}} accessToken="123" isAdmin={true} />, + ); + + await waitFor(() => { + expect(getByText("Settings")).toBeInTheDocument(); + }); + + // Click the Settings tab + fireEvent.click(getByText("Settings")); + + // Wait for the Settings panel to render + await waitFor(() => { + expect(getByText("Guardrail Settings")).toBeInTheDocument(); + }); + + // Find the info icon and hover over it + const infoIcon = container.querySelector(".anticon-info-circle"); + expect(infoIcon).toBeInTheDocument(); + + if (infoIcon) { + fireEvent.mouseEnter(infoIcon); + + // Wait for the tooltip to appear + await waitFor(() => { + expect(getByText("Guardrail is defined in the config file and cannot be edited.")).toBeInTheDocument(); + }); + } + }); +}); diff --git a/ui/litellm-dashboard/src/components/guardrails/guardrail_info.tsx b/ui/litellm-dashboard/src/components/guardrails/guardrail_info.tsx index 27292f4160..d783108de3 100644 --- a/ui/litellm-dashboard/src/components/guardrails/guardrail_info.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/guardrail_info.tsx @@ -13,7 +13,8 @@ import { TabPanels, TextInput, } from "@tremor/react"; -import { Button, Form, Input, Select, Divider } from "antd"; +import { Button, Form, Input, Select, Divider, Tooltip } from "antd"; +import { InfoCircleOutlined } from "@ant-design/icons"; import { getGuardrailInfo, updateGuardrailCall, @@ -328,6 +329,8 @@ const GuardrailInfoView: React.FC = ({ guardrailId, onClose, } }; + const isConfigGuardrail = guardrailData.guardrail_definition_location === "config"; + return (
@@ -434,7 +437,14 @@ const GuardrailInfoView: React.FC = ({ guardrailId, onClose,
Guardrail Settings - {!isEditing && setIsEditing(true)}>Edit Settings} + {isConfigGuardrail && ( + + + + )} + {!isEditing && !isConfigGuardrail && ( + setIsEditing(true)}>Edit Settings + )}
{isEditing ? (