feat(studio): Anonymizer entity picker categories and custom labels [ASTD-335]

web/packages/studio/src/routes/AnonymizerBuilderRoute/components/AnonymizerBuilderForm.tsx

@@ -1,7 +1,10 @@
 // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-import { useAnonymizerCreateRunJob } from '@nemo/sdk/generated/anonymizer/api';
+import {
+  useAnonymizerCreateRunJob,
+  useAnonymizerListEntityLabels,
+} from '@nemo/sdk/generated/anonymizer/api';
 import type { RunJob } from '@nemo/sdk/generated/anonymizer/schema';
 import {
   Banner,
@@ -47,6 +50,8 @@ export const AnonymizerBuilderForm: FC = () => {
   const [submitError, setSubmitError] = useState<string | undefined>(undefined);
 
   const { isLoading: isLoadingModels } = useDefaultRoleModels();
+  const { data: defaultEntityLabels, isLoading: isLoadingEntityLabels } =
+    useAnonymizerListEntityLabels(workspace, { query: {} });
 
   const createJob = useAnonymizerCreateRunJob({
     mutation: {
@@ -76,7 +81,10 @@ export const AnonymizerBuilderForm: FC = () => {
   const onSubmit = form.handleSubmit(
     (values) => {
       setSubmitError(undefined);
-      createJob.mutate({ workspace, data: buildAnonymizerJobRequest(values) });
+      createJob.mutate({
+        workspace,
+        data: buildAnonymizerJobRequest(values, defaultEntityLabels?.data ?? []),
+      });
     },
     (errors) => {
       const onlyModelErrors = Object.keys(errors).every((key) => key === 'roleModels');
@@ -109,7 +117,7 @@ export const AnonymizerBuilderForm: FC = () => {
                 kind="primary"
                 color="brand"
                 type="submit"
-                disabled={createJob.isPending || isLoadingModels}
+                disabled={createJob.isPending || isLoadingModels || isLoadingEntityLabels}
               >
                 Full Run
               </Button>

web/packages/studio/src/routes/AnonymizerBuilderRoute/components/EntitiesSection.tsx

@@ -5,28 +5,55 @@ import { ControlledCheckbox } from '@nemo/common/src/components/form/ControlledC
 import { ControlledCombobox } from '@nemo/common/src/components/form/ControlledCombobox';
 import { ControlledSegmentedControl } from '@nemo/common/src/components/form/ControlledSegmentedControl';
 import { useAnonymizerListEntityLabels } from '@nemo/sdk/generated/anonymizer/api';
-import { Stack, Text } from '@nvidia/foundations-react-core';
+import { Flex, Stack, Tag, Text } from '@nvidia/foundations-react-core';
 import { useWorkspaceFromPath } from '@studio/hooks/useWorkspaceFromPath';
 import {
   ENTITY_MODE_AUTO,
   ENTITY_MODE_CUSTOM,
   ENTITY_MODE_OPTIONS,
+  entityTagColor,
 } from '@studio/routes/AnonymizerBuilderRoute/constants';
+import {
+  buildEntitySections,
+  customLabelCandidate,
+} from '@studio/routes/AnonymizerBuilderRoute/entityItems';
 import type { AnonymizerFormData } from '@studio/routes/AnonymizerBuilderRoute/schema';
-import { FC } from 'react';
+import { X } from 'lucide-react';
+import { FC, useMemo, useState } from 'react';
 import { useFormContext, useWatch } from 'react-hook-form';
 
 export const EntitiesSection: FC = () => {
-  const { control } = useFormContext<AnonymizerFormData>();
+  const { control, setValue } = useFormContext<AnonymizerFormData>();
   const workspace = useWorkspaceFromPath();
   const entityMode = useWatch({ control, name: 'entityMode' });
-  const includeDefaults = useWatch({ control, name: 'includeDefaultEntities' });
+  const [inputValue, setInputValue] = useState('');
 
-  const isCustom = entityMode === ENTITY_MODE_CUSTOM;
-  const showLabelPicker = isCustom && !includeDefaults;
+  const selectedLabels = useWatch({ control, name: 'entityLabels' });
 
   const { data, isLoading } = useAnonymizerListEntityLabels(workspace, { query: {} });
-  const labels = data?.data ?? [];
+  const available = useMemo(() => data?.data ?? [], [data?.data]);
+
+  const isCustom = entityMode === ENTITY_MODE_CUSTOM;
+  const selected = useMemo(() => selectedLabels ?? [], [selectedLabels]);
+
+  const items = useMemo(() => {
+    const sections = buildEntitySections([...available]).map((section) => ({
+      kind: 'section' as const,
+      slotHeading: section.heading,
+      items: section.items,
+    }));
+    const candidate = customLabelCandidate(inputValue, [...available], selected);
+    return candidate
+      ? [{ kind: 'section' as const, slotHeading: 'Custom label', items: [candidate] }, ...sections]
+      : sections;
+  }, [available, inputValue, selected]);
+
+  const removeLabel = (label: string) =>
+    setValue(
+      'entityLabels',
+      selected.filter((value) => value !== label),
+      { shouldValidate: true }
+    );
 
   return (
     <Stack gap="density-lg">
@@ -39,27 +66,52 @@ export const EntitiesSection: FC = () => {
       />
       <Text kind="body/regular/md">
         {entityMode === ENTITY_MODE_AUTO
-          ? 'Auto-detect lets the augmenter create additional labels beyond the defaults.'
-          : 'Custom mode only outputs entities you define. Use Auto-detect to allow additional labels.'}
+          ? 'Auto-detect mode allows the augmenter to create additional labels beyond the defaults. To restrict the output entities to a defined list, use Custom.'
+          : 'Custom mode only outputs entities defined by you. To allow the augmenter to create additional labels beyond the defaults, use Auto-detect mode.'}
       </Text>
       {isCustom && (
-        <ControlledCheckbox
-          useControllerProps={{ name: 'includeDefaultEntities', control }}
-          formFieldProps={{ slotLabel: 'Include all default entities' }}
-        />
-      )}
-      {showLabelPicker && (
-        <ControlledCombobox
-          kind="multiple"
-          loading={isLoading}
-          items={labels}
-          useControllerProps={{ name: 'entityLabels', control }}
-          formFieldProps={{
-            slotLabel: 'Entity Labels',
-            slotInfo: 'Only these entity types will be detected and replaced.',
-          }}
-        />
+        <Stack gap="density-md">
+          <ControlledCombobox
+            kind="multiple"
+            aria-label="Entity labels"
+            items={items}
+            inputValue={inputValue}
+            onInputValueChange={setInputValue}
+            onChange={() => setInputValue('')}
+            placeholder="Select labels..."
+            emptyStateMessage={isLoading ? 'Loading labels...' : 'No matching labels.'}
+            multipleMode="count"
+            formatSummaryLabel={(count) => `${count} selected`}
+            useControllerProps={{ name: 'entityLabels', control }}
+            formFieldProps={{
+              slotLabel: 'Entity Labels',
+              slotInfo:
+                'Pick from the detected entity types, or type your own label and select it.',
+            }}
+          />
+          {selected.length > 0 && (
+            <Flex className="flex-wrap" gap="density-sm">
+              {selected.map((label) => (
+                <Tag
+                  key={label}
+                  color={entityTagColor(label)}
+                  kind="outline"
+                  aria-label={`Remove ${label}`}
+                  onClick={() => removeLabel(label)}
+                >
+                  {label}
+                  <X size={14} />
+                </Tag>
+              ))}
+            </Flex>
+          )}
+        </Stack>
       )}
+      <ControlledCheckbox
+        slotLabel={`Include all ${available.length} default entities`}
+        disabled={isLoading}
+        useControllerProps={{ name: 'includeDefaultEntities', control }}
+      />
     </Stack>
   );
 };

web/packages/studio/src/routes/AnonymizerBuilderRoute/constants.ts

@@ -2,6 +2,8 @@
 // SPDX-License-Identifier: Apache-2.0
 
 import { RiskTolerance } from '@nemo/sdk/generated/anonymizer/schema';
+import type { Tag } from '@nvidia/foundations-react-core';
+import type { ComponentProps } from 'react';
 
 export const SOURCE_TYPE_URL = 'url';
 export const SOURCE_TYPE_DATASET = 'dataset';
@@ -103,6 +105,132 @@ export const ENTITY_MODE_OPTIONS: { value: EntityMode; children: string }[] = [
   { value: ENTITY_MODE_AUTO, children: 'Auto-detect' },
 ];
 
+export type EntityTagColor = NonNullable<ComponentProps<typeof Tag>['color']>;
+
+interface EntityCategory {
+  readonly label: string;
+  readonly color: EntityTagColor;
+  readonly labels: readonly string[];
+}
+
+/**
+ * The entity-labels endpoint returns a flat list, so the grouping shown in the picker is
+ * curated here to match the design. Anything the API adds that isn't listed falls into Other.
+ */
+export const ENTITY_CATEGORIES: readonly EntityCategory[] = [
+  {
+    label: 'Personal Identity',
+    color: 'blue',
+    labels: [
+      'first_name',
+      'last_name',
+      'date_of_birth',
+      'age',
+      'gender',
+      'nationality',
+      'language',
+    ],
+  },
+  {
+    label: 'Demographics & Beliefs',
+    color: 'purple',
+    labels: ['race_ethnicity', 'sexuality', 'political_view', 'religious_belief'],
+  },
+  {
+    label: 'Contact & Communication',
+    color: 'teal',
+    labels: ['email', 'phone_number', 'fax_number'],
+  },
+  {
+    label: 'Location & Address',
+    color: 'green',
+    labels: [
+      'street_address',
+      'city',
+      'state',
+      'county',
+      'country',
+      'postcode',
+      'coordinate',
+      'place_name',
+      'landmark',
+    ],
+  },
+  { label: 'Date & Time', color: 'yellow', labels: ['date', 'time', 'date_time'] },
+  {
+    label: 'Government & Legal IDs',
+    color: 'red',
+    labels: ['ssn', 'national_id', 'tax_id', 'employee_id', 'certificate_license_number', 'pin'],
+  },
+  {
+    label: 'Financial',
+    color: 'blue',
+    labels: [
+      'credit_debit_card',
+      'cvv',
+      'account_number',
+      'bank_routing_number',
+      'swift_bic',
+      'monetary_amount',
+    ],
+  },
+  {
+    label: 'Medical & Health',
+    color: 'purple',
+    labels: [
+      'medical_record_number',
+      'health_plan_beneficiary_number',
+      'blood_type',
+      'biometric_identifier',
+    ],
+  },
+  {
+    label: 'Digital & Network',
+    color: 'teal',
+    labels: [
+      'ipv4',
+      'ipv6',
+      'mac_address',
+      'url',
+      'api_key',
+      'http_cookie',
+      'device_identifier',
+      'user_name',
+      'password',
+      'unique_id',
+    ],
+  },
+  { label: 'Vehicle & Transport', color: 'green', labels: ['license_plate', 'vehicle_identifier'] },
+  {
+    label: 'Employment & Organization',
+    color: 'yellow',
+    labels: ['occupation', 'employment_status', 'company_name', 'organization_name', 'customer_id'],
+  },
+  {
+    label: 'Education',
+    color: 'red',
+    labels: ['university', 'education_level', 'degree', 'field_of_study'],
+  },
+  {
+    label: 'Legal & Institutional',
+    color: 'blue',
+    labels: ['court_name', 'prison_detention_facility'],
+  },
+];
+
+export const ENTITY_CATEGORY_OTHER = 'Other';
+export const ENTITY_CUSTOM_TAG_COLOR: EntityTagColor = 'gray';
+
+const COLOR_BY_LABEL = new Map<string, EntityTagColor>(
+  ENTITY_CATEGORIES.flatMap((category) =>
+    category.labels.map((label) => [label, category.color] as const)
+  )
+);
+
+/** Custom labels have no category, so they fall back to the neutral chip colour. */
+export const entityTagColor = (label: string): EntityTagColor =>
+  COLOR_BY_LABEL.get(label) ?? ENTITY_CUSTOM_TAG_COLOR;
+
 export const DEFAULT_PREVIEW_ROWS = 1;
 
 export const MAX_COLUMN_INTROSPECTION_BYTES = 50 * 1024 * 1024;

web/packages/studio/src/routes/AnonymizerBuilderRoute/entityItems.test.ts

@@ -0,0 +1,60 @@
+// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import {
+  ENTITY_CATEGORIES,
+  ENTITY_CATEGORY_OTHER,
+  ENTITY_CUSTOM_TAG_COLOR,
+  entityTagColor,
+} from '@studio/routes/AnonymizerBuilderRoute/constants';
+import {
+  buildEntitySections,
+  customLabelCandidate,
+} from '@studio/routes/AnonymizerBuilderRoute/entityItems';
+
+const ALL_CATEGORY_LABELS = ENTITY_CATEGORIES.flatMap((category) => [...category.labels]);
+
+describe('entity categories', () => {
+  it('never lists the same label under two categories', () => {
+    expect(new Set(ALL_CATEGORY_LABELS).size).toBe(ALL_CATEGORY_LABELS.length);
+  });
+
+  it('colours labels by category and falls back for custom ones', () => {
+    expect(entityTagColor('first_name')).toBe(ENTITY_CATEGORIES[0].color);
+    expect(entityTagColor('ice_cream_flavor')).toBe(ENTITY_CUSTOM_TAG_COLOR);
+  });
+});
+
+describe('buildEntitySections', () => {
+  it('groups labels in category order and drops empty categories', () => {
+    const sections = buildEntitySections(['city', 'first_name', 'email']);
+    expect(sections).toEqual([
+      { heading: 'Personal Identity', items: ['first_name'] },
+      { heading: 'Contact & Communication', items: ['email'] },
+      { heading: 'Location & Address', items: ['city'] },
+    ]);
+  });
+
+  it('collects labels missing from the curated map under Other', () => {
+    const sections = buildEntitySections(['first_name', 'brand_new_label']);
+    expect(sections.at(-1)).toEqual({
+      heading: ENTITY_CATEGORY_OTHER,
+      items: ['brand_new_label'],
+    });
+  });
+
+  it('covers every curated label without an Other bucket', () => {
+    const sections = buildEntitySections(ALL_CATEGORY_LABELS);
+    expect(sections.map((s) => s.heading)).not.toContain(ENTITY_CATEGORY_OTHER);
+    expect(sections.flatMap((s) => s.items)).toHaveLength(ALL_CATEGORY_LABELS.length);
+  });
+});
+
+describe('customLabelCandidate', () => {
+  it('offers a trimmed candidate that is neither available nor already selected', () => {
+    expect(customLabelCandidate('  foobar ', ['email'], [])).toBe('foobar');
+    expect(customLabelCandidate('email', ['email'], [])).toBeNull();
+    expect(customLabelCandidate('foobar', [], ['foobar'])).toBeNull();
+    expect(customLabelCandidate('   ', [], [])).toBeNull();
+  });
+});