diff --git a/frontend/e2e/history.spec.ts b/frontend/e2e/history.spec.ts index 6c3db78395..d3f178ed9b 100644 --- a/frontend/e2e/history.spec.ts +++ b/frontend/e2e/history.spec.ts @@ -15,6 +15,8 @@ interface MockAttackSummary { last_message_preview?: string | null; message_count: number; related_conversation_ids: string[]; + operator?: string | null; + operation?: string | null; labels: Record; created_at: string; updated_at: string; @@ -30,6 +32,8 @@ function makeAttack(overrides: Partial & { attack_result_id: last_message_preview: null, message_count: 0, related_conversation_ids: [], + operator: null, + operation: null, labels: {}, created_at: new Date().toISOString(), updated_at: new Date().toISOString(), @@ -44,7 +48,8 @@ const ATTACKS: MockAttackSummary[] = [ attack_type: "SingleTurnAttack", target: { target_type: "OpenAIChatTarget", model_name: "gpt-4o" }, outcome: "success", - labels: { operator: "alice", operation: "test_a" }, + operator: "alice", + operation: "test_a", message_count: 3, last_message_preview: "Hello from alice", }), @@ -53,7 +58,8 @@ const ATTACKS: MockAttackSummary[] = [ attack_type: "MultiTurnAttack", target: { target_type: "OpenAIImageTarget", model_name: "dall-e-3" }, outcome: "failure", - labels: { operator: "bob", operation: "test_b" }, + operator: "bob", + operation: "test_b", message_count: 5, last_message_preview: "Hello from bob", }), @@ -62,7 +68,8 @@ const ATTACKS: MockAttackSummary[] = [ attack_type: "SingleTurnAttack", target: { target_type: "OpenAIChatTarget", model_name: "gpt-4o" }, outcome: "undetermined", - labels: { operator: "alice", operation: "test_b" }, + operator: "alice", + operation: "test_b", message_count: 1, }), makeAttack({ @@ -70,7 +77,8 @@ const ATTACKS: MockAttackSummary[] = [ attack_type: "MultiTurnAttack", target: { target_type: "OpenAIChatTarget", model_name: "gpt-4o" }, outcome: "success", - labels: { operator: "bob", operation: "test_a" }, + operator: "bob", + operation: "test_a", message_count: 2, last_message_preview: "Hello again from bob", }), @@ -83,7 +91,7 @@ function generatePaginatedAttacks(count: number): MockAttackSummary[] { attack_result_id: `atk-page-${String(i).padStart(3, "0")}`, attack_type: i % 2 === 0 ? "SingleTurnAttack" : "MultiTurnAttack", outcome: "undetermined", - labels: { operator: "paginator" }, + operator: "paginator", message_count: 1, }), ); @@ -143,10 +151,9 @@ async function mockHistoryAPIs( contentType: "application/json", body: JSON.stringify({ source: "attacks", - labels: { - operator: operatorLabels, - operation: operationLabels, - }, + labels: {}, + operators: operatorLabels, + operations: operationLabels, }), }); }); @@ -163,6 +170,8 @@ async function mockHistoryAPIs( const url = new URL(route.request().url()); const attackTypeParams = url.searchParams.getAll("attack_types"); const outcome = url.searchParams.get("outcome"); + const operatorParams = url.searchParams.getAll("operator"); + const operationParams = url.searchParams.getAll("operation"); const labelParams = url.searchParams.getAll("label"); let filtered = [...attacks]; @@ -172,6 +181,12 @@ async function mockHistoryAPIs( if (outcome) { filtered = filtered.filter((a) => a.outcome === outcome); } + if (operatorParams.length > 0) { + filtered = filtered.filter((a) => a.operator != null && operatorParams.includes(a.operator)); + } + if (operationParams.length > 0) { + filtered = filtered.filter((a) => a.operation != null && operationParams.includes(a.operation)); + } if (labelParams.length > 0) { // Group repeated label keys into OR-sets; combine across keys with AND. const grouped = new Map(); diff --git a/frontend/src/App.tsx b/frontend/src/App.tsx index 3b2f68f844..6f31c6e95c 100644 --- a/frontend/src/App.tsx +++ b/frontend/src/App.tsx @@ -110,6 +110,7 @@ interface LoadedAttack { targetSource: 'persisted' | 'active-selection' mainConversationId: string | null labels: Record | null + operator: string | null target: TargetInfo | null relatedConversationIds: string[] objective: string @@ -319,6 +320,7 @@ function App() { status: 'loading', mainConversationId: null, labels: null, + operator: null, target: null, relatedConversationIds: [], objective: '', @@ -333,6 +335,7 @@ function App() { targetSource: 'persisted', mainConversationId: attack.conversation_id, labels: attack.labels ?? {}, + operator: attack.operator ?? null, target: attack.target ?? null, relatedConversationIds: attack.related_conversation_ids ?? [], objective: attack.objective ?? '', @@ -352,6 +355,7 @@ function App() { status: isMissing ? 'not-found' : 'error', mainConversationId: null, labels: null, + operator: null, target: null, relatedConversationIds: [], objective: '', @@ -441,6 +445,7 @@ function App() { mainConversationId: convId, // New attack uses the current user's labels, so it is never operator-locked. labels: null, + operator: null, target, relatedConversationIds: [], objective: '', @@ -488,7 +493,7 @@ function App() { labels={globalLabels} onLabelsChange={handleGlobalLabelsChange} onNavigate={handleNavigate} - attackLabels={readyAttack ? readyAttack.labels : null} + attackOperator={readyAttack ? readyAttack.operator : null} attackTarget={readyAttack ? readyAttack.target : null} targetResolutionStatus={targetResolutionStatus} onRetryTargetResolution={retryTargetResolution} diff --git a/frontend/src/components/Chat/ChatWindow.test.tsx b/frontend/src/components/Chat/ChatWindow.test.tsx index b932a1e230..33bdd8d390 100644 --- a/frontend/src/components/Chat/ChatWindow.test.tsx +++ b/frontend/src/components/Chat/ChatWindow.test.tsx @@ -706,6 +706,7 @@ describe("ChatWindow Integration", () => { expect(mockedAttacksApi.createAttack).toHaveBeenCalledWith({ target_registry_name: "openai_chat_1", labels: { operator: 'testuser', operation: 'test_op' }, + system_prompt: undefined, }); expect(onConversationCreated).toHaveBeenCalledWith("ar-conv-1", "conv-1"); expect(mockedAttacksApi.addMessage).toHaveBeenCalledWith("ar-conv-1", { @@ -714,7 +715,6 @@ describe("ChatWindow Integration", () => { send: true, target_registry_name: "openai_chat_1", target_conversation_id: "conv-1", - labels: { operator: "testuser", operation: "test_op" }, }); }); @@ -2655,7 +2655,7 @@ describe("ChatWindow Integration", () => { conversationId="conv-locked" activeConversationId="conv-locked" labels={{ operator: "alice", operation: "test_op" }} - attackLabels={{ operator: "bob", operation: "test_op" }} + attackOperator="bob" /> ); @@ -3786,7 +3786,7 @@ describe("ChatWindow Integration", () => { it("allows exporting a read-only historical conversation", async () => { const user = userEvent.setup(); // Operator lock: the loaded attack belongs to a different operator. - await renderWithLoadedConversation({ attackLabels: { operator: "someone-else" } }); + await renderWithLoadedConversation({ attackOperator: "someone-else" }); const { clickSpy } = spyOnDownloadAnchor(); const exportButton = screen.getByRole("button", { name: /export conversation/i }); diff --git a/frontend/src/components/Chat/ChatWindow.tsx b/frontend/src/components/Chat/ChatWindow.tsx index d5a9129c1f..075b786cba 100644 --- a/frontend/src/components/Chat/ChatWindow.tsx +++ b/frontend/src/components/Chat/ChatWindow.tsx @@ -39,7 +39,9 @@ import { buildMessagePieces, backendMessagesToFrontend } from '../../utils/messa import { exportConversation } from '../../utils/conversationExport' import type { ExportFormat } from '../../utils/conversationExport' import type { + AddMessageRequest, AttackTargetResolutionStatus, + CreateAttackRequest, Message, MessageAttachment, TargetInstance, @@ -89,8 +91,8 @@ interface ChatWindowProps { labels?: Record onLabelsChange?: (labels: Record) => void onNavigate?: (view: ViewName) => void - /** Labels from the loaded attack (for operator locking). Null for new attacks. */ - attackLabels?: Record | null + /** Operator from the loaded attack (for operator locking). Null for new attacks. */ + attackOperator?: string | null /** Target info that the current attack was started with (for cross-target guard). */ attackTarget?: TargetInfo | null /** Result of resolving the persisted attack target against the current registry. */ @@ -118,7 +120,7 @@ export default function ChatWindow({ labels, onLabelsChange, onNavigate, - attackLabels, + attackOperator, attackTarget, targetResolutionStatus = 'idle', onRetryTargetResolution, @@ -238,10 +240,9 @@ export default function ChatWindow({ && isTargetResolutionBlocking(targetResolutionStatus), ) const currentOperator = labels?.operator - const attackOperator = attackLabels?.operator // Existing attacks are operator-locked when their operator differs from the current one. const isOperatorLocked = Boolean( - attackResultId && attackLabels && attackOperator && currentOperator && attackOperator !== currentOperator, + attackResultId && attackOperator && currentOperator && attackOperator !== currentOperator, ) // They are cross-target locked when the selected target's canonical hash differs from the persisted target. const isCrossTargetLocked = Boolean( @@ -428,11 +429,14 @@ export default function ChatWindow({ let currentConversationId = conversationId let currentActiveConversationId = activeConversationId if (!currentAttackResultId) { - const createResponse = await attacksApi.createAttack({ + const createRequest: CreateAttackRequest = { target_registry_name: activeTarget.target_registry_name, - labels: labels, + // TODO(PyRIT 1.4): Pass only dedicated attribution after legacy label aliases are removed. + // The create-attack API normalizes these aliases through _AttackAttributionInput. + labels, system_prompt: supportsSystemPrompt ? systemPrompt.trim() || undefined : undefined, - }) + } + const createResponse = await attacksApi.createAttack(createRequest) currentAttackResultId = createResponse.attack_result_id currentConversationId = createResponse.conversation_id currentActiveConversationId = currentConversationId @@ -465,15 +469,15 @@ export default function ChatWindow({ // Send message to target const converterIds = allConverterIds.length > 0 ? allConverterIds : undefined - const response = await attacksApi.addMessage(currentAttackResultId!, { + const addMessageRequest: AddMessageRequest = { role: 'user', pieces, send: true, target_registry_name: activeTarget.target_registry_name, target_conversation_id: effectiveConvId!, - labels: labels ?? undefined, converter_ids: converterIds, - }) + } + const response = await attacksApi.addMessage(currentAttackResultId!, addMessageRequest) // Clear converter state after successful send setPieceConversions({}) @@ -656,7 +660,7 @@ export default function ChatWindow({ try { const createResponse = await attacksApi.createAttack({ target_registry_name: activeTarget.target_registry_name, - labels: labels, + labels, source_conversation_id: activeConversationId, cutoff_index: messageIndex, }) @@ -707,7 +711,7 @@ export default function ChatWindow({ // Let the backend clone the conversation with new labels const createResponse = await attacksApi.createAttack({ target_registry_name: activeTarget.target_registry_name, - labels: labels, + labels, source_conversation_id: activeConversationId, cutoff_index: lastIndex, }) diff --git a/frontend/src/components/History/AttackHistory.test.tsx b/frontend/src/components/History/AttackHistory.test.tsx index a426952170..c989f5005d 100644 --- a/frontend/src/components/History/AttackHistory.test.tsx +++ b/frontend/src/components/History/AttackHistory.test.tsx @@ -687,9 +687,9 @@ describe('AttackHistory', () => { }) mockedLabelsApi.getLabels.mockResolvedValue({ source: 'attacks', + operators: ['alice', 'bob'], + operations: ['op_one'], labels: { - operator: ['alice', 'bob'], - operation: ['op_one'], custom_tag: ['val1', 'val2'], }, }) @@ -707,6 +707,39 @@ describe('AttackHistory', () => { expect(mockedLabelsApi.getLabels).toHaveBeenCalled() }) + it('should narrow arbitrary label options by selected attribution and labels', async () => { + mockedAttacksApi.listAttacks.mockResolvedValue({ + items: [], + pagination: { limit: 25, has_more: false }, + }) + mockedLabelsApi.getLabels.mockResolvedValue({ + source: 'attacks', + operators: ['alice', 'bob'], + operations: ['nightly'], + labels: { env: ['prod'] }, + }) + const activeFilters = { + ...DEFAULT_HISTORY_FILTERS, + operator: ['alice'], + operation: ['nightly'], + otherLabels: ['team:red'], + } + + render( + + + + ) + + await waitFor(() => { + expect(mockedLabelsApi.getLabels).toHaveBeenCalledWith('attacks', { + operator: ['alice'], + operation: ['nightly'], + label: ['team:red'], + }) + }) + }) + it('should show empty text with filter hint when filters active and no results', async () => { mockedAttacksApi.listAttacks.mockResolvedValue({ items: [], diff --git a/frontend/src/components/History/AttackHistory.tsx b/frontend/src/components/History/AttackHistory.tsx index 5e4e147d3f..12b7e83dd3 100644 --- a/frontend/src/components/History/AttackHistory.tsx +++ b/frontend/src/components/History/AttackHistory.tsx @@ -31,15 +31,14 @@ const PAGE_SIZE = 25 type ListParams = Parameters[0] function buildListParams(filters: HistoryFilters, pageCursor: string | undefined): ListParams { - const labelParams: string[] = [] - for (const op of filters.operator) { labelParams.push(`operator:${op}`) } - for (const op of filters.operation) { labelParams.push(`operation:${op}`) } - labelParams.push(...filters.otherLabels) + const labelParams = [...filters.otherLabels] const params: ListParams = { limit: PAGE_SIZE } if (pageCursor) params.cursor = pageCursor if (filters.attackTypes.length > 0) params.attack_types = filters.attackTypes if (filters.outcome) params.outcome = filters.outcome + if (filters.operator.length > 0) params.operator = filters.operator + if (filters.operation.length > 0) params.operation = filters.operation if (filters.converter.length > 0) params.converter_types = filters.converter // Match mode is only meaningful with >=2 converters selected. if (filters.converter.length >= 2) params.converter_types_match = filters.converterMatchMode @@ -49,6 +48,17 @@ function buildListParams(filters: HistoryFilters, pageCursor: string | undefined return params } +function buildOtherLabelOptions(labels: Record): string[] { + const options: string[] = [] + for (const [key, values] of Object.entries(labels)) { + if (key === 'operator' || key === 'operation') continue + for (const value of values) { + options.push(`${key}:${value}`) + } + } + return options.sort() +} + export default function AttackHistory({ onOpenAttack, filters, @@ -67,7 +77,11 @@ export default function AttackHistory({ const [converterOptions, setConverterOptions] = useState([]) const [operatorOptions, setOperatorOptions] = useState([]) const [operationOptions, setOperationOptions] = useState([]) - const [otherLabelOptions, setOtherLabelOptions] = useState([]) + const [allOtherLabelOptions, setAllOtherLabelOptions] = useState([]) + const [narrowedOtherLabelOptions, setNarrowedOtherLabelOptions] = useState<{ + filterKey: string + options: string[] + } | null>(null) // Pagination const [cursor, setCursor] = useState(undefined) @@ -85,6 +99,18 @@ export default function AttackHistory({ filters.otherLabels, ]) const [settledFilterKey, setSettledFilterKey] = useState(null) + const labelOptionFilterKey = JSON.stringify([ + filters.operator, + filters.operation, + filters.otherLabels, + ]) + const hasLabelOptionFilters = filters.operator.length > 0 + || filters.operation.length > 0 + || filters.otherLabels.length > 0 + const otherLabelOptions = hasLabelOptionFilters + && narrowedOtherLabelOptions?.filterKey === labelOptionFilterKey + ? narrowedOtherLabelOptions.options + : allOtherLabelOptions // Bumped from event handlers (Refresh button, pagination) to re-trigger the // fetch effect without calling setState synchronously inside it. @@ -100,7 +126,7 @@ export default function AttackHistory({ setFetchToken(prev => ({ cursor: pageCursor, filterKey, nonce: prev.nonce + 1 })) }, [filterKey]) - // Load filter options on mount + // Attack and converter options do not depend on the active history filters. useEffect(() => { attacksApi.getAttackOptions() .then(resp => setAttackTypeOptions(resp.attack_types)) @@ -110,27 +136,42 @@ export default function AttackHistory({ .catch(() => { /* ignore */ }) labelsApi.getLabels() .then(resp => { - const operators: string[] = [] - const operations: string[] = [] - const others: string[] = [] - for (const [key, values] of Object.entries(resp.labels)) { - if (key === 'operator') { - operators.push(...values) - } else if (key === 'operation') { - operations.push(...values) - } else if (key !== 'source') { - for (const val of values) { - others.push(`${key}:${val}`) - } - } - } - setOperatorOptions(operators.sort()) - setOperationOptions(operations.sort()) - setOtherLabelOptions(others.sort()) + // TODO(PyRIT 1.4): Remove the labels.* fallbacks with legacy attribution aliases. + setOperatorOptions([...(resp.operators ?? resp.labels.operator ?? [])].sort()) + setOperationOptions([...(resp.operations ?? resp.labels.operation ?? [])].sort()) + setAllOtherLabelOptions(buildOtherLabelOptions(resp.labels)) }) .catch(() => { /* ignore */ }) }, []) + // Arbitrary label options are narrowed by the active indexed attribution filters. + useEffect(() => { + if (!hasLabelOptionFilters) return + let cancelled = false + labelsApi.getLabels('attacks', { + operator: filters.operator.length > 0 ? filters.operator : undefined, + operation: filters.operation.length > 0 ? filters.operation : undefined, + label: filters.otherLabels.length > 0 ? filters.otherLabels : undefined, + }) + .then(resp => { + if (cancelled) return + setNarrowedOtherLabelOptions({ + filterKey: labelOptionFilterKey, + options: buildOtherLabelOptions(resp.labels), + }) + }) + .catch(() => { /* ignore */ }) + return () => { + cancelled = true + } + }, [ + filters.operator, + filters.operation, + filters.otherLabels, + hasLabelOptionFilters, + labelOptionFilterKey, + ]) + // Fetch attacks whenever filters change or an event handler bumps fetchToken. // All setState calls live in .then/.catch/.finally so we don't trigger // react-hooks/set-state-in-effect. diff --git a/frontend/src/components/History/AttackTable.test.tsx b/frontend/src/components/History/AttackTable.test.tsx index 62dc4e366b..07f99a8e27 100644 --- a/frontend/src/components/History/AttackTable.test.tsx +++ b/frontend/src/components/History/AttackTable.test.tsx @@ -24,7 +24,9 @@ const sampleAttacks: AttackSummary[] = [ last_message_preview: 'Hello world', message_count: 5, related_conversation_ids: ['rel-1'], - labels: { operator: 'alice', operation: 'op_one', custom: 'val' }, + operator: 'alice', + operation: 'op_one', + labels: { custom: 'val' }, created_at: '2026-01-15T10:30:00Z', updated_at: '2026-01-15T11:00:00Z', }, diff --git a/frontend/src/components/History/AttackTable.tsx b/frontend/src/components/History/AttackTable.tsx index d6826ea00e..e00983a1f6 100644 --- a/frontend/src/components/History/AttackTable.tsx +++ b/frontend/src/components/History/AttackTable.tsx @@ -104,10 +104,10 @@ export default function AttackTable({ attacks, onOpenAttack, formatDate }: Attac )} - {attack.labels.operator || '—'} + {attack.operator || '—'} - {attack.labels.operation || '—'} + {attack.operation || '—'} {attack.message_count} @@ -133,7 +133,7 @@ export default function AttackTable({ attacks, onOpenAttack, formatDate }: Attac {(() => { - const otherLabels = Object.entries(attack.labels ?? {}).filter(([k]) => k !== 'operator' && k !== 'operation' && k !== 'source') + const otherLabels = Object.entries(attack.labels ?? {}).filter(([k]) => k !== 'source') return otherLabels.length > 0 ? (
{otherLabels.slice(0, 2).map(([k, v]) => ( diff --git a/frontend/src/components/Home/Home.test.tsx b/frontend/src/components/Home/Home.test.tsx index 2acf7a3b87..84e29dca62 100644 --- a/frontend/src/components/Home/Home.test.tsx +++ b/frontend/src/components/Home/Home.test.tsx @@ -37,7 +37,9 @@ function makeAttack(overrides: Partial = {}): AttackSummary { last_message_preview: "preview", message_count: 1, related_conversation_ids: [], - labels: { operator: "alice", operation: "op_alpha" }, + operator: "alice", + operation: "op_alpha", + labels: {}, created_at: new Date().toISOString(), updated_at: new Date().toISOString(), ...overrides, @@ -125,17 +127,20 @@ describe("Home", () => { items: [ makeAttack({ attack_result_id: "ar-1", - labels: { operator: "alice", operation: "op_alpha" }, + operator: "alice", + operation: "op_alpha", updated_at: new Date(now).toISOString(), }), makeAttack({ attack_result_id: "ar-2", - labels: { operator: "alice", operation: "op_alpha" }, + operator: "alice", + operation: "op_alpha", updated_at: new Date(now - 60_000).toISOString(), }), makeAttack({ attack_result_id: "ar-3", - labels: { operator: "alice", operation: "op_beta" }, + operator: "alice", + operation: "op_beta", updated_at: new Date(now - 120_000).toISOString(), }), ], @@ -161,19 +166,22 @@ describe("Home", () => { // the group's last-activity — exercising the "newer than current" branch. makeAttack({ attack_result_id: "ar-old", - labels: { operator: "alice", operation: "op_time" }, + operator: "alice", + operation: "op_time", last_message_preview: "older than a week", updated_at: new Date(now - 10 * DAY).toISOString(), }), makeAttack({ attack_result_id: "ar-hours", - labels: { operator: "alice", operation: "op_time" }, + operator: "alice", + operation: "op_time", last_message_preview: "a few hours ago", updated_at: new Date(now - 3 * HOUR).toISOString(), }), makeAttack({ attack_result_id: "ar-days", - labels: { operator: "alice", operation: "op_time" }, + operator: "alice", + operation: "op_time", last_message_preview: "a few days ago", updated_at: new Date(now - 3 * DAY).toISOString(), }), @@ -199,21 +207,24 @@ describe("Home", () => { items: [ makeAttack({ attack_result_id: "f1", - labels: { operator: "alice", operation: "op_full" }, + operator: "alice", + operation: "op_full", outcome: "success", last_message_preview: "first preview", updated_at: new Date(now - 60_000).toISOString(), }), makeAttack({ attack_result_id: "f2", - labels: { operator: "alice", operation: "op_full" }, + operator: "alice", + operation: "op_full", outcome: null, // unknown outcome -> default icon via the ?? 'undetermined' branch last_message_preview: null, // missing preview -> falls back to attack_type updated_at: new Date(now - 120_000).toISOString(), }), makeAttack({ attack_result_id: "f3", - labels: { operator: "alice", operation: "op_full" }, + operator: "alice", + operation: "op_full", // Outcome not present in the icon map -> exercises the icon fallback branch. outcome: "mystery" as unknown as AttackSummary["outcome"], last_message_preview: "third preview", @@ -221,7 +232,8 @@ describe("Home", () => { }), makeAttack({ attack_result_id: "f4", - labels: { operator: "alice", operation: "op_full" }, + operator: "alice", + operation: "op_full", last_message_preview: "fourth preview", updated_at: new Date(now - 240_000).toISOString(), }), @@ -246,7 +258,7 @@ describe("Home", () => { items: [ makeAttack({ attack_result_id: "ar-x", - labels: { operator: "alice" }, + operation: null, }), ], pagination: { has_more: false, next_cursor: null }, diff --git a/frontend/src/components/Home/Home.tsx b/frontend/src/components/Home/Home.tsx index 2f0e4b73ab..2e700b2da0 100644 --- a/frontend/src/components/Home/Home.tsx +++ b/frontend/src/components/Home/Home.tsx @@ -56,7 +56,7 @@ function groupAttacksByOperation(attacks: AttackSummary[]): OperationGroup[] { const groups = new Map() for (const attack of attacks) { - const opLabel = attack.labels?.operation + const opLabel = attack.operation const isUnlabeled = !opLabel const key = isUnlabeled ? NO_OPERATION_KEY : opLabel const updatedAt = new Date(attack.updated_at).getTime() diff --git a/frontend/src/components/Labels/LabelsBar.test.tsx b/frontend/src/components/Labels/LabelsBar.test.tsx index a5af3163da..0f0bf37dcc 100644 --- a/frontend/src/components/Labels/LabelsBar.test.tsx +++ b/frontend/src/components/Labels/LabelsBar.test.tsx @@ -956,7 +956,9 @@ describe('LabelsBar', () => { function renderWithOperations(onChange: jest.Mock, operations: string[] = OPERATIONS) { mockedLabelsApi.getLabels.mockResolvedValue({ source: 'attacks', - labels: { operation: operations, operator: ['alice'] }, + operators: ['alice'], + operations, + labels: {}, }) render( diff --git a/frontend/src/components/Labels/LabelsBar.tsx b/frontend/src/components/Labels/LabelsBar.tsx index e83039b36f..fa9757b76f 100644 --- a/frontend/src/components/Labels/LabelsBar.tsx +++ b/frontend/src/components/Labels/LabelsBar.tsx @@ -22,6 +22,7 @@ import { useLabelsBarStyles } from './LabelsBar.styles' const validateValue = (value: string): string | null => { if (!value) return 'Value is required' + if (value.length > 128) return 'Values must be 128 characters or fewer' if (value !== value.toLowerCase()) return 'Values must be lowercase' if (!/^[a-z0-9_]+$/.test(value)) return 'Only lowercase letters, numbers, underscores' return null @@ -218,7 +219,9 @@ export default function LabelsBar({ labels, onLabelsChange }: LabelsBarProps) { // so keep anything already collected rather than replacing outright. .then(resp => setExistingLabels(prev => ({ ...resp.labels, - operation: [...new Set([...(resp.labels.operation || []), ...(prev.operation || [])])], + // TODO(PyRIT 1.4): Remove the labels.* fallbacks with legacy attribution aliases. + operator: [...new Set([...(resp.operators ?? resp.labels.operator ?? []), ...(prev.operator || [])])], + operation: [...new Set([...(resp.operations ?? resp.labels.operation ?? []), ...(prev.operation || [])])], }))) .catch(() => setLabelsFailed(true)) .finally(() => setLabelsLoading(false)) diff --git a/frontend/src/services/api.test.ts b/frontend/src/services/api.test.ts index 9de3b57847..ba6d1998be 100644 --- a/frontend/src/services/api.test.ts +++ b/frontend/src/services/api.test.ts @@ -21,6 +21,7 @@ import { configurationApi, targetsApi, attacksApi, + labelsApi, scenariosApi, } from "./api"; @@ -548,6 +549,34 @@ describe("api service", () => { }); }); + it("should get narrowed labels with repeated query parameters", async () => { + const mockResponse = { + data: { + source: "attacks", + labels: { team: ["red"] }, + }, + }; + (apiClient.get as jest.Mock).mockResolvedValueOnce(mockResponse); + + await labelsApi.getLabels("attacks", { + operator: ["alice", "bob"], + operation: ["nightly"], + label: ["team:red"], + }); + + expect(apiClient.get).toHaveBeenCalledWith("/labels", { + params: { + source: "attacks", + operator: ["alice", "bob"], + operation: ["nightly"], + label: ["team:red"], + }, + paramsSerializer: { + indexes: null, + }, + }); + }); + it("should handle add message error", async () => { const error = new Error("Target not found"); (apiClient.post as jest.Mock).mockRejectedValueOnce(error); diff --git a/frontend/src/services/api.ts b/frontend/src/services/api.ts index 9bdcac5204..b4f0a77be2 100644 --- a/frontend/src/services/api.ts +++ b/frontend/src/services/api.ts @@ -15,6 +15,7 @@ import type { CustomInitializerListResponse, RegisterInitializerRequest, CreateAttackRequest, + LabelOptionsResponse, CreateAttackResponse, AttackSummary, AttackListResponse, @@ -339,6 +340,8 @@ export const attacksApi = { has_converters?: boolean include_scenario_attacks?: boolean outcome?: string + operator?: string[] + operation?: string[] label?: string[] min_turns?: number max_turns?: number @@ -366,8 +369,18 @@ export const attacksApi = { export const labelsApi = { getLabels: async ( source: 'attacks' | 'scenarios' = 'attacks', - ): Promise<{ source: string; labels: Record }> => { - const response = await apiClient.get('/labels', { params: { source } }) + filters?: { + operator?: string[] + operation?: string[] + label?: string[] + }, + ): Promise => { + const response = await apiClient.get('/labels', { + params: { source, ...filters }, + paramsSerializer: { + indexes: null, // serialize arrays as ?key=val1&key=val2 + }, + }) return response.data }, } diff --git a/frontend/src/types/index.ts b/frontend/src/types/index.ts index 19dd9182c1..943bcff79b 100644 --- a/frontend/src/types/index.ts +++ b/frontend/src/types/index.ts @@ -306,6 +306,8 @@ export interface AttackSummary { last_message_preview?: string | null message_count: number related_conversation_ids: string[] + operator?: string | null + operation?: string | null labels: Record created_at: string updated_at: string @@ -314,6 +316,8 @@ export interface AttackSummary { export interface CreateAttackRequest { target_registry_name: string name?: string + operator?: string + operation?: string labels?: Record source_conversation_id?: string cutoff_index?: number @@ -401,7 +405,13 @@ export interface AddMessageRequest { target_registry_name?: string converter_ids?: string[] target_conversation_id: string - labels?: Record +} + +export interface LabelOptionsResponse { + source: string + operators?: string[] + operations?: string[] + labels: Record } export interface AddMessageResponse { diff --git a/pyrit/backend/mappers/attack_mappers.py b/pyrit/backend/mappers/attack_mappers.py index 9318eb9e89..c61818deda 100644 --- a/pyrit/backend/mappers/attack_mappers.py +++ b/pyrit/backend/mappers/attack_mappers.py @@ -226,6 +226,8 @@ async def attack_result_to_summary_async( """ labels = dict(ar.labels) if ar.labels else {} labels.update(stats.labels or {}) + labels.pop("operator", None) + labels.pop("operation", None) created_at, updated_at = _resolve_summary_timestamps(ar) data = {name: getattr(ar, name) for name in AttackResult.model_fields} diff --git a/pyrit/backend/models/attacks.py b/pyrit/backend/models/attacks.py index 11924fcbfb..ef1cfd5f5e 100644 --- a/pyrit/backend/models/attacks.py +++ b/pyrit/backend/models/attacks.py @@ -25,6 +25,7 @@ PromptDataType, Score, ) +from pyrit.models.results.attack_result import normalize_legacy_attack_attribution class TargetInfo(BaseModel): @@ -358,12 +359,47 @@ class PrependedMessageRequest(BaseModel): pieces: list[MessagePieceRequest] = Field(..., description="Message pieces (supports multimodal)", max_length=50) +class _AttackAttributionInput(BaseModel): + """Shared first-class attribution input with temporary legacy label aliases.""" + + operator: str | None = Field(None, max_length=128, description="Operator responsible for the attack") + operation: str | None = Field(None, max_length=128, description="Operation associated with the attack") + labels: dict[str, str] | None = Field(None, description="Arbitrary user-defined labels for filtering") + + @model_validator(mode="before") + @classmethod + def _normalize_legacy_attribution_labels(cls, data: Any) -> Any: + """ + Normalize deprecated label aliases without mutating the caller's dictionaries. + + TODO(PyRIT 1.4): Remove this validator with legacy attribution label aliases. + + Returns: + The normalized model input. + + Raises: + ValueError: If an alias is not a string or conflicts with a dedicated field. + """ + if not isinstance(data, dict) or not isinstance(data.get("labels"), dict): + return data + normalized = dict(data) + remaining, operator, operation = normalize_legacy_attack_attribution( + labels=normalized["labels"], + operator=normalized.get("operator"), + operation=normalized.get("operation"), + ) + normalized["labels"] = remaining + normalized["operator"] = operator + normalized["operation"] = operation + return normalized + + # ============================================================================ # Create Attack # ============================================================================ -class CreateAttackRequest(BaseModel): +class CreateAttackRequest(_AttackAttributionInput): """ Request to create a new attack. @@ -388,7 +424,6 @@ class CreateAttackRequest(BaseModel): prepended_conversation: list[PrependedMessageRequest] | None = Field( None, description="Messages to prepend (system prompts, branching context)", max_length=200 ) - labels: dict[str, str] | None = Field(None, description="User-defined labels for filtering") class CreateAttackResponse(BaseModel): @@ -531,11 +566,6 @@ class AddMessageRequest(BaseModel): description="The conversation_id to store and send messages under. " "Usually the attack's main conversation, but can be a related conversation.", ) - labels: dict[str, str] | None = Field( - None, - description="Request labels used for attack-level consistency checks. " - "When present, the operator must match the attack result's operator.", - ) @model_validator(mode="after") def _validate_converter_configurations(self) -> "AddMessageRequest": diff --git a/pyrit/backend/routes/attacks.py b/pyrit/backend/routes/attacks.py index 6f4cad1557..2c526d94af 100644 --- a/pyrit/backend/routes/attacks.py +++ b/pyrit/backend/routes/attacks.py @@ -9,9 +9,10 @@ """ import logging -from typing import Literal +from typing import Annotated, Literal from fastapi import APIRouter, HTTPException, Query, status +from pydantic import Field from pyrit.backend.models.attacks import ( AddMessageRequest, @@ -33,6 +34,7 @@ from pyrit.backend.models.common import ProblemDetail from pyrit.backend.routes.common import parse_label_query_params from pyrit.backend.services.attack_service import get_attack_service +from pyrit.common.deprecation import print_deprecation_message logger = logging.getLogger(__name__) @@ -75,6 +77,12 @@ async def list_attacks( # pyrit-async-suffix-exempt outcome: Literal["undetermined", "success", "failure", "error"] | None = Query( None, description="Filter by outcome" ), + operator: list[Annotated[str, Field(max_length=128)]] | None = Query( + None, description="Filter by dedicated operator values" + ), + operation: list[Annotated[str, Field(max_length=128)]] | None = Query( + None, description="Filter by dedicated operation values" + ), label: list[str] | None = Query( None, description="Filter by labels (format: key:value). May be specified multiple times; " @@ -100,8 +108,28 @@ async def list_attacks( # pyrit-async-suffix-exempt Returns: AttackListResponse: Paginated list of attack summaries. """ - service = get_attack_service() - labels = parse_label_query_params(label) + labels = parse_label_query_params(label) or {} + # TODO(PyRIT 1.4): Remove legacy attribution aliases from label query parameters. + legacy_operator = labels.pop("operator", None) + legacy_operation = labels.pop("operation", None) + if legacy_operator is not None: + print_deprecation_message( + old_item="GET /attacks?label=operator:...", + new_item="GET /attacks?operator=...", + removed_in="1.4.0", + ) + if operator is not None and operator != legacy_operator: + raise HTTPException(status_code=422, detail="operator conflicts with legacy label=operator filter") + operator = legacy_operator + if legacy_operation is not None: + print_deprecation_message( + old_item="GET /attacks?label=operation:...", + new_item="GET /attacks?operation=...", + removed_in="1.4.0", + ) + if operation is not None and operation != legacy_operation: + raise HTTPException(status_code=422, detail="operation conflicts with legacy label=operation filter") + operation = legacy_operation # Strip empty strings from the list-valued query params. The service layer # coerces an all-empty ``converter_types`` list to None ("no filter"); the # "attacks with no converters" case is expressed through ``has_converters``. @@ -109,6 +137,7 @@ async def list_attacks( # pyrit-async-suffix-exempt converter_types = [c for c in converter_types if c] if attack_types is not None: attack_types = [a for a in attack_types if a] + service = get_attack_service() return await service.list_attacks_async( attack_types=attack_types, converter_types=converter_types, @@ -116,7 +145,9 @@ async def list_attacks( # pyrit-async-suffix-exempt has_converters=has_converters, include_scenario_attacks=include_scenario_attacks, outcome=outcome, - labels=labels, + operator=operator, + operation=operation, + labels=labels or None, min_turns=min_turns, max_turns=max_turns, limit=limit, diff --git a/pyrit/backend/routes/labels.py b/pyrit/backend/routes/labels.py index 167e87a631..520192e9f4 100644 --- a/pyrit/backend/routes/labels.py +++ b/pyrit/backend/routes/labels.py @@ -7,12 +7,13 @@ Provides access to unique label values for filtering in the GUI. """ -from typing import Literal +from typing import Annotated, Literal from fastapi import APIRouter, Query from pydantic import BaseModel, Field from starlette.concurrency import run_in_threadpool +from pyrit.backend.routes.common import parse_label_query_params from pyrit.memory import CentralMemory router = APIRouter(prefix="/labels", tags=["labels"]) @@ -23,17 +24,29 @@ class LabelOptionsResponse(BaseModel): source: str = Field(..., description="Source type (e.g., 'attacks')") labels: dict[str, list[str]] = Field(..., description="Map of label keys to their unique values") + operators: list[str] | None = Field(None, description="Unique attack operators") + operations: list[str] | None = Field(None, description="Unique attack operations") @router.get( "", response_model=LabelOptionsResponse, + response_model_exclude_none=True, ) async def get_label_options( # pyrit-async-suffix-exempt source: Literal["attacks", "scenarios"] = Query( "attacks", description="Source type to get labels from.", ), + operator: list[Annotated[str, Field(max_length=128)]] | None = Query( + None, + description="Narrow attack labels by operator.", + ), + operation: list[Annotated[str, Field(max_length=128)]] | None = Query( + None, + description="Narrow attack labels by operation.", + ), + label: list[str] | None = Query(None, description="Narrow attack labels by key:value filters."), ) -> LabelOptionsResponse: """ Get unique label keys and values for filtering. @@ -43,13 +56,29 @@ async def get_label_options( # pyrit-async-suffix-exempt Args: source: The source type to query labels from. + operator: Operator values used to narrow attack rows. + operation: Operation values used to narrow attack rows. + label: Arbitrary key:value filters used to narrow attack rows. Returns: LabelOptionsResponse: Map of label keys to their unique values. """ memory = CentralMemory.get_memory_instance() - label_loader = memory.get_unique_attack_labels if source == "attacks" else memory.get_unique_scenario_labels - labels = await run_in_threadpool(label_loader) - + if source == "attacks": + label_filters = parse_label_query_params(label) + labels = await run_in_threadpool( + memory.get_unique_attack_labels, + operator=operator, + operation=operation, + labels=label_filters, + ) + attribution = ( + {} + if operator or operation or label_filters + else await run_in_threadpool(memory.get_unique_attack_attribution) + ) + return LabelOptionsResponse(source=source, labels=labels, **attribution) + + labels = await run_in_threadpool(memory.get_unique_scenario_labels) return LabelOptionsResponse(source=source, labels=labels) diff --git a/pyrit/backend/services/attack_service.py b/pyrit/backend/services/attack_service.py index ef5bb8043b..84dee78b8e 100644 --- a/pyrit/backend/services/attack_service.py +++ b/pyrit/backend/services/attack_service.py @@ -106,6 +106,8 @@ async def list_attacks_async( include_scenario_attacks: bool = True, outcome: Literal["undetermined", "success", "failure", "error"] | None = None, labels: Mapping[str, str | Sequence[str]] | None = None, + operator: Sequence[str] | None = None, + operation: Sequence[str] | None = None, min_turns: int | None = None, max_turns: int | None = None, limit: int = 20, @@ -134,6 +136,8 @@ async def list_attacks_async( include_scenario_attacks: Whether to include attacks created as part of scenario runs. Defaults to ``True`` for API compatibility. outcome: Filter by attack outcome. + operator: Filter by dedicated operator values. + operation: Filter by dedicated operation values. labels: Filter by labels. See ``MemoryInterface.get_attack_results`` for semantics (AND across label names; string equality or sequence OR within each name). @@ -162,19 +166,22 @@ async def list_attacks_async( # past the anchor, and limits in SQL, so only one page's worth of rows is materialized # instead of the full table. normalized_labels = normalize_label_filters(labels=labels) - filter_fingerprint = fingerprint_filters( - filters={ - "attack_types": effective_attack_types, - "converter_types": effective_converter_types, - "converter_types_match": converter_types_match, - "has_converters": has_converters, - "include_scenario_attacks": include_scenario_attacks, - "outcome": outcome, - "labels": normalized_labels, - "min_turns": min_turns, - "max_turns": max_turns, - } - ) + fingerprint_values: dict[str, Any] = { + "attack_types": effective_attack_types, + "converter_types": effective_converter_types, + "converter_types_match": converter_types_match, + "has_converters": has_converters, + "include_scenario_attacks": include_scenario_attacks, + "outcome": outcome, + "labels": normalized_labels, + "min_turns": min_turns, + "max_turns": max_turns, + } + if operator is not None: + fingerprint_values["operator"] = operator + if operation is not None: + fingerprint_values["operation"] = operation + filter_fingerprint = fingerprint_filters(filters=fingerprint_values) decoded_cursor = decode_keyset_cursor(cursor=cursor, fingerprint=filter_fingerprint) after = ( AttackResultKeysetCursor( @@ -186,6 +193,8 @@ async def list_attacks_async( ) results = self._memory.get_attack_results( outcome=outcome, + operator=operator, + operation=operation, labels=normalized_labels, attack_classes=effective_attack_types, converter_classes=effective_converter_types, @@ -392,6 +401,8 @@ async def create_attack_async(self, *, request: CreateAttackRequest) -> CreateAt "created_at": now.isoformat(), "target_registry_name": request.target_registry_name, }, + operator=request.operator, + operation=request.operation, labels=labels, ) @@ -649,7 +660,6 @@ async def add_message_async(self, *, attack_result_id: str, request: AddMessageR main_conversation_id = ar.conversation_id self._validate_target_match(attack_identifier=ar.get_attack_strategy_identifier(), request=request) - self._validate_operator_match(attack_result=ar, request=request) msg_conversation_id = request.target_conversation_id @@ -763,28 +773,6 @@ def _validate_target_match( f"Create a new attack to use a different target." ) - def _validate_operator_match(self, *, attack_result: AttackResult, request: AddMessageRequest) -> None: - """ - Validate that the request operator matches the attack result's operator. - - Raises: - ValueError: If the operator in the request doesn't match the attack result. - """ - if not request.labels: - return - - attack_operator = attack_result.labels.get("operator") - if not attack_operator: - return - - request_operator = request.labels.get("operator") - if request_operator and request_operator != attack_operator: - raise ValueError( - f"Operator mismatch: attack belongs to operator '{attack_operator}' " - f"but request is from '{request_operator}'. " - f"Create a new attack to continue." - ) - async def _update_attack_after_message_async( self, *, diff --git a/pyrit/memory/alembic/versions/1b3d5f7a9c2e_persist_scored_expectation.py b/pyrit/memory/alembic/versions/1b3d5f7a9c2e_persist_scored_expectation.py index bb80233ac0..3cb40abf25 100644 --- a/pyrit/memory/alembic/versions/1b3d5f7a9c2e_persist_scored_expectation.py +++ b/pyrit/memory/alembic/versions/1b3d5f7a9c2e_persist_scored_expectation.py @@ -38,16 +38,64 @@ #: Rows per page so a large score table migrates in bounded keyset batches, not one statement. _BACKFILL_BATCH_SIZE = 500 +_MSSQL_BACKFILL_SCORED_EXPECTATION_QUERY = """ +UPDATE score_entry +SET [scored_expectation] = JSON_MODIFY( + N'{"schema_version":1,"objective":null,"conditions":[]}', + N'$.objective', + score_entry.[objective] +) +FROM [ScoreEntries] AS score_entry +WHERE score_entry.[objective] IS NOT NULL + AND score_entry.[scored_expectation] IS NULL +""" +_MSSQL_RESTORE_OBJECTIVE_QUERY = """ +UPDATE score_entry +SET [objective] = objective_attribute.[value] +FROM [ScoreEntries] AS score_entry +CROSS APPLY ( + SELECT TOP (1) attribute.[value] + FROM OPENJSON( + CASE + WHEN ISJSON(score_entry.[scored_expectation]) = 1 + THEN score_entry.[scored_expectation] + ELSE N'{}' + END + ) AS attribute + WHERE attribute.[key] COLLATE Latin1_General_100_BIN2 = N'objective' + AND attribute.[type] = 1 +) AS objective_attribute +WHERE score_entry.[scored_expectation] IS NOT NULL + AND score_entry.[objective] IS NULL +""" + + +def _report_progress(message: str) -> None: + """Write migration progress to Alembic stdout, or the logger outside a migration context.""" + try: + context = op.get_context() + except (AttributeError, NameError): + logger.info(message) + return + config = context.config + if config is not None: + config.print_stdout(message) + else: + logger.info(message) + def upgrade() -> None: """Add ``scored_expectation``, fold the legacy objective into it, then drop ``objective``.""" + _report_progress("Scored expectation migration: adding scored_expectation column.") with op.batch_alter_table("ScoreEntries") as batch_op: batch_op.add_column(sa.Column("scored_expectation", sa.JSON(), nullable=True)) _backfill_scored_expectation() + _report_progress("Scored expectation migration: dropping legacy objective column.") with op.batch_alter_table("ScoreEntries") as batch_op: batch_op.drop_column("objective") + _report_progress("Scored expectation migration: upgrade completed.") def downgrade() -> None: @@ -57,13 +105,16 @@ def downgrade() -> None: This is lossy by design: only the expectation's objective survives. Typed conditions have no column in the old schema and are dropped. """ + _report_progress("Scored expectation migration: restoring legacy objective column.") with op.batch_alter_table("ScoreEntries") as batch_op: batch_op.add_column(sa.Column("objective", sa.String(), nullable=True)) _backfill_objective() + _report_progress("Scored expectation migration: dropping scored_expectation column.") with op.batch_alter_table("ScoreEntries") as batch_op: batch_op.drop_column("scored_expectation") + _report_progress("Scored expectation migration: downgrade completed.") def _backfill_scored_expectation() -> None: @@ -74,6 +125,15 @@ def _backfill_scored_expectation() -> None: into memory at once. Scores with no objective keep a NULL expectation. """ connection = op.get_bind() + if connection.dialect.name == "mssql": + _report_progress("Scored expectation backfill: applying set-based SQL Server update.") + result = connection.exec_driver_sql(_MSSQL_BACKFILL_SCORED_EXPECTATION_QUERY) + if isinstance(result.rowcount, int) and result.rowcount >= 0: + _report_progress(f"Scored expectation backfill: updated {result.rowcount} row(s).") + else: + _report_progress("Scored expectation backfill: SQL Server update completed.") + return + score_entries = sa.table( "ScoreEntries", sa.column("id"), @@ -83,6 +143,9 @@ def _backfill_scored_expectation() -> None: statement = sa.text('UPDATE "ScoreEntries" SET scored_expectation = :scored_expectation WHERE id = :score_id') last_id = None + batch_number = 0 + updated_count = 0 + _report_progress(f"Scored expectation backfill: processing rows in batches of {_BACKFILL_BATCH_SIZE}.") while True: conditions = [ score_entries.c.objective.isnot(None), @@ -97,8 +160,10 @@ def _backfill_scored_expectation() -> None: .limit(_BACKFILL_BATCH_SIZE) ).fetchall() if not rows: + _report_progress(f"Scored expectation backfill: updated {updated_count} row(s).") return last_id = rows[-1][0] + batch_number += 1 updates = [ { @@ -110,6 +175,10 @@ def _backfill_scored_expectation() -> None: for score_id, objective in rows ] connection.execute(statement, updates) + updated_count += len(updates) + _report_progress( + f"Scored expectation backfill: completed batch {batch_number}; updated {updated_count} row(s)." + ) def _backfill_objective() -> None: @@ -120,6 +189,15 @@ def _backfill_objective() -> None: Rows are read a page at a time, keyed on ``id``. """ connection = op.get_bind() + if connection.dialect.name == "mssql": + _report_progress("Objective restore: applying set-based SQL Server update.") + result = connection.exec_driver_sql(_MSSQL_RESTORE_OBJECTIVE_QUERY) + if isinstance(result.rowcount, int) and result.rowcount >= 0: + _report_progress(f"Objective restore: updated {result.rowcount} row(s).") + else: + _report_progress("Objective restore: SQL Server update completed.") + return + score_entries = sa.table( "ScoreEntries", sa.column("id"), @@ -129,6 +207,10 @@ def _backfill_objective() -> None: statement = sa.text('UPDATE "ScoreEntries" SET objective = :objective WHERE id = :score_id') last_id = None + batch_number = 0 + processed_count = 0 + updated_count = 0 + _report_progress(f"Objective restore: processing rows in batches of {_BACKFILL_BATCH_SIZE}.") while True: conditions = [ score_entries.c.scored_expectation.isnot(None), @@ -143,8 +225,11 @@ def _backfill_objective() -> None: .limit(_BACKFILL_BATCH_SIZE) ).fetchall() if not rows: + _report_progress(f"Objective restore: processed {processed_count} row(s); updated {updated_count} row(s).") return last_id = rows[-1][0] + batch_number += 1 + processed_count += len(rows) updates = [] for score_id, scored_expectation in rows: @@ -154,6 +239,11 @@ def _backfill_objective() -> None: updates.append({"score_id": score_id, "objective": objective}) if updates: connection.execute(statement, updates) + updated_count += len(updates) + _report_progress( + f"Objective restore: completed batch {batch_number}; processed {processed_count} row(s), " + f"updated {updated_count} row(s)." + ) def _extract_objective(scored_expectation: object) -> str | None: diff --git a/pyrit/memory/alembic/versions/a4c6e8f0b2d1_add_attack_attribution_and_history_indexes.py b/pyrit/memory/alembic/versions/a4c6e8f0b2d1_add_attack_attribution_and_history_indexes.py new file mode 100644 index 0000000000..8b0ae4c282 --- /dev/null +++ b/pyrit/memory/alembic/versions/a4c6e8f0b2d1_add_attack_attribution_and_history_indexes.py @@ -0,0 +1,432 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Add first-class attack attribution fields and history query indexes. + +Revision ID: a4c6e8f0b2d1 +Revises: 1b3d5f7a9c2e +Create Date: 2026-09-04 18:48:00.000000 +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any + +import sqlalchemy as sa +from alembic import op + +from pyrit.memory.memory_models import CustomUUID + +if TYPE_CHECKING: + from collections.abc import Sequence + +revision: str = "a4c6e8f0b2d1" +down_revision: str | Sequence[str] | None = "1b3d5f7a9c2e" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + +logger = logging.getLogger(__name__) + +_ATTRIBUTION_FIELDS = ("operator", "operation") +_ATTRIBUTION_MAX_LENGTH = 128 +_BATCH_SIZE = 1000 +_MSSQL_INVALID_ATTRIBUTION_QUERY = f""" +SELECT TOP (1) + attack_result.[id] AS row_id, + attribute.[key] AS field_name, + attribute.[type] AS value_type +FROM [AttackResultEntries] AS attack_result +CROSS APPLY OPENJSON(attack_result.[labels]) AS attribute +WHERE attribute.[key] COLLATE Latin1_General_100_BIN2 IN (N'operator', N'operation') + AND ( + attribute.[type] <> 1 + OR DATALENGTH(attribute.[value]) > {_ATTRIBUTION_MAX_LENGTH * 2} + ) +""" +_MSSQL_MOVE_ATTRIBUTION_QUERY = """ +UPDATE attack_result +SET + [operator] = JSON_VALUE(attack_result.[labels], N'$.operator'), + [operation] = JSON_VALUE(attack_result.[labels], N'$.operation'), + [labels] = JSON_MODIFY( + JSON_MODIFY(attack_result.[labels], N'$.operator', NULL), + N'$.operation', + NULL + ) +FROM [AttackResultEntries] AS attack_result +WHERE EXISTS ( + SELECT 1 + FROM OPENJSON(attack_result.[labels]) AS attribute + WHERE attribute.[key] COLLATE Latin1_General_100_BIN2 IN (N'operator', N'operation') +) +""" + + +def _report_progress(message: str) -> None: + """Write migration progress to Alembic stdout, or the logger outside a migration context.""" + try: + context = op.get_context() + except (AttributeError, NameError): + logger.info(message) + return + config = context.config + if config is not None: + config.print_stdout(message) + else: + logger.info(message) + + +def upgrade() -> None: + """Add attribution columns, migrate legacy labels, and replace history indexes.""" + _report_progress("Attack history migration: adding attribution columns.") + op.add_column("AttackResultEntries", sa.Column("operator", sa.Unicode(_ATTRIBUTION_MAX_LENGTH), nullable=True)) + op.add_column("AttackResultEntries", sa.Column("operation", sa.Unicode(_ATTRIBUTION_MAX_LENGTH), nullable=True)) + + _report_progress("Attack history migration: moving attribution values from labels.") + _move_attribution_from_labels() + + _report_progress("Attack history migration: validating and bounding indexed text columns.") + _bound_indexed_text_columns() + + _report_progress("Attack history migration: replacing AttackResultEntries indexes.") + op.drop_index("ix_AttackResultEntries_conversation_id", table_name="AttackResultEntries") + _report_progress("Attack history migration: creating ix_AttackResultEntries_conversation_timestamp_id.") + op.create_index( + "ix_AttackResultEntries_conversation_timestamp_id", + "AttackResultEntries", + ["conversation_id", "timestamp", "id"], + ) + _report_progress("Attack history migration: creating ix_AttackResultEntries_operator_timestamp_id.") + op.create_index( + "ix_AttackResultEntries_operator_timestamp_id", + "AttackResultEntries", + ["operator", "timestamp", "id"], + mssql_include=["conversation_id"], + ) + _report_progress("Attack history migration: creating ix_AttackResultEntries_operation_timestamp_id.") + op.create_index( + "ix_AttackResultEntries_operation_timestamp_id", + "AttackResultEntries", + ["operation", "timestamp", "id"], + mssql_include=["conversation_id"], + ) + + _report_progress("Attack history migration: replacing PromptMemoryEntries indexes.") + _drop_index_if_exists(name="idx_conversation_id", table_name="PromptMemoryEntries") + _report_progress("Attack history migration: creating ix_PromptMemoryEntries_conversation_sequence_id.") + op.create_index( + "ix_PromptMemoryEntries_conversation_sequence_id", + "PromptMemoryEntries", + ["conversation_id", "sequence", "id"], + mssql_include=["timestamp", "converted_value_data_type"], + ) + + _report_progress("Attack history migration: creating ScenarioResultEntries indexes.") + _report_progress("Attack history migration: creating ix_ScenarioResultEntries_scenario_name_timestamp_id.") + op.create_index( + "ix_ScenarioResultEntries_scenario_name_timestamp_id", + "ScenarioResultEntries", + ["scenario_name", "timestamp", "id"], + ) + _report_progress("Attack history migration: creating ix_ScenarioResultEntries_scenario_run_state_timestamp_id.") + op.create_index( + "ix_ScenarioResultEntries_scenario_run_state_timestamp_id", + "ScenarioResultEntries", + ["scenario_run_state", "timestamp", "id"], + ) + _report_progress("Attack history migration: upgrade completed.") + + +def downgrade() -> None: + """Restore legacy labels and indexes, then remove attribution columns.""" + _report_progress("Attack history migration: restoring attribution values to labels.") + _restore_attribution_to_labels() + + _report_progress("Attack history migration: restoring legacy indexes and text columns.") + op.drop_index( + "ix_ScenarioResultEntries_scenario_run_state_timestamp_id", + table_name="ScenarioResultEntries", + ) + op.drop_index( + "ix_ScenarioResultEntries_scenario_name_timestamp_id", + table_name="ScenarioResultEntries", + ) + + op.drop_index( + "ix_PromptMemoryEntries_conversation_sequence_id", + table_name="PromptMemoryEntries", + ) + + op.drop_index( + "ix_AttackResultEntries_operation_timestamp_id", + table_name="AttackResultEntries", + ) + op.drop_index( + "ix_AttackResultEntries_operator_timestamp_id", + table_name="AttackResultEntries", + ) + op.drop_index( + "ix_AttackResultEntries_conversation_timestamp_id", + table_name="AttackResultEntries", + ) + op.create_index( + "ix_AttackResultEntries_conversation_id", + "AttackResultEntries", + ["conversation_id"], + ) + + _restore_unbounded_text_columns() + op.drop_column("AttackResultEntries", "operation") + op.drop_column("AttackResultEntries", "operator") + _report_progress("Attack history migration: downgrade completed.") + + +def _attack_results_table(*, include_attribution: bool) -> sa.Table: + """ + Build a typed table for portable JSON migration reads and writes. + + Returns: + The lightweight attack-results table. + """ + columns = [ + sa.Column("id", CustomUUID(), primary_key=True), + sa.Column("labels", sa.JSON(), nullable=True), + ] + if include_attribution: + columns.extend( + [ + sa.Column("operator", sa.Unicode(_ATTRIBUTION_MAX_LENGTH), nullable=True), + sa.Column("operation", sa.Unicode(_ATTRIBUTION_MAX_LENGTH), nullable=True), + ] + ) + return sa.Table("AttackResultEntries", sa.MetaData(), *columns) + + +def _drop_index_if_exists(*, name: str, table_name: str) -> None: + """Drop an index only when it exists in the source schema.""" + bind = op.get_bind() + existing_names = {index["name"] for index in sa.inspect(bind).get_indexes(table_name)} + if name in existing_names: + op.drop_index(name, table_name=table_name) + + +def _bound_indexed_text_columns() -> None: + """Bound existing text keys before creating indexes that SQL Server accepts.""" + _validate_column_length(table_name="PromptMemoryEntries", column_name="conversation_id", max_length=128) + _validate_column_length(table_name="ScenarioResultEntries", column_name="scenario_name", max_length=256) + _validate_column_length(table_name="ScenarioResultEntries", column_name="scenario_run_state", max_length=32) + with op.batch_alter_table("PromptMemoryEntries") as batch_op: + batch_op.alter_column( + "conversation_id", + existing_type=sa.String(), + type_=sa.String(128), + existing_nullable=False, + ) + with op.batch_alter_table("ScenarioResultEntries") as batch_op: + batch_op.alter_column( + "scenario_name", + existing_type=sa.String(), + type_=sa.String(256), + existing_nullable=False, + ) + batch_op.alter_column( + "scenario_run_state", + existing_type=sa.String(), + type_=sa.String(32), + existing_nullable=False, + ) + + +def _restore_unbounded_text_columns() -> None: + """Restore the pre-migration unbounded text column types.""" + with op.batch_alter_table("ScenarioResultEntries") as batch_op: + batch_op.alter_column( + "scenario_name", + existing_type=sa.String(256), + type_=sa.String(), + existing_nullable=False, + ) + batch_op.alter_column( + "scenario_run_state", + existing_type=sa.String(32), + type_=sa.String(), + existing_nullable=False, + ) + with op.batch_alter_table("PromptMemoryEntries") as batch_op: + batch_op.alter_column( + "conversation_id", + existing_type=sa.String(128), + type_=sa.String(), + existing_nullable=False, + ) + + +def _validate_column_length(*, table_name: str, column_name: str, max_length: int) -> None: + """ + Fail before a bounded type conversion could truncate existing data. + + Raises: + ValueError: If an existing value exceeds the new bound. + """ + table = sa.Table( + table_name, + sa.MetaData(), + sa.Column(column_name, sa.String(), nullable=False), + ) + oversized_value = ( + op.get_bind() + .execute(sa.select(table.c[column_name]).where(sa.func.length(table.c[column_name]) > max_length).limit(1)) + .scalar_one_or_none() + ) + if oversized_value is not None: + raise ValueError( + f"{table_name}.{column_name} contains a value longer than {max_length} characters; " + "migration will not truncate it." + ) + + +def _move_attribution_from_labels() -> None: + """ + Move exact legacy attribution label keys into bounded scalar columns. + + Raises: + ValueError: If a legacy attribution value is invalid or too long. + """ + bind = op.get_bind() + if bind.dialect.name == "mssql": + _move_attribution_from_labels_mssql(bind=bind) + return + + _move_attribution_from_labels_portable(bind=bind) + + +def _move_attribution_from_labels_mssql(*, bind: Any) -> None: + """ + Move attribution labels with set-based SQL Server JSON operations. + + Raises: + ValueError: If a legacy attribution value is invalid or too long. + """ + _report_progress("Attack attribution backfill: validating SQL Server JSON values.") + invalid_value = bind.exec_driver_sql(_MSSQL_INVALID_ATTRIBUTION_QUERY).first() + if invalid_value is not None: + invalid = invalid_value._mapping + if invalid["value_type"] != 1: + raise ValueError( + f"AttackResultEntries row {invalid['row_id']} has non-string labels.{invalid['field_name']}; " + "cannot migrate it to a first-class string column." + ) + raise ValueError( + f"AttackResultEntries row {invalid['row_id']} has labels.{invalid['field_name']} longer than " + f"{_ATTRIBUTION_MAX_LENGTH} characters; migration will not truncate it." + ) + + _report_progress("Attack attribution backfill: applying set-based SQL Server update.") + result = bind.exec_driver_sql(_MSSQL_MOVE_ATTRIBUTION_QUERY) + if isinstance(result.rowcount, int) and result.rowcount >= 0: + _report_progress(f"Attack attribution backfill: updated {result.rowcount} row(s).") + else: + _report_progress("Attack attribution backfill: SQL Server update completed.") + + +def _move_attribution_from_labels_portable(*, bind: Any) -> None: + """ + Move attribution labels using portable SQLAlchemy operations. + + Raises: + ValueError: If a legacy attribution value is invalid or too long. + """ + table = _attack_results_table(include_attribution=True) + statement = ( + sa.update(table) + .where(table.c.id == sa.bindparam("row_id")) + .values( + labels=sa.bindparam("new_labels"), + operator=sa.bindparam("new_operator"), + operation=sa.bindparam("new_operation"), + ) + ) + rows = bind.execute(sa.select(table.c.id, table.c.labels)).all() + batch_count = (len(rows) + _BATCH_SIZE - 1) // _BATCH_SIZE + _report_progress(f"Attack attribution backfill: processing {len(rows)} row(s) in {batch_count} batch(es).") + updated_count = 0 + for batch_number, start in enumerate(range(0, len(rows), _BATCH_SIZE), start=1): + updates = [] + for row in rows[start : start + _BATCH_SIZE]: + labels = row.labels + if not isinstance(labels, dict): + continue + remaining_labels = dict(labels) + values: dict[str, Any] = {} + for field_name in _ATTRIBUTION_FIELDS: + if field_name not in remaining_labels: + continue + value = remaining_labels.pop(field_name) + if not isinstance(value, str): + raise ValueError( + f"AttackResultEntries row {row.id} has non-string labels.{field_name}; " + "cannot migrate it to a first-class string column." + ) + if len(value) > _ATTRIBUTION_MAX_LENGTH: + raise ValueError( + f"AttackResultEntries row {row.id} has labels.{field_name} longer than " + f"{_ATTRIBUTION_MAX_LENGTH} characters; migration will not truncate it." + ) + values[field_name] = value + if values: + updates.append( + { + "row_id": row.id, + "new_labels": remaining_labels, + # Both columns were just added, so writing None leaves them NULL. + "new_operator": values.get("operator"), + "new_operation": values.get("operation"), + } + ) + if updates: + bind.execute(statement, updates) + updated_count += len(updates) + _report_progress(f"Attack attribution backfill: completed batch {batch_number}/{batch_count}.") + _report_progress(f"Attack attribution backfill: updated {updated_count} row(s).") + + +def _restore_attribution_to_labels() -> None: + """ + Restore populated attribution columns to exact legacy JSON label keys. + + Raises: + ValueError: If a legacy label conflicts with its dedicated value. + """ + bind = op.get_bind() + table = _attack_results_table(include_attribution=True) + statement = sa.update(table).where(table.c.id == sa.bindparam("row_id")).values(labels=sa.bindparam("new_labels")) + rows = bind.execute(sa.select(table.c.id, table.c.labels, table.c.operator, table.c.operation)).all() + batch_count = (len(rows) + _BATCH_SIZE - 1) // _BATCH_SIZE + _report_progress(f"Attack attribution restore: processing {len(rows)} row(s) in {batch_count} batch(es).") + updated_count = 0 + for batch_number, start in enumerate(range(0, len(rows), _BATCH_SIZE), start=1): + updates = [] + for row in rows[start : start + _BATCH_SIZE]: + labels = dict(row.labels) if isinstance(row.labels, dict) else {} + changed = False + for field_name in _ATTRIBUTION_FIELDS: + value = getattr(row, field_name) + if value is None: + continue + existing = labels.get(field_name) + if existing is not None and existing != value: + raise ValueError( + f"AttackResultEntries row {row.id} has conflicting labels.{field_name} " + f"while downgrading: {existing!r} != {value!r}." + ) + labels[field_name] = value + changed = True + if changed: + updates.append({"row_id": row.id, "new_labels": labels}) + if updates: + bind.execute(statement, updates) + updated_count += len(updates) + _report_progress(f"Attack attribution restore: completed batch {batch_number}/{batch_count}.") + _report_progress(f"Attack attribution restore: updated {updated_count} row(s).") diff --git a/pyrit/memory/azure_sql_memory.py b/pyrit/memory/azure_sql_memory.py index 884b3ef832..829ed41c51 100644 --- a/pyrit/memory/azure_sql_memory.py +++ b/pyrit/memory/azure_sql_memory.py @@ -564,24 +564,28 @@ def get_conversation_stats(self, *, conversation_ids: Sequence[str]) -> dict[str sql = text( f""" SELECT - pme.conversation_id, - COUNT(DISTINCT pme.sequence) AS msg_count, - ( - SELECT TOP 1 LEFT(p2.converted_value, {ConversationStats.PREVIEW_FETCH_MAX_LEN}) - FROM "PromptMemoryEntries" p2 - WHERE p2.conversation_id = pme.conversation_id - ORDER BY p2.sequence DESC, p2.id DESC - ) AS last_preview, - ( - SELECT TOP 1 p2b.converted_value_data_type - FROM "PromptMemoryEntries" p2b - WHERE p2b.conversation_id = pme.conversation_id - ORDER BY p2b.sequence DESC, p2b.id DESC - ) AS last_data_type, - MIN(pme.timestamp) AS created_at - FROM "PromptMemoryEntries" pme - WHERE pme.conversation_id IN ({placeholders}) - GROUP BY pme.conversation_id + aggregate_rows.conversation_id, + aggregate_rows.msg_count, + latest.last_preview, + latest.last_data_type, + aggregate_rows.created_at + FROM ( + SELECT + pme.conversation_id, + COUNT(DISTINCT pme.sequence) AS msg_count, + MIN(pme.timestamp) AS created_at + FROM "PromptMemoryEntries" pme + WHERE pme.conversation_id IN ({placeholders}) + GROUP BY pme.conversation_id + ) AS aggregate_rows + OUTER APPLY ( + SELECT TOP 1 + LEFT(p2.converted_value, {ConversationStats.PREVIEW_FETCH_MAX_LEN}) AS last_preview, + p2.converted_value_data_type AS last_data_type + FROM "PromptMemoryEntries" p2 + WHERE p2.conversation_id = aggregate_rows.conversation_id + ORDER BY p2.sequence DESC, p2.id DESC + ) AS latest """ ) diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index c8d3efee95..7bbafa3e77 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -13,18 +13,20 @@ from collections.abc import Collection, Iterator, Mapping, MutableSequence, Sequence from contextlib import closing from dataclasses import dataclass -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Literal, NamedTuple, TypeVar from urllib.parse import urlparse -from sqlalchemy import MetaData, and_, case, func, literal, not_, or_, select +from sqlalchemy import MetaData, and_, case, exists, func, literal, not_, or_, select from sqlalchemy.engine.base import Engine from sqlalchemy.exc import IntegrityError, SQLAlchemyError from sqlalchemy.orm import joinedload from sqlalchemy.orm.attributes import InstrumentedAttribute, flag_modified from sqlalchemy.orm.session import Session +from pyrit.common.deprecation import print_deprecation_message + if TYPE_CHECKING: from pyrit.memory.memory_embedding import MemoryEmbedding @@ -93,6 +95,7 @@ group_conversation_message_pieces_by_sequence, sort_message_pieces, ) +from pyrit.models.results.attack_result import ATTRIBUTION_FIELDS, ATTRIBUTION_VALUE_MAX_LENGTH if TYPE_CHECKING: from sqlalchemy.sql.elements import ColumnElement @@ -104,6 +107,24 @@ IdentifierModel = TypeVar("IdentifierModel", bound=ComponentIdentifier) +def _normalize_attribution_filter_values(*, field: str, raw: str | Sequence[str]) -> tuple[str, ...]: + """ + Validate and snapshot one dedicated attribution filter. + + Returns: + tuple[str, ...]: The validated immutable filter values. + + Raises: + ValueError: If any value is not a string or exceeds the column limit. + """ + values = (raw,) if isinstance(raw, str) else tuple(raw) + if any(not isinstance(value, str) for value in values): + raise ValueError(f"{field} values must be strings") + if any(len(value) > ATTRIBUTION_VALUE_MAX_LENGTH for value in values): + raise ValueError(f"{field} values must be at most {ATTRIBUTION_VALUE_MAX_LENGTH} characters") + return values + + @dataclass(frozen=True, slots=True, kw_only=True) class _PreparedScorableContent: """A loose-content value prepared for durable database persistence.""" @@ -241,6 +262,8 @@ class _AttackResultQuery: has_converters: bool | None = None include_scenario_attacks: bool = True labels: Mapping[str, str | Sequence[str]] | None = None + operator: str | Sequence[str] | None = None + operation: str | Sequence[str] | None = None targeted_harm_categories: Sequence[str] | None = None identifier_filters: Sequence[IdentifierFilter] | None = None scenario_result_id: str | None = None @@ -250,15 +273,47 @@ class _AttackResultQuery: after: AttackResultKeysetCursor | None = None def __post_init__(self) -> None: - """Snapshot mutable sequence and mapping inputs.""" + """ + Snapshot mutable inputs and normalize legacy attribution aliases. + + TODO(PyRIT 1.4): Remove attribution handling in ``labels``. + + Raises: + ValueError: If attribution aliases conflict or exceed their maximum length. + """ for field_name in self._SEQUENCE_FIELDS: value = getattr(self, field_name) if value is not None: object.__setattr__(self, field_name, tuple(value)) + for field_name in ATTRIBUTION_FIELDS: + values = getattr(self, field_name) + if values is not None: + object.__setattr__( + self, + field_name, + _normalize_attribution_filter_values(field=field_name, raw=values), + ) + if self.labels is not None: labels = {key: value if isinstance(value, str) else tuple(value) for key, value in self.labels.items()} - object.__setattr__(self, "labels", MappingProxyType(labels)) + for field_name in ATTRIBUTION_FIELDS: + if field_name not in labels: + continue + legacy_values = _normalize_attribution_filter_values( + field=f"labels.{field_name}", + raw=labels.pop(field_name), + ) + dedicated_values = getattr(self, field_name) + if dedicated_values is not None and set(dedicated_values) != set(legacy_values): + raise ValueError(f"{field_name} conflicts with legacy labels.{field_name}") + print_deprecation_message( + old_item=f"_AttackResultQuery.labels['{field_name}']", + new_item=f"_AttackResultQuery.{field_name}", + removed_in="1.4.0", + ) + object.__setattr__(self, field_name, legacy_values) + object.__setattr__(self, "labels", MappingProxyType(labels) if labels else None) class MemoryInterface(abc.ABC): @@ -650,6 +705,14 @@ def _add_message_pieces_to_memory(self, *, message_pieces: Sequence[MessagePiece SQLAlchemyError: If the message pieces or converter identifiers cannot be persisted. """ entries = [PromptMemoryEntry(entry=piece) for piece in message_pieces] + # Sequence orders messages, so timestamp preserves the input order of pieces within one message. + latest_timestamp_by_message: dict[tuple[str, int], datetime] = {} + for entry in entries: + message_key = (entry.conversation_id, entry.sequence) + latest_timestamp = latest_timestamp_by_message.get(message_key) + if latest_timestamp is not None and entry.timestamp <= latest_timestamp: + entry.timestamp = latest_timestamp + timedelta(microseconds=1) + latest_timestamp_by_message[message_key] = entry.timestamp with closing(self.get_session()) as session: try: for piece, entry in zip(message_pieces, entries, strict=True): @@ -3470,6 +3533,8 @@ def get_attack_results( has_converters: bool | None = None, include_scenario_attacks: bool = True, labels: Mapping[str, str | Sequence[str]] | None = None, + operator: str | Sequence[str] | None = None, + operation: str | Sequence[str] | None = None, targeted_harm_categories: Sequence[str] | None = None, identifier_filters: Sequence[IdentifierFilter] | None = None, scenario_result_id: str | None = None, @@ -3513,13 +3578,14 @@ def get_attack_results( include_scenario_attacks (bool, optional): Whether to include attacks created as part of scenario runs. Defaults to ``True``. labels (Mapping[str, str | Sequence[str]] | None, optional): Filter results - by attack labels. Entries are AND-combined across label names; within a + by arbitrary attack labels. The legacy ``operator`` and ``operation`` aliases + are accepted through PyRIT 1.3 and normalized to dedicated filters. Entries + are AND-combined across label names; within a single entry, a string value is an equality match and a sequence value is an OR match over the listed values. An empty sequence applies no filter - for that label. Example: ``{"operator": "roakey", "operation": - ["roakey_op_a", "roakey_op_b"]}`` matches attacks where ``operator == - "roakey"`` AND (``operation == "roakey_op_a"`` OR ``operation == - "roakey_op_b"``). Defaults to None. + for that label. Defaults to None. + operator (str | Sequence[str] | None, optional): Filter by dedicated operator values. + operation (str | Sequence[str] | None, optional): Filter by dedicated operation values. targeted_harm_categories (Sequence[str] | None, optional): Filter results by the harm categories targeted by the attack (stored on ``AttackResultEntry.targeted_harm_categories``, auto-populated from the @@ -3542,7 +3608,7 @@ def get_attack_results( deduplication, mirroring ``min_turns``. Defaults to None. limit (int | None, optional): Maximum number of deduplicated attack results to return, ordered by recency. When either ``limit`` or ``after`` is provided, - deduplication and pagination happen in the database (via ``ROW_NUMBER()``) + deduplication and pagination happen in the database (via a ``NOT EXISTS`` anti-join) instead of loading every row into memory. Defaults to None (return all). after (AttackResultKeysetCursor | None, optional): Keyset (seek) anchor from a previous page. When provided, only results ordered strictly after the anchor @@ -3571,6 +3637,8 @@ def get_attack_results( has_converters=has_converters, include_scenario_attacks=include_scenario_attacks, labels=labels, + operator=operator, + operation=operation, targeted_harm_categories=targeted_harm_categories, identifier_filters=identifier_filters, scenario_result_id=scenario_result_id, @@ -3665,6 +3733,10 @@ def _build_attack_result_scalar_conditions(*, query: _AttackResultQuery) -> list conditions.append(AttackResultEntry.objective.contains(query.objective)) if query.outcome: conditions.append(AttackResultEntry.outcome == query.outcome) + if query.operator: + conditions.append(AttackResultEntry.operator.in_(query.operator)) + if query.operation: + conditions.append(AttackResultEntry.operation.in_(query.operation)) if query.scenario_result_id: conditions.append(AttackResultEntry.attribution_parent_id == uuid.UUID(query.scenario_result_id)) elif not query.include_scenario_attacks: @@ -3837,17 +3909,29 @@ def _query_paginated_attack_results( """ Deduplicate in SQL (filter-aware) and return one recency-ordered page of results. - Ranks rows with ``ROW_NUMBER() OVER (PARTITION BY conversation_id ORDER BY timestamp - DESC, id DESC)`` after applying ``conditions``, keeps only the newest row per - conversation (``rn == 1``) — reproducing the post-fetch Python dedup but *before* - pagination so page sizes stay correct — then applies the ``min_turns``/``max_turns`` - bounds to those winners, orders by recency, seeks past the ``after`` keyset anchor, - and applies ``limit`` in the database. The turn bounds are applied to the winners (not - inside the ranking subquery) so they never resurrect an older duplicate that happens - to fall in range. Seeking on the recency ordering tuple (rather than a numeric offset) - keeps page boundaries stable when other rows are inserted or deleted between page loads + Keeps only the newest row per ``conversation_id`` with a correlated ``NOT EXISTS`` + anti-join: a row survives when no other row that passes the same ``conditions`` + shares its conversation and sorts later on ``(timestamp, id)``. This reproduces the + post-fetch Python dedup but *before* pagination so page sizes stay correct. The + ``min_turns``/``max_turns`` bounds are applied to the surviving winners (not inside + the anti-join) so they never resurrect an older duplicate that happens to fall in + range. Seeking on the recency ordering tuple (rather than a numeric offset) keeps + page boundaries stable when other rows are inserted or deleted between page loads (offset pagination instead shifts every row after the change). + The anti-join replaces a ``ROW_NUMBER() OVER (PARTITION BY conversation_id ...)`` + window that had to rank *every* matching row on every page before a single result + could be returned, which made each page cost O(table). ``NOT EXISTS`` lets the + planner drive from the recency index, seek past the keyset anchor, probe + ``ix_AttackResultEntries_conversation_timestamp_id`` per candidate row, and stop + once ``limit`` winners are found. + + ``conditions`` are re-applied inside the anti-join through a derived table rather + than remapped onto an alias: the converter and harm-category filters are raw + ``text()`` fragments that hard-code ``"AttackResultEntries"``, so alias adaption + would silently leave them bound to the outer row and let a newer non-matching row + suppress a valid winner. + Args: conditions (list[Any]): Scalar WHERE filters applied before deduplication. min_turns (int | None): Inclusive lower bound on ``executed_turns`` for winners. @@ -3859,22 +3943,8 @@ def _query_paginated_attack_results( Returns: list[AttackResult]: The deduplicated, recency-ordered page of attack results. """ - ranked = select( - AttackResultEntry.id.label("id"), - func.row_number() - .over( - partition_by=AttackResultEntry.conversation_id, - order_by=(AttackResultEntry.timestamp.desc(), AttackResultEntry.id.desc()), - ) - .label("rn"), - ) - if conditions: - ranked = ranked.where(and_(*conditions)) - ranked_subquery = ranked.subquery() - - winner_ids = select(ranked_subquery.c.id).where(ranked_subquery.c.rn == 1) - - page_conditions: list[Any] = [AttackResultEntry.id.in_(winner_ids)] + page_conditions: list[Any] = list(conditions) + page_conditions.append(self._attack_results_not_superseded_condition(conditions=conditions)) if min_turns is not None: page_conditions.append(AttackResultEntry.executed_turns >= min_turns) if max_turns is not None: @@ -3890,6 +3960,48 @@ def _query_paginated_attack_results( ) return [entry.get_attack_result() for entry in entries] + @staticmethod + def _attack_results_not_superseded_condition(*, conditions: list[Any]) -> Any: + """ + Build the anti-join predicate keeping only the newest matching row per conversation. + + Args: + conditions (list[Any]): The same filters applied to the outer query, so dedup + picks the newest row *among matching rows*. + + Returns: + Any: A condition that is true when no later matching row shares the conversation. + """ + candidates = select( + AttackResultEntry.conversation_id.label("conversation_id"), + AttackResultEntry.timestamp.label("timestamp"), + AttackResultEntry.id.label("id"), + ) + if conditions: + candidates = candidates.where(and_(*conditions)) + # correlate(None) stops SQLAlchemy from hoisting the inner FROM onto the outer row, + # which would make every candidate trivially supersede itself. + newer = candidates.correlate(None).subquery("newer") + + return not_( + exists( + select(literal(1)) + .select_from(newer) + .where( + and_( + newer.c.conversation_id == AttackResultEntry.conversation_id, + or_( + newer.c.timestamp > AttackResultEntry.timestamp, + and_( + newer.c.timestamp == AttackResultEntry.timestamp, + newer.c.id > AttackResultEntry.id, + ), + ), + ) + ) + ) + ) + @staticmethod def _filter_attack_results_by_turns( results: list[AttackResult], *, min_turns: int | None, max_turns: int | None @@ -3935,25 +4047,43 @@ def _dedup_attack_entries(entries: Sequence[AttackResultEntry]) -> list[AttackRe seen[entry.conversation_id] = entry return [entry.get_attack_result() for entry in seen.values()] - def get_unique_attack_labels(self) -> dict[str, list[str]]: + def get_unique_attack_labels( + self, + *, + operator: Sequence[str] | None = None, + operation: Sequence[str] | None = None, + labels: Mapping[str, str | Sequence[str]] | None = None, + ) -> dict[str, list[str]]: """ - Return all unique label key-value pairs across attack results. + Return unique arbitrary labels, optionally narrowed by indexed attribution first. + + Args: + operator (Sequence[str] | None): Operator values used to narrow rows. + operation (Sequence[str] | None): Operation values used to narrow rows. + labels (Mapping[str, str | Sequence[str]] | None): Arbitrary label filters used + to narrow rows. Returns: dict[str, list[str]]: Mapping of label keys to sorted lists of unique values. """ label_values: dict[str, set[str]] = {} + filter_query = _AttackResultQuery(operator=operator, operation=operation, labels=labels) + conditions = self._build_attack_result_scalar_conditions(query=filter_query) + conditions.extend(self._build_attack_result_label_conditions(query=filter_query)) with closing(self.get_session()) as session: - are_rows = ( - session.query(AttackResultEntry.labels).filter(AttackResultEntry.labels.isnot(None)).distinct().all() - ) + query = session.query(AttackResultEntry.labels).filter(AttackResultEntry.labels.isnot(None)) + if conditions: + query = query.filter(and_(*conditions)) + are_rows = query.distinct().all() for (labels,) in are_rows: if not isinstance(labels, dict): continue for key, value in labels.items(): + if key in {"operator", "operation"}: + continue if isinstance(value, str): if key not in label_values: label_values[key] = set() @@ -3961,6 +4091,25 @@ def get_unique_attack_labels(self) -> dict[str, list[str]]: return {key: sorted(values) for key, values in sorted(label_values.items())} + def get_unique_attack_attribution(self) -> dict[str, list[str]]: + """Return unique dedicated operator and operation values from indexed columns.""" + with closing(self.get_session()) as session: + operators = [ + value + for (value,) in session.query(AttackResultEntry.operator) + .filter(AttackResultEntry.operator.isnot(None)) + .distinct() + .all() + ] + operations = [ + value + for (value,) in session.query(AttackResultEntry.operation) + .filter(AttackResultEntry.operation.isnot(None)) + .distinct() + .all() + ] + return {"operators": sorted(operators), "operations": sorted(operations)} + def add_scenario_results_to_memory(self, *, scenario_results: Sequence[ScenarioResult]) -> None: """ Insert a list of scenario results into the memory storage. diff --git a/pyrit/memory/memory_models.py b/pyrit/memory/memory_models.py index b6bbc78f78..9c533de4f2 100644 --- a/pyrit/memory/memory_models.py +++ b/pyrit/memory/memory_models.py @@ -73,6 +73,7 @@ TargetIdentifier, scorable_from_dict, ) +from pyrit.models.results.attack_result import normalize_legacy_attack_attribution logger = logging.getLogger(__name__) @@ -250,7 +251,7 @@ class PromptMemoryEntry(Base): converted_value_data_type (PromptDataType): The data type of the converted prompt (text, image) converted_value (str): The text of the converted prompt. If prompt is an image, it's a link. converted_value_sha256 (str): The SHA256 hash of the original prompt data. - idx_conversation_id (Index): The index for the conversation ID. + ix_PromptMemoryEntries_conversation_sequence_id (Index): Composite conversation ordering index. original_prompt_id (UUID): The original prompt id. It is equal to id unless it is a duplicate. scores (list[ScoreEntry]): The list of scores associated with the prompt. @@ -259,12 +260,23 @@ class PromptMemoryEntry(Base): """ __tablename__ = "PromptMemoryEntries" - __table_args__ = {"extend_existing": True} + __table_args__ = ( + Index( + "ix_PromptMemoryEntries_conversation_sequence_id", + "conversation_id", + "sequence", + "id", + mssql_include=["timestamp", "converted_value_data_type"], + ), + {"extend_existing": True}, + ) id = mapped_column(CustomUUID, nullable=False, primary_key=True) role: Mapped[Literal["system", "user", "assistant", "simulated_assistant", "tool", "developer"]] = mapped_column( String, nullable=False ) - conversation_id = mapped_column(String, nullable=False) + # Bounded so SQL Server accepts it as an index key. 128 rather than 36 because + # conversation_id is a free-form caller-supplied string, not necessarily a UUID. + conversation_id = mapped_column(String(128), nullable=False) sequence = mapped_column(INTEGER, nullable=False) timestamp = mapped_column(UTCDateTime, nullable=False) prompt_metadata: Mapped[dict[str, str | int]] = mapped_column(JSON) @@ -279,8 +291,6 @@ class PromptMemoryEntry(Base): converted_value = mapped_column(Unicode) converted_value_sha256 = mapped_column(String) - idx_conversation_id = Index("idx_conversation_id", "conversation_id") - original_prompt_id = mapped_column(CustomUUID, nullable=False) # Version of PyRIT used when this entry was created @@ -1551,6 +1561,8 @@ class AttackResultEntry(Base): outcome (AttackOutcome): The outcome of the attack, indicating success, failure, or undetermined. outcome_reason (str): Optional reason for the outcome, providing additional context. attack_metadata (dict[str, Any]): Metadata can be included as key-value pairs to provide extra context. + operator (str | None): Operator responsible for the attack. + operation (str | None): Operation associated with the attack. labels (dict[str, str]): Optional labels associated with the attack result entry. targeted_harm_categories (list[str]): Harm categories this attack targeted. pruned_conversation_ids (list[str]): List of conversation IDs that were pruned from the attack. @@ -1566,9 +1578,28 @@ class AttackResultEntry(Base): __tablename__ = "AttackResultEntries" __table_args__ = ( # Serves the PARTITION BY conversation_id dedup window in _query_paginated_attack_results. - Index("ix_AttackResultEntries_conversation_id", "conversation_id"), + Index( + "ix_AttackResultEntries_conversation_timestamp_id", + "conversation_id", + "timestamp", + "id", + ), # Serves the History recency ORDER BY timestamp DESC, id DESC and its keyset seek. Index("ix_AttackResultEntries_timestamp_id", "timestamp", "id"), + Index( + "ix_AttackResultEntries_operator_timestamp_id", + "operator", + "timestamp", + "id", + mssql_include=["conversation_id"], + ), + Index( + "ix_AttackResultEntries_operation_timestamp_id", + "operation", + "timestamp", + "id", + mssql_include=["conversation_id"], + ), # Serves scenario progress deltas scoped by parent and ordered oldest-first. Index( "ix_AttackResultEntries_attribution_parent_timestamp_id", @@ -1599,6 +1630,8 @@ class AttackResultEntry(Base): ) outcome_reason = mapped_column(String, nullable=True) attack_metadata: Mapped[dict[str, str | int | float | bool] | None] = mapped_column(JSON, nullable=True) + operator: Mapped[str | None] = mapped_column(Unicode(128), nullable=True) + operation: Mapped[str | None] = mapped_column(Unicode(128), nullable=True) labels: Mapped[dict[str, str] | None] = mapped_column(JSON, nullable=True) targeted_harm_categories: Mapped[list[str] | None] = mapped_column(JSON, nullable=True) pruned_conversation_ids: Mapped[list[str] | None] = mapped_column(JSON, nullable=True) @@ -1649,6 +1682,9 @@ def __init__(self, *, entry: AttackResult) -> None: Args: entry (AttackResult): The attack result object to convert into a database entry. + + Raises: + ValueError: If mutated legacy attribution labels are invalid or conflict. """ self.id = uuid.UUID(entry.attack_result_id) self.conversation_id = entry.conversation_id @@ -1675,7 +1711,14 @@ def __init__(self, *, entry: AttackResult) -> None: self.outcome = entry.outcome.value self.outcome_reason = entry.outcome_reason self.attack_metadata = self.filter_json_serializable_metadata(entry.metadata) - self.labels = entry.labels or {} + remaining_labels, operator, operation = normalize_legacy_attack_attribution( + labels=entry.labels or {}, + operator=entry.operator, + operation=entry.operation, + ) + self.operator = operator + self.operation = operation + self.labels = remaining_labels self.targeted_harm_categories = entry.targeted_harm_categories or None # Persist conversation references by type @@ -1807,6 +1850,8 @@ def get_attack_result(self) -> AttackResult: related_conversations=related_conversations, metadata=self.attack_metadata or {}, timestamp=self.timestamp or datetime.now(tz=timezone.utc), + operator=self.operator, + operation=self.operation, labels=self.labels or {}, targeted_harm_categories=self.targeted_harm_categories or [], error_message=self.error_message, @@ -1858,10 +1903,17 @@ class ScenarioResultEntry(Base): __tablename__ = "ScenarioResultEntries" __table_args__ = ( Index("ix_ScenarioResultEntries_timestamp_id", "timestamp", "id"), + Index("ix_ScenarioResultEntries_scenario_name_timestamp_id", "scenario_name", "timestamp", "id"), + Index( + "ix_ScenarioResultEntries_scenario_run_state_timestamp_id", + "scenario_run_state", + "timestamp", + "id", + ), {"extend_existing": True}, ) id = mapped_column(CustomUUID, nullable=False, primary_key=True) - scenario_name = mapped_column(String, nullable=False) + scenario_name = mapped_column(String(256), nullable=False) scenario_description = mapped_column(Unicode, nullable=True) scenario_version = mapped_column(INTEGER, nullable=False, default=1) pyrit_version = mapped_column(String, nullable=False) @@ -1877,7 +1929,7 @@ class ScenarioResultEntry(Base): ) objective_target_identifier: Mapped[dict[str, Any]] = mapped_column(JSON, nullable=False) objective_scorer_identifier: Mapped[dict[str, Any] | None] = mapped_column(JSON, nullable=True) - scenario_run_state: Mapped[str] = mapped_column(String, nullable=False, default="CREATED") + scenario_run_state: Mapped[str] = mapped_column(String(32), nullable=False, default="CREATED") display_group_map_json: Mapped[str | None] = mapped_column(Unicode, nullable=True) labels: Mapped[dict[str, str] | None] = mapped_column(JSON, nullable=True) number_tries: Mapped[int] = mapped_column(INTEGER, nullable=False, default=0) diff --git a/pyrit/memory/sqlite_memory.py b/pyrit/memory/sqlite_memory.py index 0d7eebab6c..52382814c6 100644 --- a/pyrit/memory/sqlite_memory.py +++ b/pyrit/memory/sqlite_memory.py @@ -433,27 +433,30 @@ def get_conversation_stats(self, *, conversation_ids: Sequence[str]) -> dict[str sql = text( f""" + WITH aggregate_rows AS ( + SELECT + conversation_id, + COUNT(DISTINCT sequence) AS msg_count, + MIN(timestamp) AS created_at + FROM "PromptMemoryEntries" + WHERE conversation_id IN ({placeholders}) + GROUP BY conversation_id + ) SELECT - pme.conversation_id, - COUNT(DISTINCT pme.sequence) AS msg_count, - ( - SELECT SUBSTR(p2.converted_value, 1, {ConversationStats.PREVIEW_FETCH_MAX_LEN}) + aggregate_rows.conversation_id, + aggregate_rows.msg_count, + SUBSTR(latest.converted_value, 1, {ConversationStats.PREVIEW_FETCH_MAX_LEN}) AS last_preview, + latest.converted_value_data_type AS last_data_type, + aggregate_rows.created_at + FROM aggregate_rows + LEFT JOIN "PromptMemoryEntries" latest + ON latest.id = ( + SELECT p2.id FROM "PromptMemoryEntries" p2 - WHERE p2.conversation_id = pme.conversation_id + WHERE p2.conversation_id = aggregate_rows.conversation_id ORDER BY p2.sequence DESC, p2.id DESC LIMIT 1 - ) AS last_preview, - ( - SELECT p2b.converted_value_data_type - FROM "PromptMemoryEntries" p2b - WHERE p2b.conversation_id = pme.conversation_id - ORDER BY p2b.sequence DESC, p2b.id DESC - LIMIT 1 - ) AS last_data_type, - MIN(pme.timestamp) AS created_at - FROM "PromptMemoryEntries" pme - WHERE pme.conversation_id IN ({placeholders}) - GROUP BY pme.conversation_id + ) """ ) diff --git a/pyrit/models/messages/message_piece.py b/pyrit/models/messages/message_piece.py index 799171a3cc..f9f807c5c8 100644 --- a/pyrit/models/messages/message_piece.py +++ b/pyrit/models/messages/message_piece.py @@ -252,10 +252,10 @@ def is_adversarial_placeholder(self) -> bool: def sort_message_pieces(message_pieces: list[MessagePiece]) -> list[MessagePiece]: """ - Group by ``conversation_id``, ordering by earliest timestamp then ``sequence``. + Group by ``conversation_id``, then order by sequence and piece timestamp. Conversations are ordered by their earliest piece's timestamp; pieces - within a conversation are ordered by ``sequence``. + within a conversation are ordered by ``sequence`` and then by creation time. Args: message_pieces: The pieces to sort. Not mutated. @@ -269,5 +269,5 @@ def sort_message_pieces(message_pieces: list[MessagePiece]) -> list[MessagePiece } return sorted( message_pieces, - key=lambda x: (earliest_timestamps[x.conversation_id], x.conversation_id or "", x.sequence), + key=lambda x: (earliest_timestamps[x.conversation_id], x.conversation_id or "", x.sequence, x.timestamp), ) diff --git a/pyrit/models/results/attack_result.py b/pyrit/models/results/attack_result.py index 054dceda4d..7b034a3095 100644 --- a/pyrit/models/results/attack_result.py +++ b/pyrit/models/results/attack_result.py @@ -6,10 +6,11 @@ import uuid from datetime import datetime, timezone from enum import Enum -from typing import Any, TypeVar +from typing import TYPE_CHECKING, Any, ClassVar, TypeVar -from pydantic import AwareDatetime, Field, field_serializer +from pydantic import AwareDatetime, Field, field_serializer, model_validator +from pyrit.common.deprecation import print_deprecation_message from pyrit.models.identifiers.component_identifier import ComponentIdentifier from pyrit.models.messages.conversation_reference import ConversationReference, ConversationType from pyrit.models.messages.message_piece import MessagePiece @@ -17,8 +18,61 @@ from pyrit.models.retry_event import RetryEvent from pyrit.models.score import Score +if TYPE_CHECKING: + from collections.abc import Mapping + AttackResultT = TypeVar("AttackResultT", bound="AttackResult") +ATTRIBUTION_FIELDS: tuple[str, str] = ("operator", "operation") +ATTRIBUTION_VALUE_MAX_LENGTH: int = 128 + + +def normalize_legacy_attack_attribution( + *, + labels: Mapping[str, Any], + operator: str | None, + operation: str | None, +) -> tuple[dict[str, Any], str | None, str | None]: + """ + Move scalar legacy label aliases to dedicated attribution fields. + + TODO(PyRIT 1.4): Remove this helper with legacy attribution label aliases. + + Args: + labels (Mapping[str, Any]): Labels that may still carry the legacy aliases. + operator (str | None): Dedicated operator value. + operation (str | None): Dedicated operation value. + + Returns: + tuple[dict[str, Any], str | None, str | None]: Arbitrary labels and resolved attribution. + + Raises: + ValueError: If an alias is invalid or disagrees with its dedicated value. + """ + remaining = dict(labels) + resolved: dict[str, Any] = {"operator": operator, "operation": operation} + for field in ATTRIBUTION_FIELDS: + current = resolved[field] + if current is not None and not isinstance(current, str): + raise ValueError(f"{field} must be a string") + if current is not None and len(current) > ATTRIBUTION_VALUE_MAX_LENGTH: + raise ValueError(f"{field} must be at most {ATTRIBUTION_VALUE_MAX_LENGTH} characters") + if field in remaining: + legacy_value = remaining.pop(field) + if not isinstance(legacy_value, str): + raise ValueError(f"labels.{field} must be a string") + if len(legacy_value) > ATTRIBUTION_VALUE_MAX_LENGTH: + raise ValueError(f"labels.{field} must be at most {ATTRIBUTION_VALUE_MAX_LENGTH} characters") + if current is not None and current != legacy_value: + raise ValueError(f"{field} conflicts with legacy labels.{field}: {current!r} != {legacy_value!r}") + print_deprecation_message( + old_item=f"labels.{field}", + new_item=field, + removed_in="1.4.0", + ) + resolved[field] = legacy_value + return remaining, resolved["operator"], resolved["operation"] + class AttackOutcome(str, Enum): """ @@ -44,6 +98,8 @@ class AttackOutcome(str, Enum): class AttackResult(StrategyResult): """Base class for all attack results.""" + ATTRIBUTION_VALUE_MAX_LENGTH: ClassVar[int] = ATTRIBUTION_VALUE_MAX_LENGTH + # Identity # Unique identifier of the conversation that produced this result conversation_id: str @@ -90,6 +146,11 @@ class AttackResult(StrategyResult): # Arbitrary metadata metadata: dict[str, Any] = Field(default_factory=dict) + # First-class attribution fields. These are deliberately separate from + # arbitrary labels so they can be indexed and queried efficiently. + operator: str | None = Field(default=None, max_length=ATTRIBUTION_VALUE_MAX_LENGTH) + operation: str | None = Field(default=None, max_length=ATTRIBUTION_VALUE_MAX_LENGTH) + # labels associated with this attack result labels: dict[str, str] = Field(default_factory=dict) @@ -115,6 +176,32 @@ class AttackResult(StrategyResult): attribution_parent_id: str | None = None attribution_data: dict[str, Any] | None = None + @model_validator(mode="before") + @classmethod + def _normalize_legacy_attribution_labels(cls, data: Any) -> Any: + """ + Move legacy attribution label aliases to their dedicated fields. + + Returns: + The normalized model input. + + Raises: + ValueError: If an alias is not a string or conflicts with a dedicated field. + """ + if not isinstance(data, dict) or not isinstance(data.get("labels"), dict): + return data + + normalized = dict(data) + remaining, operator, operation = normalize_legacy_attack_attribution( + labels=normalized["labels"], + operator=normalized.get("operator"), + operation=normalized.get("operation"), + ) + normalized["labels"] = remaining + normalized["operator"] = operator + normalized["operation"] = operation + return normalized + def get_attack_strategy_identifier(self) -> ComponentIdentifier | None: """ Return the attack strategy identifier from the composite atomic identifier. diff --git a/tests/unit/backend/test_api_routes.py b/tests/unit/backend/test_api_routes.py index 9f03d799a3..ebdbf62fd7 100644 --- a/tests/unit/backend/test_api_routes.py +++ b/tests/unit/backend/test_api_routes.py @@ -115,6 +115,8 @@ def test_list_attacks_with_filters(self, client: TestClient) -> None: has_converters=None, include_scenario_attacks=True, outcome="success", + operator=None, + operation=None, labels=None, min_turns=None, max_turns=None, @@ -579,6 +581,53 @@ def test_list_attacks_with_labels(self, client: TestClient) -> None: call_kwargs = mock_service.list_attacks_async.call_args[1] assert call_kwargs["labels"] == {"env": ["prod"], "team": ["red"]} + def test_list_attacks_with_dedicated_attribution_filters(self, client: TestClient) -> None: + with patch("pyrit.backend.routes.attacks.get_attack_service") as mock_get_service: + mock_service = MagicMock() + mock_service.list_attacks_async = AsyncMock( + return_value=AttackListResponse( + items=[], + pagination=PaginationInfo(limit=20, has_more=False, next_cursor=None, prev_cursor=None), + ) + ) + mock_get_service.return_value = mock_service + + response = client.get("/api/attacks?operator=alice&operation=nightly") + + assert response.status_code == status.HTTP_200_OK + call_kwargs = mock_service.list_attacks_async.call_args.kwargs + assert call_kwargs["operator"] == ["alice"] + assert call_kwargs["operation"] == ["nightly"] + + def test_list_attacks_legacy_attribution_label_warns_and_normalizes(self, client: TestClient) -> None: + with patch("pyrit.backend.routes.attacks.get_attack_service") as mock_get_service: + mock_service = MagicMock() + mock_service.list_attacks_async = AsyncMock( + return_value=AttackListResponse( + items=[], + pagination=PaginationInfo(limit=20, has_more=False, next_cursor=None, prev_cursor=None), + ) + ) + mock_get_service.return_value = mock_service + + with pytest.warns(DeprecationWarning, match="removed in 1.4.0"): + response = client.get("/api/attacks?label=operator:alice") + + assert response.status_code == status.HTTP_200_OK + call_kwargs = mock_service.list_attacks_async.call_args.kwargs + assert call_kwargs["operator"] == ["alice"] + assert call_kwargs["labels"] is None + + def test_list_attacks_rejects_conflicting_attribution_filters(self, client: TestClient) -> None: + response = client.get("/api/attacks?operator=alice&label=operator:bob") + + assert response.status_code == status.HTTP_422_UNPROCESSABLE_CONTENT + + def test_list_attacks_rejects_overlength_operator(self, client: TestClient) -> None: + response = client.get("/api/attacks", params={"operator": "x" * 129}) + + assert response.status_code == status.HTTP_422_UNPROCESSABLE_CONTENT + def test_get_attack_options(self, client: TestClient) -> None: """Test getting attack type options from attack results.""" with patch("pyrit.backend.routes.attacks.get_attack_service") as mock_get_service: @@ -660,8 +709,8 @@ def test_parse_labels_value_with_extra_colons(self, client: TestClient) -> None: call_kwargs = mock_service.list_attacks_async.call_args[1] assert call_kwargs["labels"] == {"url": ["http://example.com:8080"]} - def test_parse_labels_passes_keys_through_without_normalization(self, client: TestClient) -> None: - """Test that label keys are passed through as-is (DB stores canonical keys after migration).""" + def test_parse_labels_normalizes_legacy_attribution_aliases(self, client: TestClient) -> None: + """Legacy attribution label filters are routed to dedicated columns.""" with patch("pyrit.backend.routes.attacks.get_attack_service") as mock_get_service: mock_service = MagicMock() mock_service.list_attacks_async = AsyncMock( @@ -676,7 +725,9 @@ def test_parse_labels_passes_keys_through_without_normalization(self, client: Te assert response.status_code == status.HTTP_200_OK call_kwargs = mock_service.list_attacks_async.call_args[1] - assert call_kwargs["labels"] == {"operator": ["alice"], "operation": ["redteam"]} + assert call_kwargs["operator"] == ["alice"] + assert call_kwargs["operation"] == ["redteam"] + assert call_kwargs["labels"] is None def test_list_attacks_forwards_converter_types_param(self, client: TestClient) -> None: """Test that converter_types query params are forwarded to service.""" @@ -730,7 +781,8 @@ def test_list_attacks_groups_repeated_label_key_as_list(self, client: TestClient assert response.status_code == status.HTTP_200_OK call_kwargs = mock_service.list_attacks_async.call_args[1] - assert call_kwargs["labels"] == {"operator": ["alice", "bob"]} + assert call_kwargs["operator"] == ["alice", "bob"] + assert call_kwargs["labels"] is None def test_list_attacks_forwards_converter_types_match(self, client: TestClient) -> None: """converter_types_match query param is forwarded verbatim to service.""" @@ -1388,6 +1440,7 @@ def test_get_labels_for_attacks(self, client: TestClient) -> None: with patch("pyrit.backend.routes.labels.CentralMemory") as mock_memory_class: mock_memory = MagicMock() mock_memory.get_unique_attack_labels.return_value = {"env": ["prod"], "team": ["red"]} + mock_memory.get_unique_attack_attribution.return_value = {"operators": [], "operations": []} mock_memory_class.get_memory_instance.return_value = mock_memory response = client.get("/api/labels?source=attacks") @@ -1396,13 +1449,47 @@ def test_get_labels_for_attacks(self, client: TestClient) -> None: data = response.json() assert data["source"] == "attacks" assert data["labels"] == {"env": ["prod"], "team": ["red"]} - mock_memory.get_unique_attack_labels.assert_called_once() + assert data["operators"] == [] + assert data["operations"] == [] + mock_memory.get_unique_attack_labels.assert_called_once_with( + operator=None, + operation=None, + labels=None, + ) + + def test_get_labels_for_attacks_passes_narrowing_filters(self, client: TestClient) -> None: + with patch("pyrit.backend.routes.labels.CentralMemory") as mock_memory_class: + mock_memory = MagicMock() + mock_memory.get_unique_attack_labels.return_value = {"env": ["prod"]} + mock_memory.get_unique_attack_attribution.return_value = { + "operators": ["alice", "bob"], + "operations": ["nightly"], + } + mock_memory_class.get_memory_instance.return_value = mock_memory + + response = client.get( + "/api/labels", + params=[ + ("operator", "alice"), + ("operation", "nightly"), + ("label", "team:red"), + ], + ) + + assert response.status_code == status.HTTP_200_OK + mock_memory.get_unique_attack_labels.assert_called_once_with( + operator=["alice"], + operation=["nightly"], + labels={"team": ["red"]}, + ) + mock_memory.get_unique_attack_attribution.assert_not_called() def test_get_labels_empty(self, client: TestClient) -> None: """Test getting labels when no attack results exist.""" with patch("pyrit.backend.routes.labels.CentralMemory") as mock_memory_class: mock_memory = MagicMock() mock_memory.get_unique_attack_labels.return_value = {} + mock_memory.get_unique_attack_attribution.return_value = {"operators": [], "operations": []} mock_memory_class.get_memory_instance.return_value = mock_memory response = client.get("/api/labels?source=attacks") @@ -1420,6 +1507,7 @@ def test_get_labels_multiple_values(self, client: TestClient) -> None: "env": ["prod", "staging"], "team": ["blue"], } + mock_memory.get_unique_attack_attribution.return_value = {"operators": [], "operations": []} mock_memory_class.get_memory_instance.return_value = mock_memory response = client.get("/api/labels") @@ -1430,12 +1518,13 @@ def test_get_labels_multiple_values(self, client: TestClient) -> None: assert data["labels"]["team"] == ["blue"] def test_get_labels_returns_keys_without_normalization(self, client: TestClient) -> None: - """Test that label keys are returned as-is from the DB (canonical after migration).""" + """Attack attribution options are separate from arbitrary labels.""" with patch("pyrit.backend.routes.labels.CentralMemory") as mock_memory_class: mock_memory = MagicMock() - mock_memory.get_unique_attack_labels.return_value = { - "operator": ["alice", "bob"], - "operation": ["hunt", "scan"], + mock_memory.get_unique_attack_labels.return_value = {"team": ["red"]} + mock_memory.get_unique_attack_attribution.return_value = { + "operators": ["alice", "bob"], + "operations": ["hunt", "scan"], } mock_memory_class.get_memory_instance.return_value = mock_memory @@ -1443,8 +1532,9 @@ def test_get_labels_returns_keys_without_normalization(self, client: TestClient) assert response.status_code == status.HTTP_200_OK data = response.json() - assert set(data["labels"]["operator"]) == {"alice", "bob"} - assert set(data["labels"]["operation"]) == {"hunt", "scan"} + assert data["labels"] == {"team": ["red"]} + assert set(data["operators"]) == {"alice", "bob"} + assert set(data["operations"]) == {"hunt", "scan"} async def test_get_label_options_rejects_unsupported_source(self, client: TestClient) -> None: """Test that unsupported label source types are rejected.""" diff --git a/tests/unit/backend/test_attack_service.py b/tests/unit/backend/test_attack_service.py index 22310b77c7..be0714399f 100644 --- a/tests/unit/backend/test_attack_service.py +++ b/tests/unit/backend/test_attack_service.py @@ -629,22 +629,24 @@ async def test_list_attacks_formats_media_preview(self, attack_service, mock_mem assert preview == "[Image: 1780010098266691.png]" assert "C:\\" not in (preview or "") - async def test_list_attacks_filters_by_labels_directly(self, attack_service, mock_memory) -> None: - """Test that label filters are passed directly to the DB query (no legacy expansion).""" + async def test_list_attacks_filters_by_dedicated_attribution(self, attack_service, mock_memory) -> None: + """Dedicated attribution filters are passed to indexed memory columns.""" ar = make_attack_result(conversation_id="attack-canonical") + ar.operator = "alice" + ar.operation = "red" mock_memory.get_attack_results.return_value = [ar] mock_memory.get_conversation_stats.side_effect = lambda conversation_ids: { - cid: ConversationStats(message_count=1, labels={"operator": "alice", "operation": "red"}) - for cid in conversation_ids + cid: ConversationStats(message_count=1) for cid in conversation_ids } - result = await attack_service.list_attacks_async(labels={"operator": "alice", "operation": "red"}) + result = await attack_service.list_attacks_async(operator=["alice"], operation=["red"]) assert len(result.items) == 1 mock_memory.get_attack_results.assert_called_once() call_kwargs = mock_memory.get_attack_results.call_args[1] - assert call_kwargs["labels"] == {"operator": "alice", "operation": "red"} + assert call_kwargs["operator"] == ["alice"] + assert call_kwargs["operation"] == ["red"] async def test_list_attacks_forwards_min_and_max_turns(self, attack_service, mock_memory) -> None: """Both min_turns and max_turns are forwarded to the memory query.""" @@ -866,7 +868,13 @@ async def test_create_attack_stores_attack_result(self, attack_service, mock_mem mock_get_target_service.return_value = mock_target_service result = await attack_service.create_attack_async( - request=CreateAttackRequest(target_registry_name="target-1", name="My Attack") + request=CreateAttackRequest( + target_registry_name="target-1", + name="My Attack", + operator="alice", + operation="nightly", + labels={"team": "red"}, + ) ) assert result.conversation_id is not None @@ -874,6 +882,9 @@ async def test_create_attack_stores_attack_result(self, attack_service, mock_mem mock_memory.add_attack_results_to_memory.assert_called_once() stored_attack = mock_memory.add_attack_results_to_memory.call_args.kwargs["attack_results"][0] assert stored_attack.metadata["target_registry_name"] == "target-1" + assert stored_attack.operator == "alice" + assert stored_attack.operation == "nightly" + assert stored_attack.labels == {"team": "red", "source": "gui"} async def test_create_attack_stores_prepended_conversation(self, attack_service, mock_memory) -> None: """Test that create_attack stores prepended conversation messages.""" @@ -3217,40 +3228,22 @@ def test_rejects_incompatible_round_robin_target( with pytest.raises(ValueError, match="Target mismatch"): attack_service._validate_target_match(attack_identifier=attack_identifier, request=request) - async def test_rejects_mismatched_operator(self, attack_service, mock_memory) -> None: - """Should raise ValueError when request operator differs from attack operator.""" - ar = make_attack_result(conversation_id="test-id") - ar.labels["operator"] = "alice" - mock_memory.get_attack_results.return_value = [ar] - request = AddMessageRequest( - role="user", - pieces=[MessagePieceRequest(original_value="Hello")], - target_conversation_id="test-id", - send=False, - labels={"operator": "bob"}, - ) +def test_create_attack_request_normalizes_legacy_attribution_labels() -> None: + labels = {"operator": "alice", "operation": "nightly", "team": "red"} - with pytest.raises(ValueError, match="Operator mismatch"): - await attack_service.add_message_async(attack_result_id="test-id", request=request) + with pytest.warns(DeprecationWarning, match="removed in 1.4.0"): + request = CreateAttackRequest(target_registry_name="target", labels=labels) - async def test_allows_matching_operator(self, attack_service, mock_memory) -> None: - """Should NOT raise when request operator matches attack operator.""" - ar = make_attack_result(conversation_id="test-id") - ar.labels["operator"] = "alice" - mock_memory.get_attack_results.return_value = [ar] - mock_memory.get_conversation_messages.return_value = [] + assert request.operator == "alice" + assert request.operation == "nightly" + assert request.labels == {"team": "red"} + assert labels == {"operator": "alice", "operation": "nightly", "team": "red"} - request = AddMessageRequest( - role="user", - pieces=[MessagePieceRequest(original_value="Hello")], - target_conversation_id="test-id", - send=False, - labels={"operator": "alice"}, - ) - result = await attack_service.add_message_async(attack_result_id="test-id", request=request) - assert result.attack is not None +def test_create_attack_request_rejects_overlength_values() -> None: + with pytest.raises(ValueError, match="at most 128"): + CreateAttackRequest(target_registry_name="target", operator="x" * 129) class TestResolveVideoRemixMetadata: diff --git a/tests/unit/backend/test_mappers.py b/tests/unit/backend/test_mappers.py index ff79eba0e2..3587fc6421 100644 --- a/tests/unit/backend/test_mappers.py +++ b/tests/unit/backend/test_mappers.py @@ -158,6 +158,21 @@ async def test_basic_mapping(self) -> None: assert summary.target is not None assert summary.target.target_type == "TextTarget" + async def test_mapping_keeps_attribution_out_of_labels(self) -> None: + ar = _make_attack_result(name="My Attack") + ar.operator = "alice" + ar.operation = "nightly" + stats = ConversationStats( + message_count=1, + labels={"operator": "legacy", "operation": "legacy", "environment": "test"}, + ) + + summary = await attack_result_to_summary_async(ar, stats=stats) + + assert summary.operator == "alice" + assert summary.operation == "nightly" + assert summary.labels == {"test_ar_label": "test_ar_value", "environment": "test"} + async def test_round_robin_target_includes_canonical_identifier_hash(self) -> None: """Composite targets retain their full identity even when root display fields are absent.""" target_identifier = ComponentIdentifier( @@ -262,8 +277,8 @@ async def test_labels_are_mapped(self) -> None: assert summary.labels == {"env": "prod", "team": "red", "test_ar_label": "test_ar_value"} - async def test_labels_passed_through_without_normalization(self) -> None: - """Test that labels are passed through as-is (DB stores canonical keys after migration).""" + async def test_legacy_attribution_keys_are_not_merged_into_labels(self) -> None: + """Conversation-level legacy attribution keys do not leak into canonical labels.""" ar = _make_attack_result() stats = ConversationStats( message_count=1, @@ -273,8 +288,6 @@ async def test_labels_passed_through_without_normalization(self) -> None: summary = await attack_result_to_summary_async(ar, stats=stats) assert summary.labels == { - "operator": "alice", - "operation": "op_red", "env": "prod", "test_ar_label": "test_ar_value", } diff --git a/tests/unit/memory/memory_interface/test_interface_attack_results.py b/tests/unit/memory/memory_interface/test_interface_attack_results.py index 7c1f72ad9a..3331e9356e 100644 --- a/tests/unit/memory/memory_interface/test_interface_attack_results.py +++ b/tests/unit/memory/memory_interface/test_interface_attack_results.py @@ -9,6 +9,8 @@ from unittest.mock import patch import pytest +from sqlalchemy import select +from sqlalchemy.dialects import mssql from unit.mocks import get_mock_target_identifier, make_scenario_result from pyrit.common.utils import to_sha256 @@ -39,6 +41,8 @@ def create_attack_result( outcome: AttackOutcome = AttackOutcome.SUCCESS, labels: dict[str, str] | None = None, targeted_harm_categories: list[str] | None = None, + operator: str | None = None, + operation: str | None = None, ): """Helper function to create AttackResult.""" return AttackResult( @@ -46,6 +50,8 @@ def create_attack_result( objective=f"Objective {objective_num}", outcome=outcome, labels=labels or {}, + operator=operator, + operation=operation, targeted_harm_categories=targeted_harm_categories or [], ) @@ -108,14 +114,17 @@ def _drain_keyset(memory: MemoryInterface, *, page_size: int, **filters) -> list def test_attack_result_query_snapshots_mutable_inputs(): """The internal query remains stable when caller-owned containers change.""" attack_classes = ["CrescendoAttack"] - labels = {"operator": ["alice"]} - query = _AttackResultQuery(attack_classes=attack_classes, labels=labels) + operators = ["alice"] + labels = {"team": ["red"]} + query = _AttackResultQuery(attack_classes=attack_classes, operator=operators, labels=labels) attack_classes.append("ManualAttack") - labels["operator"].append("bob") + operators.append("bob") + labels["team"].append("blue") assert query.attack_classes == ("CrescendoAttack",) - assert query.labels == {"operator": ("alice",)} + assert query.operator == ("alice",) + assert query.labels == {"team": ("red",)} field_name = "limit" with pytest.raises(FrozenInstanceError): setattr(query, field_name, 10) @@ -149,7 +158,9 @@ def test_get_attack_results_forwards_all_parameters_to_query(sqlite_instance: Me converter_classes_match="any", has_converters=True, include_scenario_attacks=False, - labels={"operator": ["alice"]}, + operator=["alice"], + operation=["nightly"], + labels={"team": ["red"]}, targeted_harm_categories=["violence"], identifier_filters=[identifier_filter], scenario_result_id=str(uuid.uuid4()), @@ -172,7 +183,9 @@ def test_get_attack_results_forwards_all_parameters_to_query(sqlite_instance: Me assert query.converter_classes_match == "any" assert query.has_converters is True assert query.include_scenario_attacks is False - assert query.labels == {"operator": ("alice",)} + assert query.operator == ("alice",) + assert query.operation == ("nightly",) + assert query.labels == {"team": ("red",)} assert query.targeted_harm_categories == ("violence",) assert query.identifier_filters == (identifier_filter,) assert query.scenario_result_id is not None @@ -1270,6 +1283,84 @@ def test_get_unique_attack_labels_deduplicates_across_attacks(sqlite_instance: M assert result == {"env": ["prod"]} +def test_get_unique_attack_labels_narrows_by_attribution_and_labels(sqlite_instance: MemoryInterface): + sqlite_instance.add_attack_results_to_memory( + attack_results=[ + create_attack_result( + "conv_1", + 1, + operator="alice", + operation="nightly", + labels={"team": "red", "env": "prod"}, + ), + create_attack_result( + "conv_2", + 2, + operator="alice", + operation="daytime", + labels={"team": "blue", "env": "test"}, + ), + create_attack_result( + "conv_3", + 3, + operator="bob", + operation="nightly", + labels={"team": "red", "env": "dev"}, + ), + ] + ) + + result = sqlite_instance.get_unique_attack_labels( + operator=["alice"], + operation=["nightly"], + labels={"team": ["red"]}, + ) + + assert result == {"env": ["prod"], "team": ["red"]} + + +def test_get_attack_results_filters_dedicated_attribution_columns(sqlite_instance: MemoryInterface): + attack_results = [ + create_attack_result("conv_1", 1, operator="alice", operation="nightly"), + create_attack_result("conv_2", 2, operator="bob", operation="nightly"), + create_attack_result("conv_3", 3, operator="alice", operation="daytime"), + ] + sqlite_instance.add_attack_results_to_memory(attack_results=attack_results) + + results = sqlite_instance.get_attack_results(operator=["alice"], operation="nightly") + + assert [result.conversation_id for result in results] == ["conv_1"] + + +def test_get_attack_results_legacy_attribution_filter_warns_and_normalizes(sqlite_instance: MemoryInterface): + sqlite_instance.add_attack_results_to_memory(attack_results=[create_attack_result("conv_1", 1, operator="alice")]) + + with pytest.warns(DeprecationWarning, match="removed in 1.4.0"): + results = sqlite_instance.get_attack_results(labels={"operator": "alice"}) + + assert [result.conversation_id for result in results] == ["conv_1"] + + +def test_get_attack_results_rejects_conflicting_attribution_filters(sqlite_instance: MemoryInterface): + with pytest.raises(ValueError, match="operator conflicts"): + sqlite_instance.get_attack_results(operator="alice", labels={"operator": "bob"}) + + +def test_unique_attack_attribution_uses_dedicated_columns(sqlite_instance: MemoryInterface): + sqlite_instance.add_attack_results_to_memory( + attack_results=[ + create_attack_result("conv_1", 1, operator="bob", operation="nightly", labels={"team": "red"}), + create_attack_result("conv_2", 2, operator="alice", operation="nightly", labels={"team": "blue"}), + ] + ) + + assert sqlite_instance.get_unique_attack_attribution() == { + "operators": ["alice", "bob"], + "operations": ["nightly"], + } + assert sqlite_instance.get_unique_attack_labels() == {"team": ["blue", "red"]} + + # ============================================================================ # Attack class and converter class filtering tests # ============================================================================ @@ -1790,6 +1881,17 @@ def test_get_attack_results_pagination_returns_recency_ordered_page(sqlite_insta assert [r.conversation_id for r in page2] == ["conv-6", "conv-5", "conv-4"] +def test_get_attack_results_pagination_uses_not_exists_anti_join() -> None: + """Pagination probes for a newer duplicate instead of ranking the full result set.""" + condition = MemoryInterface._attack_results_not_superseded_condition(conditions=[]) + statement = select(AttackResultEntry.id).where(condition) + sql = str(statement.compile(dialect=mssql.dialect(), compile_kwargs={"literal_binds": True})).upper() + + assert "NOT (EXISTS" in sql + assert "ROW_NUMBER" not in sql + assert "PARTITION BY" not in sql + + def test_get_attack_results_pagination_disjoint_and_complete(sqlite_instance: MemoryInterface): """Concatenated keyset pages equal the full recency-ordered set with no gaps or duplicates.""" attack_results = [_make_attack_result(f"conv-{i}", ts_offset=i, updated_at_offset=100 + i) for i in range(25)] diff --git a/tests/unit/memory/memory_interface/test_interface_prompts.py b/tests/unit/memory/memory_interface/test_interface_prompts.py index 580529ab44..7582f01c27 100644 --- a/tests/unit/memory/memory_interface/test_interface_prompts.py +++ b/tests/unit/memory/memory_interface/test_interface_prompts.py @@ -70,6 +70,36 @@ def test_add_message_pieces_to_memory( assert len(sqlite_instance.get_message_pieces()) == num_conversations +def test_add_message_pieces_preserves_same_sequence_order(sqlite_instance: MemoryInterface): + conversation_id = str(uuid4()) + timestamp = datetime.now(tz=timezone.utc) + pieces = [ + MessagePiece( + id="00000000-0000-4000-8000-0000000000ff", + role="user", + original_value="first", + conversation_id=conversation_id, + sequence=0, + timestamp=timestamp, + ), + MessagePiece( + id="00000000-0000-4000-8000-000000000001", + role="user", + original_value="second", + conversation_id=conversation_id, + sequence=0, + timestamp=timestamp, + ), + ] + + sqlite_instance.add_message_pieces_to_memory(message_pieces=pieces) + + persisted_pieces = sqlite_instance.get_message_pieces(conversation_id=conversation_id) + assert [piece.original_value for piece in persisted_pieces] == ["first", "second"] + assert persisted_pieces[0].timestamp < persisted_pieces[1].timestamp + assert pieces[0].timestamp == pieces[1].timestamp + + def test_add_message_pieces_persists_converter_identifier_graph(sqlite_instance: MemoryInterface): target = TargetIdentifier( class_name="ConverterTarget", diff --git a/tests/unit/memory/test_azure_sql_memory.py b/tests/unit/memory/test_azure_sql_memory.py index b01fbf18ac..41d6ce5bbe 100644 --- a/tests/unit/memory/test_azure_sql_memory.py +++ b/tests/unit/memory/test_azure_sql_memory.py @@ -439,6 +439,23 @@ def test_get_attack_result_label_condition_empty_labels_dict(memory_interface: A assert not any("label_" in k for k in params) +def test_get_conversation_stats_uses_one_latest_row_apply( + uninitialized_memory_interface: AzureSQLMemory, +) -> None: + """The SQL Server query fetches preview and data type through one latest-row lookup.""" + session = MagicMock() + session.execute.return_value.fetchall.return_value = [] + + with patch.object(uninitialized_memory_interface, "get_session", return_value=session): + result = uninitialized_memory_interface.get_conversation_stats(conversation_ids=["conversation"]) + + sql = str(session.execute.call_args.args[0]) + assert result == {} + assert sql.upper().count("SELECT TOP 1") == 1 + assert "OUTER APPLY" in sql.upper() + assert "p2.converted_value_data_type AS last_data_type" in sql + + def test_scenario_history_conditions_bind_or_within_label_and_registry_values( memory_interface: AzureSQLMemory, ) -> None: diff --git a/tests/unit/memory/test_migration.py b/tests/unit/memory/test_migration.py index 8b9535b822..55c826d76e 100644 --- a/tests/unit/memory/test_migration.py +++ b/tests/unit/memory/test_migration.py @@ -2437,6 +2437,201 @@ def test_attack_recency_downgrade_restores_updated_at_and_drops_indexes(): engine.dispose() +# ============================================================================= +# First-class attack attribution and history indexes (a4c6e8f0b2d1) +# ============================================================================= + + +_ATTACK_ATTRIBUTION_REV = "a4c6e8f0b2d1" +_ATTACK_ATTRIBUTION_PREV_REV = "1b3d5f7a9c2e" + + +def _seed_attack_result_with_labels(connection, *, attack_id: str, labels: dict[str, object]) -> None: + connection.execute( + text( + 'INSERT INTO "AttackResultEntries" ' + "(id, conversation_id, objective, executed_turns, execution_time_ms, outcome, timestamp, labels) " + "VALUES (:id, :conv, 'obj', 1, 0, 'success', '2026-09-04', :labels)" + ), + {"id": attack_id, "conv": f"conv-{attack_id}", "labels": json.dumps(labels)}, + ) + + +def test_attack_attribution_migration_backfills_labels_and_indexes() -> None: + engine = create_engine("sqlite://") + attack_id = str(uuid.uuid4()) + try: + with engine.begin() as connection: + config = _config_for(connection) + command.upgrade(config, _ATTACK_ATTRIBUTION_PREV_REV) + _seed_attack_result_with_labels( + connection, + attack_id=attack_id, + labels={"operator": "alice", "operation": "nightly", "op": "keep", "team": "red"}, + ) + + command.upgrade(config, _ATTACK_ATTRIBUTION_REV) + + row = connection.execute( + text('SELECT operator, operation, labels FROM "AttackResultEntries" WHERE id = :attack_id'), + {"attack_id": attack_id}, + ).one() + attack_indexes = { + index["name"]: index["column_names"] for index in inspect(connection).get_indexes("AttackResultEntries") + } + prompt_indexes = { + index["name"]: index["column_names"] for index in inspect(connection).get_indexes("PromptMemoryEntries") + } + scenario_indexes = { + index["name"]: index["column_names"] + for index in inspect(connection).get_indexes("ScenarioResultEntries") + } + prompt_columns = { + column["name"]: column["type"] for column in inspect(connection).get_columns("PromptMemoryEntries") + } + + assert row.operator == "alice" + assert row.operation == "nightly" + assert json.loads(row.labels) == {"op": "keep", "team": "red"} + assert attack_indexes["ix_AttackResultEntries_conversation_timestamp_id"] == [ + "conversation_id", + "timestamp", + "id", + ] + assert attack_indexes["ix_AttackResultEntries_operator_timestamp_id"] == [ + "operator", + "timestamp", + "id", + ] + assert attack_indexes["ix_AttackResultEntries_operation_timestamp_id"] == [ + "operation", + "timestamp", + "id", + ] + assert prompt_indexes["ix_PromptMemoryEntries_conversation_sequence_id"] == [ + "conversation_id", + "sequence", + "id", + ] + assert prompt_columns["conversation_id"].length == 128 + assert scenario_indexes["ix_ScenarioResultEntries_scenario_name_timestamp_id"] == [ + "scenario_name", + "timestamp", + "id", + ] + assert scenario_indexes["ix_ScenarioResultEntries_scenario_run_state_timestamp_id"] == [ + "scenario_run_state", + "timestamp", + "id", + ] + finally: + engine.dispose() + + +def test_attack_attribution_migration_rejects_overlength_value() -> None: + engine = create_engine("sqlite://") + try: + with engine.begin() as connection: + config = _config_for(connection) + command.upgrade(config, _ATTACK_ATTRIBUTION_PREV_REV) + _seed_attack_result_with_labels( + connection, + attack_id=str(uuid.uuid4()), + labels={"operator": "x" * 129}, + ) + + with pytest.raises(ValueError, match="will not truncate"): + command.upgrade(config, _ATTACK_ATTRIBUTION_REV) + finally: + engine.dispose() + + +def test_attack_attribution_migration_uses_set_based_mssql_update() -> None: + import importlib + from unittest.mock import MagicMock, patch + + migration = importlib.import_module( + "pyrit.memory.alembic.versions.a4c6e8f0b2d1_add_attack_attribution_and_history_indexes" + ) + bind = MagicMock() + bind.dialect.name = "mssql" + bind.exec_driver_sql.return_value.first.return_value = None + + with patch.object(migration.op, "get_bind", return_value=bind): + migration._move_attribution_from_labels() + + assert [call.args[0] for call in bind.exec_driver_sql.call_args_list] == [ + migration._MSSQL_INVALID_ATTRIBUTION_QUERY, + migration._MSSQL_MOVE_ATTRIBUTION_QUERY, + ] + bind.execute.assert_not_called() + + +@pytest.mark.parametrize( + ("invalid_value", "error_match"), + [ + ( + {"row_id": "attack-1", "field_name": "operator", "value_type": 2}, + "non-string labels.operator", + ), + ( + {"row_id": "attack-2", "field_name": "operation", "value_type": 1}, + "labels.operation longer than 128 characters", + ), + ], +) +def test_attack_attribution_mssql_migration_rejects_invalid_value( + invalid_value: dict[str, object], error_match: str +) -> None: + import importlib + from types import SimpleNamespace + from unittest.mock import MagicMock + + migration = importlib.import_module( + "pyrit.memory.alembic.versions.a4c6e8f0b2d1_add_attack_attribution_and_history_indexes" + ) + bind = MagicMock() + bind.exec_driver_sql.return_value.first.return_value = SimpleNamespace(_mapping=invalid_value) + + with pytest.raises(ValueError, match=error_match): + migration._move_attribution_from_labels_mssql(bind=bind) + + bind.exec_driver_sql.assert_called_once_with(migration._MSSQL_INVALID_ATTRIBUTION_QUERY) + + +def test_attack_attribution_downgrade_restores_legacy_labels() -> None: + engine = create_engine("sqlite://") + attack_id = str(uuid.uuid4()) + try: + with engine.begin() as connection: + config = _config_for(connection) + command.upgrade(config, _ATTACK_ATTRIBUTION_REV) + connection.execute( + text( + 'INSERT INTO "AttackResultEntries" ' + "(id, conversation_id, objective, executed_turns, execution_time_ms, outcome, " + "timestamp, operator, operation, labels) " + "VALUES (:id, :conv, 'obj', 1, 0, 'success', '2026-09-04', " + "'alice', 'nightly', :labels)" + ), + {"id": attack_id, "conv": f"conv-{attack_id}", "labels": json.dumps({"team": "red"})}, + ) + + command.downgrade(config, _ATTACK_ATTRIBUTION_PREV_REV) + + labels = connection.execute( + text('SELECT labels FROM "AttackResultEntries" WHERE id = :attack_id'), + {"attack_id": attack_id}, + ).scalar_one() + columns = {column["name"] for column in inspect(connection).get_columns("AttackResultEntries")} + + assert json.loads(labels) == {"team": "red", "operator": "alice", "operation": "nightly"} + assert "operator" not in columns + assert "operation" not in columns + finally: + engine.dispose() + + _STRING_TYPES_REQUIRING_LENGTH = {"String", "VARCHAR", "NVARCHAR", "Unicode"} @@ -2554,6 +2749,34 @@ def test_scored_expectation_migration_script_metadata(): assert mig.depends_on is None +@pytest.mark.parametrize( + ("function_name", "query_name"), + [ + ("_backfill_scored_expectation", "_MSSQL_BACKFILL_SCORED_EXPECTATION_QUERY"), + ("_backfill_objective", "_MSSQL_RESTORE_OBJECTIVE_QUERY"), + ], +) +def test_scored_expectation_mssql_backfill_uses_set_based_update(function_name: str, query_name: str): + import importlib + from unittest.mock import MagicMock, patch + + migration = importlib.import_module("pyrit.memory.alembic.versions.1b3d5f7a9c2e_persist_scored_expectation") + connection = MagicMock() + connection.dialect.name = "mssql" + connection.exec_driver_sql.return_value.rowcount = 12 + + with ( + patch.object(migration.op, "get_bind", return_value=connection), + patch.object(migration, "_report_progress") as report_progress, + ): + getattr(migration, function_name)() + + connection.exec_driver_sql.assert_called_once_with(getattr(migration, query_name)) + connection.execute.assert_not_called() + assert report_progress.call_count == 2 + assert "updated 12 row(s)" in report_progress.call_args_list[-1].args[0] + + def test_scored_expectation_upgrade_backfills_objective_into_expectation(): """Upgrading folds a non-null objective into a versioned expectation and leaves NULLs NULL.""" id_with = str(uuid.uuid4()) diff --git a/tests/unit/memory/test_sqlite_memory.py b/tests/unit/memory/test_sqlite_memory.py index 8af4c17614..8be9ef17f5 100644 --- a/tests/unit/memory/test_sqlite_memory.py +++ b/tests/unit/memory/test_sqlite_memory.py @@ -13,7 +13,7 @@ from unittest.mock import MagicMock import pytest -from sqlalchemy import ARRAY, DateTime, Integer, String, create_engine, inspect, text +from sqlalchemy import ARRAY, DateTime, Integer, String, create_engine, event, inspect, text from sqlalchemy.dialects.postgresql import UUID from sqlalchemy.dialects.sqlite import CHAR, JSON from sqlalchemy.exc import SQLAlchemyError @@ -730,6 +730,24 @@ def test_get_conversation_stats_returns_empty_for_no_ids(sqlite_instance): assert result == {} +def test_get_conversation_stats_uses_indexed_latest_message_lookup(sqlite_instance): + statements: list[str] = [] + + def capture_statement(conn, cursor, statement, parameters, context, executemany): + statements.append(statement) + + event.listen(sqlite_instance.engine, "before_cursor_execute", capture_statement) + try: + sqlite_instance.get_conversation_stats(conversation_ids=["conversation"]) + finally: + event.remove(sqlite_instance.engine, "before_cursor_execute", capture_statement) + + sql = "\n".join(statements).upper() + assert 'LEFT JOIN "PROMPTMEMORYENTRIES" LATEST' in sql + assert "ORDER BY P2.SEQUENCE DESC, P2.ID DESC" in sql + assert "ROW_NUMBER" not in sql + + def test_get_conversation_stats_returns_empty_for_unknown_ids(sqlite_instance): """Test that get_conversation_stats omits unknown conversation IDs.""" result = sqlite_instance.get_conversation_stats(conversation_ids=["nonexistent"]) diff --git a/tests/unit/models/test_attack_result.py b/tests/unit/models/test_attack_result.py index 06482db0d6..ac7fca40f3 100644 --- a/tests/unit/models/test_attack_result.py +++ b/tests/unit/models/test_attack_result.py @@ -223,6 +223,24 @@ def test_no_error_fields_roundtrip(self) -> None: assert hydrated.retry_events == [] assert hydrated.total_retries == 0 + def test_attribution_fields_roundtrip_without_labels(self) -> None: + original = AttackResult( + conversation_id="c1", + objective="test", + operator="alice", + operation="nightly", + labels={"team": "red"}, + ) + + entry = AttackResultEntry(entry=original) + hydrated = entry.get_attack_result() + + assert entry.operator == "alice" + assert entry.operation == "nightly" + assert hydrated.operator == "alice" + assert hydrated.operation == "nightly" + assert hydrated.labels == {"team": "red"} + def test_traceback_truncation(self) -> None: """Very long tracebacks are truncated to 10KB.""" long_traceback = "x" * 20000 @@ -335,6 +353,53 @@ def test_aware_iso_string_timestamp_is_preserved(self) -> None: result = AttackResult(conversation_id="c1", objective="test", timestamp="2026-01-01T12:00:00+00:00") assert result.timestamp == datetime(2026, 1, 1, 12, 0, 0, tzinfo=timezone.utc) + def test_legacy_attribution_labels_are_normalized_without_mutation(self) -> None: + labels = {"operator": "alice", "operation": "nightly", "team": "red"} + + with pytest.warns(DeprecationWarning, match="removed in 1.4.0"): + result = AttackResult(conversation_id="c1", objective="test", labels=labels) + + assert result.operator == "alice" + assert result.operation == "nightly" + assert result.labels == {"team": "red"} + assert labels == {"operator": "alice", "operation": "nightly", "team": "red"} + assert result.model_dump(mode="json")["labels"] == {"team": "red"} + + def test_conflicting_legacy_attribution_label_is_rejected(self) -> None: + with pytest.raises(ValueError, match="operator conflicts"): + AttackResult( + conversation_id="c1", + objective="test", + operator="alice", + labels={"operator": "bob"}, + ) + + @pytest.mark.parametrize("field_name", ["operator", "operation"]) + def test_attribution_value_longer_than_128_is_rejected(self, field_name: str) -> None: + with pytest.raises(ValueError, match="at most 128"): + AttackResult(conversation_id="c1", objective="test", **{field_name: "x" * 129}) + + def test_dedicated_attribution_is_canonical(self) -> None: + result = AttackResult( + conversation_id="c1", + objective="test", + operator="alice", + operation="nightly", + labels={"team": "red"}, + ) + + dumped = result.model_dump(mode="json") + + assert dumped["operator"] == "alice" + assert dumped["operation"] == "nightly" + assert dumped["labels"] == {"team": "red"} + + result.labels["operator"] = "legacy-mutation" + assert result.model_dump(mode="json")["labels"] == { + "team": "red", + "operator": "legacy-mutation", + } + class TestAttackResultDuplicate: """duplicate() must deep-copy so mutations on the copy never touch the original.""" diff --git a/tests/unit/models/test_message_piece.py b/tests/unit/models/test_message_piece.py index e532d663b9..39c5b00ad7 100644 --- a/tests/unit/models/test_message_piece.py +++ b/tests/unit/models/test_message_piece.py @@ -644,6 +644,25 @@ def test_order_message_pieces_by_conversation_same_timestamp_different_sequences assert sort_message_pieces(pieces) == expected +def test_order_message_pieces_with_same_sequence_by_timestamp(): + earlier_piece = MessagePiece( + role="user", + original_value="first", + conversation_id="conv1", + timestamp=datetime.now(tz=timezone.utc) - timedelta(seconds=1), + sequence=1, + ) + later_piece = MessagePiece( + role="user", + original_value="second", + conversation_id="conv1", + timestamp=datetime.now(tz=timezone.utc), + sequence=1, + ) + + assert sort_message_pieces([later_piece, earlier_piece]) == [earlier_piece, later_piece] + + def test_message_piece_to_dict(): entry = MessagePiece( role="user",