pyrit.memory.MemoryInterface

Contents

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

Exports all conversations with scores to a specified file.

export_conversation_by_id(*, conversation_id)

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

export_conversation_by_orchestrator_id(*, ...)

Exports conversation data with the given orchestrator ID to a specified file.

get_all_embeddings()

Loads all EmbeddingData from the memory storage handler.

get_all_prompt_pieces()

Loads all ConversationData 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_ids_by_orchestrator(*, ...)

get_prompt_request_piece_by_memory_labels(*)

Retrieves a list of PromptRequestPiece objects that have the specified memory labels.

get_prompt_request_piece_by_orchestrator_id(*, ...)

Retrieves a list of PromptRequestPiece objects that have the specified orchestrator ID.

get_prompt_request_pieces_by_id(*, prompt_ids)

Retrieves a list of PromptRequestPiece objects that have the specified prompt ids.

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.

insert_entries(*, entries)

Inserts multiple entries into the database.

insert_entry(entry)

Inserts an entry into the Table.

query_entries(model, *[, conditions, distinct])

Fetches data from the specified table model with optional conditions.

update_entries(*, entries, update_fields)

Updates the given entries with the specified field values.

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

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_all_conversations_with_scores(*, file_path: Path | None = None, export_type: str = 'json')[source]#

Exports all conversations with scores to a specified file. :param file_path: The path to the file where the data will be exported. :type file_path: str :param If not provided: :param a default path using RESULTS_PATH will be constructed.: :param export_type: The format of the export. Defaults to “json”. :type export_type: str

export_conversation_by_id(*, conversation_id: str, file_path: Path = None, export_type: str = 'json')[source]#

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

Parameters:

  • conversation_id (str) – The ID of the conversation to export.

  • file_path (str) – The path to the file where the data will be exported.

  • provided (If not)

  • constructed. (a default path using RESULTS_PATH will be)

  • export_type (str) – The format of the export. Defaults to “json”.

export_conversation_by_orchestrator_id(*, orchestrator_id: str, file_path: Path = None, export_type: str = 'json')[source]#

Exports conversation data with the given orchestrator ID to a specified file. This will contain all conversations that were sent by the same orchestrator.

Parameters:

  • orchestrator_id (str) – The ID of the orchestrator from which to export conversations.

  • file_path (str) – The path to the file where the data will be exported.

  • provided (If not)

  • constructed. (a default path using RESULTS_PATH will be)

  • export_type (str) – The format of the export. Defaults to “json”.

abstract get_all_embeddings() → Sequence[EmbeddingDataEntry][source]#

Loads all EmbeddingData from the memory storage handler.

abstract get_all_prompt_pieces() → Sequence[PromptRequestPiece][source]#

Loads all ConversationData 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_ids_by_orchestrator(*, orchestrator_id: str) → list[str][source]#
get_prompt_request_piece_by_memory_labels(*, memory_labels: dict[str, str] = {}) → list[PromptRequestPiece][source]#

Retrieves a list of PromptRequestPiece objects that have the specified memory labels.

Parameters:

  • memory_labels (dict[str, str]) – A free-form dictionary for tagging prompts with custom labels.

  • operation (These labels can be used to track all prompts sent as part of an)

  • on (score prompts based)

  • ID (the operation)

  • dictionary. (Users can define any key-value pairs according to their needs. Defaults to an empty)

Returns:

A list of PromptRequestPiece with the specified memory labels.

Return type:

list[PromptRequestPiece]

get_prompt_request_piece_by_orchestrator_id(*, orchestrator_id: str) → list[PromptRequestPiece][source]#

Retrieves a list of PromptRequestPiece objects that have the specified orchestrator ID.

Parameters:

orchestrator_id (str) – The orchestrator ID to match.

Returns:

A list of PromptRequestPiece with the specified conversation ID.

Return type:

list[PromptRequestPiece]

abstract get_prompt_request_pieces_by_id(*, prompt_ids: list[str]) → Sequence[PromptRequestPiece][source]#

Retrieves a list of PromptRequestPiece objects that have the specified prompt ids.

Parameters:

prompt_ids (list[int]) – The prompt IDs to match.

Returns:

A list of PromptRequestPiece with the specified conversation ID.

Return type:

Sequence[PromptRequestPiece]

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]

abstract insert_entries(*, entries: list[Base]) → None[source]#

Inserts multiple entries into the database.

abstract insert_entry(entry: Base) → None[source]#

Inserts an entry into the Table.

Parameters:

entry – An instance of a SQLAlchemy model to be added to the Table.

memory_embedding: MemoryEmbedding = None#
abstract query_entries(model, *, conditions: Optional = None, distinct: bool = False) → list[Base][source]#

Fetches data from the specified table model with optional conditions.

Parameters:

  • model – The SQLAlchemy model class corresponding to the table you want to query.

  • conditions – SQLAlchemy filter conditions (Optional).

  • distinct – Whether to return distinct rows only. Defaults to False.

Returns:

List of model instances representing the rows fetched from the table.

results_path: str = None#
storage_io: StorageIO = None#
abstract update_entries(*, entries: MutableSequence[Base], update_fields: dict) → bool[source]#

Updates the given entries with the specified field values.

Parameters:

  • entries (list[Base]) – A list of SQLAlchemy model instances to be updated.

  • update_fields (dict) – A dictionary of field names and their new values.

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

Contents