pyrit.memory.MemoryInterface

pyrit.memory.MemoryInterface#

class MemoryInterface(embedding_model=None)[source]#

Bases: ABC

Represents a conversation memory that stores chat messages. This class must be overwritten with a specific implementation to store the memory objects (e.g. relational database, NoSQL database, etc.)

Parameters:

  • embedding_model (EmbeddingSupport) – If set, this includes embeddings in the memory entries

  • similarities (which are extremely useful for comparing chat messages and)

  • overhead (but also includes)

__init__(embedding_model=None)[source]#

Methods

__init__([embedding_model])

add_request_pieces_to_memory(*, request_pieces)

Inserts a list of prompt request pieces into the memory storage.

add_request_response_to_memory(*, request)

Inserts a list of prompt request pieces into the memory storage.

add_scores_to_memory(*, scores)

Inserts a list of scores into the memory storage.

add_seed_prompt_groups_to_memory(*, ...[, ...])

Inserts a list of seed prompt groups into the memory storage.

add_seed_prompts_to_memory(*, prompts[, ...])

Inserts a list of prompts into the memory storage.

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(*[, orchestrator_id, ...])

Exports conversation data with the given inputs to a specified file.

get_all_embeddings()

Loads all EmbeddingData from the memory storage handler.

get_chat_messages_with_conversation_id(*, ...)

Returns the memory for a given conversation_id.

get_conversation(*, conversation_id)

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

get_prompt_request_pieces(*[, ...])

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

get_scores_by_memory_labels(*, memory_labels)

Retrieves a list of Score objects associated with the PromptRequestPiece objects which have the specified memory labels.

get_scores_by_orchestrator_id(*, orchestrator_id)

Retrieves a list of Score objects associated with the PromptRequestPiece objects which have the specified orchestrator ID.

get_scores_by_prompt_ids(*, ...)

Gets a list of scores based on prompt_request_response_ids.

get_seed_prompt_dataset_names()

Returns a list of all seed prompt dataset names in the memory storage.

get_seed_prompt_groups(*[, dataset_name, ...])

Retrieves groups of seed prompts based on the provided filtering criteria._summary_

get_seed_prompts(*[, value, dataset_name, ...])

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

populate_prompt_piece_scores(...)

Adds scores in the database to prompt request piece objects

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

memory_embedding

results_path

results_storage_io

abstract add_request_pieces_to_memory(*, request_pieces: Sequence[PromptRequestPiece]) → None[source]#

Inserts a list of prompt request pieces into the memory storage.

add_request_response_to_memory(*, request: PromptRequestResponse) → None[source]#

Inserts a list of prompt request 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 (PromptRequestPiece) – The request piece to add to the memory.

Returns:

None

add_scores_to_memory(*, scores: list[Score]) → None[source]#

Inserts a list of scores into the memory storage.

add_seed_prompt_groups_to_memory(*, prompt_groups: list[SeedPromptGroup], added_by: str | None = None) → None[source]#

Inserts a list of seed prompt groups into the memory storage.

Parameters:

  • prompt_groups (list[SeedPromptGroup]) – 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.

add_seed_prompts_to_memory(*, prompts: list[SeedPrompt], added_by: str | None = None) → None[source]#

Inserts a list of prompts into the memory storage.

Parameters:

  • prompts (list[SeedPrompt]) – A list of prompts to insert.

  • added_by (str) – The user who added the prompts.

disable_embedding()[source]#
abstract dispose_engine()[source]#

Dispose the engine and clean up resources.

