diff --git a/web/packages/studio/src/routes/AnonymizerBuilderRoute/components/AnonymizerBuilderForm.tsx b/web/packages/studio/src/routes/AnonymizerBuilderRoute/components/AnonymizerBuilderForm.tsx index f4b84f13aa..28c16d5248 100644 --- a/web/packages/studio/src/routes/AnonymizerBuilderRoute/components/AnonymizerBuilderForm.tsx +++ b/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(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 diff --git a/web/packages/studio/src/routes/AnonymizerBuilderRoute/components/EntitiesSection.tsx b/web/packages/studio/src/routes/AnonymizerBuilderRoute/components/EntitiesSection.tsx index e700ce7cc0..cd520cfc07 100644 --- a/web/packages/studio/src/routes/AnonymizerBuilderRoute/components/EntitiesSection.tsx +++ b/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(); + const { control, setValue } = useFormContext(); 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 ( @@ -39,27 +66,52 @@ export const EntitiesSection: FC = () => { /> {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.'} {isCustom && ( - - )} - {showLabelPicker && ( - + + 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 && ( + + {selected.map((label) => ( + removeLabel(label)} + > + {label} + + + ))} + + )} + )} + ); }; diff --git a/web/packages/studio/src/routes/AnonymizerBuilderRoute/constants.ts b/web/packages/studio/src/routes/AnonymizerBuilderRoute/constants.ts index 271a29c46d..dec66eefe4 100644 --- a/web/packages/studio/src/routes/AnonymizerBuilderRoute/constants.ts +++ b/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['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( + 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; diff --git a/web/packages/studio/src/routes/AnonymizerBuilderRoute/entityItems.test.ts b/web/packages/studio/src/routes/AnonymizerBuilderRoute/entityItems.test.ts new file mode 100644 index 0000000000..0af50613ea --- /dev/null +++ b/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(); + }); +}); diff --git a/web/packages/studio/src/routes/AnonymizerBuilderRoute/entityItems.ts b/web/packages/studio/src/routes/AnonymizerBuilderRoute/entityItems.ts new file mode 100644 index 0000000000..1999b57328 --- /dev/null +++ b/web/packages/studio/src/routes/AnonymizerBuilderRoute/entityItems.ts @@ -0,0 +1,44 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { + ENTITY_CATEGORIES, + ENTITY_CATEGORY_OTHER, +} from '@studio/routes/AnonymizerBuilderRoute/constants'; + +export interface EntitySection { + readonly heading: string; + readonly items: string[]; +} + +/** Group the flat label list from the API under the curated categories, in category order. */ +export const buildEntitySections = (available: string[]): EntitySection[] => { + const remaining = new Set(available); + const sections: EntitySection[] = []; + + for (const category of ENTITY_CATEGORIES) { + const items = category.labels.filter((label) => remaining.delete(label)); + if (items.length) sections.push({ heading: category.label, items }); + } + + if (remaining.size) { + sections.push({ heading: ENTITY_CATEGORY_OTHER, items: [...remaining] }); + } + + return sections; +}; + +/** + * The typed value, when it isn't already offered or selected. Surfacing it as an item is what + * lets a custom label be added, since the underlying combobox has no create affordance. + */ +export const customLabelCandidate = ( + input: string, + available: string[], + selected: string[] +): string | null => { + const candidate = input.trim(); + if (!candidate) return null; + if (available.includes(candidate) || selected.includes(candidate)) return null; + return candidate; +}; diff --git a/web/packages/studio/src/routes/AnonymizerBuilderRoute/schema.test.ts b/web/packages/studio/src/routes/AnonymizerBuilderRoute/schema.test.ts index bc0d63cc26..47e7e40b6c 100644 --- a/web/packages/studio/src/routes/AnonymizerBuilderRoute/schema.test.ts +++ b/web/packages/studio/src/routes/AnonymizerBuilderRoute/schema.test.ts @@ -9,11 +9,13 @@ import { } from '@studio/routes/AnonymizerBuilderRoute/constants'; import { AnonymizerFormData, + anonymizerFormSchema, buildAnonymizerJobRequest, getAnonymizerFormDefaults, } from '@studio/routes/AnonymizerBuilderRoute/schema'; const ALL_ROLES = [...DETECTION_ROLES, REPLACE_ROLE, ...REWRITE_ROLES]; +const DEFAULT_LABELS = ['email', 'ssn', 'first_name']; const roleModels = (model: string, provider: string): AnonymizerFormData['roleModels'] => Object.fromEntries(ALL_ROLES.map((role) => [role, { modelId: role, model, provider }])); @@ -185,18 +187,55 @@ describe('buildAnonymizerJobRequest', () => { } }); - it('sets config.detect.entity_labels only for custom labels without defaults', () => { + it('sends only the picked labels when defaults are excluded', () => { const custom = buildAnonymizerJobRequest( - form({ entityMode: 'custom', includeDefaultEntities: false, entityLabels: ['email', 'ssn'] }) + form({ entityMode: 'custom', includeDefaultEntities: false, entityLabels: ['email', 'ssn'] }), + DEFAULT_LABELS ); expect(custom.spec.config.detect).toEqual({ entity_labels: ['email', 'ssn'] }); + }); + + it('merges defaults with custom picks when defaults are included', () => { + const merged = buildAnonymizerJobRequest( + form({ + entityMode: 'custom', + includeDefaultEntities: true, + entityLabels: ['email', 'ice_cream_flavor'], + }), + DEFAULT_LABELS + ); + expect(merged.spec.config.detect).toEqual({ + entity_labels: [...DEFAULT_LABELS, 'ice_cream_flavor'], + }); + }); + + it('omits detect when the selection adds nothing to the defaults', () => { + const defaultsOnly = buildAnonymizerJobRequest( + form({ entityMode: 'custom', includeDefaultEntities: true, entityLabels: [] }), + DEFAULT_LABELS + ); + expect(defaultsOnly.spec.config.detect).toBeUndefined(); + + const subsetOfDefaults = buildAnonymizerJobRequest( + form({ entityMode: 'custom', includeDefaultEntities: true, entityLabels: ['email'] }), + DEFAULT_LABELS + ); + expect(subsetOfDefaults.spec.config.detect).toBeUndefined(); + }); - const withDefaults = buildAnonymizerJobRequest( - form({ entityMode: 'custom', includeDefaultEntities: true, entityLabels: ['email'] }) + it('treats a duplicated default as no addition', () => { + const req = buildAnonymizerJobRequest( + form({ entityMode: 'custom', includeDefaultEntities: true, entityLabels: ['email'] }), + [...DEFAULT_LABELS, 'email'] ); - expect(withDefaults.spec.config.detect).toBeUndefined(); + expect(req.spec.config.detect).toBeUndefined(); + }); - const auto = buildAnonymizerJobRequest(form({ entityMode: 'auto', entityLabels: ['email'] })); + it('ignores entity labels in auto-detect mode', () => { + const auto = buildAnonymizerJobRequest( + form({ entityMode: 'auto', entityLabels: ['email'] }), + DEFAULT_LABELS + ); expect(auto.spec.config.detect).toBeUndefined(); }); @@ -220,3 +259,33 @@ describe('buildAnonymizerJobRequest', () => { }); }); }); + +describe('anonymizerFormSchema', () => { + const parse = (overrides: Partial) => + anonymizerFormSchema.safeParse({ + ...getAnonymizerFormDefaults(), + source: 'https://example.com/data.csv', + roleModels: roleModels('openai/gpt-oss-120b', 'default/nvidia'), + ...overrides, + }); + + it('rejects custom mode with neither labels nor defaults', () => { + const result = parse({ + entityMode: 'custom', + includeDefaultEntities: false, + entityLabels: [], + }); + expect(result.success).toBe(false); + expect(result.error?.issues.some((i) => i.path.join('.') === 'entityLabels')).toBe(true); + }); + + it('accepts custom mode with labels, or with defaults included', () => { + expect( + parse({ entityMode: 'custom', includeDefaultEntities: false, entityLabels: ['email'] }) + .success + ).toBe(true); + expect( + parse({ entityMode: 'custom', includeDefaultEntities: true, entityLabels: [] }).success + ).toBe(true); + }); +}); diff --git a/web/packages/studio/src/routes/AnonymizerBuilderRoute/schema.ts b/web/packages/studio/src/routes/AnonymizerBuilderRoute/schema.ts index 81b1ed4528..7f4fdfbf22 100644 --- a/web/packages/studio/src/routes/AnonymizerBuilderRoute/schema.ts +++ b/web/packages/studio/src/routes/AnonymizerBuilderRoute/schema.ts @@ -85,6 +85,20 @@ export const anonymizerFormSchema = z }); } } + + // Without this the request would carry no detect config and the server would fall back to + // its own defaults — the opposite of the restricted set Custom mode promises. + if ( + data.entityMode === ENTITY_MODE_CUSTOM && + !data.includeDefaultEntities && + !data.entityLabels.length + ) { + ctx.addIssue({ + code: z.ZodIssueCode.custom, + path: ['entityLabels'], + message: 'Select at least one entity label, or include the default entities', + }); + } }); export type AnonymizerFormData = z.infer; @@ -168,18 +182,41 @@ const buildRewriteConfig = (form: AnonymizerFormData): Rewrite => { return rewrite; }; -export const buildAnonymizerJobRequest = (form: AnonymizerFormData): RunJobRequest => { +/** + * entity_labels replaces the default set server-side, so "include defaults" has to send the + * defaults alongside the custom picks. Omitted entirely when the selection adds nothing, which + * leaves the server on its own defaults. + */ +const buildDetectConfig = ( + form: AnonymizerFormData, + defaultEntityLabels: string[] +): AnonymizerConfigInput['detect'] => { + if (form.entityMode !== ENTITY_MODE_CUSTOM) return undefined; + + const labels = form.includeDefaultEntities + ? [...new Set([...defaultEntityLabels, ...form.entityLabels])] + : form.entityLabels; + + if (!labels.length) return undefined; + if (form.includeDefaultEntities && labels.length === new Set(defaultEntityLabels).size) { + return undefined; + } + + return { entity_labels: labels }; +}; + +export const buildAnonymizerJobRequest = ( + form: AnonymizerFormData, + defaultEntityLabels: string[] = [] +): RunJobRequest => { const config: AnonymizerConfigInput = form.strategy === REWRITE_STRATEGY ? { rewrite: buildRewriteConfig(form) } : { replace: buildReplaceConfig(form) }; - const useCustomLabels = - form.entityMode === ENTITY_MODE_CUSTOM && - !form.includeDefaultEntities && - form.entityLabels.length > 0; - if (useCustomLabels) { - config.detect = { entity_labels: form.entityLabels }; + const detect = buildDetectConfig(form, defaultEntityLabels); + if (detect) { + config.detect = detect; } const aliasByModel = new Map();