FEAT backend attack API - revisiting attacks & conversations#1419
Open
romanlutz wants to merge 4 commits intoAzure:mainfrom
Open
FEAT backend attack API - revisiting attacks & conversations#1419romanlutz wants to merge 4 commits intoAzure:mainfrom
romanlutz wants to merge 4 commits intoAzure:mainfrom
Conversation
- Add run_initializers_async to pyrit.setup for programmatic initialization - Switch AIRTInitializer to Entra (Azure AD) auth, removing API key requirements - Add --config-file flag to pyrit_backend CLI - Use PyRIT configuration loader in FrontendCore and pyrit_backend - Update AIRTTargetInitializer with new target types Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
- Add conversation_stats model and attack_result extensions - Add get_attack_results with filtering by harm categories, labels, attack type, and converter types to memory interface - Implement SQLite-specific JSON filtering for attack results - Add memory_models field for targeted_harm_categories - Add prompt_metadata support to openai image/video/response targets - Fix missing return statements in SQLite harm_category and label filters Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
- Add attack CRUD routes with conversation management - Add message sending with target dispatch and response handling - Add attack mappers for domain-to-DTO conversion with signed blob URLs - Add attack service with video remix support and piece persistence - Expand target service and routes with registry-based target management - Add version endpoint with database info Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Builds out the PyRIT backend from a scaffold into an attack-centric API, adding multi-conversation attack support, conversation summarization, target registry interactions, and memory-layer enhancements (including Entra-auth-related initializer updates).
Changes:
- Adds attack/conversation backend endpoints keyed by
attack_result_id, plus message sending that dispatches to targets by registry name. - Introduces
ConversationStatsand new memory APIs to efficiently summarize conversations without loading full message pieces. - Updates setup/CLI and initializers (incl.
run_initializers_async+ config-file support) and expands OpenAI target/mapping behavior for richer metadata/media handling.
Reviewed changes
Copilot reviewed 43 out of 43 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/unit/target/test_video_target.py | Adds regression test for single-turn enforcement on video target requests. |
| tests/unit/target/test_image_target.py | Adds async regression test for single-turn enforcement on image target requests. |
| tests/unit/setup/test_airt_targets_initializer.py | Adds test coverage for GPT-5 high-reasoning Responses target extra body parameters. |
| tests/unit/setup/test_airt_initializer.py | Updates AIRT initializer tests to reflect Entra/token auth (no API key env vars). |
| tests/unit/memory/test_sqlite_memory.py | Adds unit tests for get_conversation_stats aggregation behavior. |
| tests/unit/memory/memory_interface/test_interface_attack_results.py | Expands tests for attack result dedupe + update semantics and renames attack/converter filters. |
| tests/unit/cli/test_pyrit_backend.py | Adds CLI tests for pyrit_backend arg parsing and initialization wiring. |
| tests/unit/cli/test_frontend_core.py | Updates FrontendCore expectations + patches to new initialization flow and defaults. |
| tests/unit/backend/test_target_service.py | Updates tests for renamed target fields (target_registry_name). |
| tests/unit/backend/test_mappers.py | Adds extensive mapper tests for media handling, blob URL signing/fetching, and summary mapping changes. |
| tests/unit/backend/test_main.py | Updates lifespan tests to reflect CLI-driven initialization instead of in-app initialization. |
| tests/unit/backend/test_api_routes.py | Updates route tests for new endpoints, request/response shapes, and renamed filters/fields. |
| pyrit/setup/initializers/airt_targets.py | Adds target config extra_kwargs and new GPT-5 high-reasoning Responses target; adjusts video target config naming/env vars. |
| pyrit/setup/initializers/airt.py | Switches AIRTInitializer to Entra/token auth and updates scorer setup accordingly. |
| pyrit/setup/initialization.py | Adds run_initializers_async helper to run initializers without reinitializing memory. |
| pyrit/setup/init.py | Re-exports run_initializers_async. |
| pyrit/prompt_target/openai/openai_video_target.py | Adds validation to enforce single-turn conversations for video target usage. |
| pyrit/prompt_target/openai/openai_target.py | Refactors OpenAITarget base inheritance/initialization behavior. |
| pyrit/prompt_target/openai/openai_response_target.py | Includes extra_body_parameters in target identifiers for richer provenance. |
| pyrit/prompt_target/openai/openai_realtime_target.py | Adjusts realtime target inheritance to ensure chat-target semantics. |
| pyrit/prompt_target/openai/openai_image_target.py | Adds validation to enforce single-turn conversations for image target usage. |
| pyrit/models/conversation_stats.py | Introduces ConversationStats dataclass for lightweight conversation summaries. |
| pyrit/models/attack_result.py | Adds attack_result_id to expose persisted identifier to callers. |
| pyrit/models/init.py | Exports ConversationStats. |
| pyrit/memory/sqlite_memory.py | Implements get_conversation_stats, improves update semantics, and renames attack/converter filter helpers. |
| pyrit/memory/memory_models.py | Populates attack_result_id when constructing AttackResult from DB entry. |
| pyrit/memory/memory_interface.py | Adds get_conversation_stats, exposes duplicate_messages, and updates attack/converter filter APIs + update helpers. |
| pyrit/memory/azure_sql_memory.py | Mirrors SQLite changes for attack/converter filters and implements get_conversation_stats. |
| pyrit/cli/pyrit_backend.py | Refactors backend CLI to use FrontendCore + adds --config-file. |
| pyrit/cli/frontend_core.py | Splits “initialize” vs “run initializers” behavior and reuses run_initializers_async. |
| pyrit/backend/services/target_service.py | Renames API surface to target_registry_name and adjusts pagination cursor semantics. |
| pyrit/backend/services/attack_service.py | Major refactor: attack_result_id-keyed operations, multi-conversation support, branching, message routing, and async DTO mapping. |
| pyrit/backend/routes/version.py | Adds optional database backend info to version response. |
| pyrit/backend/routes/targets.py | Renames target route params to target_registry_name. |
| pyrit/backend/routes/attacks.py | Adds conversation management endpoints and refactors message endpoints to support per-conversation routing. |
| pyrit/backend/models/targets.py | Updates DTO fields (registry naming) and adds supports_multiturn_chat. |
| pyrit/backend/models/attacks.py | Refactors attack DTOs for attack_result_id, target info nesting, conversation endpoints, and message routing fields. |
| pyrit/backend/models/init.py | Updates exported backend models to reflect renamed/removed DTOs. |
| pyrit/backend/mappers/target_mappers.py | Updates target mapping to registry naming and adds multiturn capability detection. |
| pyrit/backend/mappers/attack_mappers.py | Refactors summary mapping to use ConversationStats and adds media/blob URL signing/fetching logic + async message DTO mapping. |
| pyrit/backend/mappers/init.py | Exports pyrit_messages_to_dto_async instead of sync variant. |
| pyrit/backend/main.py | Removes implicit initialization in lifespan; warns when started without CLI initialization. |
Comment on lines
16
to
+22
| # Ensure emoji and other Unicode characters don't crash on Windows consoles | ||
| # that use legacy encodings like cp1252. | ||
| sys.stdout.reconfigure(errors="replace") # type: ignore[union-attr] | ||
| sys.stderr.reconfigure(errors="replace") # type: ignore[union-attr] | ||
|
|
||
| import uvicorn | ||
|
|
There was a problem hiding this comment.
import uvicorn occurs after executable statements (sys.stdout.reconfigure / sys.stderr.reconfigure). With Ruff E402 enabled for pyrit/**, this is likely to fail linting as a non-top-level import. Move the uvicorn import up with the other imports (and keep the reconfigure calls after all imports), or wrap the reconfigure logic inside main() so imports remain at the top of the module.
Comment on lines
+409
to
+417
| tb = traceback.format_exception(type(e), e, e.__traceback__) | ||
| # Include the root cause if chained | ||
| cause = e.__cause__ | ||
| if cause: | ||
| tb += traceback.format_exception(type(cause), cause, cause.__traceback__) | ||
| detail = "".join(tb) | ||
| raise HTTPException( | ||
| status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, | ||
| detail=f"Failed to add message: {str(e)}", | ||
| detail=detail, |
There was a problem hiding this comment.
The 500 error path returns a full Python traceback to the API caller. This can leak sensitive details (file paths, environment variables, internal URLs, tokens) and makes it hard to safely run the backend outside of a trusted dev environment. Gate this behavior behind an explicit debug flag (e.g., DEV_MODE / config), and otherwise return a generic error message while logging the traceback server-side.
Comment on lines
+70
to
+72
| return value.startswith("https://") and ".blob.core.windows.net/" in value | ||
|
|
||
|
|
There was a problem hiding this comment.
_is_azure_blob_url currently checks for the substring .blob.core.windows.net/ anywhere in the URL. This can misclassify non-Azure URLs (e.g., a malicious host with that string in the path) and trigger server-side requests / signing attempts, creating an SSRF surface area. Parse the URL and validate that netloc ends with blob.core.windows.net (and ideally enforce an allowlist of expected storage accounts/containers) before treating it as an Azure Blob URL.
Suggested change
| return value.startswith("https://") and ".blob.core.windows.net/" in value | |
| parsed = urlparse(value) | |
| if parsed.scheme != "https": | |
| return False | |
| if not parsed.netloc: | |
| return False | |
| # Extract hostname from possible "user@host:port" netloc forms | |
| host = parsed.netloc.split("@")[-1].split(":")[0] | |
| if not host.endswith(".blob.core.windows.net"): | |
| return False | |
| # Ensure there is a storage account name prefix before ".blob.core.windows.net" | |
| account_name = host.split(".")[0] | |
| return bool(account_name) |
Comment on lines
+1291
to
+1295
| from contextlib import closing | ||
|
|
||
| with closing(self.get_session()) as session: | ||
| from sqlalchemy.exc import SQLAlchemyError | ||
|
|
There was a problem hiding this comment.
add_attack_results_to_memory introduces local imports for closing and SQLAlchemyError even though closing is already imported at module scope. This adds noise and conflicts with the project's "imports at the top of the file" convention. Please remove the redundant local import and move SQLAlchemyError to the module imports (or reuse an existing top-level import) so this method contains only logic.
Comment on lines
+719
to
+745
| @staticmethod | ||
| def _get_last_message_preview(pieces: Sequence[PromptMemoryEntry]) -> Optional[str]: | ||
| """Return a truncated preview of the last message piece's text.""" | ||
| if not pieces: | ||
| return None | ||
| last = max(pieces, key=lambda p: p.sequence) | ||
| text = last.converted_value or "" | ||
| return text[:100] + "..." if len(text) > 100 else text | ||
|
|
||
| @staticmethod | ||
| def _count_messages(pieces: Sequence[PromptMemoryEntry]) -> int: | ||
| """ | ||
| Count distinct messages (by sequence number) in a list of pieces. | ||
|
|
||
| Returns: | ||
| The number of unique sequence values. | ||
| """ | ||
| return len(set(p.sequence for p in pieces)) | ||
|
|
||
| @staticmethod | ||
| def _get_earliest_timestamp(pieces: Sequence[PromptMemoryEntry]) -> Optional[datetime]: | ||
| """Return the earliest timestamp from a list of message pieces.""" | ||
| if not pieces: | ||
| return None | ||
| timestamps: List[datetime] = [p.timestamp for p in pieces if p.timestamp is not None] | ||
| return min(timestamps) if timestamps else None | ||
|
|
There was a problem hiding this comment.
The helper methods _get_last_message_preview, _count_messages, and _get_earliest_timestamp appear to be unused in this module (no call sites besides their definitions). Keeping unused code makes future refactors harder and increases maintenance cost; consider removing them or wiring them into the conversation-summary logic if they are intended as a fallback.
Suggested change
| @staticmethod | |
| def _get_last_message_preview(pieces: Sequence[PromptMemoryEntry]) -> Optional[str]: | |
| """Return a truncated preview of the last message piece's text.""" | |
| if not pieces: | |
| return None | |
| last = max(pieces, key=lambda p: p.sequence) | |
| text = last.converted_value or "" | |
| return text[:100] + "..." if len(text) > 100 else text | |
| @staticmethod | |
| def _count_messages(pieces: Sequence[PromptMemoryEntry]) -> int: | |
| """ | |
| Count distinct messages (by sequence number) in a list of pieces. | |
| Returns: | |
| The number of unique sequence values. | |
| """ | |
| return len(set(p.sequence for p in pieces)) | |
| @staticmethod | |
| def _get_earliest_timestamp(pieces: Sequence[PromptMemoryEntry]) -> Optional[datetime]: | |
| """Return the earliest timestamp from a list of message pieces.""" | |
| if not pieces: | |
| return None | |
| timestamps: List[datetime] = [p.timestamp for p in pieces if p.timestamp is not None] | |
| return min(timestamps) if timestamps else None | |
| # (Reserved for future conversation-info helpers; currently unused.) |
romanlutz
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR builds out the backend API from a basic scaffold into a functional attack-centric API that supports
multi-conversation attacks, conversation branching, message sending with target dispatch, and richer memory/model
support. It spans 3 logical commits stacked as a single branch.
Changes
enables pyrit_backend and FrontendCore to separate memory setup from initializer execution.
_CONTENT_SAFETY_API_KEY), switches to DefaultAzureCredential via get_azure_openai_auth / get_azure_token_provider.
method; pyrit_backend now uses FrontendCore directly instead of duplicating logic.
labels, timestamp) — avoids loading full message pieces for list views.
aggregation for conversation summaries.
implementation, and all callers.
branching logic.
(type, endpoint, model), related_conversation_ids.
- GET /{id}/conversations — list all conversations for an attack
- POST /{id}/conversations — create a related conversation (with optional branching from source_conversation_id +
cutoff_index)
- POST /{id}/change-main-conversation — swap the primary conversation
- GET /{id}/messages?conversation_id= — get messages for a specific conversation
target_conversation_id to route messages to a specific conversation, and labels for per-message labeling. Includes
full traceback in 500 error responses for debugging.
pyrit_messages_to_dto_async generates signed blob URLs for non-text content and extracts
original_filename/converted_filename.
routes/attacks.py,routes/targets.py,routes/version.py,models/attacks.py,models/targets.py,services/attack_service.py,services/target_service.py,mappers/attack_mappers.py,mappers/target_mappers.py,main.pymemory_interface.py,sqlite_memory.py,azure_sql_memory.py,memory_models.pyattack_result.py,conversation_stats.py(new),__init__.pyinitialization.py,__init__.py,airt.py,airt_targets.py,frontend_core.py,pyrit_backend.pyopenai_target.py,openai_image_target.py,openai_video_target.py,openai_response_target.py,openai_realtime_target.py