duplicate_conversation(*, conversation_id: str, new_orchestrator_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 orchestrator and conversation IDs since that would corrupt the memory. Instead, one needs to duplicate the conversation and continue with the new orchestrator ID.

Parameters:

  • conversation_id (str) – The conversation ID with existing conversations.

  • new_orchestrator_id (str, Optional) – The new orchestrator ID to assign to the duplicated conversations. If no new orchestrator ID is provided, the orchestrator ID will remain the same. Defaults to None.

Returns:

The uuid for the new conversation.

duplicate_conversation_excluding_last_turn(*, conversation_id: str, new_orchestrator_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_orchestrator_id (str, Optional) – The new orchestrator ID to assign to the duplicated conversations. If no new orchestrator ID is provided, the orchestrator ID will remain the same. Defaults to None.

Returns:

The uuid for the new conversation.

enable_embedding(embedding_model=None)[source]#
export_conversations(*, orchestrator_id: str | UUID | None = None, conversation_id: str | UUID | None = None, prompt_ids: list[str] | list[UUID] | None = None, labels: dict[str, str] | None = None, sent_after: datetime | None = None, sent_before: datetime | None = None, original_values: list[str] | None = None, converted_values: list[str] | None = None, data_type: str | None = None, not_data_type: str | None = None, converted_value_sha256: list[str] | None = None, file_path: Path | None = None, export_type: str = 'json')[source]#
Exports conversation data with the given inputs to a specified file.

Defaults to all conversations if no filters are provided.

Parameters:

  • orchestrator_id (Optional[str | uuid.UUID], optional) – The ID of the orchestrator. Defaults to None.

  • conversation_id (Optional[str | uuid.UUID], optional) – The ID of the conversation. Defaults to None.

  • prompt_ids (Optional[list[str] | list[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[list[str]], optional) – A list of original values. Defaults to None.

  • converted_values (Optional[list[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[list[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_chat_messages_with_conversation_id(*, conversation_id: str) → list[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:

list[ChatMessage]

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

Retrieves a list of PromptRequestResponse 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[PromptRequestResponse]

get_prompt_request_pieces(*, orchestrator_id: str | UUID | None = None, conversation_id: str | UUID | None = None, prompt_ids: list[str] | list[UUID] | None = None, labels: dict[str, str] | None = None, sent_after: datetime | None = None, sent_before: datetime | None = None, original_values: list[str] | None = None, converted_values: list[str] | None = None, data_type: str | None = None, not_data_type: str | None = None, converted_value_sha256: list[str] | None = None) → list[PromptRequestPiece][source]#

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

Parameters:

  • orchestrator_id (Optional[str | uuid.UUID], optional) – The ID of the orchestrator. Defaults to None.

  • conversation_id (Optional[str | uuid.UUID], optional) – The ID of the conversation. Defaults to None.

  • prompt_ids (Optional[list[str] | list[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[list[str]], optional) – A list of original values. Defaults to None.

  • converted_values (Optional[list[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[list[str]], optional) – A list of SHA256 hashes of converted values. Defaults to None.

Returns:

A list of PromptRequestPiece objects that match the specified filters.

Return type:

list[PromptRequestPiece]

Raises:

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

get_scores_by_memory_labels(*, memory_labels: dict[str, str]) → list[Score][source]#

Retrieves a list of Score objects associated with the PromptRequestPiece objects which have the specified memory labels.

Parameters:

memory_labels (dict[str, str]) – A free-form dictionary for tagging prompts with custom labels. These labels can be used to track all prompts sent as part of an operation, score prompts based on the operation ID (op_id), and tag each prompt with the relevant Responsible AI (RAI) harm category. Users can define any key-value pairs according to their needs.

Returns:

A list of Score objects associated with the PromptRequestPiece objects

which match the specified memory labels.

Return type:

list[Score]

get_scores_by_orchestrator_id(*, orchestrator_id: str) → list[Score][source]#

Retrieves a list of Score objects associated with the PromptRequestPiece objects which have the specified orchestrator ID.

Parameters:

orchestrator_id (str) – The id of the orchestrator. Can be retrieved by calling orchestrator.get_identifier()[“id”]

Returns:

A list of Score objects associated with the PromptRequestPiece objects

which match the specified orchestrator ID.

Return type:

list[Score]

get_scores_by_prompt_ids(*, prompt_request_response_ids: list[str]) → list[Score][source]#

Gets a list of scores based on prompt_request_response_ids.

get_seed_prompt_dataset_names() → list[str][source]#

Returns a list of all seed prompt dataset names in the memory storage.

get_seed_prompt_groups(*, dataset_name: str | None = None, data_types: list[str] | None = None, harm_categories: list[str] | None = None, added_by: str | None = None, authors: list[str] | None = None, groups: list[str] | None = None, source: str | None = None) → list[SeedPromptGroup][source]#

Retrieves groups of seed prompts based on the provided filtering criteria._summary_

Parameters:

  • 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 prompt groups to filter by.

  • authors (Optional[Sequence[str]], Optional) – List of authors to filter seed prompt groups by.

  • groups (Optional[Sequence[str]], Optional) – List of groups to filter seed prompt groups by.

  • source (Optional[str], Optional) – The source from which the seed prompts originated.

Returns:

A list of SeedPromptGroup objects that match the filtering criteria.

Return type:

list[SeedPromptGroup]

get_seed_prompts(*, value: str | None = None, dataset_name: str | None = None, harm_categories: list[str] | None = None, added_by: str | None = None, authors: list[str] | None = None, groups: list[str] | None = None, source: str | None = None, parameters: list[str] | None = None) → list[SeedPrompt][source]#

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.

  • dataset_name (str) – The dataset name to match. If None, all dataset names are considered.

  • harm_categories (list[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 (list[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 (list[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.

  • parameters (list[str]) – A list of parameters to filter by. Specifying parameters effectively returns prompt templates instead of prompts. If None, only prompts without parameters are returned.

Returns:

A list of prompts matching the criteria.

Return type:

list[SeedPrompt]

memory_embedding: MemoryEmbedding = None#
populate_prompt_piece_scores(prompt_request_pieces: list[PromptRequestPiece]) → list[PromptRequestPiece][source]#

Adds scores in the database to prompt request piece objects

Parameters:

prompt_request_pieces (list[PromptRequestPiece]) – The list of PromptRequestPieces to add scores to.

Returns:

None

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: str) → 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 (str) – New metadata.

Returns:

True if the update was successful, False otherwise.

Return type:

bool