Merge 85d417effe into 22b36cbcf6

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"]);
+  });
+});

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]),
+];

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<string[]>([]);
+  const [selectedActions, setSelectedActions] = useState<{ [key: string]: string }>({});
+
+  return (
+    <PiiConfiguration
+      entities={["PERSON", "EMAIL"]}
+      actions={["MASK", "BLOCK"]}
+      selectedEntities={selectedEntities}
+      selectedActions={selectedActions}
+      onEntitySelect={(entity) => {
+        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(
+      <PiiConfiguration
+        entities={[]}
+        actions={[]}
+        selectedEntities={[]}
+        selectedActions={{}}
+        onEntitySelect={() => {}}
+        onActionSelect={() => {}}
+        entityCategories={[]}
+      />,
+    );
+    expect(screen.getByText("Configure PII Protection")).toBeInTheDocument();
+  });
+
+  it("renders selected custom entities", () => {
+    render(
+      <PiiConfiguration
+        entities={["PERSON"]}
+        actions={["MASK"]}
+        selectedEntities={["NO_FODSELSNUMMER"]}
+        selectedActions={{ NO_FODSELSNUMMER: "MASK" }}
+        onEntitySelect={() => {}}
+        onActionSelect={() => {}}
+      />,
+    );
+
+    expect(screen.getByText("NO FODSELSNUMMER")).toBeInTheDocument();
+  });
+
+  it("adds a normalized custom entity with MASK", () => {
+    const onEntitySelect = vi.fn();
+    const onActionSelect = vi.fn();
+    render(
+      <PiiConfiguration
+        entities={["PERSON"]}
+        actions={["BLOCK", "MASK"]}
+        selectedEntities={[]}
+        selectedActions={{}}
+        onEntitySelect={onEntitySelect}
+        onActionSelect={onActionSelect}
+      />,
+    );
+
+    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(<StatefulPiiConfiguration />);
+
+    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(
+      <PiiConfiguration
+        entities={["PERSON"]}
+        actions={["MASK", "BLOCK"]}
+        selectedEntities={[]}
+        selectedActions={{}}
+        onEntitySelect={onEntitySelect}
+        onActionSelect={onActionSelect}
+      />,
+    );
+
+    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();
+  });
+});

ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.test.tsx

@@ -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(
-      <PiiConfiguration
-        entities={[]}
-        actions={[]}
-        selectedEntities={[]}
-        selectedActions={{}}
-        onEntitySelect={() => {}}
-        onActionSelect={() => {}}
-        entityCategories={[]}
-      />,
-    );
-    expect(screen.getByText("Configure PII Protection")).toBeInTheDocument();
-  });
-});

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<PiiConfigurationProps> = ({
   entityCategories = [],
 }) => {
   const [selectedCategories, setSelectedCategories] = useState<string[]>([]);
+  const [customEntityName, setCustomEntityName] = useState("");
+  const [customEntityError, setCustomEntityError] = useState<string | null>(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<string, string>();
-  entityCategories.forEach((category) => {
-    category.entities.forEach((entity) => {
-      entityToCategoryMap.set(entity, category.category);
-    });
-  });
+  const entityToCategoryMap = new Map<string, string>([
+    ...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<PiiConfigurationProps> = ({
     });
   };
 
+  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 (
     <div className="pii-configuration">
       <div className="flex justify-between items-center mb-5">
@@ -61,7 +91,7 @@ const PiiConfiguration: React.FC<PiiConfigurationProps> = ({
 
       <div className="mb-6">
         <CategoryFilter
-          categories={entityCategories}
+          categories={categories}
           selectedCategories={selectedCategories}
           onChange={setSelectedCategories}
         />
@@ -71,6 +101,32 @@ const PiiConfiguration: React.FC<PiiConfigurationProps> = ({
           onUnselectAll={handleUnselectAll}
           hasSelectedEntities={selectedEntities.length > 0}
         />
+        <div className="mt-4">
+          <div className="flex items-center gap-2">
+            <Input
+              value={customEntityName}
+              onChange={(event) => {
+                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"
+            />
+            <Button type="button" onClick={handleAddCustomEntity}>
+              Add entity
+            </Button>
+          </div>
+          {customEntityError && <p className="mt-1 text-sm text-destructive">{customEntityError}</p>}
+          <p className="mt-1 text-sm text-muted-foreground">
+            Add any entity your Presidio analyzer recognizes, including custom recognizers.
+          </p>
+        </div>
       </div>
 
       <PiiEntityList