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_tables(*[, export_type])	Exports all table data using the specified exporter.
export_conversations(*[, orchestrator_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_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_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.
get_session()	Provides a session for database operations.
populate_prompt_piece_scores(...)	Adds scores in the database to prompt request piece objects
print_schema()
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

DEFAULT_DB_FILE_NAME
memory_embedding
results_path
results_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_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_session() → Session

Provides a session for database operations.

print_schema()

reset_database()

Drop and recreate existing tables