pyrit.memory.MemoryInterface

pyrit.memory.MemoryInterface#

class MemoryInterface(embedding_model=None)[source]#

Bases: ABC

Abstract interface for conversation memory storage systems.

This interface defines the contract for storing and retrieving chat messages and conversation history. Implementations can use different storage backends such as files, databases, or cloud storage services.

__init__(embedding_model=None)[source]#

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__`([embedding_model])	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_scenario_results_to_memory`(*, ...)	Inserts a list of scenario results 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`()	Loads all EmbeddingData from the memory storage handler.
`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_scenario_results`(*[, ...])	Retrieves a list of ScenarioResult objects based on the specified filters.
`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.
`print_schema`()	Prints the schema of all tables in the database.
`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

`engine`
`memory_embedding`
`results_path`
`results_storage_io`

add_attack_results_to_memory(*, attack_results: Sequence[AttackResult]) → None[source]#: Inserts a list of attack results into the memory storage. The database model automatically calculates objective_sha256 for consistency.

abstract add_message_pieces_to_memory(*, message_pieces: Sequence[MessagePiece]) → None[source]#: Inserts a list of message pieces into the memory storage.

add_message_to_memory(*, request: Message) → None[source]#

Inserts a list of message pieces into the memory storage.

Automatically updates the sequence to be the next number in the conversation. If necessary, generates embedding data for applicable entries

Parameters:: request (MessagePiece) – The message piece to add to the memory.
Returns:: None

add_scenario_results_to_memory(*, scenario_results: Sequence[ScenarioResult]) → None[source]#

Inserts a list of scenario results into the memory storage.

Parameters:: scenario_results – Sequence of ScenarioResult objects to store in the database.

add_scores_to_memory(*, scores: Sequence[Score]) → None[source]#: Inserts a list of scores into the memory storage.

async add_seed_groups_to_memory(*, prompt_groups: Sequence[SeedGroup], added_by: str | None = None) → None[source]#

Inserts a list of seed groups into the memory storage.

Parameters:

prompt_groups (Sequence[SeedGroup]) – A list of prompt groups to insert.
added_by (str) – The user who added the prompt groups.

Raises:

ValueError – If a prompt group does not have at least one prompt.
ValueError – If prompt group IDs are inconsistent within the same prompt group.

async add_seeds_to_memory_async(*, prompts: Sequence[Seed], added_by: str | None = None) → None[source]#

Inserts a list of prompts into the memory storage.

Parameters:

prompts (Sequence[Seed]) – A list of prompts to insert.
added_by (str) – The user who added the prompts.

cleanup()[source]#: Ensure cleanup on process exit

disable_embedding()[source]#

abstract dispose_engine()[source]#: Dispose the engine and clean up resources.

duplicate_conversation(*, conversation_id: str, new_attack_id: str | None = None) → str[source]#

Duplicates a conversation for reuse

This can be useful when an attack strategy requires branching out from a particular point in the conversation. One cannot continue both branches with the same attack and conversation IDs since that would corrupt the memory. Instead, one needs to duplicate the conversation and continue with the new attack ID.

Parameters:

conversation_id (str) – The conversation ID with existing conversations.
new_attack_id (str, Optional) – The new attack ID to assign to the duplicated conversations. If no new attack ID is provided, the attack ID will remain the same. Defaults to None.

Returns:

The uuid for the new conversation.

duplicate_conversation_excluding_last_turn(*, conversation_id: str, new_attack_id: str | None = None) → str[source]#

Duplicate a conversation, excluding the last turn. In this case, last turn is defined as before the last user request (e.g. if there is half a turn, it just removes that half).

This can be useful when an attack strategy requires back tracking the last prompt/response pair.

Parameters:

conversation_id (str) – The conversation ID with existing conversations.
new_attack_id (str, Optional) – The new attack ID to assign to the duplicated conversations. If no new attack ID is provided, the attack ID will remain the same. Defaults to None.

Returns:

The uuid for the new conversation.

enable_embedding(embedding_model=None)[source]#

engine: Engine = None#

Exports conversation data with the given inputs to a specified file.: Defaults to all conversations if no filters are provided.

Parameters:

attack_id (Optional[str | uuid.UUID], optional) – The ID of the attack. Defaults to None.
conversation_id (Optional[str | uuid.UUID], optional) – The ID of the conversation. Defaults to None.
prompt_ids (Optional[Sequence[str] | Sequence[uuid.UUID]], optional) – A list of prompt IDs. Defaults to None.
labels (Optional[dict[str, str]], optional) – A dictionary of labels. Defaults to None.
sent_after (Optional[datetime], optional) – Filter for prompts sent after this datetime. Defaults to None.
sent_before (Optional[datetime], optional) – Filter for prompts sent before this datetime. Defaults to None.
original_values (Optional[Sequence[str]], optional) – A list of original values. Defaults to None.
converted_values (Optional[Sequence[str]], optional) – A list of converted values. Defaults to None.
data_type (Optional[str], optional) – The data type to filter by. Defaults to None.
not_data_type (Optional[str], optional) – The data type to exclude. Defaults to None.
converted_value_sha256 (Optional[Sequence[str]], optional) – A list of SHA256 hashes of converted values. Defaults to None.
file_path (Optional[Path], optional) – The path to the file where the data will be exported. Defaults to None.
export_type (str, optional) – The format of the export. Defaults to “json”.

abstract get_all_embeddings() → Sequence[EmbeddingDataEntry][source]#: Loads all EmbeddingData from the memory storage handler.

get_attack_results(*, attack_result_ids: Sequence[str] | None = None, conversation_id: str | None = None, objective: str | None = None, objective_sha256: Sequence[str] | None = None, outcome: str | None = None, targeted_harm_categories: Sequence[str] | None = None, labels: dict[str, str] | None = None) → Sequence[AttackResult][source]#

Retrieves a list of AttackResult objects based on the specified filters.

Parameters:

attack_result_ids (Optional[Sequence[str]], optional) – A list of attack result IDs. Defaults to None.
conversation_id (Optional[str], optional) – The conversation ID to filter by. Defaults to None.
objective (Optional[str], optional) – The objective to filter by (substring match). Defaults to None.
objective_sha256 (Optional[Sequence[str]], optional) – A list of objective SHA256 hashes to filter by. Defaults to None.
outcome (Optional[str], optional) – The outcome to filter by (success, failure, undetermined). Defaults to None.
targeted_harm_categories (Optional[Sequence[str]], optional) – A list of targeted harm categories to filter results by. These targeted harm categories are associated with the prompts themselves, meaning they are harm(s) we’re trying to elicit with the prompt, not necessarily one(s) that were found in the response. By providing a list, this means ALL categories in the list must be present. Defaults to None.
labels (Optional[dict[str, str]], optional) – A dictionary of memory labels to filter results by. These labels are associated with the prompts themselves, used for custom tagging and tracking. Defaults to None.

Returns:

A list of AttackResult objects that match the specified filters.

Return type:

Sequence[AttackResult]

get_chat_messages_with_conversation_id(*, conversation_id: str) → Sequence[ChatMessage][source]#

Returns the memory for a given conversation_id.

Parameters:: conversation_id (str) – The conversation ID.
Returns:: The list of chat messages.
Return type:: Sequence[ChatMessage]

get_conversation(*, conversation_id: str) → MutableSequence[Message][source]#

Retrieves a list of Message objects that have the specified conversation ID.

Parameters:: conversation_id (str) – The conversation ID to match.
Returns:: A list of chat memory entries with the specified conversation ID.
Return type:: MutableSequence[Message]

Retrieves a list of MessagePiece objects based on the specified filters.

Parameters:

attack_id (Optional[str | uuid.UUID], optional) – The ID of the attack. Defaults to None.
role (Optional[str], optional) – The role of the prompt. Defaults to None.
conversation_id (Optional[str | uuid.UUID], optional) – The ID of the conversation. Defaults to None.
prompt_ids (Optional[Sequence[str] | Sequence[uuid.UUID]], optional) – A list of prompt IDs. Defaults to None.
labels (Optional[dict[str, str]], optional) – A dictionary of labels. Defaults to None.
sent_after (Optional[datetime], optional) – Filter for prompts sent after this datetime. Defaults to None.
sent_before (Optional[datetime], optional) – Filter for prompts sent before this datetime. Defaults to None.
original_values (Optional[Sequence[str]], optional) – A list of original values. Defaults to None.
converted_values (Optional[Sequence[str]], optional) – A list of converted values. Defaults to None.
data_type (Optional[str], optional) – The data type to filter by. Defaults to None.
not_data_type (Optional[str], optional) – The data type to exclude. Defaults to None.
converted_value_sha256 (Optional[Sequence[str]], optional) – A list of SHA256 hashes of converted values. Defaults to None.

Returns:

A list of MessagePiece objects that match the specified filters.

Return type:

Sequence[MessagePiece]

Raises:

Exception – If there is an error retrieving the prompts, an exception is logged and an empty list is returned.

Retrieves scores attached to message pieces based on the specified filters.

Parameters:: get_message_pieces. (Same as)
Returns:: A list of scores extracted from the message pieces.
Return type:: Sequence[Score]

get_request_from_response(*, response: Message) → Message[source]#

Retrieves the request that produced the given response. :param request: The message object to match. :type request: Message

Returns:: The corresponding message object.
Return type:: Message

get_scenario_results(*, scenario_result_ids: Sequence[str] | None = None, scenario_name: str | None = None, scenario_version: int | None = None, pyrit_version: str | None = None, added_after: datetime | None = None, added_before: datetime | None = None, labels: dict[str, str] | None = None, objective_target_endpoint: str | None = None, objective_target_model_name: str | None = None) → Sequence[ScenarioResult][source]#

Retrieves a list of ScenarioResult objects based on the specified filters.

Parameters:

scenario_result_ids (Optional[Sequence[str]], optional) – A list of scenario result IDs. Defaults to None.
scenario_name (Optional[str], optional) – The scenario name to filter by (substring match). Defaults to None.
scenario_version (Optional[int], optional) – The scenario version to filter by. Defaults to None.
pyrit_version (Optional[str], optional) – The PyRIT version to filter by. Defaults to None.
added_after (Optional[datetime], optional) – Filter for scenarios completed after this datetime. Defaults to None.
before_time (Optional[datetime], optional) – Filter for scenarios completed before this datetime. Defaults to None.
labels (Optional[dict[str, str]], optional) – A dictionary of memory labels to filter by. Defaults to None.
objective_target_endpoint (Optional[str], optional) – Filter for scenarios where the objective_target_identifier has an endpoint attribute containing this value (case-insensitive). Defaults to None.
objective_target_model_name (Optional[str], optional) – Filter for scenarios where the objective_target_identifier has a model_name attribute containing this value (case-insensitive). Defaults to None.

Returns:

A list of ScenarioResult objects that match the specified filters.

Return type:

Sequence[ScenarioResult]

Retrieves a list of Score objects based on the specified filters.

Parameters:

score_ids (Optional[Sequence[str]]) – A list of score IDs to filter by.
score_type (Optional[str]) – The type of the score to filter by.
score_category (Optional[str]) – The category of the score to filter by.
sent_after (Optional[datetime]) – Filter for scores sent after this datetime.
sent_before (Optional[datetime]) – Filter for scores sent before this datetime.

Returns:

A list of Score objects that match the specified filters.

Return type:

Sequence[Score]

get_seed_dataset_names() → Sequence[str][source]#: Returns a list of all seed dataset names in the memory storage.

Retrieves groups of seed prompts based on the provided filtering criteria

Parameters:

value_sha256 (Optional[Sequence[str]], Optional) – SHA256 hash of value to filter seed groups by.
dataset_name (Optional[str], Optional) – Name of the dataset to filter seed prompts.
data_types (Optional[Sequence[str]], Optional) – List of data types to filter seed prompts by
(e.g.
text
image_path).
harm_categories (Optional[Sequence[str]], Optional) – List of harm categories to filter seed prompts by.
added_by (Optional[str], Optional) – The user who added the seed groups to filter by.
authors (Optional[Sequence[str]], Optional) – List of authors to filter seed groups by.
groups (Optional[Sequence[str]], Optional) – List of groups to filter seed groups by.
source (Optional[str], Optional) – The source from which the seed prompts originated.

Returns:

A list of SeedGroup objects that match the filtering criteria.

Return type:

Sequence[SeedGroup]

Retrieves a list of seed prompts based on the specified filters.

Parameters:

value (str) – The value to match by substring. If None, all values are returned.
value_sha256 (str) – The SHA256 hash of the value to match. If None, all values are returned.
dataset_name (str) – The dataset name to match. If None, all dataset names are considered.
data_types (Optional[Sequence[str], Optional) – List of data types to filter seed prompts by (e.g., text, image_path).
harm_categories (Sequence[str]) – A list of harm categories to filter by. If None,
considered. (all harm categories are) – Specifying multiple harm categories returns only prompts that are marked with all harm categories.
added_by (str) – The user who added the prompts.
authors (Sequence[str]) – A list of authors to filter by. Note that this filters by substring, so a query for “Adam Jones” may not return results if the record is “A. Jones”, “Jones, Adam”, etc. If None, all authors are considered.
groups (Sequence[str]) – A list of groups to filter by. If None, all groups are considered.
source (str) – The source to filter by. If None, all sources are considered.
is_objective (bool) – Whether to filter by prompts that are used as objectives.
parameters (Sequence[str]) – A list of parameters to filter by. Specifying parameters effectively returns prompt templates instead of prompts.
metadata (dict[str, str | int]) – A free-form dictionary for tagging prompts with custom metadata.

Returns:

A list of prompts matching the criteria.

Return type:

Sequence[SeedPrompt]

memory_embedding: MemoryEmbedding = None#

print_schema()[source]#: Prints the schema of all tables in the database.

results_path: str = None#

results_storage_io: StorageIO = None#

update_labels_by_conversation_id(*, conversation_id: str, labels: dict) → bool[source]#

Updates the labels of prompt entries in memory for a given conversation ID.

Parameters:

conversation_id (str) – The conversation ID of the entries to be updated.
labels (dict) – New dictionary of labels.

Returns:

True if the update was successful, False otherwise.

Return type:

bool

update_prompt_entries_by_conversation_id(*, conversation_id: str, update_fields: dict) → bool[source]#

Updates prompt entries for a given conversation ID with the specified field values.

Parameters:

conversation_id (str) – The conversation ID of the entries to be updated.
update_fields (dict) – A dictionary of field names and their new values (ex. {“labels”: {“test”: “value”}})

Returns:

True if the update was successful, False otherwise.

Return type:

bool

update_prompt_metadata_by_conversation_id(*, conversation_id: str, prompt_metadata: dict[str, str | int]) → bool[source]#

Updates the metadata of prompt entries in memory for a given conversation ID.

Parameters:

conversation_id (str) – The conversation ID of the entries to be updated.
metadata (dict[str, str | int]) – New metadata.

Returns:

True if the update was successful, False otherwise.

Return type:

bool

pyrit.memory.MemoryInterface

Contents

pyrit.memory.MemoryInterface#