From 9f50d18ace9b640810236b192194cca9f51c4a0e Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 15 Sep 2026 16:52:45 +0000 Subject: [PATCH 1/5] feat(ui): allow custom Presidio PII entity names in guardrail PII picker Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../_components/piiCustomEntity.test.ts | 26 ++++++++ .../guardrails/_components/piiCustomEntity.ts | 8 +++ .../_components/pii_configuration.test.tsx | 66 ++++++++++++++++++- .../_components/pii_configuration.tsx | 61 ++++++++++++++++- 4 files changed, 156 insertions(+), 5 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/piiCustomEntity.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/piiCustomEntity.ts 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..e1c04b24daa --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/piiCustomEntity.test.ts @@ -0,0 +1,26 @@ +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..4f7869d4a49 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/piiCustomEntity.ts @@ -0,0 +1,8 @@ +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.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.test.tsx index 4f822578fe8..b0f0265f067 100644 --- 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 @@ -1,5 +1,5 @@ -import { render, screen } from "@testing-library/react"; -import { describe, it, expect } from "vitest"; +import { fireEvent, render, screen } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; import PiiConfiguration from "./pii_configuration"; describe("PiiConfiguration", () => { @@ -17,4 +17,66 @@ describe("PiiConfiguration", () => { ); 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 the first action", () => { + 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("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.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.tsx index cba29ed615d..473cfc0fcb6 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,6 +19,12 @@ 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(); @@ -24,15 +33,18 @@ const PiiConfiguration: React.FC = ({ entityToCategoryMap.set(entity, category.category); }); }); + customEntities.forEach((entity) => { + entityToCategoryMap.set(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 +62,23 @@ 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; + } + + onEntitySelect(normalizedEntityName); + onActionSelect(normalizedEntityName, actions[0] ?? "MASK"); + }; + return (
@@ -61,7 +90,7 @@ const PiiConfiguration: React.FC = ({
@@ -71,6 +100,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. +

+
Date: Tue, 15 Sep 2026 16:57:19 +0000 Subject: [PATCH 2/5] style(ui): format custom Presidio PII picker files Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrails/_components/piiCustomEntity.test.ts | 9 +++------ .../guardrails/_components/piiCustomEntity.ts | 5 ++++- .../guardrails/_components/pii_configuration.tsx | 4 +++- 3 files changed, 10 insertions(+), 8 deletions(-) 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 index e1c04b24daa..0b5bfd189f5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/piiCustomEntity.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/piiCustomEntity.test.ts @@ -16,11 +16,8 @@ describe("normalizeCustomEntityName", () => { 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", - ]); + 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 index 4f7869d4a49..aae2dc2033c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/piiCustomEntity.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/piiCustomEntity.ts @@ -1,5 +1,8 @@ export const normalizeCustomEntityName = (raw: string): string | null => { - const normalized = raw.trim().toUpperCase().replace(/[\s-]+/g, "_"); + const normalized = raw + .trim() + .toUpperCase() + .replace(/[\s-]+/g, "_"); return normalized !== "" && /^[A-Z0-9_]+$/.test(normalized) ? normalized : null; }; 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 473cfc0fcb6..4efb795bb6a 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 @@ -24,7 +24,9 @@ const PiiConfiguration: React.FC = ({ const allEntities = mergeCustomEntities(entities, selectedEntities); const customEntities = allEntities.filter((entity) => !entities.includes(entity)); const categories = - customEntities.length > 0 ? [...entityCategories, { category: "Custom", entities: customEntities }] : entityCategories; + customEntities.length > 0 + ? [...entityCategories, { category: "Custom", entities: customEntities }] + : entityCategories; // Create a lookup map to quickly find an entity's category const entityToCategoryMap = new Map(); From de2cac7b672d692cd9567f409fc281c29e092407 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 15 Sep 2026 22:38:42 +0000 Subject: [PATCH 3/5] fix(ui): default custom PII entities to MASK and classify test tier Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...ration.test.tsx => pii_configuration.integration.test.tsx} | 4 ++-- .../(dashboard)/guardrails/_components/pii_configuration.tsx | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) rename ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/{pii_configuration.test.tsx => pii_configuration.integration.test.tsx} (95%) 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.integration.test.tsx similarity index 95% rename from ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.integration.test.tsx index b0f0265f067..969844c58ed 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.integration.test.tsx @@ -33,13 +33,13 @@ describe("PiiConfiguration", () => { expect(screen.getByText("NO FODSELSNUMMER")).toBeInTheDocument(); }); - it("adds a normalized custom entity with the first action", () => { + it("adds a normalized custom entity with MASK", () => { const onEntitySelect = vi.fn(); const onActionSelect = vi.fn(); render( = ({ } onEntitySelect(normalizedEntityName); - onActionSelect(normalizedEntityName, actions[0] ?? "MASK"); + onActionSelect(normalizedEntityName, "MASK"); }; return ( From 3445fb378fe8d0b7d7ab53b65526230c0b3c465c Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 15 Sep 2026 23:04:58 +0000 Subject: [PATCH 4/5] fix(ui): show newly added custom PII entity when a category filter is active Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../pii_configuration.integration.test.tsx | 53 +++++++++++++++++++ .../_components/pii_configuration.tsx | 3 ++ 2 files changed, 56 insertions(+) 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 index 969844c58ed..eb8692357fe 100644 --- 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 @@ -1,7 +1,35 @@ +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( @@ -56,6 +84,31 @@ describe("PiiConfiguration", () => { 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(); 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 e6efdf53670..609899aa95a 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 @@ -77,6 +77,9 @@ const PiiConfiguration: React.FC = ({ return; } + if (selectedCategories.length > 0 && !selectedCategories.includes("Custom")) { + setSelectedCategories([...selectedCategories, "Custom"]); + } onEntitySelect(normalizedEntityName); onActionSelect(normalizedEntityName, "MASK"); }; From 85d417effed1e5975218fc90da788b3332d2b41e Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 15 Sep 2026 23:16:21 +0000 Subject: [PATCH 5/5] refactor(ui): build PII entity category map without mutation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrails/_components/pii_configuration.tsx | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) 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 609899aa95a..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 @@ -28,16 +28,12 @@ const PiiConfiguration: React.FC = ({ ? [...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); - }); - }); - customEntities.forEach((entity) => { - entityToCategoryMap.set(entity, "Custom"); - }); + 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 = allEntities.filter((entity) => {