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."""