From c6a0110dbee8790872474fcaf459d0ef85e91508 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Fri, 4 Sep 2026 20:01:41 -0700 Subject: [PATCH 1/6] PERF: Optimize attack and scenario history Promote attack operator and operation metadata to indexed fields, add query-aligned history indexes, and reduce duplicate latest-message lookups. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 7540877a-bdf5-4309-97e1-14469bc817e7 --- frontend/src/App.tsx | 7 +- .../src/components/Chat/ChatWindow.test.tsx | 9 +- frontend/src/components/Chat/ChatWindow.tsx | 41 ++- .../components/History/AttackHistory.test.tsx | 4 +- .../src/components/History/AttackHistory.tsx | 19 +- .../components/History/AttackTable.test.tsx | 4 +- .../src/components/History/AttackTable.tsx | 6 +- frontend/src/components/Home/Home.test.tsx | 36 ++- frontend/src/components/Home/Home.tsx | 2 +- .../src/components/Labels/LabelsBar.test.tsx | 4 +- frontend/src/components/Labels/LabelsBar.tsx | 4 +- frontend/src/services/api.ts | 5 +- frontend/src/types/index.ts | 12 +- pyrit/backend/mappers/attack_mappers.py | 2 + pyrit/backend/models/attacks.py | 51 +++- pyrit/backend/routes/attacks.py | 36 ++- pyrit/backend/routes/labels.py | 10 +- pyrit/backend/services/attack_service.py | 60 ++-- ..._attack_attribution_and_history_indexes.py | 288 ++++++++++++++++++ pyrit/memory/azure_sql_memory.py | 40 +-- pyrit/memory/memory_interface.py | 138 ++++++++- pyrit/memory/memory_models.py | 83 ++++- pyrit/memory/sqlite_memory.py | 59 ++-- pyrit/models/results/attack_result.py | 67 +++- tests/unit/backend/test_api_routes.py | 79 ++++- tests/unit/backend/test_attack_service.py | 65 ++-- tests/unit/backend/test_mappers.py | 21 +- .../test_interface_attack_results.py | 65 +++- tests/unit/memory/test_azure_sql_memory.py | 17 ++ tests/unit/memory/test_migration.py | 130 ++++++++ tests/unit/models/test_attack_result.py | 62 ++++ 31 files changed, 1212 insertions(+), 214 deletions(-) create mode 100644 pyrit/memory/alembic/versions/a4c6e8f0b2d1_add_attack_attribution_and_history_indexes.py 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..52c5dd4147 100644 --- a/frontend/src/components/Chat/ChatWindow.test.tsx +++ b/frontend/src/components/Chat/ChatWindow.test.tsx @@ -705,7 +705,9 @@ describe("ChatWindow Integration", () => { await waitFor(() => { expect(mockedAttacksApi.createAttack).toHaveBeenCalledWith({ target_registry_name: "openai_chat_1", - labels: { operator: 'testuser', operation: 'test_op' }, + 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 +716,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 +2656,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 +3787,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..7dfc6c91f3 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, @@ -78,6 +80,19 @@ function matchesNarrowScreen(): boolean { && window.matchMedia(NARROW_SCREEN_QUERY).matches } +function attackAttributionFromLabels(labels?: Record): Pick< + CreateAttackRequest, + 'operator' | 'operation' | 'labels' +> { + if (!labels) return {} + const { operator, operation, ...arbitraryLabels } = labels + const attribution: Pick = {} + if (operator) attribution.operator = operator + if (operation) attribution.operation = operation + if (Object.keys(arbitraryLabels).length > 0) attribution.labels = arbitraryLabels + return attribution +} + interface ChatWindowProps { onNewAttack: () => void activeTarget: TargetInstance | null @@ -89,8 +104,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 +133,7 @@ export default function ChatWindow({ labels, onLabelsChange, onNavigate, - attackLabels, + attackOperator, attackTarget, targetResolutionStatus = 'idle', onRetryTargetResolution, @@ -238,10 +253,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 +442,12 @@ 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, + ...attackAttributionFromLabels(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 +480,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 +671,7 @@ export default function ChatWindow({ try { const createResponse = await attacksApi.createAttack({ target_registry_name: activeTarget.target_registry_name, - labels: labels, + ...attackAttributionFromLabels(labels), source_conversation_id: activeConversationId, cutoff_index: messageIndex, }) @@ -707,7 +722,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, + ...attackAttributionFromLabels(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..4e0f7e090e 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'], }, }) diff --git a/frontend/src/components/History/AttackHistory.tsx b/frontend/src/components/History/AttackHistory.tsx index 5e4e147d3f..25b6c109df 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 @@ -110,22 +109,16 @@ 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') { + if (key !== 'source') { for (const val of values) { others.push(`${key}:${val}`) } } } - setOperatorOptions(operators.sort()) - setOperationOptions(operations.sort()) + setOperatorOptions([...(resp.operators ?? resp.labels.operator ?? [])].sort()) + setOperationOptions([...(resp.operations ?? resp.labels.operation ?? [])].sort()) setOtherLabelOptions(others.sort()) }) .catch(() => { /* ignore */ }) 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..264eade0f4 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,8 @@ 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 || [])])], + 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.ts b/frontend/src/services/api.ts index 9bdcac5204..032f5b726f 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,7 +369,7 @@ export const attacksApi = { export const labelsApi = { getLabels: async ( source: 'attacks' | 'scenarios' = 'attacks', - ): Promise<{ source: string; labels: Record }> => { + ): Promise => { const response = await apiClient.get('/labels', { params: { source } }) 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..057014ce1e 100644 --- a/pyrit/backend/models/attacks.py +++ b/pyrit/backend/models/attacks.py @@ -16,6 +16,7 @@ from pyrit.backend.models._media import build_filename, infer_mime_type from pyrit.backend.models.common import PaginationInfo +from pyrit.common.deprecation import print_deprecation_message from pyrit.models import ( AttackResult, ChatMessageRole, @@ -358,12 +359,54 @@ 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. + + 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) + labels = dict(normalized["labels"]) + for field_name in ("operator", "operation"): + if field_name not in labels: + continue + legacy_value = labels.pop(field_name) + if not isinstance(legacy_value, str): + raise ValueError(f"labels.{field_name} must be a string") + dedicated_value = normalized.get(field_name) + if dedicated_value is not None and dedicated_value != legacy_value: + raise ValueError(f"{field_name} conflicts with legacy labels.{field_name}") + print_deprecation_message( + old_item=f"labels.{field_name}", + new_item=field_name, + removed_in="1.4.0", + ) + normalized[field_name] = legacy_value + normalized["labels"] = labels + return normalized + + # ============================================================================ # Create Attack # ============================================================================ -class CreateAttackRequest(BaseModel): +class CreateAttackRequest(_AttackAttributionInput): """ Request to create a new attack. @@ -388,7 +431,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 +573,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..62f135d9bc 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; " @@ -101,7 +109,27 @@ async def list_attacks( # pyrit-async-suffix-exempt AttackListResponse: Paginated list of attack summaries. """ service = get_attack_service() - labels = parse_label_query_params(label) + labels = parse_label_query_params(label) or {} + 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``. @@ -116,7 +144,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..29318d3ba9 100644 --- a/pyrit/backend/routes/labels.py +++ b/pyrit/backend/routes/labels.py @@ -23,11 +23,14 @@ 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( @@ -49,7 +52,10 @@ async def get_label_options( # pyrit-async-suffix-exempt """ 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": + labels = await run_in_threadpool(memory.get_unique_attack_labels) + attribution = 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/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..a605911a7e --- /dev/null +++ b/pyrit/memory/alembic/versions/a4c6e8f0b2d1_add_attack_attribution_and_history_indexes.py @@ -0,0 +1,288 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Add first-class attack attribution fields and history query indexes. + +Revision ID: a4c6e8f0b2d1 +Revises: 8d1e3f5a7b9c +Create Date: 2026-09-04 18:48:00.000000 +""" + +from __future__ import annotations + +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 = "8d1e3f5a7b9c" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + +_ATTRIBUTION_FIELDS = ("operator", "operation") +_ATTRIBUTION_MAX_LENGTH = 128 + + +def upgrade() -> None: + """Add attribution columns, migrate legacy labels, and replace history indexes.""" + 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)) + _move_attribution_from_labels() + _bound_indexed_text_columns() + + op.drop_index("ix_AttackResultEntries_conversation_id", table_name="AttackResultEntries") + op.create_index( + "ix_AttackResultEntries_conversation_timestamp_id", + "AttackResultEntries", + ["conversation_id", "timestamp", "id"], + ) + op.create_index( + "ix_AttackResultEntries_operator_conversation_timestamp_id", + "AttackResultEntries", + ["operator", "conversation_id", "timestamp", "id"], + ) + op.create_index( + "ix_AttackResultEntries_operation_conversation_timestamp_id", + "AttackResultEntries", + ["operation", "conversation_id", "timestamp", "id"], + ) + + _drop_index_if_exists(name="idx_conversation_id", table_name="PromptMemoryEntries") + op.create_index( + "ix_PromptMemoryEntries_conversation_sequence_id", + "PromptMemoryEntries", + ["conversation_id", "sequence", "id"], + mssql_include=["timestamp", "converted_value_data_type"], + ) + + op.create_index( + "ix_ScenarioResultEntries_scenario_name_timestamp_id", + "ScenarioResultEntries", + ["scenario_name", "timestamp", "id"], + ) + op.create_index( + "ix_ScenarioResultEntries_scenario_run_state_timestamp_id", + "ScenarioResultEntries", + ["scenario_run_state", "timestamp", "id"], + ) + + +def downgrade() -> None: + """Restore legacy labels and indexes, then remove attribution columns.""" + _restore_attribution_to_labels() + + 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_conversation_timestamp_id", + table_name="AttackResultEntries", + ) + op.drop_index( + "ix_AttackResultEntries_operator_conversation_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") + + +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=36) + _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(36), + 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(36), + 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() + table = _attack_results_table(include_attribution=True) + rows = bind.execute(sa.select(table.c.id, table.c.labels)).all() + for row in rows: + 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: + values["labels"] = remaining_labels + bind.execute(sa.update(table).where(table.c.id == row.id).values(**values)) + + +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) + rows = bind.execute(sa.select(table.c.id, table.c.labels, table.c.operator, table.c.operation)).all() + for row in rows: + 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: + bind.execute(sa.update(table).where(table.c.id == row.id).values(labels=labels)) 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..d7152dd496 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -25,6 +25,8 @@ 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 @@ -227,6 +229,8 @@ class _AttackResultQuery: "converter_classes", "targeted_harm_categories", "identifier_filters", + "operator", + "operation", ) attack_result_ids: Sequence[str] | None = None @@ -241,6 +245,8 @@ class _AttackResultQuery: has_converters: bool | None = None include_scenario_attacks: bool = True labels: Mapping[str, str | Sequence[str]] | None = None + operator: Sequence[str] | None = None + operation: 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 +256,49 @@ 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. + + 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 ("operator", "operation"): + values = getattr(self, field_name) + if values is not None and any(not isinstance(value, str) for value in values): + raise ValueError(f"{field_name} values must be strings") + if values is not None and any(len(value) > AttackResult.ATTRIBUTION_VALUE_MAX_LENGTH for value in values): + raise ValueError( + f"{field_name} values must be at most {AttackResult.ATTRIBUTION_VALUE_MAX_LENGTH} characters" + ) + 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 name in ("operator", "operation"): + if name not in labels: + continue + legacy_raw = labels.pop(name) + legacy_values = (legacy_raw,) if isinstance(legacy_raw, str) else tuple(legacy_raw) + if any(not isinstance(value, str) for value in legacy_values): + raise ValueError(f"labels.{name} values must be strings") + if any(len(value) > AttackResult.ATTRIBUTION_VALUE_MAX_LENGTH for value in legacy_values): + raise ValueError( + f"labels.{name} values must be at most {AttackResult.ATTRIBUTION_VALUE_MAX_LENGTH} characters" + ) + dedicated_values = getattr(self, name) + if dedicated_values is not None and set(dedicated_values) != set(legacy_values): + raise ValueError(f"{name} conflicts with legacy labels.{name}") + print_deprecation_message( + old_item=f"_AttackResultQuery.labels['{name}']", + new_item=f"_AttackResultQuery.{name}", + removed_in="1.4.0", + ) + object.__setattr__(self, name, legacy_values) + object.__setattr__(self, "labels", MappingProxyType(labels) if labels else None) class MemoryInterface(abc.ABC): @@ -3470,6 +3510,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 +3555,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 @@ -3558,6 +3601,11 @@ def get_attack_results( ValueError: If ``limit`` or ``after`` is combined with ``attack_result_ids`` or ``objective_sha256`` (id-batched lookups do not support SQL pagination). """ + labels, operator_values, operation_values = self._normalize_attack_attribution_filters( + labels=labels, + operator=operator, + operation=operation, + ) query = _AttackResultQuery( attack_result_ids=attack_result_ids, conversation_id=conversation_id, @@ -3571,6 +3619,8 @@ def get_attack_results( has_converters=has_converters, include_scenario_attacks=include_scenario_attacks, labels=labels, + operator=operator_values, + operation=operation_values, targeted_harm_categories=targeted_harm_categories, identifier_filters=identifier_filters, scenario_result_id=scenario_result_id, @@ -3581,6 +3631,55 @@ def get_attack_results( ) return self._query_attack_results(query=query) + @staticmethod + def _normalize_attack_attribution_filters( + *, + labels: Mapping[str, str | Sequence[str]] | None, + operator: str | Sequence[str] | None, + operation: str | Sequence[str] | None, + ) -> tuple[Mapping[str, str | Sequence[str]] | None, Sequence[str] | None, Sequence[str] | None]: + """ + Normalize deprecated attribution label aliases without mutating caller input. + + Returns: + The arbitrary labels, operator values, and operation values. + + Raises: + ValueError: If a legacy alias conflicts with its dedicated filter. + """ + operator_values = [operator] if isinstance(operator, str) else operator + operation_values = [operation] if isinstance(operation, str) else operation + for field_name, values in (("operator", operator_values), ("operation", operation_values)): + if values is not None and any(len(value) > AttackResult.ATTRIBUTION_VALUE_MAX_LENGTH for value in values): + raise ValueError( + f"{field_name} values must be at most {AttackResult.ATTRIBUTION_VALUE_MAX_LENGTH} characters" + ) + if not labels: + return labels, operator_values, operation_values + + normalized_labels = dict(labels) + normalized_dedicated = {"operator": operator_values, "operation": operation_values} + for name in ("operator", "operation"): + if name not in normalized_labels: + continue + legacy_raw = normalized_labels.pop(name) + legacy_values = [legacy_raw] if isinstance(legacy_raw, str) else list(legacy_raw) + dedicated_values = normalized_dedicated[name] + if dedicated_values is not None and set(dedicated_values) != set(legacy_values): + raise ValueError(f"{name} conflicts with legacy labels.{name}") + print_deprecation_message( + old_item=f"get_attack_results(labels={{'{name}': ...}})", + new_item=f"get_attack_results({name}=...)", + removed_in="1.4.0", + ) + normalized_dedicated[name] = legacy_values + + return ( + normalized_labels or None, + normalized_dedicated["operator"], + normalized_dedicated["operation"], + ) + def _query_attack_results(self, *, query: _AttackResultQuery) -> Sequence[AttackResult]: """ Retrieve attack results matching an immutable query. @@ -3665,6 +3764,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: @@ -3954,6 +4057,8 @@ def get_unique_attack_labels(self) -> dict[str, list[str]]: 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 +4066,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 3c8c7264a2..5f2abb8965 100644 --- a/pyrit/memory/memory_models.py +++ b/pyrit/memory/memory_models.py @@ -35,6 +35,7 @@ from typing_extensions import Self import pyrit +from pyrit.common.deprecation import print_deprecation_message from pyrit.common.utils import to_sha256 from pyrit.models import ( SEED_RESPONSE_JSON_SCHEMA_METADATA_KEY, @@ -249,7 +250,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. @@ -258,12 +259,21 @@ 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) + conversation_id = mapped_column(String(36), nullable=False) sequence = mapped_column(INTEGER, nullable=False) timestamp = mapped_column(UTCDateTime, nullable=False) prompt_metadata: Mapped[dict[str, str | int]] = mapped_column(JSON) @@ -278,8 +288,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 @@ -1543,6 +1551,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. @@ -1558,9 +1568,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_conversation_timestamp_id", + "operator", + "conversation_id", + "timestamp", + "id", + ), + Index( + "ix_AttackResultEntries_operation_conversation_timestamp_id", + "operation", + "conversation_id", + "timestamp", + "id", + ), # Serves scenario progress deltas scoped by parent and ordered oldest-first. Index( "ix_AttackResultEntries_attribution_parent_timestamp_id", @@ -1591,6 +1620,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) @@ -1641,6 +1672,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 @@ -1667,7 +1701,29 @@ 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 {} + labels = dict(entry.labels or {}) + attribution = {"operator": entry.operator, "operation": entry.operation} + for field_name in ("operator", "operation"): + if field_name not in labels: + continue + legacy_value = labels.pop(field_name) + dedicated_value = attribution[field_name] + if not isinstance(legacy_value, str): + raise ValueError(f"labels.{field_name} must be a string") + if dedicated_value is not None and dedicated_value != legacy_value: + raise ValueError(f"{field_name} conflicts with legacy labels.{field_name}") + print_deprecation_message( + old_item=f"AttackResult.labels['{field_name}']", + new_item=f"AttackResult.{field_name}", + removed_in="1.4.0", + ) + attribution[field_name] = legacy_value + for field_name, value in attribution.items(): + if value is not None and len(value) > AttackResult.ATTRIBUTION_VALUE_MAX_LENGTH: + raise ValueError(f"{field_name} must be at most {AttackResult.ATTRIBUTION_VALUE_MAX_LENGTH} characters") + self.operator = attribution["operator"] + self.operation = attribution["operation"] + self.labels = labels self.targeted_harm_categories = entry.targeted_harm_categories or None # Persist conversation references by type @@ -1799,6 +1855,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, @@ -1850,10 +1908,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) @@ -1869,7 +1934,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..4dd35c363d 100644 --- a/pyrit/memory/sqlite_memory.py +++ b/pyrit/memory/sqlite_memory.py @@ -433,27 +433,46 @@ def get_conversation_stats(self, *, conversation_ids: Sequence[str]) -> dict[str sql = text( f""" + WITH filtered AS ( + SELECT + conversation_id, + sequence, + id, + timestamp, + converted_value, + converted_value_data_type + FROM "PromptMemoryEntries" + WHERE conversation_id IN ({placeholders}) + ), + aggregate_rows AS ( + SELECT + conversation_id, + COUNT(DISTINCT sequence) AS msg_count, + MIN(timestamp) AS created_at + FROM filtered + GROUP BY conversation_id + ), + latest_rows AS ( + SELECT + conversation_id, + SUBSTR(converted_value, 1, {ConversationStats.PREVIEW_FETCH_MAX_LEN}) AS last_preview, + converted_value_data_type AS last_data_type, + ROW_NUMBER() OVER ( + PARTITION BY conversation_id + ORDER BY sequence DESC, id DESC + ) AS row_number + FROM filtered + ) SELECT - pme.conversation_id, - COUNT(DISTINCT pme.sequence) AS msg_count, - ( - SELECT SUBSTR(p2.converted_value, 1, {ConversationStats.PREVIEW_FETCH_MAX_LEN}) - FROM "PromptMemoryEntries" p2 - WHERE p2.conversation_id = pme.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 + aggregate_rows.conversation_id, + aggregate_rows.msg_count, + latest_rows.last_preview, + latest_rows.last_data_type, + aggregate_rows.created_at + FROM aggregate_rows + LEFT JOIN latest_rows + ON latest_rows.conversation_id = aggregate_rows.conversation_id + AND latest_rows.row_number = 1 """ ) diff --git a/pyrit/models/results/attack_result.py b/pyrit/models/results/attack_result.py index 054dceda4d..17da154043 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 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 @@ -44,6 +45,9 @@ class AttackOutcome(str, Enum): class AttackResult(StrategyResult): """Base class for all attack results.""" + ATTRIBUTION_VALUE_MAX_LENGTH: ClassVar[int] = 128 + _LEGACY_ATTRIBUTION_FIELDS: ClassVar[tuple[str, str]] = ("operator", "operation") + # Identity # Unique identifier of the conversation that produced this result conversation_id: str @@ -90,6 +94,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 +124,60 @@ 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): + return data + + normalized = dict(data) + labels_value = normalized.get("labels") + if labels_value is None: + return normalized + if not isinstance(labels_value, dict): + return normalized + + labels = dict(labels_value) + for field_name in cls._LEGACY_ATTRIBUTION_FIELDS: + if field_name not in labels: + continue + legacy_value = labels.pop(field_name) + if not isinstance(legacy_value, str): + raise ValueError(f"labels.{field_name} must be a string") + dedicated_value = normalized.get(field_name) + if dedicated_value is not None and dedicated_value != legacy_value: + raise ValueError( + f"{field_name} conflicts with legacy labels.{field_name}: {dedicated_value!r} != {legacy_value!r}" + ) + print_deprecation_message( + old_item=f"AttackResult.labels['{field_name}']", + new_item=f"AttackResult.{field_name}", + removed_in="1.4.0", + ) + normalized[field_name] = legacy_value + + normalized["labels"] = labels + return normalized + + @field_serializer("labels") + def _serialize_arbitrary_labels(self, labels: dict[str, str]) -> dict[str, str]: + """ + Serialize only arbitrary labels, even if the mutable mapping was modified later. + + Returns: + The labels without attribution aliases. + """ + return {key: value for key, value in labels.items() if key not in self._LEGACY_ATTRIBUTION_FIELDS} + 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..f160b84f54 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,6 +1449,8 @@ def test_get_labels_for_attacks(self, client: TestClient) -> None: data = response.json() assert data["source"] == "attacks" assert data["labels"] == {"env": ["prod"], "team": ["red"]} + assert data["operators"] == [] + assert data["operations"] == [] mock_memory.get_unique_attack_labels.assert_called_once() def test_get_labels_empty(self, client: TestClient) -> None: @@ -1403,6 +1458,7 @@ def test_get_labels_empty(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 = {} + 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 +1476,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 +1487,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 +1501,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..43748ec0b1 100644 --- a/tests/unit/memory/memory_interface/test_interface_attack_results.py +++ b/tests/unit/memory/memory_interface/test_interface_attack_results.py @@ -39,6 +39,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 +48,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 +112,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 +156,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 +181,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 +1281,48 @@ def test_get_unique_attack_labels_deduplicates_across_attacks(sqlite_instance: M assert result == {"env": ["prod"]} +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 # ============================================================================ 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 ccbbb0f500..3c562f752a 100644 --- a/tests/unit/memory/test_migration.py +++ b/tests/unit/memory/test_migration.py @@ -2437,6 +2437,136 @@ 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 = "8d1e3f5a7b9c" + + +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") + } + + 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_conversation_timestamp_id"][0] == "operator" + assert attack_indexes["ix_AttackResultEntries_operation_conversation_timestamp_id"][0] == "operation" + assert prompt_indexes["ix_PromptMemoryEntries_conversation_sequence_id"] == [ + "conversation_id", + "sequence", + "id", + ] + 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_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"} diff --git a/tests/unit/models/test_attack_result.py b/tests/unit/models/test_attack_result.py index 06482db0d6..b551e467eb 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,50 @@ 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"} + class TestAttackResultDuplicate: """duplicate() must deep-copy so mutations on the copy never touch the original.""" From 420122ed567e5a5f1d3858d31a5ed47c4c51305e Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Wed, 9 Sep 2026 13:11:40 -0700 Subject: [PATCH 2/6] PERF: Refine attack history queries Replace the full-table dedup window with an indexed anti-join, widen indexed conversation IDs, centralize attribution normalization, and narrow label discovery by active filters. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 7540877a-bdf5-4309-97e1-14469bc817e7 --- .../components/History/AttackHistory.test.tsx | 33 +++ .../src/components/History/AttackHistory.tsx | 8 +- frontend/src/services/api.ts | 7 +- pyrit/backend/models/attacks.py | 29 +-- pyrit/backend/routes/labels.py | 23 +- ..._attack_attribution_and_history_indexes.py | 130 ++++++---- pyrit/memory/memory_interface.py | 234 +++++++++--------- pyrit/memory/memory_models.py | 47 ++-- pyrit/models/results/attack_result.py | 142 ++++++++--- tests/unit/backend/test_api_routes.py | 32 ++- .../test_interface_attack_results.py | 49 ++++ tests/unit/memory/test_migration.py | 16 +- tests/unit/models/test_attack_result.py | 5 +- 13 files changed, 486 insertions(+), 269 deletions(-) diff --git a/frontend/src/components/History/AttackHistory.test.tsx b/frontend/src/components/History/AttackHistory.test.tsx index 4e0f7e090e..c989f5005d 100644 --- a/frontend/src/components/History/AttackHistory.test.tsx +++ b/frontend/src/components/History/AttackHistory.test.tsx @@ -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 25b6c109df..2b6b4da4d6 100644 --- a/frontend/src/components/History/AttackHistory.tsx +++ b/frontend/src/components/History/AttackHistory.tsx @@ -107,7 +107,11 @@ export default function AttackHistory({ attacksApi.getConverterOptions() .then(resp => setConverterOptions(resp.converter_types)) .catch(() => { /* ignore */ }) - labelsApi.getLabels() + 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 => { const others: string[] = [] for (const [key, values] of Object.entries(resp.labels)) { @@ -122,7 +126,7 @@ export default function AttackHistory({ setOtherLabelOptions(others.sort()) }) .catch(() => { /* ignore */ }) - }, []) + }, [filters.operator, filters.operation, filters.otherLabels]) // Fetch attacks whenever filters change or an event handler bumps fetchToken. // All setState calls live in .then/.catch/.finally so we don't trigger diff --git a/frontend/src/services/api.ts b/frontend/src/services/api.ts index 032f5b726f..7d5f2b28a8 100644 --- a/frontend/src/services/api.ts +++ b/frontend/src/services/api.ts @@ -369,8 +369,13 @@ export const attacksApi = { export const labelsApi = { getLabels: async ( source: 'attacks' | 'scenarios' = 'attacks', + filters?: { + operator?: string[] + operation?: string[] + label?: string[] + }, ): Promise => { - const response = await apiClient.get('/labels', { params: { source } }) + const response = await apiClient.get('/labels', { params: { source, ...filters } }) return response.data }, } diff --git a/pyrit/backend/models/attacks.py b/pyrit/backend/models/attacks.py index 057014ce1e..f3c963542b 100644 --- a/pyrit/backend/models/attacks.py +++ b/pyrit/backend/models/attacks.py @@ -16,7 +16,6 @@ from pyrit.backend.models._media import build_filename, infer_mime_type from pyrit.backend.models.common import PaginationInfo -from pyrit.common.deprecation import print_deprecation_message from pyrit.models import ( AttackResult, ChatMessageRole, @@ -26,6 +25,7 @@ PromptDataType, Score, ) +from pyrit.models.results.attack_result import ATTRIBUTION_FIELDS, pop_legacy_attribution_labels class TargetInfo(BaseModel): @@ -381,23 +381,16 @@ def _normalize_legacy_attribution_labels(cls, data: Any) -> Any: if not isinstance(data, dict) or not isinstance(data.get("labels"), dict): return data normalized = dict(data) - labels = dict(normalized["labels"]) - for field_name in ("operator", "operation"): - if field_name not in labels: - continue - legacy_value = labels.pop(field_name) - if not isinstance(legacy_value, str): - raise ValueError(f"labels.{field_name} must be a string") - dedicated_value = normalized.get(field_name) - if dedicated_value is not None and dedicated_value != legacy_value: - raise ValueError(f"{field_name} conflicts with legacy labels.{field_name}") - print_deprecation_message( - old_item=f"labels.{field_name}", - new_item=field_name, - removed_in="1.4.0", - ) - normalized[field_name] = legacy_value - normalized["labels"] = labels + remaining, resolved = pop_legacy_attribution_labels( + labels=normalized["labels"], + dedicated={field: normalized.get(field) for field in ATTRIBUTION_FIELDS}, + allow_multiple=False, + old_item="labels.{field}", + new_item="{field}", + ) + for field, values in resolved.items(): + normalized[field] = values[0] if values else None + normalized["labels"] = remaining return normalized diff --git a/pyrit/backend/routes/labels.py b/pyrit/backend/routes/labels.py index 29318d3ba9..b4806b946a 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"]) @@ -37,6 +38,15 @@ async def get_label_options( # pyrit-async-suffix-exempt "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. @@ -46,6 +56,9 @@ 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. @@ -53,7 +66,13 @@ async def get_label_options( # pyrit-async-suffix-exempt memory = CentralMemory.get_memory_instance() if source == "attacks": - labels = await run_in_threadpool(memory.get_unique_attack_labels) + 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 = await run_in_threadpool(memory.get_unique_attack_attribution) return LabelOptionsResponse(source=source, labels=labels, **attribution) 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 index a605911a7e..0f2728c3f5 100644 --- 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 @@ -28,6 +28,7 @@ _ATTRIBUTION_FIELDS = ("operator", "operation") _ATTRIBUTION_MAX_LENGTH = 128 +_BATCH_SIZE = 1000 def upgrade() -> None: @@ -44,14 +45,16 @@ def upgrade() -> None: ["conversation_id", "timestamp", "id"], ) op.create_index( - "ix_AttackResultEntries_operator_conversation_timestamp_id", + "ix_AttackResultEntries_operator_timestamp_id", "AttackResultEntries", - ["operator", "conversation_id", "timestamp", "id"], + ["operator", "timestamp", "id"], + mssql_include=["conversation_id"], ) op.create_index( - "ix_AttackResultEntries_operation_conversation_timestamp_id", + "ix_AttackResultEntries_operation_timestamp_id", "AttackResultEntries", - ["operation", "conversation_id", "timestamp", "id"], + ["operation", "timestamp", "id"], + mssql_include=["conversation_id"], ) _drop_index_if_exists(name="idx_conversation_id", table_name="PromptMemoryEntries") @@ -93,11 +96,11 @@ def downgrade() -> None: ) op.drop_index( - "ix_AttackResultEntries_operation_conversation_timestamp_id", + "ix_AttackResultEntries_operation_timestamp_id", table_name="AttackResultEntries", ) op.drop_index( - "ix_AttackResultEntries_operator_conversation_timestamp_id", + "ix_AttackResultEntries_operator_timestamp_id", table_name="AttackResultEntries", ) op.drop_index( @@ -146,14 +149,14 @@ def _drop_index_if_exists(*, name: str, table_name: str) -> None: 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=36) + _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(36), + type_=sa.String(128), existing_nullable=False, ) with op.batch_alter_table("ScenarioResultEntries") as batch_op: @@ -189,7 +192,7 @@ def _restore_unbounded_text_columns() -> None: with op.batch_alter_table("PromptMemoryEntries") as batch_op: batch_op.alter_column( "conversation_id", - existing_type=sa.String(36), + existing_type=sa.String(128), type_=sa.String(), existing_nullable=False, ) @@ -209,11 +212,7 @@ def _validate_column_length(*, table_name: str, column_name: str, max_length: in ) oversized_value = ( op.get_bind() - .execute( - sa.select(table.c[column_name]) - .where(sa.func.length(table.c[column_name]) > max_length) - .limit(1) - ) + .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: @@ -232,31 +231,51 @@ def _move_attribution_from_labels() -> None: """ 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"), + operator=sa.bindparam("new_operator"), + operation=sa.bindparam("new_operation"), + ) + ) rows = bind.execute(sa.select(table.c.id, table.c.labels)).all() - for row in rows: - 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: + for start in range(0, len(rows), _BATCH_SIZE): + updates = [] + for row in rows[start : start + _BATCH_SIZE]: + labels = row.labels + if not isinstance(labels, dict): 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." + 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 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: - values["labels"] = remaining_labels - bind.execute(sa.update(table).where(table.c.id == row.id).values(**values)) + if updates: + bind.execute(statement, updates) def _restore_attribution_to_labels() -> None: @@ -268,21 +287,26 @@ def _restore_attribution_to_labels() -> None: """ 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() - for row in rows: - 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: - bind.execute(sa.update(table).where(table.c.id == row.id).values(labels=labels)) + for start in range(0, len(rows), _BATCH_SIZE): + 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) diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index d7152dd496..d2f8c2bfc8 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -18,15 +18,13 @@ 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 @@ -95,6 +93,11 @@ group_conversation_message_pieces_by_sequence, sort_message_pieces, ) +from pyrit.models.results.attack_result import ( + ATTRIBUTION_FIELDS, + normalize_attribution_values, + pop_legacy_attribution_labels, +) if TYPE_CHECKING: from sqlalchemy.sql.elements import ColumnElement @@ -229,8 +232,6 @@ class _AttackResultQuery: "converter_classes", "targeted_harm_categories", "identifier_filters", - "operator", - "operation", ) attack_result_ids: Sequence[str] | None = None @@ -245,8 +246,8 @@ class _AttackResultQuery: has_converters: bool | None = None include_scenario_attacks: bool = True labels: Mapping[str, str | Sequence[str]] | None = None - operator: Sequence[str] | None = None - operation: 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 @@ -267,37 +268,26 @@ def __post_init__(self) -> None: if value is not None: object.__setattr__(self, field_name, tuple(value)) - for field_name in ("operator", "operation"): + for field_name in ATTRIBUTION_FIELDS: values = getattr(self, field_name) - if values is not None and any(not isinstance(value, str) for value in values): - raise ValueError(f"{field_name} values must be strings") - if values is not None and any(len(value) > AttackResult.ATTRIBUTION_VALUE_MAX_LENGTH for value in values): - raise ValueError( - f"{field_name} values must be at most {AttackResult.ATTRIBUTION_VALUE_MAX_LENGTH} characters" + if values is not None: + object.__setattr__( + self, + field_name, + normalize_attribution_values(field=field_name, raw=values, allow_multiple=True), ) if self.labels is not None: - labels = {key: value if isinstance(value, str) else tuple(value) for key, value in self.labels.items()} - for name in ("operator", "operation"): - if name not in labels: - continue - legacy_raw = labels.pop(name) - legacy_values = (legacy_raw,) if isinstance(legacy_raw, str) else tuple(legacy_raw) - if any(not isinstance(value, str) for value in legacy_values): - raise ValueError(f"labels.{name} values must be strings") - if any(len(value) > AttackResult.ATTRIBUTION_VALUE_MAX_LENGTH for value in legacy_values): - raise ValueError( - f"labels.{name} values must be at most {AttackResult.ATTRIBUTION_VALUE_MAX_LENGTH} characters" - ) - dedicated_values = getattr(self, name) - if dedicated_values is not None and set(dedicated_values) != set(legacy_values): - raise ValueError(f"{name} conflicts with legacy labels.{name}") - print_deprecation_message( - old_item=f"_AttackResultQuery.labels['{name}']", - new_item=f"_AttackResultQuery.{name}", - removed_in="1.4.0", - ) - object.__setattr__(self, name, legacy_values) + remaining, resolved = pop_legacy_attribution_labels( + labels=self.labels, + dedicated={field: getattr(self, field) for field in ATTRIBUTION_FIELDS}, + allow_multiple=True, + old_item="_AttackResultQuery.labels['{field}']", + new_item="_AttackResultQuery.{field}", + ) + for field_name, values in resolved.items(): + object.__setattr__(self, field_name, values) + labels = {key: value if isinstance(value, str) else tuple(value) for key, value in remaining.items()} object.__setattr__(self, "labels", MappingProxyType(labels) if labels else None) @@ -3585,7 +3575,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 @@ -3601,11 +3591,6 @@ def get_attack_results( ValueError: If ``limit`` or ``after`` is combined with ``attack_result_ids`` or ``objective_sha256`` (id-batched lookups do not support SQL pagination). """ - labels, operator_values, operation_values = self._normalize_attack_attribution_filters( - labels=labels, - operator=operator, - operation=operation, - ) query = _AttackResultQuery( attack_result_ids=attack_result_ids, conversation_id=conversation_id, @@ -3619,8 +3604,8 @@ def get_attack_results( has_converters=has_converters, include_scenario_attacks=include_scenario_attacks, labels=labels, - operator=operator_values, - operation=operation_values, + operator=operator, + operation=operation, targeted_harm_categories=targeted_harm_categories, identifier_filters=identifier_filters, scenario_result_id=scenario_result_id, @@ -3631,55 +3616,6 @@ def get_attack_results( ) return self._query_attack_results(query=query) - @staticmethod - def _normalize_attack_attribution_filters( - *, - labels: Mapping[str, str | Sequence[str]] | None, - operator: str | Sequence[str] | None, - operation: str | Sequence[str] | None, - ) -> tuple[Mapping[str, str | Sequence[str]] | None, Sequence[str] | None, Sequence[str] | None]: - """ - Normalize deprecated attribution label aliases without mutating caller input. - - Returns: - The arbitrary labels, operator values, and operation values. - - Raises: - ValueError: If a legacy alias conflicts with its dedicated filter. - """ - operator_values = [operator] if isinstance(operator, str) else operator - operation_values = [operation] if isinstance(operation, str) else operation - for field_name, values in (("operator", operator_values), ("operation", operation_values)): - if values is not None and any(len(value) > AttackResult.ATTRIBUTION_VALUE_MAX_LENGTH for value in values): - raise ValueError( - f"{field_name} values must be at most {AttackResult.ATTRIBUTION_VALUE_MAX_LENGTH} characters" - ) - if not labels: - return labels, operator_values, operation_values - - normalized_labels = dict(labels) - normalized_dedicated = {"operator": operator_values, "operation": operation_values} - for name in ("operator", "operation"): - if name not in normalized_labels: - continue - legacy_raw = normalized_labels.pop(name) - legacy_values = [legacy_raw] if isinstance(legacy_raw, str) else list(legacy_raw) - dedicated_values = normalized_dedicated[name] - if dedicated_values is not None and set(dedicated_values) != set(legacy_values): - raise ValueError(f"{name} conflicts with legacy labels.{name}") - print_deprecation_message( - old_item=f"get_attack_results(labels={{'{name}': ...}})", - new_item=f"get_attack_results({name}=...)", - removed_in="1.4.0", - ) - normalized_dedicated[name] = legacy_values - - return ( - normalized_labels or None, - normalized_dedicated["operator"], - normalized_dedicated["operation"], - ) - def _query_attack_results(self, *, query: _AttackResultQuery) -> Sequence[AttackResult]: """ Retrieve attack results matching an immutable query. @@ -3940,17 +3876,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. @@ -3962,22 +3910,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: @@ -3993,6 +3927,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 @@ -4038,20 +4014,36 @@ 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): diff --git a/pyrit/memory/memory_models.py b/pyrit/memory/memory_models.py index 5f2abb8965..a04092fa60 100644 --- a/pyrit/memory/memory_models.py +++ b/pyrit/memory/memory_models.py @@ -35,7 +35,6 @@ from typing_extensions import Self import pyrit -from pyrit.common.deprecation import print_deprecation_message from pyrit.common.utils import to_sha256 from pyrit.models import ( SEED_RESPONSE_JSON_SCHEMA_METADATA_KEY, @@ -73,6 +72,7 @@ TargetIdentifier, scorable_from_dict, ) +from pyrit.models.results.attack_result import pop_legacy_attribution_labels logger = logging.getLogger(__name__) @@ -273,7 +273,9 @@ class PromptMemoryEntry(Base): role: Mapped[Literal["system", "user", "assistant", "simulated_assistant", "tool", "developer"]] = mapped_column( String, nullable=False ) - conversation_id = mapped_column(String(36), 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) @@ -1577,18 +1579,18 @@ class AttackResultEntry(Base): # 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_conversation_timestamp_id", + "ix_AttackResultEntries_operator_timestamp_id", "operator", - "conversation_id", "timestamp", "id", + mssql_include=["conversation_id"], ), Index( - "ix_AttackResultEntries_operation_conversation_timestamp_id", + "ix_AttackResultEntries_operation_timestamp_id", "operation", - "conversation_id", "timestamp", "id", + mssql_include=["conversation_id"], ), # Serves scenario progress deltas scoped by parent and ordered oldest-first. Index( @@ -1701,29 +1703,16 @@ 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) - labels = dict(entry.labels or {}) - attribution = {"operator": entry.operator, "operation": entry.operation} - for field_name in ("operator", "operation"): - if field_name not in labels: - continue - legacy_value = labels.pop(field_name) - dedicated_value = attribution[field_name] - if not isinstance(legacy_value, str): - raise ValueError(f"labels.{field_name} must be a string") - if dedicated_value is not None and dedicated_value != legacy_value: - raise ValueError(f"{field_name} conflicts with legacy labels.{field_name}") - print_deprecation_message( - old_item=f"AttackResult.labels['{field_name}']", - new_item=f"AttackResult.{field_name}", - removed_in="1.4.0", - ) - attribution[field_name] = legacy_value - for field_name, value in attribution.items(): - if value is not None and len(value) > AttackResult.ATTRIBUTION_VALUE_MAX_LENGTH: - raise ValueError(f"{field_name} must be at most {AttackResult.ATTRIBUTION_VALUE_MAX_LENGTH} characters") - self.operator = attribution["operator"] - self.operation = attribution["operation"] - self.labels = labels + remaining_labels, resolved = pop_legacy_attribution_labels( + labels=entry.labels or {}, + dedicated={"operator": entry.operator, "operation": entry.operation}, + allow_multiple=False, + old_item="AttackResult.labels['{field}']", + new_item="AttackResult.{field}", + ) + self.operator = resolved["operator"][0] if resolved["operator"] else None + self.operation = resolved["operation"][0] if resolved["operation"] else None + self.labels = remaining_labels self.targeted_harm_categories = entry.targeted_harm_categories or None # Persist conversation references by type diff --git a/pyrit/models/results/attack_result.py b/pyrit/models/results/attack_result.py index 17da154043..c394378b91 100644 --- a/pyrit/models/results/attack_result.py +++ b/pyrit/models/results/attack_result.py @@ -4,6 +4,7 @@ from __future__ import annotations import uuid +from collections.abc import Mapping, Sequence from datetime import datetime, timezone from enum import Enum from typing import Any, ClassVar, TypeVar @@ -20,6 +21,96 @@ AttackResultT = TypeVar("AttackResultT", bound="AttackResult") +ATTRIBUTION_FIELDS: tuple[str, str] = ("operator", "operation") +ATTRIBUTION_VALUE_MAX_LENGTH: int = 128 + + +def normalize_attribution_values(*, field: str, raw: Any, allow_multiple: bool) -> tuple[str, ...]: + """ + Validate one attribution value, or a sequence of them, as bounded strings. + + Args: + field (str): Name used in error messages. + raw (Any): A single value or a sequence of values. + allow_multiple (bool): Whether a sequence of values is valid. + + Returns: + tuple[str, ...]: The validated values. + + Raises: + ValueError: If a value is not a string or exceeds the maximum length. + """ + if isinstance(raw, str): + values: tuple[Any, ...] = (raw,) + elif allow_multiple and isinstance(raw, Sequence): + values = tuple(raw) + else: + expected = "a string or a sequence of strings" if allow_multiple else "a string" + raise ValueError(f"{field} must be {expected}") + if any(not isinstance(value, str) for value in values): + expected = "strings" if allow_multiple else "a string" + raise ValueError(f"{field} must contain {expected}") + if any(len(value) > ATTRIBUTION_VALUE_MAX_LENGTH for value in values): + raise ValueError(f"{field} must be at most {ATTRIBUTION_VALUE_MAX_LENGTH} characters") + return values + + +def pop_legacy_attribution_labels( + *, + labels: Mapping[str, Any], + dedicated: Mapping[str, Any], + allow_multiple: bool, + old_item: str, + new_item: str, +) -> tuple[dict[str, Any], dict[str, tuple[str, ...] | None]]: + """ + Move legacy ``operator``/``operation`` label aliases onto their dedicated values. + + ``operator`` and ``operation`` used to live in the free-form ``labels`` mapping. They are + now indexed columns, so every entry point accepts the old spelling for one more release + and funnels it here. Both ``old_item`` and ``new_item`` are format strings taking a + ``field`` placeholder, so each caller reports the deprecation in its own vocabulary. + + Args: + labels (Mapping[str, Any]): Labels that may still carry the legacy aliases. + dedicated (Mapping[str, Any]): Current dedicated values, keyed by field name. + allow_multiple (bool): Whether each attribution field can contain multiple values. + old_item (str): Deprecation message template for the old spelling. + new_item (str): Deprecation message template for the replacement. + + Returns: + tuple[dict[str, Any], dict[str, tuple[str, ...] | None]]: The labels with the aliases + removed, and the resolved values per attribution field. + + Raises: + ValueError: If an alias is invalid or disagrees with its dedicated value. + """ + remaining = dict(labels) + resolved: dict[str, tuple[str, ...] | None] = {} + for field in ATTRIBUTION_FIELDS: + current = dedicated.get(field) + values = ( + None + if current is None + else normalize_attribution_values(field=field, raw=current, allow_multiple=allow_multiple) + ) + if field in remaining: + legacy_values = normalize_attribution_values( + field=f"labels.{field}", + raw=remaining.pop(field), + allow_multiple=allow_multiple, + ) + if values is not None and set(values) != set(legacy_values): + raise ValueError(f"{field} conflicts with legacy labels.{field}: {values!r} != {legacy_values!r}") + print_deprecation_message( + old_item=old_item.format(field=field), + new_item=new_item.format(field=field), + removed_in="1.4.0", + ) + values = legacy_values + resolved[field] = values + return remaining, resolved + class AttackOutcome(str, Enum): """ @@ -45,8 +136,7 @@ class AttackOutcome(str, Enum): class AttackResult(StrategyResult): """Base class for all attack results.""" - ATTRIBUTION_VALUE_MAX_LENGTH: ClassVar[int] = 128 - _LEGACY_ATTRIBUTION_FIELDS: ClassVar[tuple[str, str]] = ("operator", "operation") + ATTRIBUTION_VALUE_MAX_LENGTH: ClassVar[int] = ATTRIBUTION_VALUE_MAX_LENGTH # Identity # Unique identifier of the conversation that produced this result @@ -136,48 +226,22 @@ def _normalize_legacy_attribution_labels(cls, data: Any) -> Any: Raises: ValueError: If an alias is not a string or conflicts with a dedicated field. """ - if not isinstance(data, dict): + if not isinstance(data, dict) or not isinstance(data.get("labels"), dict): return data normalized = dict(data) - labels_value = normalized.get("labels") - if labels_value is None: - return normalized - if not isinstance(labels_value, dict): - return normalized - - labels = dict(labels_value) - for field_name in cls._LEGACY_ATTRIBUTION_FIELDS: - if field_name not in labels: - continue - legacy_value = labels.pop(field_name) - if not isinstance(legacy_value, str): - raise ValueError(f"labels.{field_name} must be a string") - dedicated_value = normalized.get(field_name) - if dedicated_value is not None and dedicated_value != legacy_value: - raise ValueError( - f"{field_name} conflicts with legacy labels.{field_name}: {dedicated_value!r} != {legacy_value!r}" - ) - print_deprecation_message( - old_item=f"AttackResult.labels['{field_name}']", - new_item=f"AttackResult.{field_name}", - removed_in="1.4.0", - ) - normalized[field_name] = legacy_value - - normalized["labels"] = labels + remaining, resolved = pop_legacy_attribution_labels( + labels=normalized["labels"], + dedicated={field: normalized.get(field) for field in ATTRIBUTION_FIELDS}, + allow_multiple=False, + old_item="AttackResult.labels['{field}']", + new_item="AttackResult.{field}", + ) + for field, values in resolved.items(): + normalized[field] = values[0] if values else None + normalized["labels"] = remaining return normalized - @field_serializer("labels") - def _serialize_arbitrary_labels(self, labels: dict[str, str]) -> dict[str, str]: - """ - Serialize only arbitrary labels, even if the mutable mapping was modified later. - - Returns: - The labels without attribution aliases. - """ - return {key: value for key, value in labels.items() if key not in self._LEGACY_ATTRIBUTION_FIELDS} - 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 f160b84f54..823fc15ca8 100644 --- a/tests/unit/backend/test_api_routes.py +++ b/tests/unit/backend/test_api_routes.py @@ -1451,7 +1451,37 @@ def test_get_labels_for_attacks(self, client: TestClient) -> None: assert data["labels"] == {"env": ["prod"], "team": ["red"]} assert data["operators"] == [] assert data["operations"] == [] - mock_memory.get_unique_attack_labels.assert_called_once() + 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"]}, + ) def test_get_labels_empty(self, client: TestClient) -> None: """Test getting labels when no attack results exist.""" 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 43748ec0b1..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 @@ -1281,6 +1283,42 @@ 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"), @@ -1843,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/test_migration.py b/tests/unit/memory/test_migration.py index 3c562f752a..454346ed9a 100644 --- a/tests/unit/memory/test_migration.py +++ b/tests/unit/memory/test_migration.py @@ -2486,6 +2486,9 @@ def test_attack_attribution_migration_backfills_labels_and_indexes() -> None: 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" @@ -2495,13 +2498,22 @@ def test_attack_attribution_migration_backfills_labels_and_indexes() -> None: "timestamp", "id", ] - assert attack_indexes["ix_AttackResultEntries_operator_conversation_timestamp_id"][0] == "operator" - assert attack_indexes["ix_AttackResultEntries_operation_conversation_timestamp_id"][0] == "operation" + 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", diff --git a/tests/unit/models/test_attack_result.py b/tests/unit/models/test_attack_result.py index b551e467eb..ac7fca40f3 100644 --- a/tests/unit/models/test_attack_result.py +++ b/tests/unit/models/test_attack_result.py @@ -395,7 +395,10 @@ def test_dedicated_attribution_is_canonical(self) -> None: assert dumped["labels"] == {"team": "red"} result.labels["operator"] = "legacy-mutation" - assert result.model_dump(mode="json")["labels"] == {"team": "red"} + assert result.model_dump(mode="json")["labels"] == { + "team": "red", + "operator": "legacy-mutation", + } class TestAttackResultDuplicate: From 8493eca9c683def95eda015a85fa384f3757585d Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Wed, 9 Sep 2026 14:04:05 -0700 Subject: [PATCH 3/6] PERF: Streamline history filter compatibility Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 7540877a-bdf5-4309-97e1-14469bc817e7 --- .../src/components/Chat/ChatWindow.test.tsx | 3 +- frontend/src/components/Chat/ChatWindow.tsx | 21 +--- .../src/components/History/AttackHistory.tsx | 72 +++++++++--- frontend/src/components/Labels/LabelsBar.tsx | 1 + frontend/src/services/api.test.ts | 29 +++++ frontend/src/services/api.ts | 7 +- pyrit/backend/models/attacks.py | 16 +-- pyrit/backend/routes/attacks.py | 1 + pyrit/backend/routes/labels.py | 6 +- pyrit/memory/memory_interface.py | 57 ++++++--- pyrit/memory/memory_models.py | 14 +-- pyrit/memory/sqlite_memory.py | 42 +++---- pyrit/models/results/attack_result.py | 110 ++++++------------ tests/unit/backend/test_api_routes.py | 1 + tests/unit/memory/test_sqlite_memory.py | 20 +++- 15 files changed, 229 insertions(+), 171 deletions(-) diff --git a/frontend/src/components/Chat/ChatWindow.test.tsx b/frontend/src/components/Chat/ChatWindow.test.tsx index 52c5dd4147..33bdd8d390 100644 --- a/frontend/src/components/Chat/ChatWindow.test.tsx +++ b/frontend/src/components/Chat/ChatWindow.test.tsx @@ -705,8 +705,7 @@ describe("ChatWindow Integration", () => { await waitFor(() => { expect(mockedAttacksApi.createAttack).toHaveBeenCalledWith({ target_registry_name: "openai_chat_1", - operator: 'testuser', - operation: 'test_op', + labels: { operator: 'testuser', operation: 'test_op' }, system_prompt: undefined, }); expect(onConversationCreated).toHaveBeenCalledWith("ar-conv-1", "conv-1"); diff --git a/frontend/src/components/Chat/ChatWindow.tsx b/frontend/src/components/Chat/ChatWindow.tsx index 7dfc6c91f3..075b786cba 100644 --- a/frontend/src/components/Chat/ChatWindow.tsx +++ b/frontend/src/components/Chat/ChatWindow.tsx @@ -80,19 +80,6 @@ function matchesNarrowScreen(): boolean { && window.matchMedia(NARROW_SCREEN_QUERY).matches } -function attackAttributionFromLabels(labels?: Record): Pick< - CreateAttackRequest, - 'operator' | 'operation' | 'labels' -> { - if (!labels) return {} - const { operator, operation, ...arbitraryLabels } = labels - const attribution: Pick = {} - if (operator) attribution.operator = operator - if (operation) attribution.operation = operation - if (Object.keys(arbitraryLabels).length > 0) attribution.labels = arbitraryLabels - return attribution -} - interface ChatWindowProps { onNewAttack: () => void activeTarget: TargetInstance | null @@ -444,7 +431,9 @@ export default function ChatWindow({ if (!currentAttackResultId) { const createRequest: CreateAttackRequest = { target_registry_name: activeTarget.target_registry_name, - ...attackAttributionFromLabels(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) @@ -671,7 +660,7 @@ export default function ChatWindow({ try { const createResponse = await attacksApi.createAttack({ target_registry_name: activeTarget.target_registry_name, - ...attackAttributionFromLabels(labels), + labels, source_conversation_id: activeConversationId, cutoff_index: messageIndex, }) @@ -722,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, - ...attackAttributionFromLabels(labels), + labels, source_conversation_id: activeConversationId, cutoff_index: lastIndex, }) diff --git a/frontend/src/components/History/AttackHistory.tsx b/frontend/src/components/History/AttackHistory.tsx index 2b6b4da4d6..12b7e83dd3 100644 --- a/frontend/src/components/History/AttackHistory.tsx +++ b/frontend/src/components/History/AttackHistory.tsx @@ -48,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, @@ -66,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) @@ -84,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. @@ -99,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)) @@ -107,26 +134,43 @@ export default function AttackHistory({ attacksApi.getConverterOptions() .then(resp => setConverterOptions(resp.converter_types)) .catch(() => { /* ignore */ }) + labelsApi.getLabels() + .then(resp => { + // 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 => { - const others: string[] = [] - for (const [key, values] of Object.entries(resp.labels)) { - if (key !== 'source') { - for (const val of values) { - others.push(`${key}:${val}`) - } - } - } - setOperatorOptions([...(resp.operators ?? resp.labels.operator ?? [])].sort()) - setOperationOptions([...(resp.operations ?? resp.labels.operation ?? [])].sort()) - setOtherLabelOptions(others.sort()) + if (cancelled) return + setNarrowedOtherLabelOptions({ + filterKey: labelOptionFilterKey, + options: buildOtherLabelOptions(resp.labels), + }) }) .catch(() => { /* ignore */ }) - }, [filters.operator, filters.operation, filters.otherLabels]) + 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 diff --git a/frontend/src/components/Labels/LabelsBar.tsx b/frontend/src/components/Labels/LabelsBar.tsx index 264eade0f4..fa9757b76f 100644 --- a/frontend/src/components/Labels/LabelsBar.tsx +++ b/frontend/src/components/Labels/LabelsBar.tsx @@ -219,6 +219,7 @@ export default function LabelsBar({ labels, onLabelsChange }: LabelsBarProps) { // so keep anything already collected rather than replacing outright. .then(resp => setExistingLabels(prev => ({ ...resp.labels, + // 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 || [])])], }))) 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 7d5f2b28a8..b4f0a77be2 100644 --- a/frontend/src/services/api.ts +++ b/frontend/src/services/api.ts @@ -375,7 +375,12 @@ export const labelsApi = { label?: string[] }, ): Promise => { - const response = await apiClient.get('/labels', { params: { source, ...filters } }) + 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/pyrit/backend/models/attacks.py b/pyrit/backend/models/attacks.py index f3c963542b..ef1cfd5f5e 100644 --- a/pyrit/backend/models/attacks.py +++ b/pyrit/backend/models/attacks.py @@ -25,7 +25,7 @@ PromptDataType, Score, ) -from pyrit.models.results.attack_result import ATTRIBUTION_FIELDS, pop_legacy_attribution_labels +from pyrit.models.results.attack_result import normalize_legacy_attack_attribution class TargetInfo(BaseModel): @@ -372,6 +372,8 @@ 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. @@ -381,16 +383,14 @@ def _normalize_legacy_attribution_labels(cls, data: Any) -> Any: if not isinstance(data, dict) or not isinstance(data.get("labels"), dict): return data normalized = dict(data) - remaining, resolved = pop_legacy_attribution_labels( + remaining, operator, operation = normalize_legacy_attack_attribution( labels=normalized["labels"], - dedicated={field: normalized.get(field) for field in ATTRIBUTION_FIELDS}, - allow_multiple=False, - old_item="labels.{field}", - new_item="{field}", + operator=normalized.get("operator"), + operation=normalized.get("operation"), ) - for field, values in resolved.items(): - normalized[field] = values[0] if values else None normalized["labels"] = remaining + normalized["operator"] = operator + normalized["operation"] = operation return normalized diff --git a/pyrit/backend/routes/attacks.py b/pyrit/backend/routes/attacks.py index 62f135d9bc..b023714074 100644 --- a/pyrit/backend/routes/attacks.py +++ b/pyrit/backend/routes/attacks.py @@ -110,6 +110,7 @@ async def list_attacks( # pyrit-async-suffix-exempt """ service = get_attack_service() 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: diff --git a/pyrit/backend/routes/labels.py b/pyrit/backend/routes/labels.py index b4806b946a..520192e9f4 100644 --- a/pyrit/backend/routes/labels.py +++ b/pyrit/backend/routes/labels.py @@ -73,7 +73,11 @@ async def get_label_options( # pyrit-async-suffix-exempt operation=operation, labels=label_filters, ) - attribution = await run_in_threadpool(memory.get_unique_attack_attribution) + 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) diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index d2f8c2bfc8..7b46e7bf1e 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -25,6 +25,8 @@ 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,11 +95,7 @@ group_conversation_message_pieces_by_sequence, sort_message_pieces, ) -from pyrit.models.results.attack_result import ( - ATTRIBUTION_FIELDS, - normalize_attribution_values, - pop_legacy_attribution_labels, -) +from pyrit.models.results.attack_result import ATTRIBUTION_FIELDS, ATTRIBUTION_VALUE_MAX_LENGTH if TYPE_CHECKING: from sqlalchemy.sql.elements import ColumnElement @@ -109,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.""" @@ -260,6 +276,8 @@ def __post_init__(self) -> None: """ 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. """ @@ -274,20 +292,27 @@ def __post_init__(self) -> None: object.__setattr__( self, field_name, - normalize_attribution_values(field=field_name, raw=values, allow_multiple=True), + _normalize_attribution_filter_values(field=field_name, raw=values), ) if self.labels is not None: - remaining, resolved = pop_legacy_attribution_labels( - labels=self.labels, - dedicated={field: getattr(self, field) for field in ATTRIBUTION_FIELDS}, - allow_multiple=True, - old_item="_AttackResultQuery.labels['{field}']", - new_item="_AttackResultQuery.{field}", - ) - for field_name, values in resolved.items(): - object.__setattr__(self, field_name, values) - labels = {key: value if isinstance(value, str) else tuple(value) for key, value in remaining.items()} + labels = {key: value if isinstance(value, str) else tuple(value) for key, value in self.labels.items()} + 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) diff --git a/pyrit/memory/memory_models.py b/pyrit/memory/memory_models.py index a04092fa60..f7314d6dec 100644 --- a/pyrit/memory/memory_models.py +++ b/pyrit/memory/memory_models.py @@ -72,7 +72,7 @@ TargetIdentifier, scorable_from_dict, ) -from pyrit.models.results.attack_result import pop_legacy_attribution_labels +from pyrit.models.results.attack_result import normalize_legacy_attack_attribution logger = logging.getLogger(__name__) @@ -1703,15 +1703,13 @@ 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) - remaining_labels, resolved = pop_legacy_attribution_labels( + remaining_labels, operator, operation = normalize_legacy_attack_attribution( labels=entry.labels or {}, - dedicated={"operator": entry.operator, "operation": entry.operation}, - allow_multiple=False, - old_item="AttackResult.labels['{field}']", - new_item="AttackResult.{field}", + operator=entry.operator, + operation=entry.operation, ) - self.operator = resolved["operator"][0] if resolved["operator"] else None - self.operation = resolved["operation"][0] if resolved["operation"] else None + self.operator = operator + self.operation = operation self.labels = remaining_labels self.targeted_harm_categories = entry.targeted_harm_categories or None diff --git a/pyrit/memory/sqlite_memory.py b/pyrit/memory/sqlite_memory.py index 4dd35c363d..52382814c6 100644 --- a/pyrit/memory/sqlite_memory.py +++ b/pyrit/memory/sqlite_memory.py @@ -433,46 +433,30 @@ def get_conversation_stats(self, *, conversation_ids: Sequence[str]) -> dict[str sql = text( f""" - WITH filtered AS ( - SELECT - conversation_id, - sequence, - id, - timestamp, - converted_value, - converted_value_data_type - FROM "PromptMemoryEntries" - WHERE conversation_id IN ({placeholders}) - ), - aggregate_rows AS ( + WITH aggregate_rows AS ( SELECT conversation_id, COUNT(DISTINCT sequence) AS msg_count, MIN(timestamp) AS created_at - FROM filtered + FROM "PromptMemoryEntries" + WHERE conversation_id IN ({placeholders}) GROUP BY conversation_id - ), - latest_rows AS ( - SELECT - conversation_id, - SUBSTR(converted_value, 1, {ConversationStats.PREVIEW_FETCH_MAX_LEN}) AS last_preview, - converted_value_data_type AS last_data_type, - ROW_NUMBER() OVER ( - PARTITION BY conversation_id - ORDER BY sequence DESC, id DESC - ) AS row_number - FROM filtered ) SELECT aggregate_rows.conversation_id, aggregate_rows.msg_count, - latest_rows.last_preview, - latest_rows.last_data_type, + 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 latest_rows - ON latest_rows.conversation_id = aggregate_rows.conversation_id - AND latest_rows.row_number = 1 + LEFT JOIN "PromptMemoryEntries" latest + ON latest.id = ( + SELECT p2.id + FROM "PromptMemoryEntries" p2 + WHERE p2.conversation_id = aggregate_rows.conversation_id + ORDER BY p2.sequence DESC, p2.id DESC + LIMIT 1 + ) """ ) diff --git a/pyrit/models/results/attack_result.py b/pyrit/models/results/attack_result.py index c394378b91..7b034a3095 100644 --- a/pyrit/models/results/attack_result.py +++ b/pyrit/models/results/attack_result.py @@ -4,10 +4,9 @@ from __future__ import annotations import uuid -from collections.abc import Mapping, Sequence from datetime import datetime, timezone from enum import Enum -from typing import Any, ClassVar, TypeVar +from typing import TYPE_CHECKING, Any, ClassVar, TypeVar from pydantic import AwareDatetime, Field, field_serializer, model_validator @@ -19,97 +18,60 @@ 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_attribution_values(*, field: str, raw: Any, allow_multiple: bool) -> tuple[str, ...]: - """ - Validate one attribution value, or a sequence of them, as bounded strings. - - Args: - field (str): Name used in error messages. - raw (Any): A single value or a sequence of values. - allow_multiple (bool): Whether a sequence of values is valid. - - Returns: - tuple[str, ...]: The validated values. - - Raises: - ValueError: If a value is not a string or exceeds the maximum length. - """ - if isinstance(raw, str): - values: tuple[Any, ...] = (raw,) - elif allow_multiple and isinstance(raw, Sequence): - values = tuple(raw) - else: - expected = "a string or a sequence of strings" if allow_multiple else "a string" - raise ValueError(f"{field} must be {expected}") - if any(not isinstance(value, str) for value in values): - expected = "strings" if allow_multiple else "a string" - raise ValueError(f"{field} must contain {expected}") - if any(len(value) > ATTRIBUTION_VALUE_MAX_LENGTH for value in values): - raise ValueError(f"{field} must be at most {ATTRIBUTION_VALUE_MAX_LENGTH} characters") - return values - - -def pop_legacy_attribution_labels( +def normalize_legacy_attack_attribution( *, labels: Mapping[str, Any], - dedicated: Mapping[str, Any], - allow_multiple: bool, - old_item: str, - new_item: str, -) -> tuple[dict[str, Any], dict[str, tuple[str, ...] | None]]: + operator: str | None, + operation: str | None, +) -> tuple[dict[str, Any], str | None, str | None]: """ - Move legacy ``operator``/``operation`` label aliases onto their dedicated values. + Move scalar legacy label aliases to dedicated attribution fields. - ``operator`` and ``operation`` used to live in the free-form ``labels`` mapping. They are - now indexed columns, so every entry point accepts the old spelling for one more release - and funnels it here. Both ``old_item`` and ``new_item`` are format strings taking a - ``field`` placeholder, so each caller reports the deprecation in its own vocabulary. + 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. - dedicated (Mapping[str, Any]): Current dedicated values, keyed by field name. - allow_multiple (bool): Whether each attribution field can contain multiple values. - old_item (str): Deprecation message template for the old spelling. - new_item (str): Deprecation message template for the replacement. + operator (str | None): Dedicated operator value. + operation (str | None): Dedicated operation value. Returns: - tuple[dict[str, Any], dict[str, tuple[str, ...] | None]]: The labels with the aliases - removed, and the resolved values per attribution field. + 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, tuple[str, ...] | None] = {} + resolved: dict[str, Any] = {"operator": operator, "operation": operation} for field in ATTRIBUTION_FIELDS: - current = dedicated.get(field) - values = ( - None - if current is None - else normalize_attribution_values(field=field, raw=current, allow_multiple=allow_multiple) - ) + 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_values = normalize_attribution_values( - field=f"labels.{field}", - raw=remaining.pop(field), - allow_multiple=allow_multiple, - ) - if values is not None and set(values) != set(legacy_values): - raise ValueError(f"{field} conflicts with legacy labels.{field}: {values!r} != {legacy_values!r}") + 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=old_item.format(field=field), - new_item=new_item.format(field=field), + old_item=f"labels.{field}", + new_item=field, removed_in="1.4.0", ) - values = legacy_values - resolved[field] = values - return remaining, resolved + resolved[field] = legacy_value + return remaining, resolved["operator"], resolved["operation"] class AttackOutcome(str, Enum): @@ -230,16 +192,14 @@ def _normalize_legacy_attribution_labels(cls, data: Any) -> Any: return data normalized = dict(data) - remaining, resolved = pop_legacy_attribution_labels( + remaining, operator, operation = normalize_legacy_attack_attribution( labels=normalized["labels"], - dedicated={field: normalized.get(field) for field in ATTRIBUTION_FIELDS}, - allow_multiple=False, - old_item="AttackResult.labels['{field}']", - new_item="AttackResult.{field}", + operator=normalized.get("operator"), + operation=normalized.get("operation"), ) - for field, values in resolved.items(): - normalized[field] = values[0] if values else None normalized["labels"] = remaining + normalized["operator"] = operator + normalized["operation"] = operation return normalized def get_attack_strategy_identifier(self) -> ComponentIdentifier | None: diff --git a/tests/unit/backend/test_api_routes.py b/tests/unit/backend/test_api_routes.py index 823fc15ca8..ebdbf62fd7 100644 --- a/tests/unit/backend/test_api_routes.py +++ b/tests/unit/backend/test_api_routes.py @@ -1482,6 +1482,7 @@ def test_get_labels_for_attacks_passes_narrowing_filters(self, client: TestClien 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.""" 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"]) From 331de21279ef991a9e6285eea84cc4fb03bc4ee0 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Wed, 9 Sep 2026 14:33:30 -0700 Subject: [PATCH 4/6] FIX: Rebase history migration on current head Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 7540877a-bdf5-4309-97e1-14469bc817e7 --- frontend/e2e/history.spec.ts | 33 ++++++++++++++----- ..._attack_attribution_and_history_indexes.py | 4 +-- tests/unit/memory/test_migration.py | 2 +- 3 files changed, 27 insertions(+), 12 deletions(-) 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/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 index 0f2728c3f5..da3f8e4588 100644 --- 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 @@ -5,7 +5,7 @@ Add first-class attack attribution fields and history query indexes. Revision ID: a4c6e8f0b2d1 -Revises: 8d1e3f5a7b9c +Revises: 1b3d5f7a9c2e Create Date: 2026-09-04 18:48:00.000000 """ @@ -22,7 +22,7 @@ from collections.abc import Sequence revision: str = "a4c6e8f0b2d1" -down_revision: str | Sequence[str] | None = "8d1e3f5a7b9c" +down_revision: str | Sequence[str] | None = "1b3d5f7a9c2e" branch_labels: str | Sequence[str] | None = None depends_on: str | Sequence[str] | None = None diff --git a/tests/unit/memory/test_migration.py b/tests/unit/memory/test_migration.py index da033ca952..13ce952497 100644 --- a/tests/unit/memory/test_migration.py +++ b/tests/unit/memory/test_migration.py @@ -2443,7 +2443,7 @@ def test_attack_recency_downgrade_restores_updated_at_and_drops_indexes(): _ATTACK_ATTRIBUTION_REV = "a4c6e8f0b2d1" -_ATTACK_ATTRIBUTION_PREV_REV = "8d1e3f5a7b9c" +_ATTACK_ATTRIBUTION_PREV_REV = "1b3d5f7a9c2e" def _seed_attack_result_with_labels(connection, *, attack_id: str, labels: dict[str, object]) -> None: From 6bfc3ae1c1fd441f1335f8e5b8e98feb1b0ce4da Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Wed, 9 Sep 2026 15:05:50 -0700 Subject: [PATCH 5/6] FIX: Preserve multimodal message ordering Validate attack filters before constructing the memory-backed service and persist a timestamp tie-breaker for pieces that share a message sequence. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 7540877a-bdf5-4309-97e1-14469bc817e7 --- pyrit/backend/routes/attacks.py | 2 +- pyrit/memory/memory_interface.py | 10 ++++++- pyrit/models/messages/message_piece.py | 6 ++-- .../test_interface_prompts.py | 30 +++++++++++++++++++ tests/unit/models/test_message_piece.py | 19 ++++++++++++ 5 files changed, 62 insertions(+), 5 deletions(-) diff --git a/pyrit/backend/routes/attacks.py b/pyrit/backend/routes/attacks.py index b023714074..2c526d94af 100644 --- a/pyrit/backend/routes/attacks.py +++ b/pyrit/backend/routes/attacks.py @@ -108,7 +108,6 @@ 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) or {} # TODO(PyRIT 1.4): Remove legacy attribution aliases from label query parameters. legacy_operator = labels.pop("operator", None) @@ -138,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, diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index 7b46e7bf1e..7bbafa3e77 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -13,7 +13,7 @@ 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 @@ -705,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): 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/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/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", From b949e06a2947fa163a5ed35fa8608b48ec262132 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Thu, 10 Sep 2026 17:31:22 -0700 Subject: [PATCH 6/6] PERF: Optimize SQL Server migration backfills Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- ...1b3d5f7a9c2e_persist_scored_expectation.py | 90 +++++++++++++ ..._attack_attribution_and_history_indexes.py | 124 +++++++++++++++++- tests/unit/memory/test_migration.py | 81 ++++++++++++ 3 files changed, 293 insertions(+), 2 deletions(-) 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 index da3f8e4588..8b0ae4c282 100644 --- 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 @@ -11,6 +11,7 @@ from __future__ import annotations +import logging from typing import TYPE_CHECKING, Any import sqlalchemy as sa @@ -26,30 +27,85 @@ 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", @@ -57,7 +113,9 @@ def upgrade() -> None: 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", @@ -65,22 +123,28 @@ def upgrade() -> None: 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", @@ -116,6 +180,7 @@ def downgrade() -> None: _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: @@ -230,6 +295,49 @@ def _move_attribution_from_labels() -> None: 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) @@ -241,7 +349,10 @@ def _move_attribution_from_labels() -> None: ) ) rows = bind.execute(sa.select(table.c.id, table.c.labels)).all() - for start in range(0, len(rows), _BATCH_SIZE): + 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 @@ -276,6 +387,9 @@ def _move_attribution_from_labels() -> None: ) 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: @@ -289,7 +403,10 @@ def _restore_attribution_to_labels() -> None: 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() - for start in range(0, len(rows), _BATCH_SIZE): + 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 {} @@ -310,3 +427,6 @@ def _restore_attribution_to_labels() -> None: 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/tests/unit/memory/test_migration.py b/tests/unit/memory/test_migration.py index 13ce952497..55c826d76e 100644 --- a/tests/unit/memory/test_migration.py +++ b/tests/unit/memory/test_migration.py @@ -2546,6 +2546,59 @@ def test_attack_attribution_migration_rejects_overlength_value() -> None: 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()) @@ -2696,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())