pyrit.memory.AzureSQLMemory

class AzureSQLMemory(*args, **kwargs)

Bases: MemoryInterface

A class to manage conversation memory using Azure SQL Server as the backend database. It leverages SQLAlchemy Base models for creating tables and provides CRUD operations to interact with the tables.

This class encapsulates the setup of the database connection, table creation based on SQLAlchemy models, and session management to perform database operations.

__init__(*, connection_string: str | None = None, results_container_url: str | None = None, results_sas_token: str | None = None, verbose: bool = False)

Initialize an Azure SQL Memory backend.

Parameters:

• connection_string (Optional[str]) – The connection string for the Azure Sql Database. If not provided, it falls back to the ‘AZURE_SQL_DB_CONNECTION_STRING’ environment variable.

• results_container_url (Optional[str]) – The URL to an Azure Storage Container. If not provided, it falls back to the ‘AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL’ environment variable.

• results_sas_token (Optional[str]) – The Shared Access Signature (SAS) token for the storage container. If not provided, falls back to the ‘AZURE_STORAGE_ACCOUNT_DB_DATA_SAS_TOKEN’ environment variable.

• verbose (bool) – Whether to enable verbose logging for the database engine. Defaults to False.

Methods

__init__(*[, connection_string, ...])	Initialize an Azure SQL Memory backend.
add_attack_results_to_memory(*, attack_results)	Insert a list of attack results into the memory storage.
add_attack_results_to_scenario(*, ...)	Add attack results to an existing scenario result in memory.
add_message_pieces_to_memory(*, message_pieces)	Insert a list of message pieces into the memory storage.
add_message_to_memory(*, request)	Insert a list of message pieces into the memory storage.
add_scenario_results_to_memory(*, ...)	Insert a list of scenario results into the memory storage.
add_scores_to_memory(*, scores)	Insert a list of scores into the memory storage.
add_seed_datasets_to_memory_async(*, ...)	Insert a list of seed datasets into the memory storage.
add_seed_groups_to_memory(*, prompt_groups)	Insert a list of seed groups into the memory storage.
add_seeds_to_memory_async(*, seeds[, added_by])	Insert a list of seeds into the memory storage.
cleanup()	Ensure cleanup on process exit.
disable_embedding()	Disable embedding functionality for the memory interface.
dispose_engine()	Dispose the engine and clean up resources.
duplicate_conversation(*, conversation_id[, ...])	Duplicate a conversation for reuse.
duplicate_conversation_excluding_last_turn(*, ...)	Duplicate a conversation, excluding the last turn.
enable_embedding([embedding_model])	Enable embedding functionality for the memory interface.
export_conversations(*[, attack_id, ...])	Export conversation data with the given inputs to a specified file.
get_all_embeddings()	Fetch all entries from the specified table and returns them as model instances.
get_attack_results(*[, attack_result_ids, ...])	Retrieve a list of AttackResult objects based on the specified filters.
get_chat_messages_with_conversation_id(*, ...)	Return the memory for a given conversation_id.
get_conversation(*, conversation_id)	Retrieve a list of Message objects that have the specified conversation ID.
get_message_pieces(*[, attack_id, role, ...])	Retrieve a list of MessagePiece objects based on the specified filters.
get_prompt_scores(*[, attack_id, role, ...])	Retrieve scores attached to message pieces based on the specified filters.
get_request_from_response(*, response)	Retrieve the request that produced the given response.
get_scenario_results(*[, ...])	Retrieve a list of ScenarioResult objects based on the specified filters.
get_scores(*[, score_ids, score_type, ...])	Retrieve a list of Score objects based on the specified filters.
get_seed_dataset_names()	Return a list of all seed dataset names in the memory storage.
get_seed_groups(*[, value, value_sha256, ...])	Retrieve groups of seed prompts based on the provided filtering criteria.
get_seeds(*[, value, value_sha256, ...])	Retrieve a list of seed prompts based on the specified filters.
get_session()	Provide a session for database operations.
print_schema()	Print the schema of all tables in the database.
reset_database()	Drop and recreate existing tables.
update_labels_by_conversation_id(*, ...)	Update the labels of prompt entries in memory for a given conversation ID.
update_prompt_entries_by_conversation_id(*, ...)	Update prompt entries for a given conversation ID with the specified field values.
update_prompt_metadata_by_conversation_id(*, ...)	Update the metadata of prompt entries in memory for a given conversation ID.
update_scenario_run_state(*, ...)	Update the run state of an existing scenario result.

Attributes

AZURE_SQL_DB_CONNECTION_STRING
AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL
AZURE_STORAGE_ACCOUNT_DB_DATA_SAS_TOKEN
SQL_COPT_SS_ACCESS_TOKEN
TOKEN_URL
engine
memory_embedding
results_path
results_storage_io

AZURE_SQL_DB_CONNECTION_STRING = 'AZURE_SQL_DB_CONNECTION_STRING'

AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL: str = 'AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL'

AZURE_STORAGE_ACCOUNT_DB_DATA_SAS_TOKEN: str = 'AZURE_STORAGE_ACCOUNT_DB_DATA_SAS_TOKEN'

SQL_COPT_SS_ACCESS_TOKEN = 1256

TOKEN_URL = 'https://database.windows.net/.default'

add_message_pieces_to_memory(*, message_pieces: Sequence[MessagePiece]) → None

Insert a list of message pieces into the memory storage.

dispose_engine()

Dispose the engine and clean up resources.

get_all_embeddings() → Sequence[EmbeddingDataEntry]

Fetch all entries from the specified table and returns them as model instances.

Returns:

A sequence of EmbeddingDataEntry instances representing all stored embeddings.

Return type:

Sequence[EmbeddingDataEntry]

get_session() → Session

Provide a session for database operations.

Returns:

A new SQLAlchemy session bound to the configured engine.

Return type:

Session

reset_database()

Drop and recreate existing tables.