[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
This commit is contained in:
yuneng-jiang 2025-10-31 16:29:24 -07:00 committed by GitHub
parent cf0d84c789
commit 725671395a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 200 additions and 4 deletions

View File

@ -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

View File

@ -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)

View File

@ -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)
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"

View File

@ -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(
<GuardrailInfoView guardrailId="123" onClose={() => {}} 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(
<GuardrailInfoView guardrailId="123" onClose={() => {}} 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();
});
}
});
});

View File

@ -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<GuardrailInfoProps> = ({ guardrailId, onClose,
}
};
const isConfigGuardrail = guardrailData.guardrail_definition_location === "config";
return (
<div className="p-4">
<div>
@ -434,7 +437,14 @@ const GuardrailInfoView: React.FC<GuardrailInfoProps> = ({ guardrailId, onClose,
<Card>
<div className="flex justify-between items-center mb-4">
<Title>Guardrail Settings</Title>
{!isEditing && <TremorButton onClick={() => setIsEditing(true)}>Edit Settings</TremorButton>}
{isConfigGuardrail && (
<Tooltip title="Guardrail is defined in the config file and cannot be edited.">
<InfoCircleOutlined />
</Tooltip>
)}
{!isEditing && !isConfigGuardrail && (
<TremorButton onClick={() => setIsEditing(true)}>Edit Settings</TremorButton>
)}
</div>
{isEditing ? (