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 the MemoryInterface.

Parameters:

embedding_model – If set, this includes embeddings in the memory entries which are extremely useful for comparing chat messages and similarities, but also includes overhead.

Methods

__init__(*[, connection_string, ...])	Initialize the MemoryInterface.
add_attack_results_to_memory(*, attack_results)	Inserts a list of attack results into the memory storage.
add_message_pieces_to_memory(*, message_pieces)	Inserts a list of message pieces into the memory storage.
add_message_to_memory(*, request)	Inserts a list of message pieces into the memory storage.
add_scores_to_memory(*, scores)	Inserts a list of scores into the memory storage.
add_seed_groups_to_memory(*, prompt_groups)	Inserts a list of seed groups into the memory storage.
add_seeds_to_memory_async(*, prompts[, added_by])	Inserts a list of prompts into the memory storage.
cleanup()	Ensure cleanup on process exit
disable_embedding()
dispose_engine()	Dispose the engine and clean up resources.
duplicate_conversation(*, conversation_id[, ...])	Duplicates a conversation for reuse
duplicate_conversation_excluding_last_turn(*, ...)	Duplicate a conversation, excluding the last turn.
enable_embedding([embedding_model])
export_conversations(*[, attack_id, ...])	Exports conversation data with the given inputs to a specified file.
get_all_embeddings()	Fetches all entries from the specified table and returns them as model instances.
get_attack_results(*[, attack_result_ids, ...])	Retrieves a list of AttackResult objects based on the specified filters.
get_chat_messages_with_conversation_id(*, ...)	Returns the memory for a given conversation_id.
get_conversation(*, conversation_id)	Retrieves a list of Message objects that have the specified conversation ID.
get_message_pieces(*[, attack_id, role, ...])	Retrieves a list of MessagePiece objects based on the specified filters.
get_prompt_scores(*[, attack_id, role, ...])	Retrieves scores attached to message pieces based on the specified filters.
get_request_from_response(*, response)	Retrieves the request that produced the given response.
get_scores(*[, score_ids, score_type, ...])	Retrieves a list of Score objects based on the specified filters.
get_seed_dataset_names()	Returns a list of all seed dataset names in the memory storage.
get_seed_groups(*[, value_sha256, ...])	Retrieves groups of seed prompts based on the provided filtering criteria
get_seeds(*[, value, value_sha256, ...])	Retrieves a list of seed prompts based on the specified filters.
get_session()	Provides a session for database operations.
print_schema()	Prints the schema of all tables in the database.
reset_database()	Drop and recreate existing tables
update_labels_by_conversation_id(*, ...)	Updates the labels of prompt entries in memory for a given conversation ID.
update_prompt_entries_by_conversation_id(*, ...)	Updates prompt entries for a given conversation ID with the specified field values.
update_prompt_metadata_by_conversation_id(*, ...)	Updates the metadata of prompt entries in memory for a given conversation ID.

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

Inserts a list of message pieces into the memory storage.

dispose_engine()

Dispose the engine and clean up resources.

get_all_embeddings() → Sequence[EmbeddingDataEntry]

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

get_session() → Session

Provides a session for database operations.

reset_database()

Drop and recreate existing tables