pyrit.memory.DuckDBMemory

class DuckDBMemory(*args, **kwargs)

Bases: MemoryInterface

A class to manage conversation memory using DuckDB 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__(*, db_path: Path | str = None, verbose: bool = False)

Methods

__init__(*[, db_path, verbose])
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_all_tables(*[, export_type])	Exports all table data using the specified exporter.
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()	Fetches all entries from the specified table and returns them as model instances.
get_all_prompt_pieces()	Fetches all entries from the specified table and returns them as model instances.
get_all_table_models()	Returns a list of all table models used in the database by inspecting the Base registry.
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.
get_session()	Provides a session for database operations.
insert_entries(*, entries)	Inserts multiple entries into the database.
insert_entry(entry)	Inserts an entry into the Table.
print_schema()
query_entries(model, *[, conditions, distinct])	Fetches data from the specified table model with optional conditions.
reset_database()	Drop and recreate existing tables
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

DEFAULT_DB_FILE_NAME
memory_embedding
results_path
storage_io

DEFAULT_DB_FILE_NAME = 'pyrit_duckdb_storage.db'

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

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

dispose_engine()

Dispose the engine and clean up resources.

export_all_tables(*, export_type: str = 'json')

Exports all table data using the specified exporter.

Iterates over all tables, retrieves their data, and exports each to a file named after the table.

Parameters:

export_type (str) – The format to export the data in (defaults to “json”).

get_all_embeddings() → list[EmbeddingDataEntry]

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

get_all_prompt_pieces() → list[PromptRequestPiece]

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

get_all_table_models() → list[Base]

Returns a list of all table models used in the database by inspecting the Base registry.

Returns:

A list of SQLAlchemy model classes.

Return type:

list[Base]

get_prompt_request_piece_by_memory_labels(*, memory_labels: dict[str, str] = {}) → list[PromptRequestPiece]

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_pieces_by_id(*, prompt_ids: list[str]) → list[PromptRequestPiece]

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

Parameters:

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

Returns:

A list of PromptRequestPiece with the specified prompt ID.

Return type:

list[PromptRequestPiece]

get_session() → Session

Provides a session for database operations.

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

Inserts multiple entries into the database.

insert_entry(entry: Base) → None

Inserts an entry into the Table.

Parameters:

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

print_schema()

query_entries(model, *, conditions: Optional = None, distinct: bool = False) → list[Base]

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 – Flag to return distinct rows (default is False).

Returns:

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

reset_database()

Drop and recreate existing tables

update_entries(*, entries: MutableSequence[Base], update_fields: dict) → bool

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.

Returns:

True if the update was successful, False otherwise.

Return type:

bool