diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/piiCustomEntity.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/piiCustomEntity.test.ts new file mode 100644 index 00000000000..0b5bfd189f5 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/piiCustomEntity.test.ts @@ -0,0 +1,23 @@ +import { describe, expect, it } from "vitest"; +import { mergeCustomEntities, normalizeCustomEntityName } from "./piiCustomEntity"; + +describe("normalizeCustomEntityName", () => { + it.each([ + ["se personnummer", "SE_PERSONNUMMER"], + [" no-fodselsnummer ", "NO_FODSELSNUMMER"], + ])("normalizes %s", (raw, expected) => { + expect(normalizeCustomEntityName(raw)).toBe(expected); + }); + + it.each(["", "bad!name"])("rejects %s", (raw) => { + expect(normalizeCustomEntityName(raw)).toBeNull(); + }); +}); + +describe("mergeCustomEntities", () => { + it("keeps supported order, dedupes, and appends unknown selected entities", () => { + expect( + mergeCustomEntities(["PERSON", "EMAIL"], ["EMAIL", "NO_FODSELSNUMMER", "PERSON", "SE_PERSONNUMMER"]), + ).toEqual(["PERSON", "EMAIL", "NO_FODSELSNUMMER", "SE_PERSONNUMMER"]); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/piiCustomEntity.ts b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/piiCustomEntity.ts new file mode 100644 index 00000000000..aae2dc2033c --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/piiCustomEntity.ts @@ -0,0 +1,11 @@ +export const normalizeCustomEntityName = (raw: string): string | null => { + const normalized = raw + .trim() + .toUpperCase() + .replace(/[\s-]+/g, "_"); + return normalized !== "" && /^[A-Z0-9_]+$/.test(normalized) ? normalized : null; +}; + +export const mergeCustomEntities = (supported: readonly string[], selected: readonly string[]): string[] => [ + ...new Set([...supported, ...selected]), +]; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.integration.test.tsx new file mode 100644 index 00000000000..eb8692357fe --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.integration.test.tsx @@ -0,0 +1,135 @@ +import { useState } from "react"; +import { fireEvent, render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { describe, expect, it, vi } from "vitest"; +import PiiConfiguration from "./pii_configuration"; + +const StatefulPiiConfiguration = () => { + const [selectedEntities, setSelectedEntities] = useState([]); + const [selectedActions, setSelectedActions] = useState<{ [key: string]: string }>({}); + + return ( + { + setSelectedEntities((previous) => + previous.includes(entity) ? previous.filter((selected) => selected !== entity) : [...previous, entity], + ); + }} + onActionSelect={(entity, action) => { + setSelectedActions((previous) => ({ ...previous, [entity]: action })); + }} + entityCategories={[ + { category: "Personal", entities: ["PERSON"] }, + { category: "Contact", entities: ["EMAIL"] }, + ]} + /> + ); +}; + +describe("PiiConfiguration", () => { + it("should render", () => { + render( + {}} + onActionSelect={() => {}} + entityCategories={[]} + />, + ); + expect(screen.getByText("Configure PII Protection")).toBeInTheDocument(); + }); + + it("renders selected custom entities", () => { + render( + {}} + onActionSelect={() => {}} + />, + ); + + expect(screen.getByText("NO FODSELSNUMMER")).toBeInTheDocument(); + }); + + it("adds a normalized custom entity with MASK", () => { + const onEntitySelect = vi.fn(); + const onActionSelect = vi.fn(); + render( + , + ); + + fireEvent.change(screen.getByRole("textbox", { name: "Custom entity name" }), { + target: { value: "se personnummer" }, + }); + fireEvent.click(screen.getByRole("button", { name: "Add entity" })); + + expect(onEntitySelect).toHaveBeenCalledWith("SE_PERSONNUMMER"); + expect(onActionSelect).toHaveBeenCalledWith("SE_PERSONNUMMER", "MASK"); + }); + + it("shows a newly added custom entity while a category filter is active", async () => { + const user = userEvent.setup(); + render(); + + const categoryInput = screen.getByPlaceholderText("Select categories to filter by"); + await user.click(categoryInput); + if (categoryInput.getAttribute("aria-expanded") !== "true") { + categoryInput.focus(); + await user.keyboard("{Enter}"); + } + await user.click(await screen.findByRole("option", { name: "Personal" })); + await user.keyboard("{Escape}"); + + expect(screen.getByText("PERSON")).toBeVisible(); + expect(screen.queryByText("EMAIL")).not.toBeInTheDocument(); + + fireEvent.change(screen.getByRole("textbox", { name: "Custom entity name" }), { + target: { value: "se personnummer" }, + }); + fireEvent.click(screen.getByRole("button", { name: "Add entity" })); + + expect(await screen.findByText("SE PERSONNUMMER")).toBeVisible(); + expect(screen.getByText("PERSON")).toBeVisible(); + }); + + it("rejects invalid custom entity names", () => { + const onEntitySelect = vi.fn(); + const onActionSelect = vi.fn(); + render( + , + ); + + fireEvent.change(screen.getByRole("textbox", { name: "Custom entity name" }), { + target: { value: "bad!name" }, + }); + fireEvent.click(screen.getByRole("button", { name: "Add entity" })); + + expect(screen.getByText("Use letters, numbers and underscores only")).toBeInTheDocument(); + expect(onEntitySelect).not.toHaveBeenCalled(); + expect(onActionSelect).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.test.tsx deleted file mode 100644 index 4f822578fe8..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.test.tsx +++ /dev/null @@ -1,20 +0,0 @@ -import { render, screen } from "@testing-library/react"; -import { describe, it, expect } from "vitest"; -import PiiConfiguration from "./pii_configuration"; - -describe("PiiConfiguration", () => { - it("should render", () => { - render( - {}} - onActionSelect={() => {}} - entityCategories={[]} - />, - ); - expect(screen.getByText("Configure PII Protection")).toBeInTheDocument(); - }); -}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.tsx index cba29ed615d..84b60f1e039 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.tsx @@ -1,6 +1,9 @@ import React, { useState } from "react"; import { CategoryFilter, PiiEntityList, QuickActions } from "./pii_components"; import { PiiConfigurationProps } from "@/components/guardrails/types"; +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import { mergeCustomEntities, normalizeCustomEntityName } from "./piiCustomEntity"; /** * A reusable component for rendering PII entity selection and action configuration @@ -16,23 +19,30 @@ const PiiConfiguration: React.FC = ({ entityCategories = [], }) => { const [selectedCategories, setSelectedCategories] = useState([]); + const [customEntityName, setCustomEntityName] = useState(""); + const [customEntityError, setCustomEntityError] = useState(null); + const allEntities = mergeCustomEntities(entities, selectedEntities); + const customEntities = allEntities.filter((entity) => !entities.includes(entity)); + const categories = + customEntities.length > 0 + ? [...entityCategories, { category: "Custom", entities: customEntities }] + : entityCategories; - // Create a lookup map to quickly find an entity's category - const entityToCategoryMap = new Map(); - entityCategories.forEach((category) => { - category.entities.forEach((entity) => { - entityToCategoryMap.set(entity, category.category); - }); - }); + const entityToCategoryMap = new Map([ + ...entityCategories.flatMap((category) => + category.entities.map((entity): [string, string] => [entity, category.category]), + ), + ...customEntities.map((entity): [string, string] => [entity, "Custom"]), + ]); // Filter entities based on selected categories - const filteredEntities = entities.filter((entity) => { + const filteredEntities = allEntities.filter((entity) => { return selectedCategories.length === 0 || selectedCategories.includes(entityToCategoryMap.get(entity) || ""); }); // Select all entities with a specified action const handleSelectAll = (action: string) => { - entities.forEach((entity) => { + allEntities.forEach((entity) => { if (!selectedEntities.includes(entity)) { onEntitySelect(entity); } @@ -50,6 +60,26 @@ const PiiConfiguration: React.FC = ({ }); }; + const handleAddCustomEntity = () => { + const normalizedEntityName = normalizeCustomEntityName(customEntityName); + if (normalizedEntityName === null) { + setCustomEntityError("Use letters, numbers and underscores only"); + return; + } + + setCustomEntityError(null); + setCustomEntityName(""); + if (allEntities.includes(normalizedEntityName)) { + return; + } + + if (selectedCategories.length > 0 && !selectedCategories.includes("Custom")) { + setSelectedCategories([...selectedCategories, "Custom"]); + } + onEntitySelect(normalizedEntityName); + onActionSelect(normalizedEntityName, "MASK"); + }; + return (
@@ -61,7 +91,7 @@ const PiiConfiguration: React.FC = ({
@@ -71,6 +101,32 @@ const PiiConfiguration: React.FC = ({ onUnselectAll={handleUnselectAll} hasSelectedEntities={selectedEntities.length > 0} /> +
+
+ { + setCustomEntityName(event.target.value); + setCustomEntityError(null); + }} + onKeyDown={(event) => { + if (event.key === "Enter") { + event.preventDefault(); + handleAddCustomEntity(); + } + }} + placeholder="Custom entity name, e.g. SE_PERSONNUMMER" + aria-label="Custom entity name" + /> + +
+ {customEntityError &&

{customEntityError}

} +

+ Add any entity your Presidio analyzer recognizes, including custom recognizers. +

+